From 11f5f0c30465176d1f4c89b60275b6ac51a57ba0 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:51:02 +0800 Subject: [PATCH 01/85] Align the doc, test, and client with the newest .env setup (image segmentation client) --- .../source/features/generative_sim/scene_engine.md | 2 +- .../scene_engine/clients/image_segmentation.py | 14 +++++++------- tests/gen_sim/scene_engine/test_clients.py | 10 +++++----- 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index 03e4b3fa4..a8792faab 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -43,7 +43,7 @@ SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL="http://host:port" SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S=30 SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS=3 SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH="/health" -SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH="/predict" +SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH="/segment_by_prompt" SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://host:port" SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 diff --git a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py index 2c56af91a..8414b9627 100644 --- a/embodichain/gen_sim/scene_engine/clients/image_segmentation.py +++ b/embodichain/gen_sim/scene_engine/clients/image_segmentation.py @@ -36,14 +36,14 @@ def __init__( timeout_s: int, max_attempts: int, health_path: str, - segment_single_object_path: str, + segment_by_prompt_path: str, session: requests.Session | None = None, ) -> None: self._base_url = base_url.rstrip("/") self._timeout_s = timeout_s self._max_attempts = max_attempts self._health_path = health_path - self._segment_single_object_path = segment_single_object_path + self._segment_by_prompt_path = segment_by_prompt_path self._session = session or requests.Session() @classmethod @@ -96,7 +96,7 @@ def segment_single_object( try: with resolved_image_path.open("rb") as image_file: response = self._session.post( - self._url(self._segment_single_object_path), + self._url(self._segment_by_prompt_path), data={"prompt": prompt}, files={"image": (resolved_image_path.name, image_file)}, timeout=self._timeout_s, @@ -138,7 +138,7 @@ def _load_dotenv_config() -> dict[str, Any]: "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S", "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH", ) try: timeout_s = int(values["SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S"]) @@ -165,7 +165,7 @@ def _load_dotenv_config() -> dict[str, Any]: string_keys = ( "SCENE_ENGINE_IMAGE_SEGMENTATION_BASE_URL", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH", ) for key in string_keys: if not values[key].strip(): @@ -176,8 +176,8 @@ def _load_dotenv_config() -> dict[str, Any]: "timeout_s": timeout_s, "max_attempts": max_attempts, "health_path": values["SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH"].strip(), - "segment_single_object_path": values[ - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH" + "segment_by_prompt_path": values[ + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH" ].strip(), } diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index 513f0c5c0..a2ccacb21 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -79,7 +79,7 @@ def test_clients_load_their_required_dotenv_values( "SCENE_ENGINE_IMAGE_SEGMENTATION_TIMEOUT_S": "30", "SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS": "2", "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH": "/health", - "SCENE_ENGINE_IMAGE_SEGMENTATION_SINGLE_OBJECT_PATH": "/predict", + "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH": "/segment_by_prompt", } llm_values = { "OPENAI_API_KEY": "test-key", @@ -107,7 +107,7 @@ def test_clients_load_their_required_dotenv_values( assert geometry_client._base_url == "http://geometry" assert geometry_client._generate_objects_path == "/objects" assert segmentation_client._base_url == "http://segment" - assert segmentation_client._segment_single_object_path == "/predict" + assert segmentation_client._segment_by_prompt_path == "/segment_by_prompt" assert llm_client_config.default_query == {"api-version": "1"} assert llm_client_config.base_url == "http://llm/v1" @@ -146,7 +146,7 @@ def test_service_health_checks_use_the_configured_health_path() -> None: timeout_s=30, max_attempts=1, health_path="/health", - segment_single_object_path="/predict", + segment_by_prompt_path="/segment_by_prompt", session=segmentation_session, ) @@ -170,7 +170,7 @@ def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) timeout_s=30, max_attempts=1, health_path="/health", - segment_single_object_path="/predict", + segment_by_prompt_path="/segment_by_prompt", session=session, ) @@ -178,7 +178,7 @@ def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) rle_mask ] assert session.post_call is not None - assert session.post_call["url"] == "http://segment/predict" + assert session.post_call["url"] == "http://segment/segment_by_prompt" assert session.post_call["data"] == {"prompt": "table"} From 4c68adf2002dcbe27a41a707c1ad7be05b9502a1 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:09:49 +0800 Subject: [PATCH 02/85] Add image generation client + update .env + update doc + test files --- .../features/generative_sim/scene_engine.md | 10 +- .../scene_engine/clients/image_generation.py | 172 ++++++++++++++++++ tests/gen_sim/scene_engine/test_clients.py | 106 ++++++++++- 3 files changed, 285 insertions(+), 3 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/clients/image_generation.py diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index a8792faab..c9f569ee2 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -29,8 +29,8 @@ python -m embodichain scene-engine \ ## Configuration -Scene Engine reads the LLM, segmentation, and geometry-generation settings -from `embodichain/gen_sim/.env`: +Scene Engine reads the LLM, segmentation, image-generation, and +geometry-generation settings from `embodichain/gen_sim/.env`: ```bash OPENAI_API_KEY="your-api-key" @@ -45,6 +45,12 @@ SCENE_ENGINE_IMAGE_SEGMENTATION_MAX_ATTEMPTS=3 SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH="/health" SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH="/segment_by_prompt" +SCENE_ENGINE_IMAGE_GENERATION_BASE_URL="http://host:port" +SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S=120 +SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH="/health" +SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH="/generate_image_by_prompt" + SCENE_ENGINE_GEOMETRY_GENERATION_BASE_URL="http://host:port" SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS=3 diff --git a/embodichain/gen_sim/scene_engine/clients/image_generation.py b/embodichain/gen_sim/scene_engine/clients/image_generation.py new file mode 100644 index 000000000..4286e26f8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/image_generation.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 + +from pathlib import Path +from typing import Any + +import requests + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +class ImageGenerationClient: + """Manage the Image Generation Server connection.""" + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + generate_image_by_prompt_path: str, + session: requests.Session | None = None, + ) -> None: + self._base_url = base_url.rstrip("/") + self._timeout_s = timeout_s + self._max_attempts = max_attempts + self._health_path = health_path + self._generate_image_by_prompt_path = generate_image_by_prompt_path + self._session = session or requests.Session() + + @classmethod + def from_dotenv(cls) -> "ImageGenerationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) + + def check_health(self) -> None: + last_error: requests.RequestException | RuntimeError | None = None + for _ in range(self._max_attempts): + try: + response = self._session.get( + self._url(self._health_path), + timeout=10, # Use a shorter timeout for avoiding long waits. + ) + response.raise_for_status() + response_data = response.json() + if ( + not isinstance(response_data, dict) + or response_data.get("ok") is not True + ): + raise RuntimeError( + "Image Generation Server health response does not contain ok=true." + ) + return + except (requests.RequestException, ValueError, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Image Generation Server health check failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def close(self) -> None: + self._session.close() + + def generate_image_by_prompt( + self, + *, + prompt: str, + output_path: str | Path, + ) -> Path: + """Generate one PNG image from ``prompt`` and save it to ``output_path``.""" + prompt = prompt.strip() + if not prompt: + raise ValueError("Image generation prompt must not be empty.") + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = self._session.post( + self._url(self._generate_image_by_prompt_path), + json={"prompt": prompt}, + timeout=self._timeout_s, + ) + response.raise_for_status() + content_type = response.headers.get("content-type", "").split(";")[0] + if content_type != "image/png": + raise RuntimeError( + "Image Generation Server response is not a PNG image." + ) + resolved_output_path.write_bytes(response.content) + return resolved_output_path + except (requests.RequestException, RuntimeError) as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Image Generation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _url(self, path: str) -> str: + return f"{self._base_url}/{path.lstrip('/')}" + + +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL", + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S", + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH", + ) + try: + timeout_s = int(values["SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S must be an integer." + ) from exc + if timeout_s < 1: + raise ValueError("SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S must be at least 1.") + + try: + max_attempts = int(values["SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS"]) + except (TypeError, ValueError) as exc: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS must be an integer." + ) from exc + if max_attempts < 1: + raise ValueError( + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS must be at least 1." + ) + + string_keys = ( + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH", + ) + for key in string_keys: + if not values[key].strip(): + raise ValueError(f"Scene Engine .env key {key} must be a non-empty string.") + + return { + "base_url": values["SCENE_ENGINE_IMAGE_GENERATION_BASE_URL"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": values["SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH"].strip(), + "generate_image_by_prompt_path": values[ + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH" + ].strip(), + } diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index a2ccacb21..948db9c73 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -23,6 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.clients import geometry_generation +from embodichain.gen_sim.scene_engine.clients import image_generation from embodichain.gen_sim.scene_engine.clients import image_segmentation from embodichain.gen_sim.scene_engine.llms import load_config @@ -30,9 +31,16 @@ class _Response: """Minimal successful HTTP response used by client unit tests.""" - def __init__(self, payload: object, *, content: bytes = b"") -> None: + def __init__( + self, + payload: object, + *, + content: bytes = b"", + headers: dict[str, str] | None = None, + ) -> None: self._payload = payload self.content = content + self.headers = headers or {} def raise_for_status(self) -> None: return None @@ -81,6 +89,13 @@ def test_clients_load_their_required_dotenv_values( "SCENE_ENGINE_IMAGE_SEGMENTATION_HEALTH_PATH": "/health", "SCENE_ENGINE_IMAGE_SEGMENTATION_BY_PROMPT_PATH": "/segment_by_prompt", } + image_generation_values = { + "SCENE_ENGINE_IMAGE_GENERATION_BASE_URL": "http://image-generation/", + "SCENE_ENGINE_IMAGE_GENERATION_TIMEOUT_S": "120", + "SCENE_ENGINE_IMAGE_GENERATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH": "/generate_image_by_prompt", + } llm_values = { "OPENAI_API_KEY": "test-key", "OPENAI_MODEL": "test-model", @@ -96,18 +111,29 @@ def test_clients_load_their_required_dotenv_values( "read_scene_engine_env_values", lambda *_: segmentation_values, ) + monkeypatch.setattr( + image_generation, + "read_scene_engine_env_values", + lambda *_: image_generation_values, + ) monkeypatch.setattr( load_config, "read_scene_engine_env_values", lambda *_: llm_values ) geometry_client = geometry_generation.GeometryGenerationClient.from_dotenv() segmentation_client = image_segmentation.ImageSegmentationClient.from_dotenv() + image_generation_client = image_generation.ImageGenerationClient.from_dotenv() llm_client_config = load_config.load_llm_config() assert geometry_client._base_url == "http://geometry" assert geometry_client._generate_objects_path == "/objects" assert segmentation_client._base_url == "http://segment" assert segmentation_client._segment_by_prompt_path == "/segment_by_prompt" + assert image_generation_client._base_url == "http://image-generation" + assert ( + image_generation_client._generate_image_by_prompt_path + == "/generate_image_by_prompt" + ) assert llm_client_config.default_query == {"api-version": "1"} assert llm_client_config.base_url == "http://llm/v1" @@ -149,12 +175,90 @@ def test_service_health_checks_use_the_configured_health_path() -> None: segment_by_prompt_path="/segment_by_prompt", session=segmentation_session, ) + image_generation_session = _Session(get_payload={"ok": True}) + image_generation_client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=image_generation_session, + ) geometry_client.check_health() segmentation_client.check_health() + image_generation_client.check_health() assert geometry_session.get_calls == [("http://geometry/health", 10)] assert segmentation_session.get_calls == [("http://segment/health", 30)] + assert image_generation_session.get_calls == [ + ("http://image-generation/health", 10) + ] + + +def test_image_generation_client_posts_prompt_and_writes_png( + tmp_path: Path, +) -> None: + png_bytes = b"\x89PNG\r\n\x1a\nimage" + + class ImageGenerationSession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + {}, + content=png_bytes, + headers={"content-type": "image/png"}, + ) + + session = ImageGenerationSession(get_payload={"ok": True}) + client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=session, + ) + + output_path = client.generate_image_by_prompt( + prompt="a red mug on a wooden table", + output_path=tmp_path / "generated.png", + ) + + assert output_path.read_bytes() == png_bytes + assert session.post_call is not None + assert session.post_call["url"] == ( + "http://image-generation/generate_image_by_prompt" + ) + assert session.post_call["json"] == {"prompt": "a red mug on a wooden table"} + + +def test_image_generation_client_rejects_non_png_response(tmp_path: Path) -> None: + class ImageGenerationSession(_Session): + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response( + {"ok": False, "error": "failed"}, + content=b'{"ok": false}', + headers={"content-type": "application/json"}, + ) + + session = ImageGenerationSession(get_payload={"ok": True}) + client = image_generation.ImageGenerationClient( + base_url="http://image-generation", + timeout_s=120, + max_attempts=1, + health_path="/health", + generate_image_by_prompt_path="/generate_image_by_prompt", + session=session, + ) + + with pytest.raises(RuntimeError, match="request failed after 1 attempts") as exc: + client.generate_image_by_prompt( + prompt="a red mug on a wooden table", + output_path=tmp_path / "generated.png", + ) + assert "response is not a PNG image" in str(exc.value.__cause__) def test_segmentation_client_posts_prompt_and_returns_rle_masks(tmp_path: Path) -> None: From 88de161799aa8e65949ff958dd73dcd348a5e364 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:21:29 +0800 Subject: [PATCH 03/85] Reformat the code, add scene edit basic framework, notice that the docs and tests may not be updated --- embodichain/gen_sim/scene_engine/cli/start.py | 26 +- .../gen_sim/scene_engine/pipeline/edit.py | 89 +++++++ .../scene_engine/pipeline/editing/__init__.py | 19 ++ .../editing/scene_edit_understanding.py | 43 ++++ .../gen_sim/scene_engine/pipeline/generate.py | 4 +- .../pipeline/generation/__init__.py | 19 ++ .../{ => generation}/scene_generation.py | 0 .../{ => generation}/scene_understanding.py | 0 .../pipeline/utils/scene_importer.py | 233 ++++++++++++++++++ tests/gen_sim/scene_engine/test_scene_edit.py | 133 ++++++++++ .../scene_engine/test_scene_understanding.py | 2 +- 11 files changed, 564 insertions(+), 4 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/edit.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py rename embodichain/gen_sim/scene_engine/pipeline/{ => generation}/scene_generation.py (100%) rename embodichain/gen_sim/scene_engine/pipeline/{ => generation}/scene_understanding.py (100%) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py create mode 100644 tests/gen_sim/scene_engine/test_scene_edit.py diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 59454b09f..edced21c4 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -21,6 +21,7 @@ from pathlib import Path from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image +from embodichain.gen_sim.scene_engine.pipeline.edit import edit_scene _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} @@ -28,6 +29,8 @@ def cli_scene_engine( image: str | Path, output_root: str | Path, + *, + edit_prompt: str | None = None, ) -> None: """Generate one scene using the required ``gen_sim/.env`` settings.""" resolved_image_path = Path(image).expanduser().resolve() @@ -41,12 +44,27 @@ def cli_scene_engine( ) resolved_output_root = Path(output_root).expanduser().resolve() + # If this scene needs editing. + if edit_prompt is not None: + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + if not resolved_output_root.is_dir() or not any(resolved_output_root.iterdir()): + raise ValueError( + "Output root must exist and contain files when edit_prompt is provided." + ) + resolved_output_root.mkdir(parents=True, exist_ok=True) generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, ) + if edit_prompt is not None: + edit_scene( + output_root=resolved_output_root, + edit_prompt=edit_prompt, + ) print("Successfully completed!") @@ -68,9 +86,15 @@ def main(argv: Sequence[str] | None = None) -> None: required=True, help="Path to the output directory", ) + parser.add_argument( + "--edit_prompt", + type=str, + default=None, + help="Optional text instruction for editing an existing output root", + ) args = parser.parse_args(argv) - cli_scene_engine(args.image, args.output_root) + cli_scene_engine(args.image, args.output_root, edit_prompt=args.edit_prompt) if __name__ == "__main__": diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py new file mode 100644 index 000000000..c055decf2 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -0,0 +1,89 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + understand_scene_edit, +) +from embodichain.utils.logger import log_info + + +def edit_scene( + *, + output_root: str | Path, + edit_prompt: str, +) -> None: + """Apply one text edit instruction to an existing Scene Engine output.""" + + # Initialize the VLM client that will interpret the edit instruction. + vlm_client = OpenAICompatibleVLM.from_dotenv() + scene_importer = SceneExportImporter(output_root=output_root) + # Validate scene_export, write scene.json, and return Scene; failures raise before editing. + scene = scene_importer.import_scene() + + # 1. Edit Understanding + log_info("Starting Edit Understanding") + edit_plan = understand_scene_edit( + scene=scene, + edit_prompt=edit_prompt, + output_root=output_root, # Has already been resolved. + vlm_client=vlm_client, + ) + log_info("Completed Edit Understanding") + + # 2. Scene Graph Update(or Initialization) + log_info("Starting Scene Graph Update") + # updated_scene_graph = update_scene_graph( + # scene=scene, + # edit_plan=edit_plan, + # output_root=output_root, + # ) + log_info("Completed Scene Graph Update") + + # 3. Prepare Objects. + log_info("Preparing Objects if necessary") + # scene = prepare_objects( + # scene=scene, + # output_root=output_root, + # ) + log_info("Completed Preparing Objects") + + # 4. Layout Editing + log_info("Starting Layout Editing") + # scene = edit_layout( + # scene=scene, + # edit_plan=edit_plan, + # scene_graph=updated_scene_graph, + # output_root=output_root, + # ) + log_info("Completed Layout Editing") + + # 5. Scene Export + # Re export the scene to the same output format, + # and delete some temporary files or folders. + log_info("Starting Scene Export") + log_info("Completed Scene Export") + + raise NotImplementedError("Scene editing is not implemented yet.") diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py new file mode 100644 index 000000000..c5db9531b --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -0,0 +1,43 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) + + +def understand_scene_edit( + *, + scene: Scene, + edit_prompt: str, + output_root: str | Path, + vlm_client: OpenAICompatibleVLM, +) -> dict[str, object]: + """Understand one text edit instruction for an existing scene.""" + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + + return { + "edit_prompt": edit_prompt, + "operations": [], + } diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 773a897bf..17e0a69dc 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -26,12 +26,12 @@ GeometryGenerationClient, ) -from embodichain.gen_sim.scene_engine.pipeline.scene_understanding import ( +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_understanding import ( understand_scene, ) from embodichain.utils.logger import log_info -from embodichain.gen_sim.scene_engine.pipeline.scene_generation import ( +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( generate_scene_and_refine, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py new file mode 100644 index 000000000..015c41510 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py similarity index 100% rename from embodichain/gen_sim/scene_engine/pipeline/scene_generation.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py similarity index 100% rename from embodichain/gen_sim/scene_engine/pipeline/scene_understanding.py rename to embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py new file mode 100644 index 000000000..716ce036d --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.utils.logger import log_info + +_Y_UP_TO_Z_UP_ROTATION = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0], + ], + dtype=float, +) +_Z_UP_TO_Y_UP_ROTATION = _Y_UP_TO_Z_UP_ROTATION.T + + +class SceneExportImporter: + """Import an editable ``Scene`` from an exported Scene Engine directory.""" + + def __init__( + self, + *, + output_root: str | Path, + ) -> None: + self.output_root = Path(output_root).expanduser().resolve() + self.scene_export_root = self.output_root / "scene_export" + self.mesh_assets_root = self.scene_export_root / "mesh_assets" + self.scene_config_path = self.scene_export_root / "scene_config.json" + self.scene_json_path = self.scene_export_root / "scene.json" + + def import_scene(self) -> Scene: + """Validate the scene export, write ``scene.json``, and return a ``Scene``.""" + # Editing only runs on an existing Scene Engine output directory. + if not self.output_root.is_dir() or not any(self.output_root.iterdir()): + raise ValueError( + "Output root must exist and contain files when edit_prompt is provided." + ) + + # The editor consumes the portable scene export and its copied GLB assets. + if not self.scene_export_root.is_dir(): + raise FileNotFoundError( + f"Scene export directory not found: {self.scene_export_root}" + ) + if not self.mesh_assets_root.is_dir(): + raise FileNotFoundError( + f"Scene mesh assets directory not found: {self.mesh_assets_root}" + ) + if not self.scene_config_path.is_file(): + raise FileNotFoundError(f"Scene config not found: {self.scene_config_path}") + + try: + scene_config = json.loads( + self.scene_config_path.read_text(encoding="utf-8") + ) + except json.JSONDecodeError as exc: + raise ValueError( + f"Scene config is not valid JSON: {self.scene_config_path}" + ) from exc + if not isinstance(scene_config, dict): + raise ValueError("Scene config must be a JSON object.") + + scene = self._scene_from_config(scene_config) + if self.scene_json_path.exists(): + self.scene_json_path.unlink() + self.scene_json_path.write_text( + json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Imported scene JSON: {self.scene_json_path}") + return scene + + def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: + """Build a y-up ``Scene`` from the z-up scene-export config.""" + # The table is the required support object for scene-edit operations. + background = scene_config.get("background", []) + if not isinstance(background, list): + raise ValueError("Scene config background must be a list.") + table_entry = next( + ( + scene_object + for scene_object in background + if isinstance(scene_object, dict) and scene_object.get("uid") == "table" + ), + None, + ) + if table_entry is None: + raise ValueError("Scene config background must contain a table entry.") + + rigid_object_entries = scene_config.get("rigid_object", []) + if not isinstance(rigid_object_entries, list): + raise ValueError("Scene config rigid_object must be a list.") + + return Scene( + objects=[ + self._scene_object_from_export_entry(table_entry, kind="table"), + *[ + self._scene_object_from_export_entry(entry, kind="asset") + for entry in rigid_object_entries + ], + ] + ) + + def _scene_object_from_export_entry( + self, + entry: object, + *, + kind: str, + ) -> SceneObject: + """Convert one z-up scene-export entry back to a y-up ``SceneObject``.""" + if not isinstance(entry, dict): + raise ValueError("Scene config entries must be objects.") + uid = entry.get("uid") + if not isinstance(uid, str) or not uid: + raise ValueError("Scene config entries must contain a valid uid.") + + glb_path = self._resolve_export_glb_path(entry, uid=uid) + pos_z_up = self._vector3( + entry.get("init_pos", [0.0, 0.0, 0.0]), + field_name=f"{uid}.init_pos", + ) + rot_z_up = self._vector3( + entry.get("init_rot", [0.0, 0.0, 0.0]), + field_name=f"{uid}.init_rot", + ) + scale = self._vector3( + entry.get("body_scale", [1.0, 1.0, 1.0]), + field_name=f"{uid}.body_scale", + ) + + pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) + rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() + rotation_y_up = ( + _Z_UP_TO_Y_UP_ROTATION @ rotation_z_up @ _Z_UP_TO_Y_UP_ROTATION.T + ) + rot_y_up = Rotation.from_matrix(rotation_y_up).as_euler("xyz", degrees=True) + + return SceneObject( + id=uid, + kind=kind, # type: ignore[arg-type] + category=uid, + name=uid, + description=str(entry.get("description") or uid), + simready_glb_path=str(glb_path), + rot=rot_y_up.tolist(), + pos=pos_y_up.tolist(), + scale=scale, + physics=ObjectPhysics( + body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] + attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), + max_convex_hull_num=max(1, int(entry.get("max_convex_hull_num", 32))), + ), + ) + + def _resolve_export_glb_path( + self, + entry: dict[str, Any], + *, + uid: str, + ) -> Path: + """Validate one exported mesh reference and return its absolute GLB path.""" + shape = entry.get("shape") + if not isinstance(shape, dict) or not isinstance(shape.get("fpath"), str): + raise ValueError(f"Scene object {uid!r} must contain shape.fpath.") + fpath = Path(shape["fpath"]) + if fpath.is_absolute(): + raise ValueError(f"Scene object {uid!r} shape.fpath must be relative.") + if fpath.suffix.lower() != ".glb": + raise ValueError(f"Scene object {uid!r} shape.fpath must point to a GLB.") + glb_path = (self.scene_export_root / fpath).resolve() + if self.scene_export_root.resolve() not in glb_path.parents: + raise ValueError( + f"Scene object {uid!r} shape.fpath must stay within " + f"{self.scene_export_root.resolve()}." + ) + if not glb_path.is_file(): + raise FileNotFoundError(f"Scene object {uid!r} GLB not found: {glb_path}") + return glb_path + + @staticmethod + def _vector3(value: object, *, field_name: str) -> list[float]: + """Validate one length-3 numeric vector.""" + if not isinstance(value, list) or len(value) != 3: + raise ValueError( + f"Scene config field {field_name!r} must be a length-3 list." + ) + vector = [float(item) for item in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"Scene config field {field_name!r} must be finite.") + return vector + + @staticmethod + def _physics_attrs(value: object) -> dict[str, float | int]: + """Validate exported physics attributes.""" + if not isinstance(value, dict) or not value: + raise ValueError("Scene object attrs must be a non-empty object.") + attrs: dict[str, float | int] = {} + for key, item in value.items(): + if not isinstance(key, str) or not isinstance(item, (float, int)): + raise ValueError("Scene object attrs must map strings to numbers.") + attrs[key] = item + return attrs + + +def import_scene_from_output_root(output_root: str | Path) -> Scene: + """Import an editable ``Scene`` from ``scene_export/scene_config.json``.""" + return SceneExportImporter(output_root=output_root).import_scene() diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py new file mode 100644 index 000000000..aef5a1462 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -0,0 +1,133 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + import_scene_from_output_root, +) + + +def _write_scene_export( + output_root: Path, + *, + include_table: bool = True, + include_asset_mesh: bool = True, +) -> None: + scene_export_root = output_root / "scene_export" + table_mesh_path = scene_export_root / "mesh_assets" / "table" / "table.glb" + asset_mesh_path = scene_export_root / "mesh_assets" / "cup" / "cup.glb" + table_mesh_path.parent.mkdir(parents=True) + asset_mesh_path.parent.mkdir(parents=True) + table_mesh_path.write_bytes(b"glTF-table") + if include_asset_mesh: + asset_mesh_path.write_bytes(b"glTF-cup") + + background = [] + if include_table: + background.append( + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/table/table.glb", + }, + "attrs": {"mass": 1.0}, + "body_type": "kinematic", + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "max_convex_hull_num": 16, + } + ) + scene_config = { + "background": background, + "rigid_object": [ + { + "uid": "cup", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/cup/cup.glb", + }, + "attrs": {"mass": 1.0}, + "body_type": "dynamic", + "init_pos": [1.0, -3.0, 2.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 2.0, 3.0], + "max_convex_hull_num": 32, + } + ], + } + (scene_export_root / "scene_config.json").write_text( + json.dumps(scene_config), + encoding="utf-8", + ) + + +def test_import_scene_from_output_root_writes_y_up_scene_json(tmp_path: Path) -> None: + _write_scene_export(tmp_path) + (tmp_path / "scene_export" / "scene.json").write_text( + '{"old": true}', + encoding="utf-8", + ) + + scene = import_scene_from_output_root(tmp_path) + scene_json = json.loads( + (tmp_path / "scene_export" / "scene.json").read_text(encoding="utf-8") + ) + + assert scene.table is not None + assert scene.table.id == "table" + assert scene.assets[0].id == "cup" + assert scene.assets[0].pos == [1.0, 2.0, 3.0] + assert scene.assets[0].scale == [1.0, 2.0, 3.0] + assert scene.assets[0].simready_glb_path == str( + (tmp_path / "scene_export" / "mesh_assets" / "cup" / "cup.glb").resolve() + ) + assert scene_json["objects"][1]["id"] == "cup" + assert scene_json["objects"][1]["pos"] == [1.0, 2.0, 3.0] + + +def test_check_scene_export_for_edit_requires_export_directories( + tmp_path: Path, +) -> None: + with pytest.raises(ValueError, match="Output root"): + import_scene_from_output_root(tmp_path) + + (tmp_path / "scene_export").mkdir() + with pytest.raises(FileNotFoundError, match="mesh assets"): + import_scene_from_output_root(tmp_path) + + +def test_check_scene_export_for_edit_requires_table(tmp_path: Path) -> None: + _write_scene_export(tmp_path, include_table=False) + + with pytest.raises(ValueError, match="table"): + import_scene_from_output_root(tmp_path) + + +def test_check_scene_export_for_edit_requires_rigid_object_glb( + tmp_path: Path, +) -> None: + _write_scene_export(tmp_path, include_asset_mesh=False) + + with pytest.raises(FileNotFoundError, match="cup"): + import_scene_from_output_root(tmp_path) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 1fdb51f7b..5d9a285c3 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -23,7 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.pipeline import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding def _response(*, asset_name: str = "cup") -> str: From 139850b8f9d05d4a5cc83587a9ec42d96010f3cf Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 12:25:19 +0800 Subject: [PATCH 04/85] Basic design of scene graph --- .../gen_sim/scene_engine/core/scene_graph.py | 352 ++++++++++++++++++ .../gen_sim/scene_engine/test_scene_graph.py | 309 +++++++++++++++ 2 files changed, 661 insertions(+) create mode 100644 embodichain/gen_sim/scene_engine/core/scene_graph.py create mode 100644 tests/gen_sim/scene_engine/test_scene_graph.py diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py new file mode 100644 index 000000000..806c5a0ee --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -0,0 +1,352 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +TABLE_OBJECT_ID = "table" + +# 9-grid table regions, treat the table as a 3x3 grid. +TableRegion = Literal[ + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", +] + +# A on B, then B is the parent node of A. +SupportRelationType = Literal["on"] + +# A PlanarRelation with B, then A and B must have the same parent node. +PlanarRelationType = Literal["left_of", "right_of", "in_front_of", "behind"] +SceneConstraintType = SupportRelationType | PlanarRelationType + + +@dataclass +class SceneGraphNode: + """One object node in the edit-time scene hierarchy.""" + + object_id: str + parent_id: str | None + parent_relation: SupportRelationType | None = None + table_region: TableRegion | None = None + + def __post_init__(self) -> None: + """Validate local node fields before graph-level checks.""" + if not self.object_id: + raise ValueError("object_id must be non-empty.") + # If the node is the table. + if self.object_id == TABLE_OBJECT_ID: + if self.parent_id is not None: + raise ValueError("table must not have a parent.") + if self.parent_relation is not None: + raise ValueError("table must not have a parent relation.") + # If the node is not the table. + elif self.parent_id is None: + raise ValueError("non-table nodes must have a parent.") + elif self.parent_relation not in {None, "on"}: + raise ValueError("non-table nodes must be on their parent.") + + def to_dict(self) -> dict[str, object]: + """Serialize this node for scene graph debugging artifacts.""" + return { + "object_id": self.object_id, + "parent_id": self.parent_id, + "parent_relation": self.parent_relation, + "table_region": self.table_region, + } + + +@dataclass +class SceneGraphRelation: + """One edit-time spatial relation between two non-table objects.""" + + source_id: str + relation: PlanarRelationType + target_id: str + + def __post_init__(self) -> None: + """Validate local relation fields before graph-level checks.""" + if not self.source_id or not self.target_id: + raise ValueError("relation endpoints must be non-empty.") + if self.source_id == self.target_id: + raise ValueError("relation endpoints must be different.") + + def to_dict(self) -> dict[str, object]: + """Serialize this planar relation for scene graph debugging artifacts.""" + return { + "source_id": self.source_id, + "relation": self.relation, + "target_id": self.target_id, + } + + +@dataclass +class SceneGraph: + """Layered support graph plus planar relations for scene editing.""" + + nodes: list[SceneGraphNode] = field(default_factory=list) + relations: list[SceneGraphRelation] = field(default_factory=list) + validate_on_refresh: bool = True # Validate after each automatic refresh. + + def __post_init__(self) -> None: + """Normalize new graphs so downstream stages see canonical constraints.""" + self.refresh() + + def refresh(self) -> None: + """Normalize the graph and optionally validate semantic constraints.""" + # First normalize then validate (if applicable). + self.normalize() + if self.validate_on_refresh: + self.validate() + + def node_by_id(self) -> dict[str, SceneGraphNode]: + """Return nodes keyed by object id, raising on duplicate ids.""" + nodes_by_id: dict[str, SceneGraphNode] = {} + for node in self.nodes: + if node.object_id in nodes_by_id: + raise ValueError(f"Duplicate scene graph node: {node.object_id}") + nodes_by_id[node.object_id] = node + return nodes_by_id + + def normalize(self) -> None: + """Materialize inverse planar relations and remove duplicates.""" + self._materialize_inverse_planar_relations() + self._deduplicate_relations() + + def layer_by_id(self) -> dict[str, int]: + """Return the layer depth of each node inferred from parent links.""" + # Build fast lookup tables before walking the table-rooted tree. + nodes_by_id = self.node_by_id() + children_by_parent = self._children_by_parent() + layers: dict[str, int] = {} + visiting: set[str] = set() + visited: set[str] = set() + + def visit(node_id: str, layer: int) -> None: + # A node already on the recursion path means the parent chain loops. + if node_id in visiting: + raise ValueError(f"Parent cycle detected at node: {node_id}") + if node_id in visited: + return + # Missing parent nodes cannot contribute a valid table-rooted layer. + if node_id not in nodes_by_id: + raise ValueError(f"Parent node does not exist: {node_id}") + + visiting.add(node_id) + layers[node_id] = layer + # Children are exactly one support level above their parent. + for child in children_by_parent.get(node_id, []): + visit(child.object_id, layer + 1) + visiting.remove(node_id) + visited.add(node_id) + + if TABLE_OBJECT_ID not in nodes_by_id: + raise ValueError("Scene graph must contain a table node.") + visit(TABLE_OBJECT_ID, 0) + return layers + + def derive_constraints(self) -> list[dict[str, str]]: + """Return support constraints plus materialized planar relations.""" + self.refresh() + constraints: list[dict[str, str]] = [] + for node in self.nodes: + if node.parent_id is None: + continue + # Parent links become direct support constraints. + constraints.append( + self._constraint_dict( + source_id=node.object_id, + relation=node.parent_relation, + target_id=node.parent_id, + ), + ) + for relation in self.relations: + # Inverse planar relations are already stored during normalization. + constraints.append( + self._constraint_dict( + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ), + ) + return self._deduplicate_constraints(constraints) + + def validate(self) -> None: + """Validate hierarchy, table regions, and planar relation constraints.""" + # id -> Node mapping. + nodes_by_id = self.node_by_id() + # Table node must exist. + if TABLE_OBJECT_ID not in nodes_by_id: + raise ValueError("Scene graph must contain a table node.") + + # Validate the table-rooted support tree before checking sibling relations. + for node in self.nodes: + if node.object_id == TABLE_OBJECT_ID: + if node.table_region is not None: + raise ValueError("table must not have a table_region.") + continue + parent = nodes_by_id.get(node.parent_id) + # Parent must exist, except for the table. (root node) + if parent is None: + raise ValueError(f"Parent node does not exist: {node.parent_id}") + if node.parent_relation is None: + raise ValueError( + f"Node {node.object_id} must define its parent relation." + ) + # Table regions are only valid for objects directly on the table. + if node.table_region is not None and node.parent_id != TABLE_OBJECT_ID: + raise ValueError("table_region is only valid for objects on the table.") + # Get id -> layer mapping. + layers = self.layer_by_id() + if len(layers) != len(nodes_by_id): + raise ValueError("All scene graph nodes must be reachable from the table.") + # Validate planar relations between nodes with the same parent. + for relation in self.relations: + source = nodes_by_id.get(relation.source_id) + target = nodes_by_id.get(relation.target_id) + if source is None or target is None: + raise ValueError("Planar relation endpoint does not exist.") + if source.parent_id != target.parent_id: + raise ValueError("Planar relation endpoints must share one parent.") + if source.parent_relation != "on" or target.parent_relation != "on": + raise ValueError("Planar relation endpoints must be on their parent.") + # Validate and imply planar relation. + self._validate_planar_relation_conflicts() + + def to_dict(self) -> dict[str, object]: + """Serialize the normalized graph state.""" + self.refresh() + return { + "nodes": [node.to_dict() for node in self.nodes], + "relations": [relation.to_dict() for relation in self.relations], + } + + def _children_by_parent(self) -> dict[str, list[SceneGraphNode]]: + children_by_parent: dict[str, list[SceneGraphNode]] = {} + for node in self.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node) + return children_by_parent + + def _deduplicate_relations(self) -> None: + """Remove duplicate planar relations while preserving the first occurrence.""" + deduplicated: list[SceneGraphRelation] = [] + seen: set[tuple[str, PlanarRelationType, str]] = set() + for relation in self.relations: + key = (relation.source_id, relation.relation, relation.target_id) + # Only identical triples are duplicates; inverse relations are both retained. + if key in seen: + continue + seen.add(key) + deduplicated.append(relation) + self.relations = deduplicated + + def _materialize_inverse_planar_relations(self) -> None: + """Add the inverse of every planar relation to the graph.""" + inverse_relations = [ + SceneGraphRelation( + source_id=relation.target_id, + relation=self._inverse_planar_relation(relation.relation), + target_id=relation.source_id, + ) + for relation in self.relations + ] + self.relations.extend(inverse_relations) + + def _deduplicate_constraints( + self, + constraints: list[dict[str, str]], + ) -> list[dict[str, str]]: + deduplicated: list[dict[str, str]] = [] + seen: set[tuple[str, SceneConstraintType, str]] = set() + for constraint in constraints: + key = ( + constraint["source_id"], + constraint["relation"], + constraint["target_id"], + ) + if key in seen: + continue + seen.add(key) + deduplicated.append(constraint) + return deduplicated + + def _constraint_dict( + self, + *, + source_id: str, + relation: SceneConstraintType, + target_id: str, + ) -> dict[str, str]: + return { + "source_id": source_id, + "relation": relation, + "target_id": target_id, + } + + def _validate_planar_relation_conflicts(self) -> None: + implied_relations: dict[tuple[str, str], PlanarRelationType] = {} + for relation in self.relations: + self._add_implied_planar_relation( + implied_relations, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ) + self._add_implied_planar_relation( + implied_relations, + source_id=relation.target_id, + relation=self._inverse_planar_relation(relation.relation), + target_id=relation.source_id, + ) + + def _add_implied_planar_relation( + self, + implied_relations: dict[tuple[str, str], PlanarRelationType], + *, + source_id: str, + relation: PlanarRelationType, + target_id: str, + ) -> None: + key = (source_id, target_id) + existing_relation = implied_relations.get(key) + if existing_relation is not None and existing_relation != relation: + raise ValueError( + f"Conflicting planar relations: {source_id} " + f"{existing_relation} and {relation} {target_id}" + ) + implied_relations[key] = relation + + @classmethod + def _inverse_planar_relation( + cls, + relation: PlanarRelationType, + ) -> PlanarRelationType: + if relation == "left_of": + return "right_of" + if relation == "right_of": + return "left_of" + if relation == "in_front_of": + return "behind" + return "in_front_of" diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py new file mode 100644 index 000000000..970a3fc68 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -0,0 +1,309 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) + + +def test_scene_graph_accepts_layered_on_relations() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + table_region="center", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + table_region="right_center", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + graph.validate() + assert graph.layer_by_id()["spoon"] == 2 + + +def test_scene_graph_rejects_planar_relations_without_common_parent() -> None: + with pytest.raises(ValueError, match="share one parent"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="right_of", + target_id="spoon", + ), + ], + ) + + +def test_scene_graph_rejects_conflicting_planar_relations() -> None: + with pytest.raises(ValueError, match="Conflicting planar relations"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + SceneGraphRelation( + source_id="cup", + relation="left_of", + target_id="plate", + ), + ], + ) + + +def test_scene_graph_requires_explicit_parent_relation() -> None: + with pytest.raises(ValueError, match="parent relation"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + ), + ], + ) + + +def test_scene_graph_can_skip_validation_during_refresh() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + ), + ], + validate_on_refresh=False, + ) + + with pytest.raises(ValueError, match="parent relation"): + graph.validate() + + +def test_scene_graph_rejects_unsupported_parent_relation() -> None: + with pytest.raises(ValueError, match="must be on their parent"): + SceneGraphNode( + object_id="orange", + parent_id="box", + parent_relation="inside", + ) + + +def test_scene_graph_derives_layers_from_parent_links() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + ), + ], + ) + + assert graph.layer_by_id() == { + "table": 0, + "plate": 1, + "spoon": 2, + } + + +def test_scene_graph_layer_by_id_requires_table_root() -> None: + graph = SceneGraph(nodes=[], validate_on_refresh=False) + + with pytest.raises(ValueError, match="table node"): + graph.layer_by_id() + + +def test_scene_graph_rejects_table_region_for_non_table_parent() -> None: + with pytest.raises(ValueError, match="only valid for objects on the table"): + SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="plate", + parent_relation="on", + table_region="center", + ), + ], + ) + + +def test_scene_graph_derives_support_and_inverse_planar_constraints() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + constraints = graph.derive_constraints() + + assert constraints == [ + {"source_id": "plate", "relation": "on", "target_id": "table"}, + {"source_id": "cup", "relation": "on", "target_id": "table"}, + {"source_id": "plate", "relation": "left_of", "target_id": "cup"}, + {"source_id": "cup", "relation": "right_of", "target_id": "plate"}, + ] + + +def test_scene_graph_materializes_inverse_planar_relations() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + ], + relations=[ + SceneGraphRelation( + source_id="plate", + relation="left_of", + target_id="cup", + ), + ], + ) + + assert [relation.to_dict() for relation in graph.relations] == [ + {"source_id": "plate", "relation": "left_of", "target_id": "cup"}, + {"source_id": "cup", "relation": "right_of", "target_id": "plate"}, + ] + + +def test_scene_graph_to_dict_serializes_graph_state() -> None: + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + table_region="center", + ), + ], + ) + + graph_dict = graph.to_dict() + + assert graph_dict == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + }, + { + "object_id": "plate", + "parent_id": "table", + "parent_relation": "on", + "table_region": "center", + }, + ], + "relations": [], + } From ecaacda1f8ef0bac9b1aba417ca1316c7529b2ef Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 15:41:40 +0800 Subject: [PATCH 05/85] i 1. modify cli/start logic 2. add xy position info in scene object data sturcture 3. add some test files 4. modify scene export and scene import, thus the scene graph will be export and import automatically --- embodichain/gen_sim/scene_engine/cli/start.py | 37 ++--- .../gen_sim/scene_engine/core/scene_object.py | 2 + .../gen_sim/scene_engine/pipeline/edit.py | 28 ++-- .../editing/scene_edit_understanding.py | 5 + .../gen_sim/scene_engine/pipeline/generate.py | 4 +- .../pipeline/generation/scene_generation.py | 26 ++- .../generation/scene_understanding.py | 37 ++++- .../pipeline/utils/scene_exporter.py | 14 ++ .../pipeline/utils/scene_importer.py | 148 +++++++++++++++++- .../test_scene_core_and_export.py | 66 +++++++- tests/gen_sim/scene_engine/test_scene_edit.py | 2 + .../scene_engine/test_scene_engine_config.py | 52 ++++++ .../scene_engine/test_scene_understanding.py | 44 ++++++ 13 files changed, 419 insertions(+), 46 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index edced21c4..6d46f32dc 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -27,12 +27,25 @@ def cli_scene_engine( - image: str | Path, + image: str | Path | None, output_root: str | Path, *, edit_prompt: str | None = None, ) -> None: - """Generate one scene using the required ``gen_sim/.env`` settings.""" + """Generate a scene from an image, edit an export, or do both in sequence.""" + resolved_output_root = Path(output_root).expanduser().resolve() + if edit_prompt is not None: + edit_prompt = edit_prompt.strip() + if not edit_prompt: + raise ValueError("Edit prompt must not be empty.") + + if image is None: + if edit_prompt is None: + raise ValueError("Provide --image, --edit_prompt, or both.") + edit_scene(output_root=resolved_output_root, edit_prompt=edit_prompt) + print("Successfully completed!") + return + resolved_image_path = Path(image).expanduser().resolve() if not resolved_image_path.exists(): raise FileNotFoundError(f"Image input not found: {resolved_image_path}") @@ -43,19 +56,7 @@ def cli_scene_engine( "Image input must have one of these extensions: .jpg, .jpeg, .png" ) - resolved_output_root = Path(output_root).expanduser().resolve() - # If this scene needs editing. - if edit_prompt is not None: - edit_prompt = edit_prompt.strip() - if not edit_prompt: - raise ValueError("Edit prompt must not be empty.") - if not resolved_output_root.is_dir() or not any(resolved_output_root.iterdir()): - raise ValueError( - "Output root must exist and contain files when edit_prompt is provided." - ) - resolved_output_root.mkdir(parents=True, exist_ok=True) - generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, @@ -71,14 +72,14 @@ def cli_scene_engine( def main(argv: Sequence[str] | None = None) -> None: parser = argparse.ArgumentParser( prog="embodichain scene-engine", - description="Generate a Scene Engine export from one input image.", + description="Generate a Scene Engine export, edit one, or do both.", epilog="Service settings are read from embodichain/gen_sim/.env.", ) parser.add_argument( "--image", type=str, - required=True, - help="Path to the required input image file (.jpg, .jpeg, or .png)", + required=False, + help="Optional input image file (.jpg, .jpeg, or .png)", ) parser.add_argument( "--output_root", @@ -90,7 +91,7 @@ def main(argv: Sequence[str] | None = None) -> None: "--edit_prompt", type=str, default=None, - help="Optional text instruction for editing an existing output root", + help="Text instruction for editing an existing or newly generated output root", ) args = parser.parse_args(argv) diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index d9a837e87..0a8ed3aea 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -66,6 +66,7 @@ class SceneObject: rot: list[float] | None = None # Final y-up Euler XYZ rotation in degrees. pos: list[float] | None = None # Final y-up world position in metres. scale: list[float] | None = None # Final y-up object scale. + center_xy: list[float] | None = None # Z-up table-frame XY AABB center. physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. def to_dict(self) -> dict[str, object]: @@ -81,5 +82,6 @@ def to_dict(self) -> dict[str, object]: "rot": self.rot, "pos": self.pos, "scale": self.scale, + "center_xy": self.center_xy, "physics": self.physics.to_dict() if self.physics is not None else None, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index c055decf2..3b6c95e02 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -16,6 +16,7 @@ from __future__ import annotations +import json from pathlib import Path from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( @@ -41,11 +42,19 @@ def edit_scene( vlm_client = OpenAICompatibleVLM.from_dotenv() scene_importer = SceneExportImporter(output_root=output_root) # Validate scene_export, write scene.json, and return Scene; failures raise before editing. - scene = scene_importer.import_scene() + scene, scene_graph = scene_importer.import_scene_and_graph() + print( + json.dumps( + {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, + indent=2, + ensure_ascii=False, + ) + ) # 1. Edit Understanding + # And update or initialize the scene graph log_info("Starting Edit Understanding") - edit_plan = understand_scene_edit( + updated_scene_graph = understand_scene_edit( scene=scene, edit_prompt=edit_prompt, output_root=output_root, # Has already been resolved. @@ -53,16 +62,7 @@ def edit_scene( ) log_info("Completed Edit Understanding") - # 2. Scene Graph Update(or Initialization) - log_info("Starting Scene Graph Update") - # updated_scene_graph = update_scene_graph( - # scene=scene, - # edit_plan=edit_plan, - # output_root=output_root, - # ) - log_info("Completed Scene Graph Update") - - # 3. Prepare Objects. + # 2. Prepare Objects. log_info("Preparing Objects if necessary") # scene = prepare_objects( # scene=scene, @@ -70,7 +70,7 @@ def edit_scene( # ) log_info("Completed Preparing Objects") - # 4. Layout Editing + # 3. Layout Editing log_info("Starting Layout Editing") # scene = edit_layout( # scene=scene, @@ -80,7 +80,7 @@ def edit_scene( # ) log_info("Completed Layout Editing") - # 5. Scene Export + # 4. Scene Export # Re export the scene to the same output format, # and delete some temporary files or folders. log_info("Starting Scene Export") diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index c5db9531b..48dda64bf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -20,6 +20,11 @@ from pathlib import Path from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + TABLE_OBJECT_ID, +) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 17e0a69dc..32aa6550f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -51,7 +51,7 @@ def generate_scene_from_image( # 1. Scene Understanding log_info("Starting Scene Understanding") - scene = understand_scene( + scene, scene_graph = understand_scene( scene=scene, image_path=image_path, output_root=resolved_output_root, @@ -69,6 +69,7 @@ def generate_scene_from_image( image_path=image_path, output_root=resolved_output_root, scene=scene, + scene_graph=scene_graph, geometry_generation_client=geometry_generation_client, ) finally: @@ -79,6 +80,7 @@ def generate_scene_from_image( log_info("Starting Scene Export") scene_exporter = SceneExporter( scene=scene, + scene_graph=scene_graph, output_root=resolved_output_root, ) scene_exporter.export() diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index d14ce364f..1cfe3f592 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -28,6 +28,7 @@ GeometryGenerationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( AssetsGroupSupportClamp, @@ -62,11 +63,14 @@ def generate_scene_and_refine( image_path: str | Path, output_root: str | Path, scene: Scene, + scene_graph: SceneGraph, *, geometry_generation_client: GeometryGenerationClient, ) -> Scene: resolved_image_path = _validate_image_path(image_path) + # Validate the scene graph before layout refinement consumes it. + scene_graph.validate() # Create stage output directory. stage_output_root = Path(output_root).expanduser().resolve() / "scene_generation" if stage_output_root.exists(): @@ -195,16 +199,18 @@ def _generate_coarse_results_from_masks( return None -def _update_scene_final_y_up_layout( +def _update_scene_final_y_up_layout_and_z_up_centers( *, scene: Scene, table_layout: dict[str, object], assets_layout: list[dict[str, object]], + geometry_root: str | Path, ) -> None: - """Copy final y-up layout values into the matching table and asset objects.""" + """Write final y-up layouts and z-up XY centers into the scene.""" if scene.table is None: raise ValueError("Cannot update a final layout without a table.") + # Keep final poses in the y-up layout convention used by exported GLBs. _copy_y_up_layout_to_scene_object(scene.table, table_layout) assets_by_id = {asset.id: asset for asset in scene.assets} layout_ids = set() @@ -223,6 +229,17 @@ def _update_scene_final_y_up_layout( f"Final layout is missing scene assets: {sorted(missing_assets)}." ) + # Measure final geometry in z-up so scene edits can compare tabletop XY positions. + table_mesh, assets_aabb_corners_by_id = _measure_table_and_assets_in_z_up_world( + table_layout=table_layout, + assets_layout=assets_layout, + geometry_root=geometry_root, + ) + # Persist AABB centers for future scene-edit object disambiguation. + scene.table.center_xy = table_mesh.bounds[:, :2].mean(axis=0).tolist() + for asset in scene.assets: + asset.center_xy = assets_aabb_corners_by_id[asset.id].mean(axis=0).tolist() + def _copy_y_up_layout_to_scene_object( scene_object: SceneObject, @@ -407,11 +424,12 @@ def _layout_refinement( ) refined_assets_layout = gravity_settler.settle() - # Update the scene data structure with the final y-up layout values. - _update_scene_final_y_up_layout( + # Update the scene data structure with the final layout and spatial metadata. + _update_scene_final_y_up_layout_and_z_up_centers( scene=scene, table_layout=refined_table_layout, assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, ) return refined_table_layout, refined_assets_layout diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f91c665e6..00a9ea250 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -29,6 +29,11 @@ ImageSegmentationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + TABLE_OBJECT_ID, +) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, @@ -150,7 +155,7 @@ def understand_scene( *, vlm_client: OpenAICompatibleVLM, json_max_attempts: int = 3, -) -> Scene: +) -> tuple[Scene, SceneGraph]: resolved_image_path = _validate_image_path(image_path) # The output in this stage will keep a JSON which contains @@ -181,12 +186,40 @@ def understand_scene( finally: image_segmentation_client.close() # Kill the session to avoid resource leaks. + # Use the segmented image to initialize the scene graph + # with the help of the VLM client. + # But at here, we do with the simplest way (hard code). + scene_graph = _initialize_scene_graph_from_segmented_scene(scene) + # Write the Updated scene JSON for debugging. (stage_output_root / "scene.json").write_text( json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) - return scene + (stage_output_root / "scene_graph.json").write_text( + json.dumps(scene_graph.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return scene, scene_graph + + +def _initialize_scene_graph_from_segmented_scene(scene: Scene) -> SceneGraph: + """Build the initial graph assuming every segmented asset rests on the table.""" + if scene.table is None: + raise ValueError("Cannot initialize a scene graph without a table.") + return SceneGraph( + nodes=[ + SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), + *[ + SceneGraphNode( + object_id=asset.id, + parent_id=TABLE_OBJECT_ID, + parent_relation="on", + ) + for asset in scene.assets + ], + ], + ) def _analyze_image_objects( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index fb66c30c4..391ccd156 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -26,6 +26,7 @@ from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.utils.logger import log_info @@ -46,12 +47,15 @@ def __init__( self, *, scene: Scene, + scene_graph: SceneGraph, output_root: str | Path, ) -> None: self.scene = scene + self.scene_graph = scene_graph self.output_root = Path(output_root).expanduser().resolve() self.export_root = self.output_root / "scene_export" self.scene_config_path: Path | None = None + self.scene_graph_path: Path | None = None def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -72,6 +76,9 @@ def export(self) -> Path: object_ids = [scene_object.id for scene_object in scene_objects] if len(set(object_ids)) != len(object_ids): raise ValueError("Scene export requires unique table and asset ids.") + self.scene_graph.validate() + if set(self.scene_graph.node_by_id()) != set(object_ids): + raise ValueError("Scene graph nodes must match exported scene object ids.") exported_entries = { scene_object.id: self._copy_scene_object_to_assets( @@ -106,6 +113,12 @@ def export(self) -> Path: encoding="utf-8", ) log_info(f"Exported scene config: {self.scene_config_path}") + self.scene_graph_path = self.export_root / "scene_graph.json" + self.scene_graph_path.write_text( + json.dumps(self.scene_graph.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene graph: {self.scene_graph_path}") return self.scene_config_path @staticmethod @@ -179,6 +192,7 @@ def _scene_object_config( # Do not permute this scale: it belongs to the original y-up GLB, # which SimulationManager itself converts to z-up. "body_scale": scale_y_up, + "center_xy": scene_object.center_xy, "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 716ce036d..1ae6cdb1b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -24,6 +24,11 @@ from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, @@ -53,10 +58,28 @@ def __init__( self.scene_export_root = self.output_root / "scene_export" self.mesh_assets_root = self.scene_export_root / "mesh_assets" self.scene_config_path = self.scene_export_root / "scene_config.json" + self.scene_graph_path = self.scene_export_root / "scene_graph.json" self.scene_json_path = self.scene_export_root / "scene.json" def import_scene(self) -> Scene: """Validate the scene export, write ``scene.json``, and return a ``Scene``.""" + scene = self._load_scene() + self._write_scene_json(scene) + return scene + + def import_scene_and_graph(self) -> tuple[Scene, SceneGraph]: + """Import a scene and graph after validating the complete edit input.""" + scene = self._load_scene() + scene_graph = self._load_scene_graph() + if set(scene_graph.node_by_id()) != { + scene_object.id for scene_object in scene.objects + }: + raise ValueError("Scene graph nodes must match imported scene object ids.") + self._write_scene_json(scene) + return scene, scene_graph + + def _load_scene(self) -> Scene: + """Validate the exported scene files and restore the ``Scene`` data.""" # Editing only runs on an existing Scene Engine output directory. if not self.output_root.is_dir() or not any(self.output_root.iterdir()): raise ValueError( @@ -86,15 +109,29 @@ def import_scene(self) -> Scene: if not isinstance(scene_config, dict): raise ValueError("Scene config must be a JSON object.") - scene = self._scene_from_config(scene_config) - if self.scene_json_path.exists(): - self.scene_json_path.unlink() + return self._scene_from_config(scene_config) + + def _load_scene_graph(self) -> SceneGraph: + """Read and validate the exported scene graph.""" + if not self.scene_graph_path.is_file(): + raise FileNotFoundError(f"Scene graph not found: {self.scene_graph_path}") + try: + scene_graph_data = json.loads( + self.scene_graph_path.read_text(encoding="utf-8") + ) + except json.JSONDecodeError as exc: + raise ValueError( + f"Scene graph is not valid JSON: {self.scene_graph_path}" + ) from exc + return self._scene_graph_from_data(scene_graph_data) + + def _write_scene_json(self, scene: Scene) -> None: + """Write the restored scene debugging artifact after validation succeeds.""" self.scene_json_path.write_text( json.dumps(scene.to_dict(), indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) log_info(f"Imported scene JSON: {self.scene_json_path}") - return scene def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: """Build a y-up ``Scene`` from the z-up scene-export config.""" @@ -127,6 +164,88 @@ def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: ] ) + @staticmethod + def _scene_graph_from_data(value: object) -> SceneGraph: + """Build a validated ``SceneGraph`` from exported graph JSON.""" + if not isinstance(value, dict) or set(value) != {"nodes", "relations"}: + raise ValueError("Scene graph must contain exactly nodes and relations.") + nodes_value = value["nodes"] + relations_value = value["relations"] + if not isinstance(nodes_value, list) or not isinstance(relations_value, list): + raise ValueError("Scene graph nodes and relations must be lists.") + + nodes = [ + SceneExportImporter._scene_graph_node_from_data(node) + for node in nodes_value + ] + relations = [ + SceneExportImporter._scene_graph_relation_from_data(relation) + for relation in relations_value + ] + return SceneGraph(nodes=nodes, relations=relations) + + @staticmethod + def _scene_graph_node_from_data(value: object) -> SceneGraphNode: + if not isinstance(value, dict) or set(value) != { + "object_id", + "parent_id", + "parent_relation", + "table_region", + }: + raise ValueError("Scene graph nodes must use the serialized node schema.") + object_id = value["object_id"] + parent_id = value["parent_id"] + parent_relation = value["parent_relation"] + table_region = value["table_region"] + if not isinstance(object_id, str) or not isinstance( + parent_id, (str, type(None)) + ): + raise ValueError("Scene graph node ids must be strings or null.") + if parent_relation not in {None, "on"}: + raise ValueError("Scene graph parent_relation must be 'on' or null.") + if table_region not in { + None, + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", + }: + raise ValueError("Scene graph table_region is invalid.") + return SceneGraphNode( + object_id=object_id, + parent_id=parent_id, + parent_relation=parent_relation, + table_region=table_region, + ) + + @staticmethod + def _scene_graph_relation_from_data(value: object) -> SceneGraphRelation: + if not isinstance(value, dict) or set(value) != { + "source_id", + "relation", + "target_id", + }: + raise ValueError( + "Scene graph relations must use the serialized relation schema." + ) + source_id = value["source_id"] + relation = value["relation"] + target_id = value["target_id"] + if not isinstance(source_id, str) or not isinstance(target_id, str): + raise ValueError("Scene graph relation ids must be strings.") + if relation not in {"left_of", "right_of", "in_front_of", "behind"}: + raise ValueError("Scene graph relation is invalid.") + return SceneGraphRelation( + source_id=source_id, + relation=relation, + target_id=target_id, + ) + def _scene_object_from_export_entry( self, entry: object, @@ -153,6 +272,9 @@ def _scene_object_from_export_entry( entry.get("body_scale", [1.0, 1.0, 1.0]), field_name=f"{uid}.body_scale", ) + center_xy = entry.get("center_xy") + if center_xy is not None: + center_xy = self._vector2(center_xy, field_name=f"{uid}.center_xy") pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() @@ -171,6 +293,7 @@ def _scene_object_from_export_entry( rot=rot_y_up.tolist(), pos=pos_y_up.tolist(), scale=scale, + center_xy=center_xy, physics=ObjectPhysics( body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), @@ -193,6 +316,11 @@ def _resolve_export_glb_path( raise ValueError(f"Scene object {uid!r} shape.fpath must be relative.") if fpath.suffix.lower() != ".glb": raise ValueError(f"Scene object {uid!r} shape.fpath must point to a GLB.") + expected_fpath = Path("mesh_assets") / uid / f"{uid}.glb" + if fpath != expected_fpath: + raise ValueError( + f"Scene object {uid!r} shape.fpath must be {expected_fpath.as_posix()!r}." + ) glb_path = (self.scene_export_root / fpath).resolve() if self.scene_export_root.resolve() not in glb_path.parents: raise ValueError( @@ -215,6 +343,18 @@ def _vector3(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be finite.") return vector + @staticmethod + def _vector2(value: object, *, field_name: str) -> list[float]: + """Validate one length-2 numeric vector.""" + if not isinstance(value, list) or len(value) != 2: + raise ValueError( + f"Scene config field {field_name!r} must be a length-2 list." + ) + vector = [float(item) for item in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"Scene config field {field_name!r} must be finite.") + return vector + @staticmethod def _physics_attrs(value: object) -> dict[str, float | int]: """Validate exported physics attributes.""" diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 56c12af1c..b7b614d49 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -24,11 +24,18 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) def _scene_object( @@ -60,6 +67,24 @@ def _physics(body_type: str) -> ObjectPhysics: ) +def _scene_graph(scene: Scene) -> SceneGraph: + if scene.table is None: + raise ValueError("Test scene must contain a table.") + return SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + *[ + SceneGraphNode( + object_id=asset.id, + parent_id="table", + parent_relation="on", + ) + for asset in scene.assets + ], + ] + ) + + def test_scene_returns_one_table_and_ordered_assets() -> None: table = _scene_object(object_id="table", kind="table") asset = _scene_object(object_id="cup", kind="asset") @@ -120,9 +145,12 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No glb_path=asset_glb, physics=_physics("dynamic"), ) + asset.center_xy = [0.25, -0.5] + scene = Scene(objects=[table, asset]) export_path = SceneExporter( - scene=Scene(objects=[table, asset]), + scene=scene, + scene_graph=_scene_graph(scene), output_root=tmp_path / "output", ).export() exported = json.loads(export_path.read_text(encoding="utf-8")) @@ -136,7 +164,31 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert entry["body_type"] == "dynamic" assert entry["init_pos"] == [1.0, -3.0, 2.0] assert entry["body_scale"] == [1.0, 2.0, 3.0] + assert entry["center_xy"] == [0.25, -0.5] assert np.allclose(entry["init_rot"], [0.0, 0.0, 0.0]) + assert json.loads((export_path.parent / "scene_graph.json").read_text()) == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + }, + { + "object_id": "cup", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + }, + ], + "relations": [], + } + + imported_scene, imported_graph = SceneExportImporter( + output_root=tmp_path / "output" + ).import_scene_and_graph() + assert [asset.id for asset in imported_scene.assets] == ["cup"] + assert imported_graph.to_dict() == _scene_graph(scene).to_dict() def test_scene_export_requires_final_physics(tmp_path: Path) -> None: @@ -145,7 +197,13 @@ def test_scene_export_requires_final_physics(tmp_path: Path) -> None: table = _scene_object(object_id="table", kind="table", glb_path=glb_path) with pytest.raises(ValueError, match="no SimReady physics"): - SceneExporter(scene=Scene(objects=[table]), output_root=tmp_path).export() + SceneExporter( + scene=Scene(objects=[table]), + scene_graph=SceneGraph( + nodes=[SceneGraphNode(object_id="table", parent_id=None)] + ), + output_root=tmp_path, + ).export() def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: @@ -165,7 +223,9 @@ def test_scene_export_rejects_backslash_in_object_id(tmp_path: Path) -> None: ) with pytest.raises(ValueError, match="not safe for a GLB filename"): + scene = Scene(objects=[table, unsafe_asset]) SceneExporter( - scene=Scene(objects=[table, unsafe_asset]), + scene=scene, + scene_graph=_scene_graph(scene), output_root=tmp_path / "output", ).export() diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py index aef5a1462..fd1a9e933 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit.py +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -72,6 +72,7 @@ def _write_scene_export( "init_pos": [1.0, -3.0, 2.0], "init_rot": [0.0, 0.0, 0.0], "body_scale": [1.0, 2.0, 3.0], + "center_xy": [1.0, -3.0], "max_convex_hull_num": 32, } ], @@ -99,6 +100,7 @@ def test_import_scene_from_output_root_writes_y_up_scene_json(tmp_path: Path) -> assert scene.assets[0].id == "cup" assert scene.assets[0].pos == [1.0, 2.0, 3.0] assert scene.assets[0].scale == [1.0, 2.0, 3.0] + assert scene.assets[0].center_xy == [1.0, -3.0] assert scene.assets[0].simready_glb_path == str( (tmp_path / "scene_export" / "mesh_assets" / "cup" / "cup.glb").resolve() ) diff --git a/tests/gen_sim/scene_engine/test_scene_engine_config.py b/tests/gen_sim/scene_engine/test_scene_engine_config.py index 01914e144..3754210d9 100644 --- a/tests/gen_sim/scene_engine/test_scene_engine_config.py +++ b/tests/gen_sim/scene_engine/test_scene_engine_config.py @@ -87,6 +87,58 @@ def generate_scene(*, image_path: Path, output_root: Path) -> None: } +def test_scene_engine_cli_edits_existing_output_without_an_image( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + captured: dict[str, object] = {} + + def edit_scene(*, output_root: Path, edit_prompt: str) -> None: + captured["output_root"] = output_root + captured["edit_prompt"] = edit_prompt + + monkeypatch.setattr(start, "edit_scene", edit_scene) + output_root = tmp_path / "existing_output" + + start.cli_scene_engine(None, output_root, edit_prompt="move the cup right") + + assert captured == { + "output_root": output_root.resolve(), + "edit_prompt": "move the cup right", + } + + +def test_scene_engine_cli_generates_then_edits_when_both_inputs_exist( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + call_order: list[str] = [] + + def generate_scene(*, image_path: Path, output_root: Path) -> None: + call_order.append("generate") + + def edit_scene(*, output_root: Path, edit_prompt: str) -> None: + call_order.append("edit") + + monkeypatch.setattr(start, "generate_scene_from_image", generate_scene) + monkeypatch.setattr(start, "edit_scene", edit_scene) + + start.cli_scene_engine( + image_path, + tmp_path / "output", + edit_prompt="move the cup right", + ) + + assert call_order == ["generate", "edit"] + + +def test_scene_engine_cli_requires_an_image_or_edit_prompt(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="--image, --edit_prompt, or both"): + start.cli_scene_engine(None, tmp_path / "output") + + @pytest.mark.parametrize("image_name", ["missing.png", "scene.gif"]) def test_scene_engine_cli_rejects_invalid_image_inputs( tmp_path: Path, diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 5d9a285c3..9f5ae3884 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -23,6 +23,7 @@ import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding @@ -83,3 +84,46 @@ def complete(self, **_: object) -> str: assert scene.table is not None assert [asset.id for asset in scene.assets] == ["cup_001"] + + +def test_initial_scene_graph_places_every_asset_on_table() -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="cup_001", + kind="asset", + category="cup", + name="blue cup", + description="A blue cup.", + ), + ], + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene + ) + + assert scene_graph.to_dict() == { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + }, + { + "object_id": "cup_001", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + }, + ], + "relations": [], + } From 51e1c64d131fa3b322e035e44f941a222ee1a115 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:25:40 +0800 Subject: [PATCH 06/85] Fixed the prompt of scne understanding: delete the location in description, and add check code --- .../generation/scene_understanding.py | 29 +++++++++++-------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index 00a9ea250..c69ec3719 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -67,8 +67,8 @@ 3. Do not merge objects merely resting on another object. A mug on a table and the table are separate entries. 4. List every visible physical instance separately. If two objects look alike, - keep the same category and name, but distinguish them in description using - location. Do not add location to name. + keep the same category and name. Do not encode location or spatial context + in any semantic field. 5. category is a lower-case singular snake_case class, such as mug, book, potted_plant, or coffee_table. It must not contain color or material. 6. name contains only color, material, texture, shape, and object description. @@ -78,8 +78,9 @@ shape, and visible structural details. Do not mention image coverage, image position, camera framing, or viewpoint. For example, do not write "occupying most of the image" or "at the center of the image". -8. For assets, description may include all visible details, including location - and spatial context. +8. For assets, description contains only visible category, material, color, + texture, shape, and structural details. Do not mention location, the table, + or any relationship to another object. Return JSON only: no Markdown, comments, or prose outside this exact schema: { @@ -92,14 +93,14 @@ { "category": "mug", "name": "blue ceramic mug", - "description": "small blue ceramic mug on the left side of the table" + "description": "small blue ceramic mug with a curved handle" } ] } -For two identical blue mugs, output two asset entries with the same category and -name, and use their descriptions to state left/right or front/back. Do not -infer objects that are not visible. Use an empty assets array when no objects -are visible. Every field must be a non-empty string.""" +For two identical blue mugs, output two asset entries with the same category, +name, and description. Do not infer objects that are not visible. Use an empty +assets array when no objects are visible. Every field must be a non-empty +string.""" _USER_PROMPT = "Analyze the provided image and return only the required JSON object." @@ -363,13 +364,17 @@ def _parse_scene_object_fields( f"VLM JSON key {field_name}.category must be a lower-case snake_case " "class name." ) - if _LOCATION_WORD_PATTERN.search( - fields["name"] - ): # Check whether the name contains location. + # Check whether the name and description contain location or relationship words. + if _LOCATION_WORD_PATTERN.search(fields["name"]): raise ValueError( f"VLM JSON key {field_name}.name must not contain location or " "relationship words." ) + if _LOCATION_WORD_PATTERN.search(fields["description"]): + raise ValueError( + f"VLM JSON key {field_name}.description must not contain location or " + "relationship words." + ) return fields From 64ca9f12b5cf22ec8fa2bc0a468ce4083b10dadb Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:53:06 +0800 Subject: [PATCH 07/85] modify the prompt of scene understanding, to avoid a very long description in objects' names --- .../scene_engine/pipeline/generation/scene_understanding.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index c69ec3719..cae1c612f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -71,9 +71,9 @@ in any semantic field. 5. category is a lower-case singular snake_case class, such as mug, book, potted_plant, or coffee_table. It must not contain color or material. -6. name contains only color, material, texture, shape, and object description. - It must not contain position or relations, such as left, right, on, in, or - near. +6. name is a concise human-readable phrase containing only color, material, + texture, shape, and object details. It may contain spaces, but must not + contain position or relations, such as left, right, on, in, or near. 7. For table, description contains only its category, material, color, texture, shape, and visible structural details. Do not mention image coverage, image position, camera framing, or viewpoint. For example, do not write "occupying From 4236f9dc11d049914aa1ca1e8fed353482d31c20 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:31:14 +0800 Subject: [PATCH 08/85] finish the llm understand scene edit --- .../scene_engine/core/scene_edit_plan.py | 239 +++++++++++++++ .../gen_sim/scene_engine/pipeline/edit.py | 7 +- .../editing/scene_edit_understanding.py | 280 +++++++++++++++++- .../scene_engine/test_scene_edit_plan.py | 204 +++++++++++++ .../scene_engine/test_scene_understanding.py | 8 + 5 files changed, 725 insertions(+), 13 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/core/scene_edit_plan.py create mode 100644 tests/gen_sim/scene_engine/test_scene_edit_plan.py diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py new file mode 100644 index 000000000..26fe96bc9 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -0,0 +1,239 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Literal + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneConstraintType, + SceneGraph, + TABLE_OBJECT_ID, +) + +__all__ = ["SceneEditOperation", "SceneEditPlan"] + +SceneEditOperationType = Literal["add", "move", "delete"] + + +@dataclass(frozen=True) +class SceneEditOperation: + """One normalized edit operation produced from an LLM edit draft.""" + + op: SceneEditOperationType + object_id: str | None = None + target_id: str | None = None + relation: SceneConstraintType | None = None + category: str | None = None + name: str | None = None + description: str | None = None + + def to_dict(self) -> dict[str, object]: + """Serialize one normalized edit operation.""" + return { + "op": self.op, + "object_id": self.object_id, + "target_id": self.target_id, + "relation": self.relation, + "category": self.category, + "name": self.name, + "description": self.description, + } + + +@dataclass +class SceneEditPlan: + """Validated operations against one immutable pre-edit scene state.""" + + scene: Scene + scene_graph: SceneGraph + operations: list[SceneEditOperation] = field(default_factory=list) + + def __post_init__(self) -> None: + """Validate the plan before later stages prepare assets or edit layouts.""" + self.validate() + + def to_dict(self) -> dict[str, object]: + """Serialize the input scene state and normalized edit operations.""" + return { + "scene": self.scene.to_dict(), + "scene_graph": self.scene_graph.to_dict(), + "operations": [operation.to_dict() for operation in self.operations], + } + + def validate(self) -> None: + """Validate object references and edit conflicts against the input scene.""" + # Scene object IDs must remain a one-to-one lookup key for edit operations. + scene_object_ids = {scene_object.id for scene_object in self.scene.objects} + if len(scene_object_ids) != len(self.scene.objects): + raise ValueError("Scene edit input must contain unique object ids.") + # Parent and child checks require the graph to describe this exact scene. + if set(self.scene_graph.node_by_id()) != scene_object_ids: + raise ValueError("Scene edit plan graph nodes must match scene object ids.") + + existing_object_ids = set(scene_object_ids) + added_object_ids: set[str] = set() + # Collect deletions first so other operations cannot target removed objects. + deleted_object_ids = { + operation.object_id + for operation in self.operations + if operation.op == "delete" + } + if None in deleted_object_ids: + raise ValueError("Delete operations must identify an existing object.") + + edited_object_ids: set[str] = set() + # Validate each operation against the unchanged input scene and graph. + for operation in self.operations: + self._validate_operation( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + edited_object_ids=edited_object_ids, + added_object_ids=added_object_ids, + ) + # A removed support object must not leave any child objects orphaned. + self._validate_deleted_subtrees(deleted_object_ids) + + def _validate_operation( + self, + *, + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + edited_object_ids: set[str], + added_object_ids: set[str], + ) -> None: + if operation.op == "add": + self._validate_add_operation( + operation, + existing_object_ids, + deleted_object_ids, + added_object_ids, + ) + return + if operation.op not in {"move", "delete"}: + raise ValueError(f"Unsupported scene edit operation: {operation.op!r}") + if operation.object_id not in existing_object_ids: + raise ValueError( + "Move and delete operations must reference existing objects." + ) + if operation.object_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot be moved or deleted.") + # Existing objects accept only one move or delete instruction per plan. + if operation.object_id in edited_object_ids: + raise ValueError("An existing object may have only one edit operation.") + edited_object_ids.add(operation.object_id) + + if operation.op == "delete": + # Delete carries no new metadata or spatial placement. + if any( + value is not None + for value in ( + operation.target_id, + operation.relation, + operation.category, + operation.name, + operation.description, + ) + ): + raise ValueError("Delete operations may only specify object_id.") + return + + self._validate_position_reference( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + ) + if any( + value is not None + for value in (operation.category, operation.name, operation.description) + ): + raise ValueError("Move operations must not declare a new object.") + + @staticmethod + def _validate_add_operation( + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + added_object_ids: set[str], + ) -> None: + if not operation.object_id: + raise ValueError("Add operations must have a generated object_id.") + # Generated IDs must not collide with the input scene or this add batch. + if ( + operation.object_id in existing_object_ids + or operation.object_id in added_object_ids + ): + raise ValueError("Add operations must use unique new object ids.") + added_object_ids.add(operation.object_id) + if not all( + isinstance(value, str) and value.strip() + for value in (operation.category, operation.name, operation.description) + ): + raise ValueError("Add operations require category, name, and description.") + SceneEditPlan._validate_position_reference( + operation=operation, + existing_object_ids=existing_object_ids, + deleted_object_ids=deleted_object_ids, + ) + + @staticmethod + def _validate_position_reference( + *, + operation: SceneEditOperation, + existing_object_ids: set[str], + deleted_object_ids: set[str], + ) -> None: + if (operation.target_id is None) != (operation.relation is None): + raise ValueError("target_id and relation must be specified together.") + if operation.target_id is None: + return + if operation.target_id not in existing_object_ids: + raise ValueError("Edit targets must reference existing scene objects.") + # One edit may not position an object relative to a deleted target. + if operation.target_id in deleted_object_ids: + raise ValueError("Edit targets must not reference deleted objects.") + + def _validate_deleted_subtrees(self, deleted_object_ids: set[str]) -> None: + # Index the support graph once before checking every deleted parent. + children_by_parent: dict[str, list[str]] = {} + for node in self.scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + for object_id in deleted_object_ids: + descendants = self._descendant_ids(object_id, children_by_parent) + if not descendants.issubset(deleted_object_ids): + raise ValueError( + "Deleting a parent requires deleting all of its children." + ) + + @staticmethod + def _descendant_ids( + object_id: str, + children_by_parent: dict[str, list[str]], + ) -> set[str]: + descendants: set[str] = set() + # Traverse every support descendant, not only direct children. + pending = list(children_by_parent.get(object_id, [])) + while pending: + child_id = pending.pop() + descendants.add(child_id) + pending.extend(children_by_parent.get(child_id, [])) + return descendants diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index 3b6c95e02..93001afb6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -43,6 +43,7 @@ def edit_scene( scene_importer = SceneExportImporter(output_root=output_root) # Validate scene_export, write scene.json, and return Scene; failures raise before editing. scene, scene_graph = scene_importer.import_scene_and_graph() + # Only for debug. print( json.dumps( {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, @@ -52,12 +53,12 @@ def edit_scene( ) # 1. Edit Understanding - # And update or initialize the scene graph + # Will return an already checked scene edit plan. log_info("Starting Edit Understanding") - updated_scene_graph = understand_scene_edit( + scene_edit_plan = understand_scene_edit( scene=scene, + scene_graph=scene_graph, edit_prompt=edit_prompt, - output_root=output_root, # Has already been resolved. vlm_client=vlm_client, ) log_info("Completed Edit Understanding") diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 48dda64bf..8099eda4b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -17,32 +17,292 @@ from __future__ import annotations -from pathlib import Path +import json -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import ( - SceneGraph, - SceneGraphNode, - TABLE_OBJECT_ID, +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( + SceneEditOperation, + SceneEditPlan, ) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) +_EDIT_SYSTEM_PROMPT = """You convert one user instruction into edits for an existing tabletop scene. + +Use an existing object ID only when it appears in the supplied Existing object +IDs list. IDs identify existing objects exactly; never invent, correct, or +renumber them. The table ID is "table" and cannot be moved or deleted. + +Each operation is one of: +1. move: move one existing object. object_id identifies it. target_id and + relation are either both provided or both null. +2. delete: delete one existing object. Only object_id is provided. +3. add: create one new object. object_id must be null. Provide a lower-case + singular snake_case category, name, and description. Multiple add operations + may have the same category and name; their final IDs are assigned by the + program in operation order. target_id and relation are either both provided + or both null. + +For a positioned move or add, target_id must be an Existing object ID and +relation must be one of on, left_of, right_of, in_front_of, or behind. Do not +position a new object relative to another newly added object. + +Each existing object's center_xy is its center position [x, y] in the +table-frame Z-up world coordinate system. Smaller x is left, larger x is right, +larger y is in front, and smaller y is behind. Use center_xy only to disambiguate +references such as "the bottle on the left"; do not output coordinates. Express +the requested position using target_id and one allowed relation instead. + +For every newly added object, category is its lower-case singular snake_case +class. name contains only color, material, texture, shape, and object details. +description contains only visible category, material, color, texture, shape, +and structural details. name and description must not mention position, the +table, or relations to any object. + +Return JSON only: no Markdown, comments, or prose. Every operation must contain +exactly these fields: op, object_id, target_id, relation, category, name, and +description. Use null for every field that does not apply to an operation: +{ + "operations": [ + { + "op": "move", + "object_id": "bottle_001", + "target_id": "book_001", + "relation": "right_of", + "category": null, + "name": null, + "description": null + }, + { + "op": "delete", + "object_id": "cup_001", + "target_id": null, + "relation": null, + "category": null, + "name": null, + "description": null + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel" + }, + { + "op": "add", + "object_id": null, + "target_id": "book_001", + "relation": "right_of", + "category": "orange", + "name": "small orange", + "description": "small round orange with a textured peel" + } + ] +} +The two orange additions intentionally share category and name. Do not add +fields beyond the required schema.""" + def understand_scene_edit( *, scene: Scene, + scene_graph: SceneGraph, edit_prompt: str, - output_root: str | Path, vlm_client: OpenAICompatibleVLM, -) -> dict[str, object]: + json_max_attempts: int = 3, +) -> SceneEditPlan: """Understand one text edit instruction for an existing scene.""" edit_prompt = edit_prompt.strip() if not edit_prompt: raise ValueError("Edit prompt must not be empty.") + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + # Give the VLM only the scene metadata needed to identify existing objects. + simplified_scene_info = _simplify_scene_info(scene=scene) + operations = _vlm_understand_scene_edit( + scene=scene, + edit_prompt=edit_prompt, + simplified_scene_info=simplified_scene_info, + vlm_client=vlm_client, + json_max_attempts=json_max_attempts, + ) + + # SceneEditPlan validates all references against the immutable input scene graph. + return SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=operations, + ) + +def _simplify_scene_info(scene: Scene) -> dict[str, object]: + """Return the object metadata needed for edit instruction resolution.""" return { - "edit_prompt": edit_prompt, - "operations": [], + "existing_object_ids": [scene_object.id for scene_object in scene.objects], + "objects": [ + { + "id": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, + "description": scene_object.description, + "center_xy": scene_object.center_xy, + } + for scene_object in scene.objects + ], } + + +def _vlm_understand_scene_edit( + *, + scene: Scene, + edit_prompt: str, + simplified_scene_info: dict[str, object], + vlm_client: OpenAICompatibleVLM, + json_max_attempts: int, +) -> list[SceneEditOperation]: + """Return parsed edit operations from the VLM with assigned add IDs.""" + # Construct user prompt. + user_prompt = ( + f"User edit instruction:\n{edit_prompt}\n\n" + "Existing scene metadata:\n" + f"{json.dumps(simplified_scene_info, indent=2, ensure_ascii=False)}" + ) + last_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_EDIT_SYSTEM_PROMPT, + user_prompt=user_prompt, + ) + try: + value = json.loads(_strip_json_code_fence(response_text)) + return _parse_scene_edit_operations(value, scene=scene) + except (json.JSONDecodeError, ValueError) as exc: + last_error = ValueError(f"VLM returned invalid scene edit JSON: {exc}") + continue + + assert last_error is not None + raise ValueError( + "VLM returned invalid scene edit JSON after " + f"{json_max_attempts} attempts: {last_error}" + ) from last_error + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if len(lines) < 3 or not lines[-1].strip().startswith("```"): + raise ValueError("VLM response contains an incomplete JSON code fence.") + return "\n".join(lines[1:-1]).strip() + + +def _parse_scene_edit_operations( + value: object, + *, + scene: Scene, +) -> list[SceneEditOperation]: + """Parse the strict VLM edit-draft schema into typed operations with add IDs.""" + if not isinstance(value, dict) or set(value) != {"operations"}: + raise ValueError("Scene edit draft must contain exactly operations.") + # Get and validate a list operation value. + operations_value = value["operations"] + if not isinstance(operations_value, list): + raise ValueError("Scene edit draft operations must be a list.") + + expected_keys = { + "op", + "object_id", + "target_id", + "relation", + "category", + "name", + "description", + } + # Get ids and counts of existing objects to assign new add IDs. + assigned_object_ids = {scene_object.id for scene_object in scene.objects} + category_counts = { + category: sum( + scene_object.category == category for scene_object in scene.objects + ) + for category in {scene_object.category for scene_object in scene.objects} + } + operations: list[SceneEditOperation] = [] + for value in operations_value: + if not isinstance(value, dict) or not isinstance(value.get("op"), str): + raise ValueError("Scene edit operations must contain a string op.") + op = value["op"] + if op not in {"add", "move", "delete"}: + raise ValueError("Scene edit operation op is invalid.") + if set(value) != expected_keys: + raise ValueError("Scene edit operations must use the required schema.") + object_id = _optional_string(value.get("object_id"), field_name="object_id") + category = _optional_string(value.get("category"), field_name="category") + if op == "add": + if object_id is not None: + raise ValueError("VLM add operations must set object_id to null.") + if category is None: + raise ValueError("VLM add operations must provide a category.") + # Add operation should generate new id here. + # Never believe the LLM could always generate a valid id. + object_id = _next_add_object_id( + category=category, + category_counts=category_counts, + assigned_object_ids=assigned_object_ids, + ) + operations.append( + SceneEditOperation( + op=op, + object_id=object_id, + target_id=_optional_string( + value.get("target_id"), field_name="target_id" + ), + relation=_optional_relation(value.get("relation")), + category=category, + name=_optional_string(value.get("name"), field_name="name"), + description=_optional_string( + value.get("description"), field_name="description" + ), + ) + ) + return operations + + +def _next_add_object_id( + *, + category: str, + category_counts: dict[str, int], + assigned_object_ids: set[str], +) -> str: + """Assign the next available ID for one new object category.""" + index = category_counts.get(category, 0) + 1 + object_id = f"{category}_{index:03d}" + # In case the scene have orange_001 and orange_003. + while object_id in assigned_object_ids: + index += 1 + object_id = f"{category}_{index:03d}" + category_counts[category] = index + assigned_object_ids.add(object_id) + return object_id + + +def _optional_string(value: object, *, field_name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Scene edit operation {field_name} must be a string or null.") + return value.strip() + + +def _optional_relation(value: object) -> str | None: + if value is None: + return None + if value not in {"on", "left_of", "right_of", "in_front_of", "behind"}: + raise ValueError("Scene edit operation relation is invalid.") + return value diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py new file mode 100644 index 000000000..c01049eb6 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -0,0 +1,204 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( + SceneEditOperation, + SceneEditPlan, +) +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + _parse_scene_edit_operations, +) + + +def _scene_and_graph() -> tuple[Scene, SceneGraph]: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="blue book", + description="A blue book.", + ), + SceneObject( + id="orange_001", + kind="asset", + category="orange", + name="orange", + description="An orange.", + ), + ] + ) + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="orange_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + return scene, scene_graph + + +def test_scene_edit_plan_accepts_add_without_a_position() -> None: + scene, scene_graph = _scene_and_graph() + + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + assert len(plan.operations) == 1 + assert plan.to_dict()["operations"] == [ + { + "op": "add", + "object_id": "cup_001", + "target_id": None, + "relation": None, + "category": "cup", + "name": "green cup", + "description": "A small green ceramic cup.", + } + ] + + +def test_scene_edit_plan_accepts_multiple_new_objects_with_the_same_category() -> None: + scene, scene_graph = _scene_and_graph() + + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="orange_002", + category="orange", + name="small orange", + description="A small round orange with a textured peel.", + ), + SceneEditOperation( + op="add", + object_id="orange_003", + category="orange", + name="large orange", + description="A large round orange with a textured peel.", + ), + ], + ) + + assert [operation.category for operation in plan.operations] == ["orange", "orange"] + + +def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: + scene, _ = _scene_and_graph() + draft = { + "operations": [ + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + }, + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "category": "orange", + "name": "small_orange", + "description": "A small round orange with a textured peel.", + }, + ] + } + + operations = _parse_scene_edit_operations( + json.loads(json.dumps(draft)), scene=scene + ) + + assert [operation.object_id for operation in operations] == [ + "orange_002", + "orange_003", + ] + + +def test_scene_edit_plan_rejects_targets_outside_the_input_scene() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="existing scene objects"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="spoon_001", + target_id="new_orange_001", + relation="left_of", + category="spoon", + name="metal spoon", + description="A metal spoon.", + ) + ], + ) + + +def test_scene_edit_plan_requires_deleting_all_children_of_a_deleted_parent() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="all of its children"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[SceneEditOperation(op="delete", object_id="book_001")], + ) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 9f5ae3884..8b9357f33 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -63,6 +63,14 @@ def test_image_object_analysis_rejects_location_words_in_object_names() -> None: ) +def test_image_object_analysis_rejects_location_words_in_object_descriptions() -> None: + response = json.loads(_response()) + response["assets"][0]["description"] = "A small ceramic cup on the table." + + with pytest.raises(ValueError, match="description must not contain location"): + scene_understanding._parse_image_object_analysis_response(json.dumps(response)) + + def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: class VLM: def __init__(self) -> None: From e52db87f1352f99574ddbcf36957624ff5f3fb98 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:54:11 +0800 Subject: [PATCH 09/85] finish the scene edit plan -> updated scene graph --- .../scene_engine/core/scene_edit_plan.py | 14 +- .../gen_sim/scene_engine/core/scene_graph.py | 152 +++++++++++++++ .../gen_sim/scene_engine/pipeline/edit.py | 35 ++-- .../editing/scene_edit_understanding.py | 97 +++++++++- .../scene_engine/test_scene_edit_plan.py | 183 ++++++++++++++++++ 5 files changed, 459 insertions(+), 22 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 26fe96bc9..a7ac152d7 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -78,6 +78,15 @@ def to_dict(self) -> dict[str, object]: def validate(self) -> None: """Validate object references and edit conflicts against the input scene.""" + # Edit-plan rules: + # - move and delete identify one existing non-table object with object_id. + # - add carries generated object_id plus non-empty category, name, and description. + # - move always supplies target_id and relation; add may omit both. + # - target_id and relation are otherwise supplied together or both absent. + # - every target is from the pre-edit scene; new and deleted objects are invalid targets. + # - an existing object has at most one move or delete operation in one plan. + # - delete carries no placement or new-object metadata and must delete every descendant. + # - these checks validate intent only; they do not mutate Scene or SceneGraph. # Scene object IDs must remain a one-to-one lookup key for edit operations. scene_object_ids = {scene_object.id for scene_object in self.scene.objects} if len(scene_object_ids) != len(self.scene.objects): @@ -88,7 +97,7 @@ def validate(self) -> None: existing_object_ids = set(scene_object_ids) added_object_ids: set[str] = set() - # Collect deletions first so other operations cannot target removed objects. + # Collect deletion intents first so move and add cannot target them regardless of order. deleted_object_ids = { operation.object_id for operation in self.operations @@ -155,6 +164,8 @@ def _validate_operation( raise ValueError("Delete operations may only specify object_id.") return + if operation.target_id is None or operation.relation is None: + raise ValueError("Move operations must specify target_id and relation.") self._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, @@ -204,6 +215,7 @@ def _validate_position_reference( raise ValueError("target_id and relation must be specified together.") if operation.target_id is None: return + # Targets come only from the pre-edit scene, so new objects cannot be targets. if operation.target_id not in existing_object_ids: raise ValueError("Edit targets must reference existing scene objects.") # One edit may not position an object relative to a deleted target. diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index 806c5a0ee..514176211 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -129,6 +129,158 @@ def node_by_id(self) -> dict[str, SceneGraphNode]: nodes_by_id[node.object_id] = node return nodes_by_id + def remove_nodes(self, object_ids: set[str]) -> None: + """Remove nodes and their incident planar relations, then validate.""" + # If no node to be removed, return directly. + if not object_ids: + return + if TABLE_OBJECT_ID in object_ids: + raise ValueError("The table cannot be removed from a scene graph.") + unknown_object_ids = object_ids - set(self.node_by_id()) + if unknown_object_ids: + raise ValueError( + f"Cannot remove unknown scene graph nodes: {sorted(unknown_object_ids)}" + ) + + # Removing every incident relation prevents dangling planar endpoints. + self.nodes = [node for node in self.nodes if node.object_id not in object_ids] + self.relations = [ + relation + for relation in self.relations + if relation.source_id not in object_ids + and relation.target_id not in object_ids + ] + # Refresh. + self.refresh() + + def add_node(self, node: SceneGraphNode) -> None: + """Add one node and validate the resulting graph.""" + if node.object_id in self.node_by_id(): + raise ValueError(f"Duplicate scene graph node: {node.object_id}") + self.nodes.append(node) + self.refresh() + + def apply_updates( + self, + *, + deleted_object_ids: set[str], + added_object_ids: list[str], + on_parent_updates: list[tuple[str, str]], + planar_relation_updates: list[tuple[str, PlanarRelationType, str]], + ) -> None: + """Apply one atomic batch of node and relationship updates.""" + if TABLE_OBJECT_ID in deleted_object_ids: + raise ValueError("The table cannot be removed from a scene graph.") + + existing_object_ids = set(self.node_by_id()) + unknown_object_ids = deleted_object_ids - existing_object_ids + if unknown_object_ids: + raise ValueError( + f"Cannot remove unknown scene graph nodes: {sorted(unknown_object_ids)}" + ) + + # Delete all requested nodes before resolving new parents and relations. + self.nodes = [ + node for node in self.nodes if node.object_id not in deleted_object_ids + ] + self.relations = [ + relation + for relation in self.relations + if relation.source_id not in deleted_object_ids + and relation.target_id not in deleted_object_ids + ] + + remaining_object_ids = set(self.node_by_id()) + if len(added_object_ids) != len(set(added_object_ids)): + raise ValueError("Added scene graph node ids must be unique.") + duplicate_object_ids = set(added_object_ids) & remaining_object_ids + if duplicate_object_ids: + raise ValueError( + f"Duplicate scene graph nodes: {sorted(duplicate_object_ids)}" + ) + + # New nodes default to the table; later updates replace that parent when needed. + self.nodes.extend( + SceneGraphNode( + object_id=object_id, + parent_id=TABLE_OBJECT_ID, + parent_relation="on", + ) + for object_id in added_object_ids + ) + + # Apply support-parent changes before planar updates need the final parent. + for object_id, parent_id in on_parent_updates: + self._set_on_parent(object_id=object_id, parent_id=parent_id) + + # Resolve chained planar parent inheritance before adding final relations. + self._resolve_planar_parent_updates(planar_relation_updates) + for source_id, relation, target_id in planar_relation_updates: + self._clear_incident_planar_relations(source_id) + self.relations.append( + SceneGraphRelation( + source_id=source_id, + relation=relation, + target_id=target_id, + ) + ) + + # Normalize inverse relations and reject invalid final graph constraints. + self.refresh() + + def _set_on_parent(self, *, object_id: str, parent_id: str) -> None: + """Replace one node's support parent and stale planar constraints.""" + nodes_by_id = self.node_by_id() + if object_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot be moved onto another object.") + if object_id not in nodes_by_id or parent_id not in nodes_by_id: + raise ValueError( + "Parent updates must reference existing scene graph nodes." + ) + if object_id == parent_id: + raise ValueError("A scene graph node cannot be its own parent.") + + node = nodes_by_id[object_id] + node.parent_id = parent_id + node.parent_relation = "on" + node.table_region = None + self._clear_incident_planar_relations(object_id) + + def _resolve_planar_parent_updates( + self, + planar_relation_updates: list[tuple[str, PlanarRelationType, str]], + ) -> None: + """Make every planar source share its target's final support parent.""" + for _ in range(len(planar_relation_updates)): + changed = False + for source_id, _, target_id in planar_relation_updates: + nodes_by_id = self.node_by_id() + if source_id == TABLE_OBJECT_ID: + raise ValueError("The table cannot have a planar relation.") + if source_id not in nodes_by_id or target_id not in nodes_by_id: + raise ValueError( + "Planar updates must reference existing scene graph nodes." + ) + target_parent_id = nodes_by_id[target_id].parent_id + if target_parent_id is None: + raise ValueError("Planar relation targets must have a parent.") + source = nodes_by_id[source_id] + if source.parent_id != target_parent_id: + source.parent_id = target_parent_id + source.parent_relation = "on" + source.table_region = None + changed = True + if not changed: + return + + def _clear_incident_planar_relations(self, object_id: str) -> None: + """Remove planar constraints invalidated when one node changes parent.""" + self.relations = [ + relation + for relation in self.relations + if relation.source_id != object_id and relation.target_id != object_id + ] + def normalize(self) -> None: """Materialize inverse planar relations and remove duplicates.""" self._materialize_inverse_planar_relations() diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index 93001afb6..a13ef17fd 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -44,31 +44,40 @@ def edit_scene( # Validate scene_export, write scene.json, and return Scene; failures raise before editing. scene, scene_graph = scene_importer.import_scene_and_graph() # Only for debug. - print( - json.dumps( - {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, - indent=2, - ensure_ascii=False, - ) - ) + # print( + # json.dumps( + # {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, + # indent=2, + # ensure_ascii=False, + # ) + # ) # 1. Edit Understanding - # Will return an already checked scene edit plan. + # Will return an already checked scene edit plan + # and a validated updated scene graph. log_info("Starting Edit Understanding") - scene_edit_plan = understand_scene_edit( + scene_edit_plan, updated_scene_graph = understand_scene_edit( scene=scene, scene_graph=scene_graph, edit_prompt=edit_prompt, vlm_client=vlm_client, ) log_info("Completed Edit Understanding") + # Only for debug. + # print( + # json.dumps( + # {"updated scene graph": updated_scene_graph.to_dict()}, + # indent=2, + # ensure_ascii=False, + # ) + # ) # 2. Prepare Objects. log_info("Preparing Objects if necessary") - # scene = prepare_objects( - # scene=scene, - # output_root=output_root, - # ) + scene = prepare_scene_edit_assets( + scene=scene, + scene_edit_plan=scene_edit_plan, + ) log_info("Completed Preparing Objects") # 3. Layout Editing diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 8099eda4b..1eb3a857f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -24,7 +24,12 @@ SceneEditPlan, ) from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + PlanarRelationType, + SceneGraph, + SceneGraphNode, + SceneGraphRelation, +) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) @@ -36,8 +41,8 @@ renumber them. The table ID is "table" and cannot be moved or deleted. Each operation is one of: -1. move: move one existing object. object_id identifies it. target_id and - relation are either both provided or both null. +1. move: move one existing object. object_id, target_id, and relation must all + be provided. 2. delete: delete one existing object. Only object_id is provided. 3. add: create one new object. object_id must be null. Provide a lower-case singular snake_case category, name, and description. Multiple add operations @@ -45,9 +50,9 @@ program in operation order. target_id and relation are either both provided or both null. -For a positioned move or add, target_id must be an Existing object ID and -relation must be one of on, left_of, right_of, in_front_of, or behind. Do not -position a new object relative to another newly added object. +For every move and every positioned add, target_id must be an Existing object +ID and relation must be one of on, left_of, right_of, in_front_of, or behind. +Do not position a new object relative to another newly added object. Each existing object's center_xy is its center position [x, y] in the table-frame Z-up world coordinate system. Smaller x is left, larger x is right, @@ -115,7 +120,7 @@ def understand_scene_edit( edit_prompt: str, vlm_client: OpenAICompatibleVLM, json_max_attempts: int = 3, -) -> SceneEditPlan: +) -> tuple[SceneEditPlan, SceneGraph]: """Understand one text edit instruction for an existing scene.""" edit_prompt = edit_prompt.strip() if not edit_prompt: @@ -133,11 +138,87 @@ def understand_scene_edit( ) # SceneEditPlan validates all references against the immutable input scene graph. - return SceneEditPlan( + scene_edit_plan = SceneEditPlan( scene=scene, scene_graph=scene_graph, operations=operations, ) + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=scene_edit_plan, + ) + return scene_edit_plan, updated_scene_graph + + +def _build_updated_scene_graph( + *, + scene_graph: SceneGraph, + scene_edit_plan: SceneEditPlan, +) -> SceneGraph: + """Build and validate the target graph implied by one edit plan.""" + # Copy every mutable graph value so the pre-edit graph remains unchanged. + updated_scene_graph = SceneGraph( + nodes=[ + SceneGraphNode( + object_id=node.object_id, + parent_id=node.parent_id, + parent_relation=node.parent_relation, + table_region=node.table_region, + ) + for node in scene_graph.nodes + ], + relations=[ + SceneGraphRelation( + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + ) + for relation in scene_graph.relations + ], + validate_on_refresh=scene_graph.validate_on_refresh, + ) + _apply_scene_edit_plan_to_scene_graph( + scene_graph=updated_scene_graph, + scene_edit_plan=scene_edit_plan, + ) + return updated_scene_graph + + +def _apply_scene_edit_plan_to_scene_graph( + *, + scene_graph: SceneGraph, + scene_edit_plan: SceneEditPlan, +) -> None: + """Apply the target graph updates implied by add and move operations.""" + deleted_object_ids: set[str] = set() + added_object_ids: list[str] = [] + on_parent_updates: list[tuple[str, str]] = [] + planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] + for operation in scene_edit_plan.operations: + if operation.op == "delete": + if operation.object_id is not None: + deleted_object_ids.add(operation.object_id) + continue + if operation.object_id is None: + raise ValueError("Add and move operations must have an object_id.") + if operation.op == "add": + added_object_ids.append(operation.object_id) + if operation.target_id is None or operation.relation is None: + continue + if operation.relation == "on": + on_parent_updates.append((operation.object_id, operation.target_id)) + continue + planar_relation_updates.append( + (operation.object_id, operation.relation, operation.target_id) + ) + + # Apply all graph changes atomically so intermediate edit states need not be valid. + scene_graph.apply_updates( + deleted_object_ids=deleted_object_ids, + added_object_ids=added_object_ids, + on_parent_updates=on_parent_updates, + planar_relation_updates=planar_relation_updates, + ) def _simplify_scene_info(scene: Scene) -> dict[str, object]: diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index c01049eb6..30c11c995 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -31,8 +31,13 @@ ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + _apply_scene_edit_plan_to_scene_graph, + _build_updated_scene_graph, _parse_scene_edit_operations, ) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) def _scene_and_graph() -> tuple[Scene, SceneGraph]: @@ -202,3 +207,181 @@ def test_scene_edit_plan_requires_deleting_all_children_of_a_deleted_parent() -> scene_graph=scene_graph, operations=[SceneEditOperation(op="delete", object_id="book_001")], ) + + +def test_scene_edit_plan_requires_a_position_for_move_operations() -> None: + scene, scene_graph = _scene_and_graph() + + with pytest.raises(ValueError, match="must specify target_id and relation"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[SceneEditOperation(op="move", object_id="book_001")], + ) + + +def test_scene_edit_asset_preparation_skips_plans_without_adds() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="book_001", + target_id="table", + relation="on", + ) + ], + ) + + prepared_scene = prepare_scene_edit_assets( + scene=scene, + scene_edit_plan=plan, + ) + + assert prepared_scene is scene + + +def test_scene_edit_graph_builder_copies_the_pre_edit_graph() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan(scene=scene, scene_graph=scene_graph) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph is not scene_graph + assert updated_scene_graph.nodes is not scene_graph.nodes + assert updated_scene_graph.relations is not scene_graph.relations + assert updated_scene_graph.to_dict() == scene_graph.to_dict() + + +def test_scene_edit_graph_builder_removes_deleted_nodes() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation(op="delete", object_id="orange_001"), + SceneEditOperation(op="delete", object_id="book_001"), + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert set(updated_scene_graph.node_by_id()) == {"table"} + assert set(scene_graph.node_by_id()) == {"table", "book_001", "orange_001"} + + +def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + added_node = updated_scene_graph.node_by_id()["cup_001"] + assert added_node.parent_id == "table" + assert added_node.parent_relation == "on" + + +def test_scene_edit_graph_builder_updates_move_on_parent() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="orange_001", + target_id="table", + relation="on", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph.node_by_id()["orange_001"].parent_id == "table" + + +def test_scene_edit_graph_builder_adds_planar_relation_with_target_parent() -> None: + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + target_id="book_001", + relation="right_of", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + updated_scene_graph = _build_updated_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert updated_scene_graph.node_by_id()["cup_001"].parent_id == "table" + assert any( + relation.source_id == "cup_001" + and relation.relation == "right_of" + and relation.target_id == "book_001" + for relation in updated_scene_graph.relations + ) + + +def test_scene_edit_plan_application_adds_new_nodes_before_relationship_updates() -> ( + None +): + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + target_id="book_001", + relation="right_of", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + + _apply_scene_edit_plan_to_scene_graph( + scene_graph=scene_graph, + scene_edit_plan=plan, + ) + + assert scene_graph.node_by_id()["cup_001"].parent_id == "table" From 95f475979f18aae5105def735242aafa4f592683 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 12 Aug 2026 20:50:46 +0800 Subject: [PATCH 10/85] move client try catch outside the understand_scene --- .../gen_sim/scene_engine/pipeline/generate.py | 22 ++++++++++++++----- .../generation/scene_understanding.py | 21 +++++++----------- 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 32aa6550f..1df9027f9 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -25,6 +25,9 @@ from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, ) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) from embodichain.gen_sim.scene_engine.pipeline.generation.scene_understanding import ( understand_scene, @@ -51,12 +54,19 @@ def generate_scene_from_image( # 1. Scene Understanding log_info("Starting Scene Understanding") - scene, scene_graph = understand_scene( - scene=scene, - image_path=image_path, - output_root=resolved_output_root, - vlm_client=vlm_client, - ) + # Load .env settings and fail if the Image Segmentation Server is unavailable. + image_segmentation_client = ImageSegmentationClient.from_dotenv() + try: + image_segmentation_client.check_health() + scene, scene_graph = understand_scene( + scene=scene, + image_path=image_path, + output_root=resolved_output_root, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) + finally: + image_segmentation_client.close() # Close the session after scene understanding. log_info("Completed Scene Understanding") # 2. Objects + Coarse Layout Generation diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index cae1c612f..f2103a945 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -155,6 +155,7 @@ def understand_scene( output_root: str | Path, *, vlm_client: OpenAICompatibleVLM, + image_segmentation_client: ImageSegmentationClient, json_max_attempts: int = 3, ) -> tuple[Scene, SceneGraph]: @@ -173,19 +174,13 @@ def understand_scene( json_max_attempts=json_max_attempts, ) - # Load .env settings and fail if the Image Segmentation Server is unavailable. - image_segmentation_client = ImageSegmentationClient.from_dotenv() - try: - image_segmentation_client.check_health() # Error raising will happen internally. - _segment_scene( - image_path=resolved_image_path, - stage_output_root=stage_output_root, - scene=scene, - vlm_client=vlm_client, - image_segmentation_client=image_segmentation_client, - ) - finally: - image_segmentation_client.close() # Kill the session to avoid resource leaks. + _segment_scene( + image_path=resolved_image_path, + stage_output_root=stage_output_root, + scene=scene, + vlm_client=vlm_client, + image_segmentation_client=image_segmentation_client, + ) # Use the segmented image to initialize the scene graph # with the help of the VLM client. From a72f1011698b5c2f45c123bfb6bdb785cee1f441 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 13 Aug 2026 17:07:06 +0800 Subject: [PATCH 11/85] finish part of the scene edit: from user prompt to simreadyed asset, but no real-world scale... --- .../scene_engine/core/scene_edit_plan.py | 11 + .../gen_sim/scene_engine/core/scene_graph.py | 38 ++- .../gen_sim/scene_engine/pipeline/edit.py | 64 ++-- .../editing/scene_edit_asset_preparation.py | 309 ++++++++++++++++++ .../editing/scene_edit_understanding.py | 70 +++- .../pipeline/generation/scene_generation.py | 6 +- .../utils/image_segmentation_utils.py | 42 +++ .../pipeline/utils/scene_importer.py | 14 +- ...ene_processor.py => simready_processor.py} | 14 +- .../scene_engine/test_scene_edit_plan.py | 142 +++++++- 10 files changed, 646 insertions(+), 64 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py rename embodichain/gen_sim/scene_engine/pipeline/utils/{simready_scene_processor.py => simready_processor.py} (97%) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index a7ac152d7..83ad5bd79 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -23,6 +23,7 @@ from embodichain.gen_sim.scene_engine.core.scene_graph import ( SceneConstraintType, SceneGraph, + TableRegion, TABLE_OBJECT_ID, ) @@ -39,6 +40,7 @@ class SceneEditOperation: object_id: str | None = None target_id: str | None = None relation: SceneConstraintType | None = None + table_region: TableRegion | None = None category: str | None = None name: str | None = None description: str | None = None @@ -50,6 +52,7 @@ def to_dict(self) -> dict[str, object]: "object_id": self.object_id, "target_id": self.target_id, "relation": self.relation, + "table_region": self.table_region, "category": self.category, "name": self.name, "description": self.description, @@ -82,6 +85,7 @@ def validate(self) -> None: # - move and delete identify one existing non-table object with object_id. # - add carries generated object_id plus non-empty category, name, and description. # - move always supplies target_id and relation; add may omit both. + # - table_region is only valid with target_id=table and relation=on. # - target_id and relation are otherwise supplied together or both absent. # - every target is from the pre-edit scene; new and deleted objects are invalid targets. # - an existing object has at most one move or delete operation in one plan. @@ -156,6 +160,7 @@ def _validate_operation( for value in ( operation.target_id, operation.relation, + operation.table_region, operation.category, operation.name, operation.description, @@ -213,6 +218,12 @@ def _validate_position_reference( ) -> None: if (operation.target_id is None) != (operation.relation is None): raise ValueError("target_id and relation must be specified together.") + if operation.table_region is not None and ( + operation.target_id != TABLE_OBJECT_ID or operation.relation != "on" + ): + raise ValueError( + "table_region requires target_id='table' and relation='on'." + ) if operation.target_id is None: return # Targets come only from the pre-edit scene, so new objects cannot be targets. diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index 514176211..c5c49c625 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -21,7 +21,7 @@ TABLE_OBJECT_ID = "table" -# 9-grid table regions, treat the table as a 3x3 grid. +# Static type constraint for the nine regions of the tabletop 3x3 grid. TableRegion = Literal[ "left_back", "back_center", @@ -33,6 +33,20 @@ "front_center", "right_front", ] +# Runtime membership set for validating serialized and user-provided regions. +TABLE_REGIONS = frozenset( + { + "left_back", + "back_center", + "right_back", + "left_center", + "center", + "right_center", + "left_front", + "front_center", + "right_front", + } +) # A on B, then B is the parent node of A. SupportRelationType = Literal["on"] @@ -55,6 +69,8 @@ def __post_init__(self) -> None: """Validate local node fields before graph-level checks.""" if not self.object_id: raise ValueError("object_id must be non-empty.") + if self.table_region not in {None, *TABLE_REGIONS}: + raise ValueError("table_region is invalid.") # If the node is the table. if self.object_id == TABLE_OBJECT_ID: if self.parent_id is not None: @@ -165,7 +181,7 @@ def apply_updates( *, deleted_object_ids: set[str], added_object_ids: list[str], - on_parent_updates: list[tuple[str, str]], + on_parent_updates: list[tuple[str, str, TableRegion | None]], planar_relation_updates: list[tuple[str, PlanarRelationType, str]], ) -> None: """Apply one atomic batch of node and relationship updates.""" @@ -210,8 +226,12 @@ def apply_updates( ) # Apply support-parent changes before planar updates need the final parent. - for object_id, parent_id in on_parent_updates: - self._set_on_parent(object_id=object_id, parent_id=parent_id) + for object_id, parent_id, table_region in on_parent_updates: + self._set_on_parent( + object_id=object_id, + parent_id=parent_id, + table_region=table_region, + ) # Resolve chained planar parent inheritance before adding final relations. self._resolve_planar_parent_updates(planar_relation_updates) @@ -228,7 +248,13 @@ def apply_updates( # Normalize inverse relations and reject invalid final graph constraints. self.refresh() - def _set_on_parent(self, *, object_id: str, parent_id: str) -> None: + def _set_on_parent( + self, + *, + object_id: str, + parent_id: str, + table_region: TableRegion | None = None, + ) -> None: """Replace one node's support parent and stale planar constraints.""" nodes_by_id = self.node_by_id() if object_id == TABLE_OBJECT_ID: @@ -243,7 +269,7 @@ def _set_on_parent(self, *, object_id: str, parent_id: str) -> None: node = nodes_by_id[object_id] node.parent_id = parent_id node.parent_relation = "on" - node.table_region = None + node.table_region = table_region self._clear_incident_planar_relations(object_id) def _resolve_planar_parent_updates( diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index a13ef17fd..6ba7898d1 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -16,9 +16,17 @@ from __future__ import annotations -import json from pathlib import Path +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, ) @@ -28,6 +36,9 @@ from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( understand_scene_edit, ) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) from embodichain.utils.logger import log_info @@ -37,20 +48,14 @@ def edit_scene( edit_prompt: str, ) -> None: """Apply one text edit instruction to an existing Scene Engine output.""" + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) # Initialize the VLM client that will interpret the edit instruction. vlm_client = OpenAICompatibleVLM.from_dotenv() scene_importer = SceneExportImporter(output_root=output_root) # Validate scene_export, write scene.json, and return Scene; failures raise before editing. scene, scene_graph = scene_importer.import_scene_and_graph() - # Only for debug. - # print( - # json.dumps( - # {"scene": scene.to_dict(), "scene_graph": scene_graph.to_dict()}, - # indent=2, - # ensure_ascii=False, - # ) - # ) # 1. Edit Understanding # Will return an already checked scene edit plan @@ -63,22 +68,31 @@ def edit_scene( vlm_client=vlm_client, ) log_info("Completed Edit Understanding") - # Only for debug. - # print( - # json.dumps( - # {"updated scene graph": updated_scene_graph.to_dict()}, - # indent=2, - # ensure_ascii=False, - # ) - # ) - # 2. Prepare Objects. - log_info("Preparing Objects if necessary") - scene = prepare_scene_edit_assets( - scene=scene, - scene_edit_plan=scene_edit_plan, - ) - log_info("Completed Preparing Objects") + # 2. Prepare Objects + log_info("Starting Objects Preparation") + # Initialize all the clients and then check. + image_generation_client = ImageGenerationClient.from_dotenv() + geometry_generation_client = GeometryGenerationClient.from_dotenv() + image_segmentation_client = ImageSegmentationClient.from_dotenv() + try: + image_generation_client.check_health() + geometry_generation_client.check_health() + image_segmentation_client.check_health() + # Return a list of added SceneObjects assets. + # Now do not support editing the table. + added_assets = prepare_scene_edit_assets( + scene_edit_plan=scene_edit_plan, + output_root=resolved_output_root, + image_generation_client=image_generation_client, + geometry_generation_client=geometry_generation_client, + image_segmentation_client=image_segmentation_client, + ) + finally: + image_generation_client.close() + geometry_generation_client.close() + image_segmentation_client.close() + log_info("Completed Objects Preparation") # 3. Layout Editing log_info("Starting Layout Editing") @@ -96,4 +110,4 @@ def edit_scene( log_info("Starting Scene Export") log_info("Completed Scene Export") - raise NotImplementedError("Scene editing is not implemented yet.") + return None diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py new file mode 100644 index 000000000..70c02d1a7 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -0,0 +1,309 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import shutil + +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + MaskCandidate, + build_mask_candidates, + invert_mask_if_foreground_is_off_center, + save_binary_mask, + union_overlapping_mask_candidates, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, +) + +__all__ = ["prepare_scene_edit_assets"] + + +@dataclass(frozen=True) +class _AddedAssetInfo: + """Semantic information needed while preparing one newly added asset.""" + + object_id: str + category: str + name: str + description: str + + +def prepare_scene_edit_assets( + *, + scene_edit_plan: SceneEditPlan, + output_root: str | Path, + image_generation_client: ImageGenerationClient, + geometry_generation_client: GeometryGenerationClient, + image_segmentation_client: ImageSegmentationClient, +) -> list[SceneObject]: + """Prepare and return SimReady assets required by add operations.""" + # Prepare descriptions for all newly added objects. + added_asset_descriptions = _collect_added_asset_descriptions(scene_edit_plan) + # Skip asset generation when the edit plan only moves or deletes existing objects. + if not added_asset_descriptions: + return [] + + # Recreate this stage only when new assets need image, segmentation, and geometry outputs. + stage_output_root = ( + Path(output_root).expanduser().resolve() / "scene_editing" / "asset_preparation" + ) + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + generated_asset_images = _generate_added_asset_images( + added_asset_descriptions=added_asset_descriptions, + stage_output_root=stage_output_root, + image_generation_client=image_generation_client, + ) + generated_asset_masks = _segment_generated_added_asset_images( + added_asset_descriptions=added_asset_descriptions, + generated_asset_images=generated_asset_images, + stage_output_root=stage_output_root, + image_segmentation_client=image_segmentation_client, + ) + + generated_asset_glbs = _generate_added_assets_coarse_geometry( + generated_asset_images=generated_asset_images, + generated_asset_masks=generated_asset_masks, + stage_output_root=stage_output_root, + geometry_generation_client=geometry_generation_client, + ) + # Build a list of added SceneObjects. + added_assets = _build_added_scene_objects( + added_asset_descriptions=added_asset_descriptions, + generated_asset_glbs=generated_asset_glbs, + ) + # The temporary scene contains only new assets because the existing table is reused. + tmp_scene = Scene(objects=added_assets) + simready_processor = SimReadyProcessor( + scene=tmp_scene, + coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), + coarse_geometry_root=stage_output_root / "coarse_geometry", + simready_geometry_root=stage_output_root / "simready_geometry", + ) + # process_assets() validates and processes assets only; it does not require a table. + simready_processor.process_assets() + # Canonical GLBs use identity edit-time poses; layout editing sets them later. + _reset_added_asset_layouts(added_assets) + return added_assets + + +def _build_added_scene_objects( + *, + added_asset_descriptions: list[_AddedAssetInfo], + generated_asset_glbs: list[tuple[str, Path]], +) -> list[SceneObject]: + """Build temporary SceneObjects from generated coarse GLBs.""" + glbs_by_id = dict(generated_asset_glbs) + if len(glbs_by_id) != len(generated_asset_glbs): + raise ValueError("Generated asset GLBs must use unique object ids.") + assets: list[SceneObject] = [] + for asset_info in added_asset_descriptions: + glb_path = glbs_by_id.get(asset_info.object_id) + if glb_path is None: + raise ValueError(f"Generated asset {asset_info.object_id!r} has no GLB.") + assets.append( + SceneObject( + id=asset_info.object_id, + kind="asset", + category=asset_info.category, + name=asset_info.name, + description=asset_info.description, + simready_glb_path=str(glb_path), + ) + ) + return assets + + +def _reset_added_asset_layouts(added_assets: list[SceneObject]) -> None: + """Reset added asset poses after SimReady canonicalization.""" + for asset in added_assets: + asset.rot = [0.0, 0.0, 0.0] + asset.pos = [0.0, 0.0, 0.0] + asset.scale = [1.0, 1.0, 1.0] + + +def _coarse_layouts_by_id( + generated_asset_glbs: list[tuple[str, Path]], +) -> dict[str, dict[str, object]]: + """Build edit-time layouts with fixed identity poses and scale.""" + return { + object_id: { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + for object_id, _ in generated_asset_glbs + } + + +def _collect_added_asset_descriptions( + scene_edit_plan: SceneEditPlan, +) -> list[_AddedAssetInfo]: + """Return complete semantic information for add operations in plan order.""" + # Existing assets already have SimReady GLBs, so only add operations need assets. + added_asset_descriptions: list[_AddedAssetInfo] = [] + for operation in scene_edit_plan.operations: + if operation.op != "add": + continue + if ( + operation.object_id is None + or operation.category is None + or operation.name is None + or operation.description is None + ): + raise ValueError( + "Add operations must have an object_id, category, name, and description." + ) + added_asset_descriptions.append( + _AddedAssetInfo( + object_id=operation.object_id, + category=operation.category, + name=operation.name, + description=operation.description, + ) + ) + return added_asset_descriptions + + +def _generate_added_asset_images( + *, + added_asset_descriptions: list[_AddedAssetInfo], + stage_output_root: Path, + image_generation_client: ImageGenerationClient, +) -> list[tuple[str, Path]]: + """Generate one stable PNG for each new object description.""" + # Prepare a list. + generated_asset_images: list[tuple[str, Path]] = [] + # Create a subdir. + image_output_root = stage_output_root / "generated_images" + image_output_root.mkdir(parents=True, exist_ok=True) + + for asset_info in added_asset_descriptions: + object_id = asset_info.object_id + # Stable object IDs preserve the image-to-asset mapping across later stages. + image_path = image_generation_client.generate_image_by_prompt( + prompt=asset_info.description, + output_path=image_output_root / f"{object_id}.png", + ) + generated_asset_images.append((object_id, image_path)) + return generated_asset_images + + +def _segment_generated_added_asset_images( + *, + added_asset_descriptions: list[_AddedAssetInfo], + generated_asset_images: list[tuple[str, Path]], + stage_output_root: Path, + image_segmentation_client: ImageSegmentationClient, +) -> list[tuple[str, Path]]: + """Segment each generated image with its description and return binary masks.""" + asset_info_by_id = { + asset_info.object_id: asset_info for asset_info in added_asset_descriptions + } + if len(asset_info_by_id) != len(added_asset_descriptions): + raise ValueError("Added asset descriptions must use unique object ids.") + + masks_output_root = stage_output_root / "generated_masks" + masks_output_root.mkdir(parents=True, exist_ok=True) + generated_asset_masks: list[tuple[str, Path]] = [] + for object_id, image_path in generated_asset_images: + asset_info = asset_info_by_id.get(object_id) + if asset_info is None: + raise ValueError(f"Generated image {object_id!r} has no description.") + + candidates: list[MaskCandidate] = [] + # Retry with simpler semantic prompts when the detailed description is not found. + for prompt in (asset_info.description, asset_info.name, asset_info.category): + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=prompt, + ) + ), + min_iou=0.8, + ) + if candidates: + break + # A single generated object may still yield multiple SAM3 candidates; use the first one. + if not candidates: + raise ValueError( + f"Generated asset {object_id!r} produced no segmentation candidates." + ) + with Image.open(image_path) as image: + image_size = image.size + mask_path = save_binary_mask( + invert_mask_if_foreground_is_off_center(candidates[0]), + image_size=image_size, + output_path=masks_output_root / f"{object_id}_mask.png", + ) + generated_asset_masks.append((object_id, mask_path)) + return generated_asset_masks + + +def _generate_added_assets_coarse_geometry( + *, + generated_asset_images: list[tuple[str, Path]], + generated_asset_masks: list[tuple[str, Path]], + stage_output_root: Path, + geometry_generation_client: GeometryGenerationClient, +) -> list[tuple[str, Path]]: + """Generate one coarse GLB for each generated image and binary mask.""" + masks_by_id = dict(generated_asset_masks) + if len(masks_by_id) != len(generated_asset_masks): + raise ValueError("Generated asset masks must use unique object ids.") + if set(masks_by_id) != {object_id for object_id, _ in generated_asset_images}: + raise ValueError( + "Generated asset images and masks must have matching object ids." + ) + + geometry_output_root = stage_output_root / "coarse_geometry" + geometry_output_root.mkdir(parents=True, exist_ok=True) + generated_asset_glbs: list[tuple[str, Path]] = [] + for object_id, image_path in generated_asset_images: + # Each generated object has its own color image, so it needs an individual request. + geometry_generation_client.generate_objects( + image_path=image_path, + object_masks=[(object_id, masks_by_id[object_id])], + output_root=geometry_output_root, + ) + glb_path = geometry_output_root / f"{object_id}.glb" + if not glb_path.is_file(): + raise FileNotFoundError( + f"Geometry generation did not produce a GLB for {object_id!r}: {glb_path}" + ) + generated_asset_glbs.append((object_id, glb_path)) + return generated_asset_glbs diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 1eb3a857f..e66fd8653 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -29,6 +29,8 @@ SceneGraph, SceneGraphNode, SceneGraphRelation, + TABLE_REGIONS, + TableRegion, ) from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, @@ -52,6 +54,17 @@ For every move and every positioned add, target_id must be an Existing object ID and relation must be one of on, left_of, right_of, in_front_of, or behind. +When the target is the tabletop, use target_id "table", relation "on", and set +table_region to one of left_back, back_center, right_back, left_center, center, +right_center, left_front, front_center, or right_front. Do not use a planar +relation with the table. For non-table placement, table_region must be null. +If an add operation has no target_id and relation, it is placed on the table by +default and table_region must be null. +In the tabletop 9-grid, smaller x means left, larger x means right, smaller y +means back, and larger y means front: left_back is the upper-left/back cell, +back_center is the upper-center/back cell, right_back is the upper-right/back +cell, left_center/center/right_center are the middle row, and +left_front/front_center/right_front are the lower/front row. Do not position a new object relative to another newly added object. Each existing object's center_xy is its center position [x, y] in the @@ -67,8 +80,8 @@ table, or relations to any object. Return JSON only: no Markdown, comments, or prose. Every operation must contain -exactly these fields: op, object_id, target_id, relation, category, name, and -description. Use null for every field that does not apply to an operation: +exactly these fields: op, object_id, target_id, relation, table_region, category, +name, and description. Use null for every field that does not apply: { "operations": [ { @@ -76,6 +89,7 @@ "object_id": "bottle_001", "target_id": "book_001", "relation": "right_of", + "table_region": null, "category": null, "name": null, "description": null @@ -85,6 +99,7 @@ "object_id": "cup_001", "target_id": null, "relation": null, + "table_region": null, "category": null, "name": null, "description": null @@ -92,8 +107,9 @@ { "op": "add", "object_id": null, - "target_id": null, - "relation": null, + "target_id": "table", + "relation": "on", + "table_region": "back_center", "category": "orange", "name": "small orange", "description": "small round orange with a textured peel" @@ -103,9 +119,20 @@ "object_id": null, "target_id": "book_001", "relation": "right_of", + "table_region": null, "category": "orange", "name": "small orange", "description": "small round orange with a textured peel" + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "category": "banana", + "name": "yellow banana", + "description": "curved yellow banana with a green stem" } ] } @@ -128,7 +155,10 @@ def understand_scene_edit( if json_max_attempts < 1: raise ValueError("json_max_attempts must be at least 1.") # Give the VLM only the scene metadata needed to identify existing objects. - simplified_scene_info = _simplify_scene_info(scene=scene) + simplified_scene_info = _simplify_scene_info( + scene=scene, + scene_graph=scene_graph, + ) operations = _vlm_understand_scene_edit( scene=scene, edit_prompt=edit_prompt, @@ -192,7 +222,7 @@ def _apply_scene_edit_plan_to_scene_graph( """Apply the target graph updates implied by add and move operations.""" deleted_object_ids: set[str] = set() added_object_ids: list[str] = [] - on_parent_updates: list[tuple[str, str]] = [] + on_parent_updates: list[tuple[str, str, TableRegion | None]] = [] planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] for operation in scene_edit_plan.operations: if operation.op == "delete": @@ -206,7 +236,13 @@ def _apply_scene_edit_plan_to_scene_graph( if operation.target_id is None or operation.relation is None: continue if operation.relation == "on": - on_parent_updates.append((operation.object_id, operation.target_id)) + on_parent_updates.append( + ( + operation.object_id, + operation.target_id, + operation.table_region, + ) + ) continue planar_relation_updates.append( (operation.object_id, operation.relation, operation.target_id) @@ -221,8 +257,15 @@ def _apply_scene_edit_plan_to_scene_graph( ) -def _simplify_scene_info(scene: Scene) -> dict[str, object]: +def _simplify_scene_info( + *, + scene: Scene, + scene_graph: SceneGraph, +) -> dict[str, object]: """Return the object metadata needed for edit instruction resolution.""" + table_regions_by_id = { + node.object_id: node.table_region for node in scene_graph.nodes + } return { "existing_object_ids": [scene_object.id for scene_object in scene.objects], "objects": [ @@ -232,6 +275,7 @@ def _simplify_scene_info(scene: Scene) -> dict[str, object]: "name": scene_object.name, "description": scene_object.description, "center_xy": scene_object.center_xy, + "table_region": table_regions_by_id.get(scene_object.id), } for scene_object in scene.objects ], @@ -302,6 +346,7 @@ def _parse_scene_edit_operations( "object_id", "target_id", "relation", + "table_region", "category", "name", "description", @@ -345,6 +390,7 @@ def _parse_scene_edit_operations( value.get("target_id"), field_name="target_id" ), relation=_optional_relation(value.get("relation")), + table_region=_optional_table_region(value.get("table_region")), category=category, name=_optional_string(value.get("name"), field_name="name"), description=_optional_string( @@ -387,3 +433,11 @@ def _optional_relation(value: object) -> str | None: if value not in {"on", "left_of", "right_of", "in_front_of", "behind"}: raise ValueError("Scene edit operation relation is invalid.") return value + + +def _optional_table_region(value: object) -> TableRegion | None: + if value is None: + return None + if value not in TABLE_REGIONS: + raise ValueError("Scene edit operation table_region is invalid.") + return value diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 1cfe3f592..d8d82b19b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -48,8 +48,8 @@ quaternion_wxyz_to_euler_xyz_degrees, transform_matrix_to_layout_object, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.simready_scene_processor import ( - SimReadySceneProcessor, +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, ) from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( TableSupportSurfaceDetector, @@ -105,7 +105,7 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - simready_processor = SimReadySceneProcessor( + simready_processor = SimReadyProcessor( scene=scene, coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index e0c3f772a..ccf9af949 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -94,6 +94,35 @@ def decode_rle_mask(mask_rle: dict[str, Any]) -> Image.Image: return Image.frombytes("L", (width, height), bytes(pixels)) +def invert_mask_if_foreground_is_off_center(candidate: MaskCandidate) -> MaskCandidate: + """Invert a mask when its foreground is less concentrated at image center. + + This heuristic is intended for generated single-object images, where the + object is expected near the center and SAM3 may return its background. + """ + mask = decode_rle_mask(candidate.mask_rle) + width, height = mask.size + left, top = width // 6, height // 6 + right, bottom = width - left, height - top + center_mask = mask.crop((left, top, right, bottom)) + + center_foreground_ratio = _foreground_ratio(center_mask) + total_foreground_pixels = _foreground_pixel_count(mask) + outside_foreground_pixels = total_foreground_pixels - _foreground_pixel_count( + center_mask + ) + outside_pixel_count = width * height - center_mask.width * center_mask.height + outside_foreground_ratio = outside_foreground_pixels / outside_pixel_count + if center_foreground_ratio >= outside_foreground_ratio: + return candidate + + inverted_mask = mask.point(lambda value: 0 if value else 255) + return MaskCandidate( + index=candidate.index, + mask_rle=_encode_binary_mask_rle(inverted_mask), + ) + + def union_overlapping_mask_candidates( candidates: list[MaskCandidate], *, @@ -361,6 +390,19 @@ def _mask_iou(first_mask: Image.Image, second_mask: Image.Image) -> float: return intersection.histogram()[255] / union_pixels +def _foreground_pixel_count(mask: Image.Image) -> int: + """Return the number of white pixels in one binary mask.""" + return mask.convert("L").histogram()[255] + + +def _foreground_ratio(mask: Image.Image) -> float: + """Return the white-pixel ratio in one non-empty image region.""" + pixel_count = mask.width * mask.height + if pixel_count == 0: + raise ValueError("Mask region must contain at least one pixel.") + return _foreground_pixel_count(mask) / pixel_count + + def _encode_binary_mask_rle(mask: Image.Image) -> dict[str, Any]: binary_mask = mask.convert("L").point( lambda value: 255 if value else 0 diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 1ae6cdb1b..a59bc7bdf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -28,6 +28,7 @@ SceneGraph, SceneGraphNode, SceneGraphRelation, + TABLE_REGIONS, ) from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, @@ -203,18 +204,7 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: raise ValueError("Scene graph node ids must be strings or null.") if parent_relation not in {None, "on"}: raise ValueError("Scene graph parent_relation must be 'on' or null.") - if table_region not in { - None, - "left_back", - "back_center", - "right_back", - "left_center", - "center", - "right_center", - "left_front", - "front_center", - "right_front", - }: + if table_region is not None and table_region not in TABLE_REGIONS: raise ValueError("Scene graph table_region is invalid.") return SceneGraphNode( object_id=object_id, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py similarity index 97% rename from embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py rename to embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 40b4a6f94..6b5a09b0b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_scene_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -53,7 +53,7 @@ @dataclass(frozen=True) -class SimReadySceneProcessorConfig: +class SimReadyProcessorConfig: """Object-category policy for SimReady mesh canonicalization.""" upright_container_id_tokens: frozenset[str] = frozenset( @@ -61,8 +61,8 @@ class SimReadySceneProcessorConfig: ) # Object-id tokens that enable upright-container standardization. -class SimReadySceneProcessor: - """Create SimReady GLBs and layouts for one table and its scene assets.""" +class SimReadyProcessor: + """Create SimReady GLBs and layouts for scene objects.""" def __init__( self, @@ -71,7 +71,7 @@ def __init__( coarse_layout_by_id: dict[str, dict[str, object]], coarse_geometry_root: str | Path, simready_geometry_root: str | Path, - config: SimReadySceneProcessorConfig | None = None, + config: SimReadyProcessorConfig | None = None, ) -> None: self.scene = scene self.coarse_layout_by_id = coarse_layout_by_id @@ -81,7 +81,7 @@ def __init__( ) self.simready_table_layout: dict[str, object] | None = None self.simready_assets_layout: list[dict[str, object]] | None = None - self.config = config if config is not None else SimReadySceneProcessorConfig() + self.config = config if config is not None else SimReadyProcessorConfig() if not self.config.upright_container_id_tokens: raise ValueError("upright_container_id_tokens must not be empty.") @@ -317,8 +317,8 @@ def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: lower_points = standardized_points[ standardized_points[:, 2] < axis_min + axis_range * 0.2 ] - upper_volume = SimReadySceneProcessor._convex_hull_volume(upper_points) - lower_volume = SimReadySceneProcessor._convex_hull_volume(lower_points) + upper_volume = SimReadyProcessor._convex_hull_volume(upper_points) + lower_volume = SimReadyProcessor._convex_hull_volume(lower_points) # Bottles usually have a smaller top (neck) than bottom; flip if necessary. if upper_volume > lower_volume: diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index 30c11c995..5579e8d55 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -17,8 +17,11 @@ from __future__ import annotations import json +from pathlib import Path import pytest +from PIL import Image +import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( @@ -108,6 +111,7 @@ def test_scene_edit_plan_accepts_add_without_a_position() -> None: "object_id": "cup_001", "target_id": None, "relation": None, + "table_region": None, "category": "cup", "name": "green cup", "description": "A small green ceramic cup.", @@ -151,6 +155,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "object_id": None, "target_id": None, "relation": None, + "table_region": None, "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", @@ -160,6 +165,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "object_id": None, "target_id": None, "relation": None, + "table_region": None, "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", @@ -220,7 +226,9 @@ def test_scene_edit_plan_requires_a_position_for_move_operations() -> None: ) -def test_scene_edit_asset_preparation_skips_plans_without_adds() -> None: +def test_scene_edit_asset_preparation_skips_plans_without_adds( + tmp_path: Path, +) -> None: scene, scene_graph = _scene_and_graph() plan = SceneEditPlan( scene=scene, @@ -234,13 +242,141 @@ def test_scene_edit_asset_preparation_skips_plans_without_adds() -> None: ) ], ) + previous_asset_output = tmp_path / "scene_editing" / "asset_preparation" + previous_asset_output.mkdir(parents=True) + (previous_asset_output / "previous.txt").write_text("keep", encoding="utf-8") - prepared_scene = prepare_scene_edit_assets( + prepared_assets = prepare_scene_edit_assets( + scene_edit_plan=plan, + output_root=tmp_path, + image_generation_client=object(), # type: ignore[arg-type] + geometry_generation_client=object(), # type: ignore[arg-type] + image_segmentation_client=object(), # type: ignore[arg-type] + ) + + assert prepared_assets == [] + assert (previous_asset_output / "previous.txt").is_file() + + +def test_scene_edit_asset_preparation_generates_one_image_per_add( + tmp_path: Path, +) -> None: + class ImageGenerationClient: + def __init__(self) -> None: + self.requests: list[tuple[str, Path]] = [] + + def generate_image_by_prompt(self, *, prompt: str, output_path: Path) -> Path: + self.requests.append((prompt, output_path)) + Image.new("RGB", (6, 6), "white").save(output_path) + return output_path + + class ImageSegmentationClient: + def __init__(self) -> None: + self.requests: list[tuple[Path, str]] = [] + + def segment_single_object( + self, + *, + image_path: Path, + prompt: str, + ) -> list[dict[str, object]]: + self.requests.append((image_path, prompt)) + return [ + { + "size": [6, 6], + "counts": [7, 4, 2, 4, 2, 4, 2, 4, 7], + "starts_with": 1, + } + ] + + class GeometryGenerationClient: + def __init__(self) -> None: + self.requests: list[tuple[Path, list[tuple[str, Path]], Path]] = [] + + def generate_objects( + self, + *, + image_path: Path, + object_masks: list[tuple[str, Path]], + output_root: Path, + ) -> tuple[dict[str, object], list[dict[str, object]]]: + self.requests.append((image_path, object_masks, output_root)) + output_root.mkdir(parents=True, exist_ok=True) + for object_id, _ in object_masks: + trimesh.creation.box().export(output_root / f"{object_id}.glb") + return {}, [{"scale": [1.25, 1.5, 1.75]}] + + scene, scene_graph = _scene_and_graph() + plan = SceneEditPlan( scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="add", + object_id="cup_001", + category="cup", + name="green cup", + description="A small green ceramic cup.", + ) + ], + ) + image_generation_client = ImageGenerationClient() + image_segmentation_client = ImageSegmentationClient() + geometry_generation_client = GeometryGenerationClient() + + prepared_assets = prepare_scene_edit_assets( scene_edit_plan=plan, + output_root=tmp_path, + image_generation_client=image_generation_client, # type: ignore[arg-type] + geometry_generation_client=geometry_generation_client, # type: ignore[arg-type] + image_segmentation_client=image_segmentation_client, # type: ignore[arg-type] ) - assert prepared_scene is scene + expected_image_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "generated_images" + / "cup_001.png" + ) + assert [asset.id for asset in prepared_assets] == ["cup_001"] + assert prepared_assets[0].simready_glb_path is not None + assert prepared_assets[0].rot == [0.0, 0.0, 0.0] + assert prepared_assets[0].pos == [0.0, 0.0, 0.0] + assert prepared_assets[0].scale == [1.0, 1.0, 1.0] + assert image_generation_client.requests == [ + ("A small green ceramic cup.", expected_image_path) + ] + assert expected_image_path.is_file() + assert image_segmentation_client.requests == [ + (expected_image_path, "A small green ceramic cup.") + ] + generated_mask_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "generated_masks" + / "cup_001_mask.png" + ) + assert generated_mask_path.is_file() + with Image.open(generated_mask_path) as mask: + assert mask.getpixel((3, 3)) == 255 + assert mask.getpixel((0, 0)) == 0 + generated_glb_path = ( + tmp_path + / "scene_editing" + / "asset_preparation" + / "coarse_geometry" + / "cup_001.glb" + ) + assert geometry_generation_client.requests == [ + ( + expected_image_path, + [("cup_001", generated_mask_path)], + generated_glb_path.parent, + ) + ] + assert generated_glb_path.read_bytes().startswith(b"glTF") def test_scene_edit_graph_builder_copies_the_pre_edit_graph() -> None: From a2398ebdee2cc2a2d5fc7179521c8e6410e6c9d5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:27:18 +0800 Subject: [PATCH 12/85] finish basic simready real world scale + semantic-based rotation --- .../gen_sim/scene_engine/pipeline/edit.py | 1 + .../editing/scene_edit_asset_preparation.py | 11 + .../gen_sim/scene_engine/pipeline/generate.py | 1 + .../pipeline/generation/scene_generation.py | 11 + .../pipeline/utils/simready_processor.py | 96 ++++- .../utils/simready_processor_utils.py | 370 ++++++++++++++++++ 6 files changed, 486 insertions(+), 4 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index 6ba7898d1..fac2b6769 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -87,6 +87,7 @@ def edit_scene( image_generation_client=image_generation_client, geometry_generation_client=geometry_generation_client, image_segmentation_client=image_segmentation_client, + vlm_client=vlm_client, ) finally: image_generation_client.close() diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index 70c02d1a7..c6ec26fb3 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -35,6 +35,9 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( MaskCandidate, build_mask_candidates, @@ -44,6 +47,7 @@ ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( SimReadyProcessor, + SimReadyProcessorConfig, ) __all__ = ["prepare_scene_edit_assets"] @@ -66,6 +70,7 @@ def prepare_scene_edit_assets( image_generation_client: ImageGenerationClient, geometry_generation_client: GeometryGenerationClient, image_segmentation_client: ImageSegmentationClient, + vlm_client: OpenAICompatibleVLM | None = None, ) -> list[SceneObject]: """Prepare and return SimReady assets required by add operations.""" # Prepare descriptions for all newly added objects. @@ -112,6 +117,12 @@ def prepare_scene_edit_assets( coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), coarse_geometry_root=stage_output_root / "coarse_geometry", simready_geometry_root=stage_output_root / "simready_geometry", + # Scene editing will later provide the VLM-selected scale and rotation. + config=SimReadyProcessorConfig( + use_vlm_scale=vlm_client is not None, + use_vlm_rotation=vlm_client is not None, + ), + vlm_client=vlm_client, ) # process_assets() validates and processes assets only; it does not require a table. simready_processor.process_assets() diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 1df9027f9..144551c29 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -81,6 +81,7 @@ def generate_scene_from_image( scene=scene, scene_graph=scene_graph, geometry_generation_client=geometry_generation_client, + vlm_client=vlm_client, ) finally: geometry_generation_client.close() # Kill the session to avoid resource leaks. diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index d8d82b19b..0d736da76 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -30,6 +30,9 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_support_clamp import ( AssetsGroupSupportClamp, ) @@ -50,6 +53,7 @@ ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( SimReadyProcessor, + SimReadyProcessorConfig, ) from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( TableSupportSurfaceDetector, @@ -66,6 +70,7 @@ def generate_scene_and_refine( scene_graph: SceneGraph, *, geometry_generation_client: GeometryGenerationClient, + vlm_client: OpenAICompatibleVLM, ) -> Scene: resolved_image_path = _validate_image_path(image_path) @@ -110,6 +115,12 @@ def generate_scene_and_refine( coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, + # Image-to-scene uses the geometry service's coarse scale directly. + config=SimReadyProcessorConfig( + use_vlm_scale=False, + use_vlm_rotation=False, + ), + vlm_client=vlm_client, ) simready_assets_layout = simready_processor.process_assets() simready_table_layout = simready_processor.process_table() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 6b5a09b0b..b22e24f6b 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -32,6 +32,15 @@ ObjectPhysics, SceneObject, ) +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + query_vlm_object_rotation_and_target_size, + compute_uniform_xy_scale_for_target, + render_object_front_top_views, + rotate_glb_about_x_axis, +) from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { @@ -56,6 +65,9 @@ class SimReadyProcessorConfig: """Object-category policy for SimReady mesh canonicalization.""" + use_vlm_scale: bool = False # Use the VLM-selected asset scale. + use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. + upright_container_id_tokens: frozenset[str] = frozenset( {"bottle", "can", "jar", "flask", "thermos"} ) # Object-id tokens that enable upright-container standardization. @@ -72,6 +84,7 @@ def __init__( coarse_geometry_root: str | Path, simready_geometry_root: str | Path, config: SimReadyProcessorConfig | None = None, + vlm_client: OpenAICompatibleVLM | None = None, ) -> None: self.scene = scene self.coarse_layout_by_id = coarse_layout_by_id @@ -82,8 +95,13 @@ def __init__( self.simready_table_layout: dict[str, object] | None = None self.simready_assets_layout: list[dict[str, object]] | None = None self.config = config if config is not None else SimReadyProcessorConfig() + self.vlm_client = vlm_client if not self.config.upright_container_id_tokens: raise ValueError("upright_container_id_tokens must not be empty.") + if ( + self.config.use_vlm_scale or self.config.use_vlm_rotation + ) and vlm_client is None: + raise ValueError("vlm_client is required when VLM transforms are enabled.") def process_table(self) -> dict[str, object]: """Process the required scene table and return its SimReady layout.""" @@ -104,7 +122,13 @@ def process_assets(self) -> list[dict[str, object]]: self.simready_assets_layout = processed_assets return self.simready_assets_layout - def _process_object(self, scene_object: SceneObject) -> dict[str, object]: + def _process_object( + self, + scene_object: SceneObject, + *, + scale: object | None = None, + rot: object | None = None, + ) -> dict[str, object]: """Canonicalize one coarse object and write its SimReady GLB.""" object_id = scene_object.id object_role = scene_object.kind @@ -113,12 +137,18 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, object]: coarse_layout = self.coarse_layout_by_id.get(object_id) if coarse_layout is None: raise ValueError(f"Coarse layout does not contain object {object_id!r}.") + prepared_glb_path, vlm_scale = self._prepare_vlm_rotated_glb(scene_object) + selected_scale = scale + if selected_scale is None: + selected_scale = vlm_scale or coarse_layout.get("scale") simready_mesh, simready_transform = self._canonicalize_object_mesh( - coarse_glb_path=self.coarse_geometry_root / f"{object_id}.glb", + coarse_glb_path=prepared_glb_path, object_id=object_id, - rot=coarse_layout.get("rot"), + # An enabled external rotation replaces the coarse-layout rotation. + rot=coarse_layout.get("rot") if rot is None else rot, pos=coarse_layout.get("pos"), - scale=coarse_layout.get("scale"), + # An enabled VLM scale replaces the coarse-layout scale. + scale=selected_scale, ) output_path = self.simready_geometry_root / f"{object_id}.glb" output_path.parent.mkdir(parents=True, exist_ok=True) @@ -132,6 +162,64 @@ def _process_object(self, scene_object: SceneObject) -> dict[str, object]: log_info(f"Created SimReady {object_role}: {object_id!r}.") return {"id": object_id, **simready_transform} + def _prepare_vlm_rotated_glb( + self, scene_object: SceneObject + ) -> tuple[Path, list[float] | None]: + """Render, query, and optionally bake the VLM-selected x-axis rotation.""" + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + if not (self.config.use_vlm_scale or self.config.use_vlm_rotation): + return coarse_path, None + decision = self._vlm_transform_for_object( + scene_object, + use_scale=self.config.use_vlm_scale, + use_rotation=self.config.use_vlm_rotation, + ) + rotate_about_x = bool(decision["rotate_about_x"]) + vlm_scale = compute_uniform_xy_scale_for_target( + glb_path=coarse_path, + target_xy_size_cm=decision["target_xy_size_cm"], + rotate_about_x=rotate_about_x, + ) + rotated_path = rotate_glb_about_x_axis( + input_path=coarse_path, + output_path=self.simready_geometry_root + / "vlm_rotated" + / f"{scene_object.id}.glb", + rotate=rotate_about_x, + ) + # The scale flag controls whether this VLM-derived isotropic scale is used. + # Apply the same factor on x, y, and z to preserve the asset's proportions. + return ( + rotated_path, + [vlm_scale, vlm_scale, vlm_scale] if self.config.use_vlm_scale else None, + ) + + def _vlm_transform_for_object( + self, + scene_object: SceneObject, + *, + use_scale: bool, + use_rotation: bool, + ) -> dict[str, object]: + """Render the object and return the validated VLM pose decision.""" + del use_scale, use_rotation + assert self.vlm_client is not None + coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" + needed_layout = "This asset needs to be place on the table that will not move a lot after simulation." + debug_root = self.simready_geometry_root.parent / "debug" + rendered_path = render_object_front_top_views( + glb_path=coarse_path, + output_path=debug_root / "vlm_views" / f"{scene_object.id}.png", + ) + # Both semantic questions are always answered in one multimodal call. + return query_vlm_object_rotation_and_target_size( + scene_object_description=scene_object.description, + needed_layout=needed_layout, + rendered_views_path=rendered_path, + vlm_client=self.vlm_client, + debug_output_path=debug_root / "vlm_outputs" / f"{scene_object.id}.json", + ) + @staticmethod def _fixed_physics_for_kind(kind: str) -> ObjectPhysics: """Create the fixed initial physics profile for one SimReady object.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py new file mode 100644 index 000000000..b7f718374 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -0,0 +1,370 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw, ImageFont +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) + +_VLM_SYSTEM_PROMPT = """You inspect one isolated 3D object from front and top views. +Use the object description and the rendered views together. + +Use the rendered views and the needed layout to decide whether the object should +be rotated around its own center by +90 degrees around the z-up world's x axis. +The z-up world is right-handed: x is left-right, y is front-back, and z is up. +In the composed image, FRONT VIEW is the left panel: x is horizontal and z is +vertical; the upper-right marker shows the positive z direction. TOP VIEW is +the right panel: x is horizontal and y is vertical; the upper-right markers +show the positive x and y directions. +Do not confuse the top view with looking at the object from above in the image +description: it is a projection along the z axis onto the x-y plane. +After deciding and applying that rotation, estimate the object's desired AABB +footprint on the x-y plane in real-world centimetres. The first value is the x +size and the second value is the y size. + +Return JSON only with exactly this schema: +{ + "rotate_about_x": false, + "target_xy_size_cm": [12.0, 5.0] +} + +Examples: +- Fork lying flat on a table: in FRONT VIEW the fork is mostly a thin + horizontal line; in TOP VIEW its length is visible. Keep it flat with + rotate_about_x=false, and use the tabletop footprint, for example + target_xy_size_cm=[15.0, 3.0]. +- Fork placed in a pen holder: the desired fork is upright, so its long axis is + approximately z. If the input coarse fork is lying in the x-y plane, set + rotate_about_x=true; if the input coarse fork is already upright, set it to + false. The target is the footprint inside the holder, not the fork's full + length, for example target_xy_size_cm=[3.0, 3.0]. +- Fork requested to lie flat on a table even when the input coarse fork is + upright: set rotate_about_x=true and estimate the final flat footprint, for + example target_xy_size_cm=[15.0, 3.0]. +- Bottle already standing on its flat base: keep it upright with + rotate_about_x=false and use target_xy_size_cm=[8.0, 8.0]. +""" + + +def render_object_front_top_views( + *, + glb_path: str | Path, + output_path: str | Path, + resolution: int = 512, +) -> Path: + """Render fixed z-up front/top views and compose them horizontally.""" + if resolution <= 0: + raise ValueError("resolution must be positive.") + source_path = Path(glb_path).expanduser().resolve() + if not source_path.is_file(): + raise FileNotFoundError(f"GLB for VLM rendering not found: {source_path}") + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + front_path = output_path.with_name(f"{output_path.stem}_front.png") + top_path = output_path.with_name(f"{output_path.stem}_top.png") + try: + import bpy + from mathutils import Vector + except ImportError as exc: + raise RuntimeError( + "Blender's bpy is required for SimReady VLM view rendering." + ) from exc + + bpy.ops.wm.read_factory_settings(use_empty=True) + bpy.ops.import_scene.gltf(filepath=str(source_path)) + if not any(obj.type == "MESH" for obj in bpy.context.scene.objects): + raise ValueError(f"GLB contains no mesh objects: {source_path}") + scene = bpy.context.scene + # Eevee renders imported GLB materials and textures instead of Workbench previews. + try: + scene.render.engine = "BLENDER_EEVEE_NEXT" + except TypeError: + scene.render.engine = "BLENDER_EEVEE" + scene.render.resolution_x = resolution + scene.render.resolution_y = resolution + scene.render.resolution_percentage = 100 + scene.render.image_settings.file_format = "PNG" + scene.render.film_transparent = False + if scene.world is None: + scene.world = bpy.data.worlds.new("VLM_World") + scene.world.color = (0.08, 0.08, 0.08) + for name, location, energy in ( + ("VLM_Key", (2.0, -2.0, 3.0), 700.0), + ("VLM_Fill", (-2.0, 1.0, 2.0), 400.0), + ): + light_data = bpy.data.lights.new(name, type="AREA") + light_data.energy = energy + light_data.shape = "DISK" + light_data.size = 4.0 + light = bpy.data.objects.new(name, light_data) + light.location = location + light.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - light.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.collection.objects.link(light) + camera_data = bpy.data.cameras.new("VLM_Camera") + camera = bpy.data.objects.new("VLM_Camera", camera_data) + scene.collection.objects.link(camera) + scene.camera = camera + camera.data.type = "ORTHO" + camera.data.ortho_scale = 1.25 + + def render_view(path: Path, location: tuple[float, float, float]) -> None: + camera.location = location + camera.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - camera.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.render.filepath = str(path) + bpy.ops.render.render(write_still=True) + + # Blender uses a right-handed z-up world; front is viewed along +y. + render_view(front_path, (0.0, -3.0, 0.0)) + render_view(top_path, (0.0, 0.0, 3.0)) + with Image.open(front_path) as front, Image.open(top_path) as top: + composed = Image.new("RGB", (resolution * 2, resolution), "white") + composed.paste(front.convert("RGB"), (0, 0)) + composed.paste(top.convert("RGB"), (resolution, 0)) + draw = ImageDraw.Draw(composed) + # Use a readable scaled font for the panel labels when available. + try: + font = ImageFont.truetype( + "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", + max(24, resolution // 16), + ) + except OSError: + font = ImageFont.load_default() + # Label each panel so the VLM and manual debugging can distinguish views. + for label, origin in (("FRONT VIEW", (0, 0)), ("TOP VIEW", (resolution, 0))): + x, y = origin + text_box = draw.textbbox((x + 16, y + 16), label, font=font) + draw.rectangle( + (text_box[0] - 8, text_box[1] - 6, text_box[2] + 8, text_box[3] + 6), + fill="white", + ) + draw.text((x + 16, y + 16), label, fill="black", font=font) + # Mark the positive axes used by each projection for VLM interpretation. + _draw_arrow( + draw, + (resolution - 62, 62), + (resolution - 62, 20), + "+Z", + font, + color="blue", + ) + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 42, 62), + "+X", + font, + color="red", + ) + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 92, 20), + "+Y", + font, + color="green", + ) + composed.save(output_path) + return output_path + + +def _draw_arrow( + draw: ImageDraw.ImageDraw, + start: tuple[int, int], + end: tuple[int, int], + label: str, + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, + color: str, +) -> None: + """Draw one labeled positive-axis arrow on a rendered view.""" + dx, dy = end[0] - start[0], end[1] - start[1] + length = max(abs(dx), abs(dy)) + if length == 0: + raise ValueError("Axis arrow start and end must differ.") + unit_x, unit_y = dx / length, dy / length + perpendicular_x, perpendicular_y = -unit_y, unit_x + head_length = 14.0 + head_width = 8.0 + tip_x, tip_y = end + base_x = tip_x - unit_x * head_length + base_y = tip_y - unit_y * head_length + arrowhead = ( + (tip_x, tip_y), + ( + base_x + perpendicular_x * head_width, + base_y + perpendicular_y * head_width, + ), + ( + base_x - perpendicular_x * head_width, + base_y - perpendicular_y * head_width, + ), + ) + draw.line((*start, *end), fill=color, width=4) + draw.polygon(arrowhead, fill=color) + # Put each axis label beside its arrowhead so it does not cover the arrow. + draw.text((int(tip_x + 8), int(tip_y - 8)), label, fill=color, font=font) + + +def query_vlm_object_rotation_and_target_size( + *, + scene_object_description: str, + needed_layout: str, + rendered_views_path: str | Path, + vlm_client: OpenAICompatibleVLM, + debug_output_path: str | Path | None = None, +) -> dict[str, object]: + """Ask the VLM for rotation and post-rotation tabletop footprint.""" + response_text = vlm_client.complete( + system_prompt=_VLM_SYSTEM_PROMPT, + user_prompt=( + f"Object description:\n{scene_object_description}\n\n" + f"Needed layout:\n{needed_layout}\n\n" + "The image contains front view on the left and top view on the right." + ), + image_path=rendered_views_path, + ) + try: + value = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM transform response is not valid JSON: {exc}") from exc + if not isinstance(value, dict) or set(value) != { + "rotate_about_x", + "target_xy_size_cm", + }: + raise ValueError( + "VLM transform response must contain exactly rotate_about_x and " + "target_xy_size_cm." + ) + if not isinstance(value["rotate_about_x"], bool): + raise ValueError("VLM rotate_about_x must be boolean.") + target_size = value["target_xy_size_cm"] + if ( + not isinstance(target_size, list) + or len(target_size) != 2 + or not all(isinstance(item, (int, float)) for item in target_size) + or not all(np.isfinite(item) and item > 0 for item in target_size) + ): + raise ValueError("VLM target_xy_size_cm must contain two positive numbers.") + if debug_output_path is not None: + output_path = Path(debug_output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + { + "description": scene_object_description, + "needed_layout": needed_layout, + "rendered_views_path": str( + Path(rendered_views_path).expanduser().resolve() + ), + "vlm_output": value, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return value + + +def compute_uniform_xy_scale_for_target( + *, + glb_path: str | Path, + target_xy_size_cm: list[float], + rotate_about_x: bool, +) -> float: + """Compute an isotropic scale from the rotated mesh XY AABB and target size.""" + loaded = trimesh.load(Path(glb_path).expanduser().resolve(), process=False) + mesh = ( + loaded.dump(concatenate=True) if isinstance(loaded, trimesh.Scene) else loaded + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"GLB is not a mesh: {glb_path}") + if len(target_xy_size_cm) != 2 or any(value <= 0 for value in target_xy_size_cm): + raise ValueError("target_xy_size_cm must contain two positive values.") + if rotate_about_x: + center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-center) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(transform) + mesh.apply_translation(center) + actual_xy_size = mesh.bounds[1, :2] - mesh.bounds[0, :2] + if np.any(actual_xy_size <= 0): + raise ValueError("Rotated mesh must have a positive XY AABB.") + # Convert the VLM's centimetres to metres before comparing with the GLB AABB. + target_xy_size_m = np.asarray(target_xy_size_cm, dtype=float) / 100.0 + axis_scales = target_xy_size_m / actual_xy_size + # Use sqrt(target XY area / actual XY area) as one uniform scale on all axes. + return float(np.sqrt(axis_scales[0] * axis_scales[1])) + + +def rotate_glb_about_x_axis( + *, + input_path: str | Path, + output_path: str | Path, + rotate: bool, +) -> Path: + """Bake an optional +90-degree x-axis rotation around the mesh centre.""" + # Current coarse layouts are either flat on xy with possible random z rotation, + # or upright with almost no random y rotation, so this x-axis toggle is enough. + source_path = Path(input_path).expanduser().resolve() + destination_path = Path(output_path).expanduser().resolve() + destination_path.parent.mkdir(parents=True, exist_ok=True) + loaded = trimesh.load(source_path, process=False) + mesh = ( + loaded.dump(concatenate=True) if isinstance(loaded, trimesh.Scene) else loaded + ) + if not isinstance(mesh, trimesh.Trimesh): + raise ValueError(f"GLB is not a mesh: {source_path}") + if rotate: + center = mesh.bounds.mean(axis=0) + mesh.apply_translation(-center) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(transform) + mesh.apply_translation(center) + mesh.export(destination_path, file_type="glb") + return destination_path + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].strip() == "```": + lines = lines[:-1] + return "\n".join(lines).strip() From ccab0f3dd93a826375220d1fd35e7de32aa41510 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 15:34:00 +0800 Subject: [PATCH 13/85] Moved the table support infos into simready. --- .../gen_sim/scene_engine/core/scene_object.py | 6 ++ .../pipeline/generation/scene_generation.py | 67 +++++++++++-------- .../pipeline/utils/scene_exporter.py | 3 + .../pipeline/utils/scene_importer.py | 27 ++++++++ .../pipeline/utils/simready_processor.py | 39 +++++++++++ .../pipeline/utils/table_support_surface.py | 53 ++++++++++++++- 6 files changed, 167 insertions(+), 28 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index 0a8ed3aea..2b868e3c3 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -67,6 +67,9 @@ class SceneObject: pos: list[float] | None = None # Final y-up world position in metres. scale: list[float] | None = None # Final y-up object scale. center_xy: list[float] | None = None # Z-up table-frame XY AABB center. + support_surface_z: float | None = None # Detected tabletop height in z-up. + support_contour_xy: list[list[float]] | None = None # Outer support contour. + support_optimization_rect_xy: list[list[float]] | None = None # Safe XY rectangle. physics: ObjectPhysics | None = None # Assigned when SimReady processing succeeds. def to_dict(self) -> dict[str, object]: @@ -83,5 +86,8 @@ def to_dict(self) -> dict[str, object]: "pos": self.pos, "scale": self.scale, "center_xy": self.center_xy, + "support_surface_z": self.support_surface_z, + "support_contour_xy": self.support_contour_xy, + "support_optimization_rect_xy": self.support_optimization_rect_xy, "physics": self.physics.to_dict() if self.physics is not None else None, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 0d736da76..60f91ca81 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -23,6 +23,7 @@ import numpy as np import trimesh +from shapely.geometry import Polygon from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, @@ -55,9 +56,6 @@ SimReadyProcessor, SimReadyProcessorConfig, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( - TableSupportSurfaceDetector, -) from embodichain.utils.logger import log_info _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} @@ -115,6 +113,7 @@ def generate_scene_and_refine( coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, + debug_output_root=debug_output_root, # Image-to-scene uses the geometry service's coarse scale directly. config=SimReadyProcessorConfig( use_vlm_scale=False, @@ -248,6 +247,18 @@ def _update_scene_final_y_up_layout_and_z_up_centers( ) # Persist AABB centers for future scene-edit object disambiguation. scene.table.center_xy = table_mesh.bounds[:, :2].mean(axis=0).tolist() + if scene.table.support_contour_xy is not None: + # Move SimReady-local support geometry into the final table-frame position. + table_center_xy = np.asarray(scene.table.center_xy, dtype=float) + scene.table.support_contour_xy = [ + (np.asarray(point, dtype=float) + table_center_xy).tolist() + for point in scene.table.support_contour_xy + ] + if scene.table.support_optimization_rect_xy is not None: + scene.table.support_optimization_rect_xy = [ + (np.asarray(point, dtype=float) + table_center_xy).tolist() + for point in scene.table.support_optimization_rect_xy + ] for asset in scene.assets: asset.center_xy = assets_aabb_corners_by_id[asset.id].mean(axis=0).tolist() @@ -360,31 +371,36 @@ def _layout_refinement( log_info("Scene has no movable assets; skipping support-region clamping.") return refined_table_layout, [] - # 4. Detect the actual upward support triangles instead of projecting the - # entire table mesh to one convex hull. The result retains concavities - # (for example, an L-shaped tabletop) and is the only boundary used for - # placement below. - ( - table_world_mesh_z_up, - assets_aabb_2d_z_up_world_corners_by_id, - ) = _measure_table_and_assets_in_z_up_world( - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, - geometry_root=simready_geometry_output_root, - ) - support_detector = TableSupportSurfaceDetector( - table_world_mesh=table_world_mesh_z_up, - debug_output_root=debug_output_root, + # 4. Reuse support geometry detected during SimReady processing. + if ( + scene.table is None + or scene.table.support_contour_xy is None + or scene.table.support_optimization_rect_xy is None + ): + raise ValueError("Scene table has no persisted support geometry.") + table_support_polygon = Polygon(scene.table.support_contour_xy) + table_optimization_rectangle = Polygon(scene.table.support_optimization_rect_xy) + if not table_support_polygon.is_valid or table_support_polygon.is_empty: + raise ValueError("Scene table support contour is not a valid polygon.") + if ( + not table_optimization_rectangle.is_valid + or table_optimization_rectangle.is_empty + ): + raise ValueError("Scene table optimization rectangle is not valid.") + _, assets_aabb_2d_z_up_world_corners_by_id = ( + _measure_table_and_assets_in_z_up_world( + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + geometry_root=simready_geometry_output_root, + ) ) - table_support_region = support_detector.detect() - support_detector.save_support_surface_debug_images() # 5. Keep the complete clutter rigid in the table plane. A successful # result applies one shared z-up XY delta to every AABB, so it preserves # all existing asset-to-asset relations. It is *not* an asset packing # pass: pre-existing overlap is deliberately left to a later optimizer. group_clamp = AssetsGroupSupportClamp( - support_region=table_support_region.support_polygon, + support_region=table_support_polygon, assets_aabb_2d_z_up_world_corners_by_id=( assets_aabb_2d_z_up_world_corners_by_id ), @@ -405,13 +421,10 @@ def _layout_refinement( ) ) - # 6. Restore the previous pairwise AABB separation stage, but constrain - # every candidate with the actual support polygon rather than the legacy - # largest internal rectangle. Assets may now move independently only as - # much as needed to remove overlap; every resulting AABB remains on the - # L-shaped, circular, or otherwise non-convex support region. + # 6. Optimize independent asset positions inside the conservative rectangle. + # The clamp above already used the exact outer contour for the shared shift. overlap_optimizer = AssetsSupportLayoutOptimizer( - support_region=table_support_region.support_polygon, + support_region=table_optimization_rectangle, assets_aabb_2d_z_up_world_corners_by_id=( clamped_assets_aabb_2d_z_up_world_corners_by_id ), diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index 391ccd156..4063bf92c 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -193,6 +193,9 @@ def _scene_object_config( # which SimulationManager itself converts to z-up. "body_scale": scale_y_up, "center_xy": scene_object.center_xy, + "support_surface_z": scene_object.support_surface_z, + "support_contour_xy": scene_object.support_contour_xy, + "support_optimization_rect_xy": scene_object.support_optimization_rect_xy, "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index a59bc7bdf..eefeee6b6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -265,6 +265,16 @@ def _scene_object_from_export_entry( center_xy = entry.get("center_xy") if center_xy is not None: center_xy = self._vector2(center_xy, field_name=f"{uid}.center_xy") + support_surface_z = entry.get("support_surface_z") + if support_surface_z is not None: + support_surface_z = float(support_surface_z) + support_contour_xy = self._points2( + entry.get("support_contour_xy"), field_name=f"{uid}.support_contour_xy" + ) + support_optimization_rect_xy = self._points2( + entry.get("support_optimization_rect_xy"), + field_name=f"{uid}.support_optimization_rect_xy", + ) pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() @@ -284,6 +294,9 @@ def _scene_object_from_export_entry( pos=pos_y_up.tolist(), scale=scale, center_xy=center_xy, + support_surface_z=support_surface_z, + support_contour_xy=support_contour_xy, + support_optimization_rect_xy=support_optimization_rect_xy, physics=ObjectPhysics( body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), @@ -345,6 +358,20 @@ def _vector2(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be finite.") return vector + @classmethod + def _points2(cls, value: object, *, field_name: str) -> list[list[float]] | None: + """Validate an optional list of XY points from the scene export.""" + if value is None: + return None + if not isinstance(value, list) or len(value) < 3: + raise ValueError( + f"Scene config field {field_name!r} must contain 3 points." + ) + return [ + cls._vector2(point, field_name=f"{field_name}[{index}]") + for index, point in enumerate(value) + ] + @staticmethod def _physics_attrs(value: object) -> dict[str, float | int]: """Validate exported physics attributes.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index b22e24f6b..55860be7f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -41,6 +41,9 @@ render_object_front_top_views, rotate_glb_about_x_axis, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( + TableSupportSurfaceDetector, +) from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { @@ -83,6 +86,7 @@ def __init__( coarse_layout_by_id: dict[str, dict[str, object]], coarse_geometry_root: str | Path, simready_geometry_root: str | Path, + debug_output_root: str | Path | None = None, config: SimReadyProcessorConfig | None = None, vlm_client: OpenAICompatibleVLM | None = None, ) -> None: @@ -92,6 +96,12 @@ def __init__( self.simready_geometry_root = ( Path(simready_geometry_root).expanduser().resolve() ) + # Save rendered debug images. + self.debug_output_root = ( + Path(debug_output_root).expanduser().resolve() + if debug_output_root is not None + else None + ) self.simready_table_layout: dict[str, object] | None = None self.simready_assets_layout: list[dict[str, object]] | None = None self.config = config if config is not None else SimReadyProcessorConfig() @@ -159,9 +169,38 @@ def _process_object( ) scene_object.simready_glb_path = str(output_path) scene_object.physics = self._fixed_physics_for_kind(object_role) + # For table. (currently the id is fixed into table) + if object_role == "table": + # Detect and persist all reusable tabletop support geometry at SimReady time. + support_detector = TableSupportSurfaceDetector( + table_world_mesh=self._z_up_table_mesh(simready_mesh), + debug_output_root=self.debug_output_root, + ) + support_region = support_detector.detect() + scene_object.support_surface_z = support_region.top_z + scene_object.support_contour_xy = [ + [float(x), float(y)] + for x, y in support_region.support_polygon.exterior.coords[:-1] + ] + scene_object.support_optimization_rect_xy = [ + [float(x), float(y)] + for x, y in support_region.optimization_rectangle.exterior.coords[:-1] + ] + if self.debug_output_root is not None: + # Keep the 3D selected surface and 2D contour diagnostics beside SimReady output. + support_detector.save_support_surface_debug_images() log_info(f"Created SimReady {object_role}: {object_id!r}.") return {"id": object_id, **simready_transform} + @staticmethod + def _z_up_table_mesh(mesh: trimesh.Trimesh) -> trimesh.Trimesh: + """Convert one canonical y-up GLB mesh into the detector's z-up frame.""" + y_up_to_z_up = np.eye(4) + y_up_to_z_up[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + z_up_mesh = mesh.copy() + z_up_mesh.apply_transform(y_up_to_z_up) + return z_up_mesh + def _prepare_vlm_rotated_glb( self, scene_object: SceneObject ) -> tuple[Path, list[float] | None]: diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py index 6c3ad3e94..983c16098 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_support_surface.py @@ -55,6 +55,7 @@ class TableSupportRegion: vertices: np.ndarray # Full z-up table vertex array referenced by ``faces``. faces: np.ndarray # Indices of triangles selected as the main support surface. support_polygon: Polygon # Largest valid outer support contour in z-up XY. + optimization_rectangle: Polygon # Axis-aligned rectangle fully inside the contour. class TableSupportSurfaceDetector: @@ -131,11 +132,13 @@ def detect(self) -> TableSupportRegion: selected_vertices = face_vertices[selected_faces] vertices = mesh.vertices.copy() faces = mesh.faces[selected_faces].copy() + support_polygon = self._extract_largest_support_polygon(vertices[faces, :2]) self.support_region = TableSupportRegion( top_z=float(selected_vertices[:, :, 2].max()), vertices=vertices, faces=faces, - support_polygon=self._extract_largest_support_polygon(vertices[faces, :2]), + support_polygon=support_polygon, + optimization_rectangle=self._largest_inscribed_rectangle(support_polygon), ) return self.support_region @@ -367,6 +370,46 @@ def _extract_largest_support_polygon(cls, triangles_xy: np.ndarray) -> Polygon: raise ValueError("The merged 2D support contour is degenerate.") return Polygon(boundary_xy) + @staticmethod + def _largest_inscribed_rectangle(polygon: Polygon) -> Polygon: + """Find a conservative axis-aligned rectangle contained by the support contour.""" + coordinates = np.asarray(polygon.exterior.coords[:-1], dtype=float) + x_values = np.unique(coordinates[:, 0]) + y_values = np.unique(coordinates[:, 1]) + # Keep the search bounded for highly tessellated support contours. + if len(x_values) > 48: + x_values = x_values[np.linspace(0, len(x_values) - 1, 48, dtype=int)] + if len(y_values) > 48: + y_values = y_values[np.linspace(0, len(y_values) - 1, 48, dtype=int)] + + best_rectangle: Polygon | None = None + best_area = 0.0 + for x_index, minimum_x in enumerate(x_values[:-1]): + for maximum_x in x_values[x_index + 1 :]: + if maximum_x <= minimum_x: + continue + for y_index, minimum_y in enumerate(y_values[:-1]): + for maximum_y in y_values[y_index + 1 :]: + if maximum_y <= minimum_y: + continue + rectangle = Polygon( + [ + (minimum_x, minimum_y), + (maximum_x, minimum_y), + (maximum_x, maximum_y), + (minimum_x, maximum_y), + ] + ) + area = rectangle.area + if area > best_area and polygon.covers(rectangle): + best_rectangle = rectangle + best_area = area + if best_rectangle is None: + raise ValueError( + "Support contour has no non-degenerate inscribed rectangle." + ) + return best_rectangle + @staticmethod def _face_adjacency(mesh: trimesh.Trimesh) -> dict[int, set[int]]: """Build a face adjacency dictionary for the mesh.""" @@ -476,6 +519,14 @@ def _save_support_region_2d_image( linewidth=2.0, label="outer support contour", ) + rectangle_xy = np.asarray(support_region.optimization_rectangle.exterior.coords) + axis.plot( + rectangle_xy[:, 0], + rectangle_xy[:, 1], + color="seagreen", + linewidth=2.0, + label="optimization rectangle", + ) axis.autoscale_view() axis.set_aspect("equal", adjustable="box") axis.set_xlabel("x (z-up world)") From 9b27be2541a38e66a50944cfddc1fe4da01bf4b6 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 19:16:56 +0800 Subject: [PATCH 14/85] Finished scene edit --- .../scene_engine/cli/test_text_to_simready.py | 183 ++++ .../gen_sim/scene_engine/pipeline/edit.py | 33 +- .../editing/scene_edit_layout_generation.py | 77 ++ .../pipeline/utils/scene_exporter.py | 35 +- .../pipeline/utils/scene_importer.py | 26 +- .../utils/scene_layout_constructor.py | 399 +++++++++ .../pipeline/utils/scene_layout_optimizer.py | 785 ++++++++++++++++++ .../test_scene_core_and_export.py | 69 ++ 8 files changed, 1593 insertions(+), 14 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py b/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py new file mode 100644 index 000000000..e3db15b16 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py @@ -0,0 +1,183 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import argparse +import json +from collections.abc import Sequence +from pathlib import Path +import shutil + +from PIL import Image + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + build_mask_candidates, + invert_mask_if_foreground_is_off_center, + save_binary_mask, + union_overlapping_mask_candidates, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, +) + +__all__ = ["main", "run_text_to_simready"] + + +def run_text_to_simready(*, text: str, output_root: str | Path) -> SceneObject: + """Run one manual text-to-SimReady asset pipeline for debugging.""" + text = text.strip() + if not text: + raise ValueError("Text prompt must not be empty.") + + root = Path(output_root).expanduser().resolve() + if root.exists(): + shutil.rmtree(root) + debug_root = root / "debug" + image_root = root / "generated_images" + mask_root = root / "masks" + coarse_root = root / "coarse_geometry" + simready_root = root / "simready_geometry" + for directory in (debug_root, image_root, mask_root, coarse_root, simready_root): + directory.mkdir(parents=True, exist_ok=True) + + object_id = "asset_001" + scene_object = SceneObject( + id=object_id, + kind="asset", + category="asset", + name=text, + description=text, + ) + + image_generation_client = ImageGenerationClient.from_dotenv() + image_segmentation_client = ImageSegmentationClient.from_dotenv() + geometry_generation_client = GeometryGenerationClient.from_dotenv() + vlm_client = OpenAICompatibleVLM.from_dotenv() + try: + image_generation_client.check_health() + image_segmentation_client.check_health() + geometry_generation_client.check_health() + + # Generate a centered single-object image from the semantic text prompt. + image_path = image_generation_client.generate_image_by_prompt( + prompt=text, + output_path=image_root / f"{object_id}.png", + ) + with Image.open(image_path) as image: + image_size = image.size + + # Segment the generated object and apply the single-object foreground heuristic. + candidates = union_overlapping_mask_candidates( + build_mask_candidates( + image_segmentation_client.segment_single_object( + image_path=image_path, + prompt=text, + ) + ), + min_iou=0.8, + ) + if not candidates: + raise ValueError("Image segmentation returned no mask candidates.") + mask_path = save_binary_mask( + invert_mask_if_foreground_is_off_center(candidates[0]), + image_size=image_size, + output_path=mask_root / f"{object_id}.png", + ) + + # Generate one coarse GLB using the generated image and its binary mask. + geometry_generation_client.generate_objects( + image_path=image_path, + object_masks=[(object_id, mask_path)], + output_root=coarse_root, + ) + coarse_glb_path = coarse_root / f"{object_id}.glb" + if not coarse_glb_path.is_file(): + raise FileNotFoundError(f"Coarse GLB was not generated: {coarse_glb_path}") + + # Use identity coarse layout; VLM determines rotation and real-world size. + processor = SimReadyProcessor( + scene=Scene(objects=[scene_object]), + coarse_layout_by_id={ + object_id: { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + }, + coarse_geometry_root=coarse_root, + simready_geometry_root=simready_root, + config=SimReadyProcessorConfig( + use_vlm_scale=True, + use_vlm_rotation=True, + ), + vlm_client=vlm_client, + ) + simready_layout = processor.process_assets() + (root / "result.json").write_text( + json.dumps( + { + "input_text": text, + "scene_object": scene_object.to_dict(), + "simready_layout": simready_layout, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return scene_object + finally: + image_generation_client.close() + image_segmentation_client.close() + geometry_generation_client.close() + + +def main(argv: Sequence[str] | None = None) -> None: + """Run the manual text-to-SimReady CLI.""" + parser = argparse.ArgumentParser( + prog="embodichain test-text-to-simready", + description="Debug text-to-image-to-segmentation-to-SimReady generation.", + ) + parser.add_argument("--text", required=True, help="Description of one object.") + parser.add_argument( + "--output_root", + required=True, + help="Directory for all intermediate and final artifacts.", + ) + args = parser.parse_args(argv) + scene_object = run_text_to_simready(text=args.text, output_root=args.output_root) + print(f"Generated SimReady asset: {scene_object.simready_glb_path}") + + +if __name__ == "__main__": + main() diff --git a/embodichain/gen_sim/scene_engine/pipeline/edit.py b/embodichain/gen_sim/scene_engine/pipeline/edit.py index fac2b6769..5b71af556 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/edit.py +++ b/embodichain/gen_sim/scene_engine/pipeline/edit.py @@ -33,12 +33,18 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( SceneExportImporter, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import ( + SceneExporter, +) from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( understand_scene_edit, ) from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( prepare_scene_edit_assets, ) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_layout_generation import ( + edit_layout, +) from embodichain.utils.logger import log_info @@ -95,20 +101,25 @@ def edit_scene( image_segmentation_client.close() log_info("Completed Objects Preparation") - # 3. Layout Editing - log_info("Starting Layout Editing") - # scene = edit_layout( - # scene=scene, - # edit_plan=edit_plan, - # scene_graph=updated_scene_graph, - # output_root=output_root, - # ) - log_info("Completed Layout Editing") + # 3. Layout Generation + log_info("Starting Layout Generation") + post_edit_scene = edit_layout( + scene=scene, + scene_edit_plan=scene_edit_plan, + updated_scene_graph=updated_scene_graph, + added_assets=added_assets, + output_root=resolved_output_root, + ) + log_info("Completed Layout Generation") # 4. Scene Export - # Re export the scene to the same output format, - # and delete some temporary files or folders. log_info("Starting Scene Export") + scene_exporter = SceneExporter( + scene=post_edit_scene, + scene_graph=updated_scene_graph, + output_root=resolved_output_root, + ) + scene_exporter.export() log_info("Completed Scene Export") return None diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py new file mode 100644 index 000000000..9c6b06410 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_layout_generation.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import shutil +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutConstructor, +) + + +def edit_layout( + *, + scene: Scene, + scene_edit_plan: SceneEditPlan, + updated_scene_graph: SceneGraph, + added_assets: list[SceneObject], + output_root: str | Path, +) -> Scene: + """Dispatch one edit-layout optimization from the goal scene graph.""" + formal_scene = scene + goal_scene_graph = updated_scene_graph + generated_scene_objects = added_assets + # Recreate this stage only when new assets need image, segmentation, and geometry outputs. + stage_output_root = ( + Path(output_root).expanduser().resolve() + / "scene_editing" + / "layout_optimization" + ) + if stage_output_root.exists(): + shutil.rmtree(stage_output_root) + stage_output_root.mkdir(parents=True, exist_ok=True) + + # Add and move operations are the only layout variables in this edit pass. + layout_variable_ids = { + operation.object_id + for operation in scene_edit_plan.operations + if operation.op in {"add", "move"} + } + if None in layout_variable_ids: + raise ValueError("Add and move operations must identify an object.") + layout_variable_ids = { + object_id for object_id in layout_variable_ids if object_id is not None + } + + # Optimize the layout constrained by the goal scene graph. + layout_constructor = SceneLayoutConstructor( + formal_scene=formal_scene, + goal_scene_graph=goal_scene_graph, + layout_variable_ids=layout_variable_ids, + generated_scene_objects=generated_scene_objects, + output_root=stage_output_root, + ) + # Optimize. + post_edit_scene = layout_constructor.construct() + + return post_edit_scene diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index 4063bf92c..7451296b4 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -56,6 +56,7 @@ def __init__( self.export_root = self.output_root / "scene_export" self.scene_config_path: Path | None = None self.scene_graph_path: Path | None = None + self.scene_json_path: Path | None = None def export(self) -> Path: """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. @@ -119,6 +120,17 @@ def export(self) -> Path: encoding="utf-8", ) log_info(f"Exported scene graph: {self.scene_graph_path}") + self.scene_json_path = self.export_root / "scene.json" + self.scene_json_path.write_text( + json.dumps(self.scene.to_dict(), indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + log_info(f"Exported scene JSON: {self.scene_json_path}") + # Remove only assets absent from the completed scene export. + self._remove_stale_mesh_assets( + mesh_assets_root=mesh_assets_root, + object_ids=set(object_ids), + ) return self.scene_config_path @staticmethod @@ -148,9 +160,28 @@ def _copy_scene_object_to_assets( ) destination_glb_path = mesh_assets_root / object_id / f"{object_id}.glb" destination_glb_path.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(source_glb_path, destination_glb_path) + # Imported assets already live at their export destination. + if not destination_glb_path.is_file() or not source_glb_path.samefile( + destination_glb_path + ): + shutil.copy2(source_glb_path, destination_glb_path) return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() + @staticmethod + def _remove_stale_mesh_assets( + *, + mesh_assets_root: Path, + object_ids: set[str], + ) -> None: + """Remove copied asset directories that no longer belong to the scene.""" + for asset_root in mesh_assets_root.iterdir(): + if asset_root.name in object_ids: + continue + if asset_root.is_dir(): + shutil.rmtree(asset_root) + else: + asset_root.unlink() + @staticmethod def _scene_object_config( *, @@ -179,6 +210,8 @@ def _scene_object_config( return { "uid": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, "description": scene_object.description, "shape": { "shape_type": "Mesh", diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index eefeee6b6..9ea558a2a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -286,8 +286,16 @@ def _scene_object_from_export_entry( return SceneObject( id=uid, kind=kind, # type: ignore[arg-type] - category=uid, - name=uid, + category=self._semantic_text( + entry.get("category"), + field_name=f"{uid}.category", + default=uid, + ), + name=self._semantic_text( + entry.get("name"), + field_name=f"{uid}.name", + default=uid, + ), description=str(entry.get("description") or uid), simready_glb_path=str(glb_path), rot=rot_y_up.tolist(), @@ -346,6 +354,20 @@ def _vector3(value: object, *, field_name: str) -> list[float]: raise ValueError(f"Scene config field {field_name!r} must be finite.") return vector + @staticmethod + def _semantic_text( + value: object, + *, + field_name: str, + default: str, + ) -> str: + """Read one non-empty semantic label with a legacy-export fallback.""" + if value is None: + return default + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Scene config field {field_name!r} must be non-empty.") + return value + @staticmethod def _vector2(value: object, *, field_name: str) -> list[float]: """Validate one length-2 numeric vector.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py new file mode 100644 index 000000000..38706ac47 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -0,0 +1,399 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + TABLE_OBJECT_ID, + SceneGraph, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( + SceneLayoutOptimizerConfig, + SceneLayoutOptimizer, +) + + +@dataclass(frozen=True) +class SceneLayoutGroup: + """One parent and its direct on-children handled in one layout pass.""" + + parent_id: str + child_ids: list[str] + + +@dataclass(frozen=True) +class SceneLayoutProblem: + """Prepared graph-constrained inputs for one scene-layout construction.""" + + post_edit_scene: Scene + goal_scene_graph: SceneGraph + layout_variable_ids: set[str] + initial_xy_by_id: dict[str, list[float] | None] + groups: list[SceneLayoutGroup] + + +class SceneLayoutConstructor: + """Construct a scene layout from its goal graph. + + ``formal_scene`` may be empty for text-to-scene. In that case every table + and asset object must be supplied through ``generated_scene_objects``. + """ + + def __init__( + self, + *, + formal_scene: Scene, + goal_scene_graph: SceneGraph, + layout_variable_ids: set[str], + generated_scene_objects: list[SceneObject], + output_root: str | Path, + config: SceneLayoutOptimizerConfig | None = None, + ) -> None: + self.formal_scene = formal_scene + self.goal_scene_graph = goal_scene_graph + self.layout_variable_ids = layout_variable_ids + self.generated_scene_objects = generated_scene_objects + self.output_root = Path(output_root).expanduser().resolve() + self.layout_optimizer = SceneLayoutOptimizer(config=config) + self._current_xy_by_id: dict[str, list[float] | None] = {} + self._solved_delta_xy_by_id: dict[str, list[float]] = {} + self._updated_object_ids: set[str] = set() + + def construct(self) -> Scene: + """Construct table-root layouts before later stacked-group refinement.""" + layout_problem = self._build_problem() + self._current_xy_by_id = { + object_id: list(initial_xy) if initial_xy is not None else None + for object_id, initial_xy in layout_problem.initial_xy_by_id.items() + } + self._solved_delta_xy_by_id = {} + self._updated_object_ids = set() + if ( + layout_problem.groups + and layout_problem.groups[0].parent_id != TABLE_OBJECT_ID + ): + raise ValueError("The first layout group must be rooted at the table.") + + for group in layout_problem.groups: + if group.parent_id == TABLE_OBJECT_ID: + self._optimize_table_group( + layout_problem=layout_problem, + group=group, + ) + continue + self._optimize_parent_group( + layout_problem=layout_problem, + group=group, + ) + + return layout_problem.post_edit_scene + + def _optimize_table_group( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + ) -> None: + """Optimize all direct on-table children before any stacked child groups.""" + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table group optimization requires a table.") + if table.support_optimization_rect_xy is None: + raise ValueError( + "Table group optimization requires a table support optimization rectangle." + ) + + root_ids = set(group.child_ids) + root_relations = [ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in root_ids and relation.target_id in root_ids + ] + root_seed_xy_by_id: dict[str, list[float]] = {} + for root_id in group.child_ids: + inherited_xy = self._current_xy_by_id[root_id] + # New roots start from the table-local origin; imported roots keep their pose. + root_seed_xy_by_id[root_id] = ( + [0.0, 0.0] if inherited_xy is None else list(inherited_xy) + ) + self._current_xy_by_id[root_id] = root_seed_xy_by_id[root_id] + + nodes_by_id = layout_problem.goal_scene_graph.node_by_id() + solved_root_xy_by_id = self.layout_optimizer.optimize_table_root_xy( + assets_by_id={ + asset.id: asset for asset in layout_problem.post_edit_scene.assets + }, + root_ids=group.child_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids={ + root_id + for root_id in group.child_ids + if layout_problem.initial_xy_by_id[root_id] is not None + }, + fixed_root_xy_by_id={ + root_id: ( + None + if root_id in layout_problem.layout_variable_ids + else self._current_xy_by_id[root_id] + ) + for root_id in group.child_ids + }, + root_table_regions_by_id={ + root_id: nodes_by_id[root_id].table_region + for root_id in group.child_ids + }, + table_optimization_rect_xy=table.support_optimization_rect_xy, + root_relations=root_relations, + ) + if table.support_surface_z is None and any( + root_id in layout_problem.layout_variable_ids for root_id in group.child_ids + ): + raise ValueError("Table group optimization requires support_surface_z.") + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + for root_id, solved_xy in solved_root_xy_by_id.items(): + seed_xy = root_seed_xy_by_id[root_id] + delta_xy = [ + solved_xy[0] - seed_xy[0], + solved_xy[1] - seed_xy[1], + ] + self._current_xy_by_id[root_id] = list(solved_xy) + self._solved_delta_xy_by_id[root_id] = delta_xy + if root_id in layout_problem.layout_variable_ids: + # Direct add/move roots receive a new pose on the table support. + assert table.support_surface_z is not None + self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( + scene_object=assets_by_id[root_id], + support_region_z=table.support_surface_z, + center_xy=solved_xy, + ) + self._updated_object_ids.add(root_id) + self._propagate_descendant_delta( + scene=layout_problem.post_edit_scene, + root_id=root_id, + delta_xy=delta_xy, + ) + + def _propagate_descendant_delta( + self, + *, + scene: Scene, + root_id: str, + delta_xy: list[float], + ) -> None: + """Move every positioned descendant by one solved ancestor XY delta.""" + if delta_xy == [0.0, 0.0]: + return + assets_by_id = {asset.id: asset for asset in scene.assets} + children_by_parent: dict[str, list[str]] = {} + for node in self.goal_scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + pending = list(children_by_parent.get(root_id, [])) + while pending: + descendant_id = pending.pop(0) + descendant_xy = self._current_xy_by_id[descendant_id] + if descendant_xy is not None: + self._current_xy_by_id[descendant_id] = [ + descendant_xy[0] + delta_xy[0], + descendant_xy[1] + delta_xy[1], + ] + self.layout_optimizer.translate_scene_object_y_up_by_z_up_delta( + scene_object=assets_by_id[descendant_id], + delta_xy=delta_xy, + ) + self._updated_object_ids.add(descendant_id) + pending.extend(children_by_parent.get(descendant_id, [])) + + def _optimize_parent_group( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + ) -> None: + """Optimize one settled parent's direct on-children in local XY coordinates.""" + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + parent = assets_by_id.get(group.parent_id) + if parent is None: + raise ValueError(f"Parent {group.parent_id!r} is not an asset.") + parent_aabb = self.layout_optimizer.scene_object_z_up_world_aabb( + scene_object=parent + ) + parent_aabb_xy = [ + [parent_aabb[0][0], parent_aabb[0][1]], + [parent_aabb[1][0], parent_aabb[1][1]], + ] + parent_center_xy = [ + (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, + (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, + ] + child_seed_xy_by_id: dict[str, list[float]] = {} + for child_id in group.child_ids: + inherited_xy = self._current_xy_by_id[child_id] + # New children start at their parent's current AABB center. + child_seed_xy_by_id[child_id] = ( + parent_center_xy if inherited_xy is None else list(inherited_xy) + ) + self._current_xy_by_id[child_id] = child_seed_xy_by_id[child_id] + + solved_child_xy_by_id = self.layout_optimizer.optimize_parent_child_xy( + assets_by_id=assets_by_id, + child_ids=group.child_ids, + child_seed_xy_by_id=child_seed_xy_by_id, + imported_child_ids={ + child_id + for child_id in group.child_ids + if layout_problem.initial_xy_by_id[child_id] is not None + }, + fixed_child_xy_by_id={ + child_id: ( + None + if child_id in layout_problem.layout_variable_ids + else self._current_xy_by_id[child_id] + ) + for child_id in group.child_ids + }, + parent_aabb_xy=parent_aabb_xy, + ) + parent_top_z = parent_aabb[1][2] + for child_id, solved_xy in solved_child_xy_by_id.items(): + seed_xy = child_seed_xy_by_id[child_id] + delta_xy = [ + solved_xy[0] - seed_xy[0], + solved_xy[1] - seed_xy[1], + ] + self._current_xy_by_id[child_id] = list(solved_xy) + self._solved_delta_xy_by_id[child_id] = delta_xy + if child_id in layout_problem.layout_variable_ids: + # Variable children are placed directly above the parent's current top. + self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( + scene_object=assets_by_id[child_id], + support_region_z=parent_top_z, + center_xy=solved_xy, + ) + self._updated_object_ids.add(child_id) + self._propagate_descendant_delta( + scene=layout_problem.post_edit_scene, + root_id=child_id, + delta_xy=delta_xy, + ) + + def _build_problem(self) -> SceneLayoutProblem: + """Build post-edit objects and preserve formal-scene centers as seeds.""" + self.goal_scene_graph.validate() + graph_object_ids = set(self.goal_scene_graph.node_by_id()) + generated_objects_by_id = self._generated_scene_objects_by_id() + + # The goal graph removes deleted formal-scene objects from the layout input. + post_edit_objects = [ + scene_object + for scene_object in self.formal_scene.objects + if scene_object.id in graph_object_ids + ] + imported_object_ids = {scene_object.id for scene_object in post_edit_objects} + if imported_object_ids.intersection(generated_objects_by_id): + raise ValueError( + "Generated scene objects must not reuse formal scene object ids." + ) + post_edit_objects.extend(generated_objects_by_id.values()) + + post_edit_scene = Scene(objects=post_edit_objects) + post_edit_object_ids = { + scene_object.id for scene_object in post_edit_scene.objects + } + if post_edit_object_ids != graph_object_ids: + raise ValueError("Goal scene graph and post-edit scene have different ids.") + if not self.layout_variable_ids.issubset(post_edit_object_ids - {"table"}): + raise ValueError( + "Only post-edit assets may participate in layout optimization." + ) + + initial_xy_by_id = { + asset.id: self._initial_xy( + asset, + is_generated=asset.id in generated_objects_by_id, + ) + for asset in post_edit_scene.assets + } + for object_id, initial_xy in initial_xy_by_id.items(): + if initial_xy is None and object_id not in self.layout_variable_ids: + raise ValueError( + f"New asset {object_id!r} must participate in layout optimization." + ) + + return SceneLayoutProblem( + post_edit_scene=post_edit_scene, + goal_scene_graph=self.goal_scene_graph, + layout_variable_ids=set(self.layout_variable_ids), + initial_xy_by_id=initial_xy_by_id, + groups=self._build_groups(), + ) + + def _build_groups(self) -> list[SceneLayoutGroup]: + """Build table-rooted BFS groups of direct on-children.""" + children_by_parent: dict[str, list[str]] = {} + for node in self.goal_scene_graph.nodes: + if node.parent_id is None: + continue + if node.parent_relation != "on": + raise ValueError(f"Node {node.object_id!r} must be on its parent.") + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + groups: list[SceneLayoutGroup] = [] + pending = [TABLE_OBJECT_ID] + while pending: + parent_id = pending.pop(0) + child_ids = children_by_parent.get(parent_id, []) + if not child_ids: + continue + groups.append(SceneLayoutGroup(parent_id=parent_id, child_ids=child_ids)) + pending.extend(child_ids) + return groups + + def _generated_scene_objects_by_id(self) -> dict[str, SceneObject]: + """Index generated scene objects before merging them into the formal scene.""" + generated_objects_by_id = { + scene_object.id: scene_object + for scene_object in self.generated_scene_objects + } + if len(generated_objects_by_id) != len(self.generated_scene_objects): + raise ValueError("Generated scene objects must use unique object ids.") + return generated_objects_by_id + + @staticmethod + def _initial_xy( + asset: SceneObject, + *, + is_generated: bool, + ) -> list[float] | None: + """Retain formal-scene centers while generated assets await initialization.""" + if is_generated: + return None + if asset.center_xy is None or len(asset.center_xy) != 2: + raise ValueError( + f"Formal-scene asset {asset.id!r} must have a 2D center_xy." + ) + return [float(value) for value in asset.center_xy] diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py new file mode 100644 index 000000000..85e98f0e8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py @@ -0,0 +1,785 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) + + +@dataclass(frozen=True) +class SceneLayoutOptimizerConfig: + """Numerical controls shared by each graph-layout solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid numerical controls before assembling a layout problem.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class SceneLayoutOptimizer: + """Solve graph-constrained XY layouts and apply resulting poses.""" + + def __init__(self, *, config: SceneLayoutOptimizerConfig | None = None) -> None: + self.config = config if config is not None else SceneLayoutOptimizerConfig() + + def optimize_table_root_xy( + self, + *, + assets_by_id: dict[str, SceneObject], + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + fixed_root_xy_by_id: dict[str, list[float] | None], + root_table_regions_by_id: dict[str, str | None], + table_optimization_rect_xy: list[list[float]], + root_relations: list[SceneGraphRelation], + ) -> dict[str, list[float]]: + """Solve direct table-child centers with graph and AABB constraints.""" + return _optimize_table_root_xy( + assets_by_id=assets_by_id, + root_ids=root_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids=imported_root_ids, + fixed_root_xy_by_id=fixed_root_xy_by_id, + root_table_regions_by_id=root_table_regions_by_id, + table_optimization_rect_xy=table_optimization_rect_xy, + root_relations=root_relations, + config=self.config, + ) + + def optimize_parent_child_xy( + self, + *, + assets_by_id: dict[str, SceneObject], + child_ids: list[str], + child_seed_xy_by_id: dict[str, list[float]], + imported_child_ids: set[str], + fixed_child_xy_by_id: dict[str, list[float] | None], + parent_aabb_xy: list[list[float]], + ) -> dict[str, list[float]]: + """Solve direct on-children inside one parent's current XY AABB.""" + child_half_extents_xy = _asset_half_extents_xy( + assets_by_id=assets_by_id, + object_ids=child_ids, + ) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + child_index = {child_id: index for index, child_id in enumerate(child_ids)} + parent_bounds = _bounds_from_points(parent_aabb_xy) + for child_id in child_ids: + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=child_index, + root_id=child_id, + bounds=parent_bounds, + half_extents_xy=child_half_extents_xy[child_id], + ) + fixed_xy = fixed_child_xy_by_id[child_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=child_index, + root_id=child_id, + fixed_xy=fixed_xy, + ) + + solved_child_xy_by_id = _solve_root_xy( + root_ids=child_ids, + root_seed_xy_by_id=child_seed_xy_by_id, + imported_root_ids=imported_child_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=child_ids, + root_seed_xy_by_id=child_seed_xy_by_id, + imported_root_ids=imported_child_ids, + root_half_extents_xy=child_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + solved_root_xy_by_id=solved_child_xy_by_id, + config=self.config, + ) + + @staticmethod + def scene_object_z_up_world_aabb( + *, + scene_object: SceneObject, + ) -> list[list[float]]: + """Return one object's current z-up world AABB as [min, max].""" + return _scene_object_z_up_world_aabb(scene_object=scene_object) + + @staticmethod + def update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, + ) -> None: + """Place one SimReady asset on a horizontal z-up support region.""" + _update_scene_object_y_up_pose_from_z_up_support( + scene_object=scene_object, + support_region_z=support_region_z, + center_xy=center_xy, + clearance_m=clearance_m, + ) + + @staticmethod + def translate_scene_object_y_up_by_z_up_delta( + *, + scene_object: SceneObject, + delta_xy: list[float], + ) -> None: + """Translate one existing y-up pose by a solved z-up XY delta.""" + _translate_scene_object_y_up_by_z_up_delta( + scene_object=scene_object, + delta_xy=delta_xy, + ) + + +def _optimize_table_root_xy( + *, + assets_by_id: dict[str, SceneObject], + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + fixed_root_xy_by_id: dict[str, list[float] | None], + root_table_regions_by_id: dict[str, str | None], + table_optimization_rect_xy: list[list[float]], + root_relations: list[SceneGraphRelation], + config: SceneLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve direct table-child centers with graph and AABB constraints.""" + root_half_extents_xy = _asset_half_extents_xy( + assets_by_id=assets_by_id, + object_ids=root_ids, + ) + inequality_constraints, equality_constraints = _build_table_root_constraints( + root_ids=root_ids, + root_half_extents_xy=root_half_extents_xy, + root_relations=root_relations, + root_table_regions_by_id=root_table_regions_by_id, + table_optimization_rect_xy=table_optimization_rect_xy, + fixed_root_xy_by_id=fixed_root_xy_by_id, + config=config, + ) + solved_root_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + return _refine_root_collisions( + root_ids=root_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids=imported_root_ids, + root_half_extents_xy=root_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + solved_root_xy_by_id=solved_root_xy_by_id, + config=config, + ) + + +def _update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, +) -> None: + """Place one SimReady asset on a horizontal z-up support region. + + ``SceneObject`` stores poses in y-up before export. The target center and + support height are z-up values because layout optimization uses that frame. + """ + if not np.isfinite(support_region_z): + raise ValueError("support_region_z must be finite.") + if clearance_m < 0.0 or not np.isfinite(clearance_m): + raise ValueError("clearance_m must be finite and non-negative.") + target_xy = _two_floats(center_xy, field_name="center_xy") + rotation_y_up = _three_floats_or_default( + scene_object.rot, + field_name="rot", + default=[0.0, 0.0, 0.0], + ) + mesh = _asset_z_up_mesh_at_zero_translation( + scene_object=scene_object, + rotation_y_up=rotation_y_up, + ) + target_position_z_up = np.array( + [ + target_xy[0] - float(mesh.bounds[:, 0].mean()), + target_xy[1] - float(mesh.bounds[:, 1].mean()), + float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), + ] + ) + z_up_to_y_up = np.linalg.inv(_y_up_to_z_up_matrix()) + # Persist the y-up pose that SceneExporter later converts back to z-up. + scene_object.pos = (z_up_to_y_up[:3, :3] @ target_position_z_up).tolist() + scene_object.rot = rotation_y_up + scene_object.center_xy = target_xy + + +def _translate_scene_object_y_up_by_z_up_delta( + *, + scene_object: SceneObject, + delta_xy: list[float], +) -> None: + """Translate one existing y-up pose by a solved z-up XY delta.""" + dx, dy = _two_floats(delta_xy, field_name="delta_xy") + current_pos = _three_floats_or_default( + scene_object.pos, + field_name="pos", + default=None, + ) + # z-up x maps to y-up x, while z-up y maps to negative y-up z. + scene_object.pos = [ + current_pos[0] + dx, + current_pos[1], + current_pos[2] - dy, + ] + if scene_object.center_xy is not None: + scene_object.center_xy = [ + scene_object.center_xy[0] + dx, + scene_object.center_xy[1] + dy, + ] + + +def _scene_object_z_up_world_aabb( + *, + scene_object: SceneObject, +) -> list[list[float]]: + """Measure one current SceneObject pose in z-up world coordinates.""" + position_y_up = _three_floats_or_default( + scene_object.pos, + field_name="pos", + default=None, + ) + mesh = _asset_z_up_mesh_at_zero_translation(scene_object=scene_object) + position_z_up = _y_up_to_z_up_matrix()[:3, :3] @ np.asarray( + position_y_up, + dtype=float, + ) + mesh.apply_translation(position_z_up) + return mesh.bounds.tolist() + + +def _build_table_root_constraints( + *, + root_ids: list[str], + root_half_extents_xy: dict[str, np.ndarray], + root_relations: list[SceneGraphRelation], + root_table_regions_by_id: dict[str, str | None], + table_optimization_rect_xy: list[list[float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + config: SceneLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard table, region, planar, and fixed-root constraints.""" + root_index = {root_id: index for index, root_id in enumerate(root_ids)} + table_bounds = _bounds_from_points(table_optimization_rect_xy) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + + for root_id in root_ids: + region_bounds = _table_region_bounds( + table_bounds=table_bounds, + table_region=root_table_regions_by_id[root_id], + ) + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=root_id, + bounds=region_bounds, + half_extents_xy=root_half_extents_xy[root_id], + ) + fixed_xy = fixed_root_xy_by_id[root_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=root_id, + fixed_xy=fixed_xy, + ) + + for relation in root_relations: + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=root_half_extents_xy[relation.source_id], + target_half_extents_xy=root_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: SceneLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve one root-group center model with the legacy SLSQP settings.""" + root_index = {root_id: index for index, root_id in enumerate(root_ids)} + initial_xy = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + x0 = initial_xy.reshape(-1) + + def unpack(values: np.ndarray) -> dict[str, list[float]]: + return { + root_id: [float(values[2 * index]), float(values[2 * index + 1])] + for root_id, index in root_index.items() + } + + def objective(values: np.ndarray) -> float: + coordinates = values.reshape(-1, 2) + loss = 0.0 + for root_id, index in root_index.items(): + if root_id in imported_root_ids: + delta = coordinates[index] - initial_xy[index] + loss += config.imported_seed_weight * float(delta @ delta) + for first_index in range(len(root_ids)): + for second_index in range(first_index + 1, len(root_ids)): + distance = float( + np.linalg.norm(coordinates[first_index] - coordinates[second_index]) + ) + shortfall = max(0.0, config.min_center_distance_m - distance) + loss += config.min_center_distance_weight * shortfall**2 + return loss + + constraints: list[dict[str, object]] = [] + for row, bound in inequality_constraints: + constraints.append( + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + ) + for row, bound in equality_constraints: + constraints.append( + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + ) + + result = minimize( + objective, + x0, + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise ValueError(f"Table layout optimization failed: {result.message}") + return unpack(np.asarray(result.x, dtype=float)) + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + solved_root_xy_by_id: dict[str, list[float]], + config: SceneLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Add AABB separation constraints until the table roots no longer overlap.""" + seen_pairs: set[tuple[str, str]] = set() + current_xy_by_id = solved_root_xy_by_id + for _ in range(config.max_collision_rounds): + overlaps = _root_aabb_overlaps( + root_ids=root_ids, + root_half_extents_xy=root_half_extents_xy, + xy_by_id=current_xy_by_id, + ) + if not overlaps: + return current_xy_by_id + added_constraint_count = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + pair_key = tuple(sorted((first_id, second_id))) + if pair_key in seen_pairs: + continue + inequality_constraints.append( + _aabb_separation_constraint( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + first_half_extents_xy=root_half_extents_xy[first_id], + second_half_extents_xy=root_half_extents_xy[second_id], + first_xy=current_xy_by_id[first_id], + second_xy=current_xy_by_id[second_id], + collision_margin_m=config.collision_margin_m, + ) + ) + seen_pairs.add(pair_key) + added_constraint_count += 1 + if added_constraint_count == 0: + break + current_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current_xy_by_id, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + + remaining_pairs = [ + f"{first_id}/{second_id}" + for _, first_id, second_id in _root_aabb_overlaps( + root_ids=root_ids, + root_half_extents_xy=root_half_extents_xy, + xy_by_id=current_xy_by_id, + ) + ] + raise ValueError( + "Table-root AABB collisions remain after layout refinement: " + f"{remaining_pairs}." + ) + + +def _asset_half_extents_xy( + *, + assets_by_id: dict[str, SceneObject], + object_ids: list[str], +) -> dict[str, np.ndarray]: + """Measure each asset's oriented z-up footprint around its XY center.""" + half_extents_xy: dict[str, np.ndarray] = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Table root {object_id!r} is not an asset.") + half_extents_xy[object_id] = _asset_half_extent_xy(asset) + return half_extents_xy + + +def _asset_half_extent_xy(asset: SceneObject) -> np.ndarray: + """Measure one SimReady GLB with its current orientation and scale.""" + mesh = _asset_z_up_mesh_at_zero_translation(scene_object=asset) + return (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + + +def _asset_z_up_mesh_at_zero_translation( + *, + scene_object: SceneObject, + rotation_y_up: list[float] | None = None, +): + """Load one SimReady GLB in z-up with orientation and scale but no position.""" + asset = scene_object + if asset.simready_glb_path is None: + raise ValueError(f"Asset {asset.id!r} has no SimReady GLB path.") + y_up_layout = { + "id": asset.id, + "rot": ( + rotation_y_up + if rotation_y_up is not None + else _three_floats_or_default( + asset.rot, + field_name="rot", + default=[0.0, 0.0, 0.0], + ) + ), + "pos": [0.0, 0.0, 0.0], + "scale": _three_floats_or_default( + asset.scale, + field_name="scale", + default=[1.0, 1.0, 1.0], + ), + } + y_up_to_z_up = _y_up_to_z_up_matrix() + z_up_layout = transform_matrix_to_layout_object( + asset.id, + y_up_to_z_up + @ layout_object_to_transform_matrix(y_up_layout) + @ np.linalg.inv(y_up_to_z_up), + ) + mesh = load_glb_mesh(asset.simready_glb_path) + mesh.apply_transform(y_up_to_z_up) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + +def _y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate transform used by SceneExporter and layout stages.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix + + +def _two_floats(value: object, *, field_name: str) -> list[float]: + """Validate one finite two-value vector.""" + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError(f"{field_name} must contain two values.") + vector = [float(component) for component in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"{field_name} must contain finite values.") + return vector + + +def _three_floats_or_default( + value: object, + *, + field_name: str, + default: list[float] | None, +) -> list[float]: + """Return a finite three-value vector or the canonical SimReady default.""" + if value is None: + if default is None: + raise ValueError(f"{field_name} must contain three values.") + return list(default) + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"{field_name} must contain three values.") + vector = [float(component) for component in value] + if not np.all(np.isfinite(vector)): + raise ValueError(f"{field_name} must contain finite values.") + return vector + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + """Return [[min_x, min_y], [max_x, max_y]] from finite XY points.""" + coordinates = np.asarray(points, dtype=float) + if coordinates.ndim != 2 or coordinates.shape[1] != 2 or len(coordinates) < 2: + raise ValueError("XY bounds must contain at least two points.") + if not np.all(np.isfinite(coordinates)): + raise ValueError("XY bounds must contain finite values.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _table_region_bounds( + *, + table_bounds: np.ndarray, + table_region: str | None, +) -> np.ndarray: + """Return the requested 3x3 table region, with y increasing toward front.""" + if table_region is None: + return table_bounds.copy() + column_by_region = { + "left_back": 0, + "left_center": 0, + "left_front": 0, + "back_center": 1, + "center": 1, + "front_center": 1, + "right_back": 2, + "right_center": 2, + "right_front": 2, + } + row_by_region = { + "left_front": 0, + "front_center": 0, + "right_front": 0, + "left_center": 1, + "center": 1, + "right_center": 1, + "left_back": 2, + "back_center": 2, + "right_back": 2, + } + if table_region not in column_by_region: + raise ValueError(f"Unsupported table region {table_region!r}.") + minimum, maximum = table_bounds + cell_size = (maximum - minimum) / 3.0 + region_minimum = minimum + cell_size * np.array( + [column_by_region[table_region], row_by_region[table_region]] + ) + return np.stack([region_minimum, region_minimum + cell_size]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + """Keep one root's complete AABB inside the given rectangular bounds.""" + minimum = bounds[0] + half_extents_xy + maximum = bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError( + f"Asset {root_id!r} cannot fit inside its assigned table region." + ) + variable_count = 2 * len(root_index) + root_offset = 2 * root_index[root_id] + for axis in range(2): + upper_row = np.zeros(variable_count) + upper_row[root_offset + axis] = 1.0 + constraints.append((upper_row, float(maximum[axis]))) + lower_row = np.zeros(variable_count) + lower_row[root_offset + axis] = -1.0 + constraints.append((lower_row, -float(minimum[axis]))) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + """Use equality constraints so unchanged formal objects remain fixed.""" + variable_count = 2 * len(root_index) + root_offset = 2 * root_index[root_id] + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(variable_count) + row[root_offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + """Require directional relations to clear both sibling AABB footprints.""" + if source_id not in root_index or target_id not in root_index: + raise ValueError("Table-root planar relations must reference table roots.") + axis, source_sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_sign is None: + raise ValueError(f"Unsupported planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = source_sign + row[2 * root_index[target_id] + axis] = -source_sign + required_distance = ( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ) + constraints.append((row, -float(required_distance))) + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + root_half_extents_xy: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return root pairs whose current XY AABBs overlap without a margin.""" + overlaps: list[tuple[float, str, str]] = [] + for first_index, first_id in enumerate(root_ids): + first_xy = np.asarray(xy_by_id[first_id], dtype=float) + first_half_extents = root_half_extents_xy[first_id] + for second_id in root_ids[first_index + 1 :]: + second_xy = np.asarray(xy_by_id[second_id], dtype=float) + second_half_extents = root_half_extents_xy[second_id] + overlap_xy = np.minimum( + first_xy + first_half_extents, + second_xy + second_half_extents, + ) - np.maximum( + first_xy - first_half_extents, + second_xy - second_half_extents, + ) + if np.all(overlap_xy > 1e-9): + overlaps.append((float(np.min(overlap_xy)), first_id, second_id)) + return sorted(overlaps, reverse=True) + + +def _aabb_separation_constraint( + *, + root_ids: list[str], + first_id: str, + second_id: str, + first_half_extents_xy: np.ndarray, + second_half_extents_xy: np.ndarray, + first_xy: list[float], + second_xy: list[float], + collision_margin_m: float, +) -> tuple[np.ndarray, float]: + """Separate one overlapping pair along its shallowest penetration axis.""" + root_index = {root_id: index for index, root_id in enumerate(root_ids)} + first_xy_array = np.asarray(first_xy, dtype=float) + second_xy_array = np.asarray(second_xy, dtype=float) + overlap_xy = np.minimum( + first_xy_array + first_half_extents_xy, + second_xy_array + second_half_extents_xy, + ) - np.maximum( + first_xy_array - first_half_extents_xy, + second_xy_array - second_half_extents_xy, + ) + axis = int(np.argmin(overlap_xy)) + first_is_lower = first_xy_array[axis] < second_xy_array[axis] or ( + first_xy_array[axis] == second_xy_array[axis] and first_id < second_id + ) + row = np.zeros(2 * len(root_ids)) + first_coefficient = 1.0 if first_is_lower else -1.0 + row[2 * root_index[first_id] + axis] = first_coefficient + row[2 * root_index[second_id] + axis] = -first_coefficient + required_distance = ( + first_half_extents_xy[axis] + second_half_extents_xy[axis] + collision_margin_m + ) + return row, -float(required_distance) diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index b7b614d49..b50126751 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -161,6 +161,8 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert (export_path.parent / "mesh_assets/cup/cup.glb").read_bytes() == b"glTF-cup" entry = exported["rigid_object"][0] assert entry["uid"] == "cup" + assert entry["category"] == "asset" + assert entry["name"] == "cup" assert entry["body_type"] == "dynamic" assert entry["init_pos"] == [1.0, -3.0, 2.0] assert entry["body_scale"] == [1.0, 2.0, 3.0] @@ -188,9 +190,76 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No output_root=tmp_path / "output" ).import_scene_and_graph() assert [asset.id for asset in imported_scene.assets] == ["cup"] + assert imported_scene.assets[0].category == "asset" + assert imported_scene.assets[0].name == "cup" assert imported_graph.to_dict() == _scene_graph(scene).to_dict() +def test_scene_export_overwrites_an_existing_scene_export(tmp_path: Path) -> None: + table_glb = tmp_path / "table.glb" + cup_glb = tmp_path / "cup.glb" + banana_glb = tmp_path / "banana.glb" + table_glb.write_bytes(b"glTF-table") + cup_glb.write_bytes(b"glTF-cup") + banana_glb.write_bytes(b"glTF-banana") + output_root = tmp_path / "output" + + initial_table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + initial_cup = _scene_object( + object_id="cup", + kind="asset", + glb_path=cup_glb, + physics=_physics("dynamic"), + ) + initial_scene = Scene(objects=[initial_table, initial_cup]) + SceneExporter( + scene=initial_scene, + scene_graph=_scene_graph(initial_scene), + output_root=output_root, + ).export() + + # The imported table mesh already occupies its final export location. + exported_table_glb = ( + output_root / "scene_export" / "mesh_assets" / "table" / "table.glb" + ) + updated_table = _scene_object( + object_id="table", + kind="table", + glb_path=exported_table_glb, + physics=_physics("kinematic"), + ) + banana = _scene_object( + object_id="banana", + kind="asset", + glb_path=banana_glb, + physics=_physics("dynamic"), + ) + updated_scene = Scene(objects=[updated_table, banana]) + SceneExporter( + scene=updated_scene, + scene_graph=_scene_graph(updated_scene), + output_root=output_root, + ).export() + + scene_export_root = output_root / "scene_export" + assert exported_table_glb.read_bytes() == b"glTF-table" + assert ( + scene_export_root / "mesh_assets" / "banana" / "banana.glb" + ).read_bytes() == b"glTF-banana" + assert not (scene_export_root / "mesh_assets" / "cup").exists() + assert ( + json.loads((scene_export_root / "scene.json").read_text(encoding="utf-8"))[ + "objects" + ][1]["id"] + == "banana" + ) + + def test_scene_export_requires_final_physics(tmp_path: Path) -> None: glb_path = tmp_path / "table.glb" glb_path.write_bytes(b"glTF") From 78e01a3fe9fe5dfe2fd0df753709ff1b79ea7235 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 19:28:27 +0800 Subject: [PATCH 15/85] fix some bug before pr --- .../features/generative_sim/scene_engine.md | 22 ++- .../scene_engine/cli/test_text_to_simready.py | 183 ------------------ .../pipeline/utils/scene_layout_optimizer.py | 12 +- .../test_scene_layout_optimizer.py | 138 +++++++++++++ 4 files changed, 165 insertions(+), 190 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py create mode 100644 tests/gen_sim/scene_engine/test_scene_layout_optimizer.py diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index c9f569ee2..e4255163e 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -27,6 +27,23 @@ python -m embodichain scene-engine \ --output_root /path/to/scene_output ``` +## Scene Editing + +Edit an existing valid Scene Engine output with an instruction: + +```bash +embodichain scene-engine \ + --output_root /path/to/scene_output \ + --edit_prompt "add a red cup to the front-center of the tabletop" +``` + +`--image` and `--edit_prompt` may also be provided together. Scene Engine then +generates the image-based scene first and applies the edit to that export. An +edit-only invocation requires an existing `scene_export` directory. The edit +overwrites its `scene_config.json`, `scene_graph.json`, `scene.json`, and final +`mesh_assets`; intermediate generation and edit artifacts remain available for +debugging. + ## Configuration Scene Engine reads the LLM, segmentation, image-generation, and @@ -72,9 +89,12 @@ The important final outputs are: scene_output/ |-- scene_understanding/ # Object analysis, masks, and stage JSON |-- scene_generation/ # Generated, SimReady, and layout-debug artifacts +|-- scene_editing/ # Present after edits; generated asset/debug artifacts `-- scene_export/ |-- mesh_assets/ # Final GLBs - `-- scene_config.json # Exported scene description + |-- scene_config.json # Exported z-up scene description + |-- scene_graph.json # Table support and planar relation graph + `-- scene.json # Scene Engine object metadata and y-up poses ``` Validate the export without opening a window: diff --git a/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py b/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py deleted file mode 100644 index e3db15b16..000000000 --- a/embodichain/gen_sim/scene_engine/cli/test_text_to_simready.py +++ /dev/null @@ -1,183 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import argparse -import json -from collections.abc import Sequence -from pathlib import Path -import shutil - -from PIL import Image - -from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( - GeometryGenerationClient, -) -from embodichain.gen_sim.scene_engine.clients.image_generation import ( - ImageGenerationClient, -) -from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( - ImageSegmentationClient, -) -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject -from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( - OpenAICompatibleVLM, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( - build_mask_candidates, - invert_mask_if_foreground_is_off_center, - save_binary_mask, - union_overlapping_mask_candidates, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( - SimReadyProcessor, - SimReadyProcessorConfig, -) - -__all__ = ["main", "run_text_to_simready"] - - -def run_text_to_simready(*, text: str, output_root: str | Path) -> SceneObject: - """Run one manual text-to-SimReady asset pipeline for debugging.""" - text = text.strip() - if not text: - raise ValueError("Text prompt must not be empty.") - - root = Path(output_root).expanduser().resolve() - if root.exists(): - shutil.rmtree(root) - debug_root = root / "debug" - image_root = root / "generated_images" - mask_root = root / "masks" - coarse_root = root / "coarse_geometry" - simready_root = root / "simready_geometry" - for directory in (debug_root, image_root, mask_root, coarse_root, simready_root): - directory.mkdir(parents=True, exist_ok=True) - - object_id = "asset_001" - scene_object = SceneObject( - id=object_id, - kind="asset", - category="asset", - name=text, - description=text, - ) - - image_generation_client = ImageGenerationClient.from_dotenv() - image_segmentation_client = ImageSegmentationClient.from_dotenv() - geometry_generation_client = GeometryGenerationClient.from_dotenv() - vlm_client = OpenAICompatibleVLM.from_dotenv() - try: - image_generation_client.check_health() - image_segmentation_client.check_health() - geometry_generation_client.check_health() - - # Generate a centered single-object image from the semantic text prompt. - image_path = image_generation_client.generate_image_by_prompt( - prompt=text, - output_path=image_root / f"{object_id}.png", - ) - with Image.open(image_path) as image: - image_size = image.size - - # Segment the generated object and apply the single-object foreground heuristic. - candidates = union_overlapping_mask_candidates( - build_mask_candidates( - image_segmentation_client.segment_single_object( - image_path=image_path, - prompt=text, - ) - ), - min_iou=0.8, - ) - if not candidates: - raise ValueError("Image segmentation returned no mask candidates.") - mask_path = save_binary_mask( - invert_mask_if_foreground_is_off_center(candidates[0]), - image_size=image_size, - output_path=mask_root / f"{object_id}.png", - ) - - # Generate one coarse GLB using the generated image and its binary mask. - geometry_generation_client.generate_objects( - image_path=image_path, - object_masks=[(object_id, mask_path)], - output_root=coarse_root, - ) - coarse_glb_path = coarse_root / f"{object_id}.glb" - if not coarse_glb_path.is_file(): - raise FileNotFoundError(f"Coarse GLB was not generated: {coarse_glb_path}") - - # Use identity coarse layout; VLM determines rotation and real-world size. - processor = SimReadyProcessor( - scene=Scene(objects=[scene_object]), - coarse_layout_by_id={ - object_id: { - "rot": [0.0, 0.0, 0.0], - "pos": [0.0, 0.0, 0.0], - "scale": [1.0, 1.0, 1.0], - } - }, - coarse_geometry_root=coarse_root, - simready_geometry_root=simready_root, - config=SimReadyProcessorConfig( - use_vlm_scale=True, - use_vlm_rotation=True, - ), - vlm_client=vlm_client, - ) - simready_layout = processor.process_assets() - (root / "result.json").write_text( - json.dumps( - { - "input_text": text, - "scene_object": scene_object.to_dict(), - "simready_layout": simready_layout, - }, - indent=2, - ensure_ascii=False, - ) - + "\n", - encoding="utf-8", - ) - return scene_object - finally: - image_generation_client.close() - image_segmentation_client.close() - geometry_generation_client.close() - - -def main(argv: Sequence[str] | None = None) -> None: - """Run the manual text-to-SimReady CLI.""" - parser = argparse.ArgumentParser( - prog="embodichain test-text-to-simready", - description="Debug text-to-image-to-segmentation-to-SimReady generation.", - ) - parser.add_argument("--text", required=True, help="Description of one object.") - parser.add_argument( - "--output_root", - required=True, - help="Directory for all intermediate and final artifacts.", - ) - args = parser.parse_args(argv) - scene_object = run_text_to_simready(text=args.text, output_root=args.output_root) - print(f"Generated SimReady asset: {scene_object.simready_glb_path}") - - -if __name__ == "__main__": - main() diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py index 85e98f0e8..ea6387cd2 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py @@ -628,15 +628,15 @@ def _table_region_bounds( "right_front": 2, } row_by_region = { - "left_front": 0, - "front_center": 0, - "right_front": 0, + "left_back": 0, + "back_center": 0, + "right_back": 0, "left_center": 1, "center": 1, "right_center": 1, - "left_back": 2, - "back_center": 2, - "right_back": 2, + "left_front": 2, + "front_center": 2, + "right_front": 2, } if table_region not in column_by_region: raise ValueError(f"Unsupported table region {table_region!r}.") diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py new file mode 100644 index 000000000..36b4bfa79 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -0,0 +1,138 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import trimesh + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutConstructor, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( + _table_region_bounds, +) + + +def _asset( + *, + object_id: str, + glb_path: Path, + center_xy: list[float] | None = None, + pos: list[float] | None = None, +) -> SceneObject: + return SceneObject( + id=object_id, + kind="asset", + category=object_id, + name=object_id, + description=object_id, + simready_glb_path=str(glb_path), + rot=[0.0, 0.0, 0.0], + pos=pos or [0.0, 0.0, 0.0], + scale=[1.0, 1.0, 1.0], + center_xy=center_xy, + ) + + +def test_table_regions_put_front_at_larger_y() -> None: + table_bounds = np.asarray([[0.0, 0.0], [3.0, 3.0]]) + + assert np.allclose( + _table_region_bounds( + table_bounds=table_bounds, + table_region="back_center", + ), + [[1.0, 0.0], [2.0, 1.0]], + ) + assert np.allclose( + _table_region_bounds( + table_bounds=table_bounds, + table_region="front_center", + ), + [[1.0, 2.0], [2.0, 3.0]], + ) + + +def test_layout_constructor_places_new_child_on_parent_top( + tmp_path: Path, +) -> None: + book_glb = tmp_path / "book.glb" + cup_glb = tmp_path / "cup.glb" + # SimReady GLBs are y-up, so the book's short vertical axis is y. + trimesh.creation.box(extents=[1.0, 0.2, 1.0]).export(book_glb) + trimesh.creation.box(extents=[0.2, 0.2, 0.2]).export(cup_glb) + + table = SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + support_surface_z=0.0, + support_optimization_rect_xy=[ + [-2.0, -2.0], + [2.0, -2.0], + [2.0, 2.0], + [-2.0, 2.0], + ], + ) + book = _asset( + object_id="book_001", + glb_path=book_glb, + center_xy=[0.0, 0.0], + # This y-up position maps to a z-up center at z=0.52 m. + pos=[0.0, 0.52, 0.0], + ) + cup = _asset(object_id="cup_001", glb_path=cup_glb) + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + + post_edit_scene = SceneLayoutConstructor( + formal_scene=Scene(objects=[table, book]), + goal_scene_graph=graph, + layout_variable_ids={"cup_001"}, + generated_scene_objects=[cup], + output_root=tmp_path, + ).construct() + + placed_cup = next( + asset for asset in post_edit_scene.assets if asset.id == "cup_001" + ) + assert placed_cup.center_xy == [0.0, 0.0] + # book top is z=0.62 m; cup half-height is 0.1 m and clearance is 0.02 m. + assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) From 3cd914776975b97498c4b213e2ac183ab6ef24ef Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:05:55 +0800 Subject: [PATCH 16/85] fix a real-world size bug, using z-up mesh --- .../utils/simready_processor_utils.py | 36 ++++++++++++++-- .../test_simready_processor_utils.py | 41 +++++++++++++++++++ 2 files changed, 74 insertions(+), 3 deletions(-) create mode 100644 tests/gen_sim/scene_engine/test_simready_processor_utils.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py index b7f718374..539bace70 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -17,7 +17,10 @@ from __future__ import annotations import json +import os from pathlib import Path +import sys +from typing import Callable import numpy as np from PIL import Image, ImageDraw, ImageFont @@ -92,8 +95,12 @@ def render_object_front_top_views( "Blender's bpy is required for SimReady VLM view rendering." ) from exc - bpy.ops.wm.read_factory_settings(use_empty=True) - bpy.ops.import_scene.gltf(filepath=str(source_path)) + _run_blender_operation_silently( + lambda: bpy.ops.wm.read_factory_settings(use_empty=True) + ) + _run_blender_operation_silently( + lambda: bpy.ops.import_scene.gltf(filepath=str(source_path)) + ) if not any(obj.type == "MESH" for obj in bpy.context.scene.objects): raise ValueError(f"GLB contains no mesh objects: {source_path}") scene = bpy.context.scene @@ -141,7 +148,7 @@ def render_view(path: Path, location: tuple[float, float, float]) -> None: .to_euler() ) scene.render.filepath = str(path) - bpy.ops.render.render(write_still=True) + _run_blender_operation_silently(lambda: bpy.ops.render.render(write_still=True)) # Blender uses a right-handed z-up world; front is viewed along +y. render_view(front_path, (0.0, -3.0, 0.0)) @@ -197,6 +204,25 @@ def render_view(path: Path, location: tuple[float, float, float]) -> None: return output_path +def _run_blender_operation_silently(operation: Callable[[], object]) -> object: + """Run one bpy operation without forwarding Blender-native console output.""" + # bpy writes render progress directly to process file descriptors, not Python streams. + sys.stdout.flush() + sys.stderr.flush() + saved_stdout_fd = os.dup(1) + saved_stderr_fd = os.dup(2) + try: + with open(os.devnull, "w", encoding="utf-8") as null_output: + os.dup2(null_output.fileno(), 1) + os.dup2(null_output.fileno(), 2) + return operation() + finally: + os.dup2(saved_stdout_fd, 1) + os.dup2(saved_stderr_fd, 2) + os.close(saved_stdout_fd) + os.close(saved_stderr_fd) + + def _draw_arrow( draw: ImageDraw.ImageDraw, start: tuple[int, int], @@ -311,6 +337,10 @@ def compute_uniform_xy_scale_for_target( raise ValueError(f"GLB is not a mesh: {glb_path}") if len(target_xy_size_cm) != 2 or any(value <= 0 for value in target_xy_size_cm): raise ValueError("target_xy_size_cm must contain two positive values.") + # GLB geometry is y-up, while the target footprint is defined on z-up table XY. + y_up_to_z_up = np.eye(4) + y_up_to_z_up[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + mesh.apply_transform(y_up_to_z_up) if rotate_about_x: center = mesh.bounds.mean(axis=0) mesh.apply_translation(-center) diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py new file mode 100644 index 000000000..51459b56b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +import pytest +import trimesh + +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + compute_uniform_xy_scale_for_target, +) + + +def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: + """Measure y-up GLBs against the VLM's z-up XY target footprint.""" + glb_path = tmp_path / "flat_fork.glb" + # In y-up, the thin vertical axis is y; in z-up it becomes the z axis. + trimesh.creation.box(extents=[2.0, 0.01, 0.5]).export(glb_path) + + scale = compute_uniform_xy_scale_for_target( + glb_path=glb_path, + target_xy_size_cm=[200.0, 50.0], + rotate_about_x=False, + ) + + assert scale == pytest.approx(1.0) From 69431c2a49b19778a29b1e112a299bda78a56edb Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 17:41:36 +0800 Subject: [PATCH 17/85] update image-to-scene scene understanding: let vlm give orientation state for assets who need to be calibrated --- .../gen_sim/scene_engine/core/scene_graph.py | 14 +- .../editing/scene_edit_understanding.py | 1 + .../generation/scene_understanding.py | 153 +++++++++++++++- .../utils/image_segmentation_utils.py | 61 +++++++ .../pipeline/utils/scene_importer.py | 5 + .../test_scene_core_and_export.py | 28 +++ .../gen_sim/scene_engine/test_scene_graph.py | 4 + .../scene_engine/test_scene_understanding.py | 170 +++++++++++++++++- 8 files changed, 427 insertions(+), 9 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index c5c49c625..e13f65aba 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -54,16 +54,23 @@ # A PlanarRelation with B, then A and B must have the same parent node. PlanarRelationType = Literal["left_of", "right_of", "in_front_of", "behind"] SceneConstraintType = SupportRelationType | PlanarRelationType +OrientationState = Literal["standing", "lying"] @dataclass class SceneGraphNode: - """One object node in the edit-time scene hierarchy.""" + """One object node in the edit-time scene hierarchy. + + ``orientation_state`` is an image-derived placement semantic, rather than + an edge to the node itself or an exact three-dimensional transform. + """ object_id: str parent_id: str | None parent_relation: SupportRelationType | None = None table_region: TableRegion | None = None + # Preserves image-observed placement semantics for later pose refinement. + orientation_state: OrientationState | None = None def __post_init__(self) -> None: """Validate local node fields before graph-level checks.""" @@ -71,12 +78,16 @@ def __post_init__(self) -> None: raise ValueError("object_id must be non-empty.") if self.table_region not in {None, *TABLE_REGIONS}: raise ValueError("table_region is invalid.") + if self.orientation_state not in {None, "standing", "lying"}: + raise ValueError("orientation_state is invalid.") # If the node is the table. if self.object_id == TABLE_OBJECT_ID: if self.parent_id is not None: raise ValueError("table must not have a parent.") if self.parent_relation is not None: raise ValueError("table must not have a parent relation.") + if self.orientation_state is not None: + raise ValueError("table must not have an orientation state.") # If the node is not the table. elif self.parent_id is None: raise ValueError("non-table nodes must have a parent.") @@ -90,6 +101,7 @@ def to_dict(self) -> dict[str, object]: "parent_id": self.parent_id, "parent_relation": self.parent_relation, "table_region": self.table_region, + "orientation_state": self.orientation_state, } diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index e66fd8653..a648981e5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -194,6 +194,7 @@ def _build_updated_scene_graph( parent_id=node.parent_id, parent_relation=node.parent_relation, table_region=node.table_region, + orientation_state=node.orientation_state, ) for node in scene_graph.nodes ], diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f2103a945..c5f340714 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -41,6 +41,7 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( MaskCandidate, build_mask_candidates, + render_asset_mask_id_overlay, render_image_without_masks, render_numbered_mask_candidates, save_binary_mask, @@ -147,6 +148,24 @@ Return JSON only, with exactly one key: assignments. It must be null or an array of asset_id and mask_index objects. Do not include Markdown or any other text.""" +_ORIENTATION_STATE_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. +Each visible asset has an outline and an ID label. Determine whether each listed +object is standing, lying, or unknown in the image. + +Use "standing" only when an upright container, such as a bottle, can, jar, +flask, or thermos, is resting vertically on its base. Use "lying" only when +such a container rests on its side. Use null for the table, every other object +type, or any uncertain case. + +Return JSON only, with exactly this schema. Include every supplied object ID +exactly once and do not add IDs: +{ + "orientation_states": [ + {"object_id": "bottle_001", "orientation_state": "standing"}, + {"object_id": "table", "orientation_state": null} + ] +}""" +_UPRIGHT_CONTAINER_ID_TOKENS = frozenset({"bottle", "can", "jar", "flask", "thermos"}) def understand_scene( @@ -174,7 +193,8 @@ def understand_scene( json_max_attempts=json_max_attempts, ) - _segment_scene( + # Receive the validated whole-scene mask, for VLM output the scene graph. + asset_mask_id_overlay_path = _segment_scene( image_path=resolved_image_path, stage_output_root=stage_output_root, scene=scene, @@ -185,7 +205,11 @@ def understand_scene( # Use the segmented image to initialize the scene graph # with the help of the VLM client. # But at here, we do with the simplest way (hard code). - scene_graph = _initialize_scene_graph_from_segmented_scene(scene) + scene_graph = _initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=asset_mask_id_overlay_path, + vlm_client=vlm_client, + ) # Write the Updated scene JSON for debugging. (stage_output_root / "scene.json").write_text( @@ -199,18 +223,38 @@ def understand_scene( return scene, scene_graph -def _initialize_scene_graph_from_segmented_scene(scene: Scene) -> SceneGraph: +def _initialize_scene_graph_from_segmented_scene( + scene: Scene, + *, + asset_mask_id_overlay_path: str | Path, + vlm_client: OpenAICompatibleVLM, +) -> SceneGraph: """Build the initial graph assuming every segmented asset rests on the table.""" + # Get simplified scene info for VLM. + scene_info = _simplify_scene_info_for_graph_initialization(scene=scene) + resolved_asset_mask_id_overlay_path = _validate_image_path( + asset_mask_id_overlay_path + ) if scene.table is None: raise ValueError("Cannot initialize a scene graph without a table.") + orientation_states_by_id = _query_orientation_states( + scene_info=scene_info, + asset_mask_id_overlay_path=resolved_asset_mask_id_overlay_path, + vlm_client=vlm_client, + ) return SceneGraph( nodes=[ SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), *[ - SceneGraphNode( + SceneGraphNode( # semi-hard code. object_id=asset.id, parent_id=TABLE_OBJECT_ID, parent_relation="on", + orientation_state=( + orientation_states_by_id[asset.id] + if _is_upright_container_id(asset.id) + else None + ), ) for asset in scene.assets ], @@ -218,6 +262,93 @@ def _initialize_scene_graph_from_segmented_scene(scene: Scene) -> SceneGraph: ) +def _simplify_scene_info_for_graph_initialization( + *, + scene: Scene, +) -> dict[str, object]: + """Return the object metadata needed to initialize an image-based graph.""" + return { + "existing_object_ids": [scene_object.id for scene_object in scene.objects], + } + + +def _query_orientation_states( + *, + scene_info: dict[str, object], + asset_mask_id_overlay_path: Path, + vlm_client: OpenAICompatibleVLM, +) -> dict[str, str | None]: + """Return validated image-observed orientation states keyed by object ID.""" + response_text = vlm_client.complete( + image_path=asset_mask_id_overlay_path, + system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, + user_prompt=json.dumps(scene_info, ensure_ascii=False), + ) + return _parse_orientation_states_response( + response_text=response_text, + existing_object_ids=scene_info["existing_object_ids"], + ) + + +def _parse_orientation_states_response( + *, + response_text: str, + existing_object_ids: object, +) -> dict[str, str | None]: + """Parse a complete VLM orientation-state response for known object IDs.""" + if not isinstance(existing_object_ids, list) or not all( + isinstance(object_id, str) for object_id in existing_object_ids + ): + raise ValueError("Scene graph initialization requires string object IDs.") + json_text = _strip_json_code_fence(response_text) + try: + payload = json.loads(json_text) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM response is not valid JSON: {exc.msg}") from exc + if not isinstance(payload, dict) or set(payload) != {"orientation_states"}: + raise ValueError("VLM JSON must contain exactly the key: orientation_states.") + states_value = payload["orientation_states"] + if not isinstance(states_value, list): + raise ValueError("VLM JSON key orientation_states must be an array.") + + orientation_states_by_id: dict[str, str | None] = {} + for index, state_value in enumerate(states_value): + if not isinstance(state_value, dict) or set(state_value) != { + "object_id", + "orientation_state", + }: + raise ValueError( + "VLM JSON orientation_states[" + f"{index}] must contain exactly object_id and orientation_state." + ) + object_id = state_value["object_id"] + orientation_state = state_value["orientation_state"] + if not isinstance(object_id, str) or not object_id: + raise ValueError( + f"VLM JSON orientation_states[{index}].object_id is invalid." + ) + if orientation_state not in {None, "standing", "lying"}: + raise ValueError( + f"VLM JSON orientation_states[{index}].orientation_state is invalid." + ) + if object_id in orientation_states_by_id: + raise ValueError(f"VLM JSON repeats orientation state for {object_id!r}.") + orientation_states_by_id[object_id] = orientation_state + + if set(orientation_states_by_id) != set(existing_object_ids): + raise ValueError( + "VLM JSON orientation states must match all existing object IDs." + ) + return orientation_states_by_id + + +def _is_upright_container_id(object_id: str) -> bool: + """Return whether an object ID identifies a standardized upright container.""" + return bool( + set(re.findall(r"[a-z0-9]+", object_id.lower())) & _UPRIGHT_CONTAINER_ID_TOKENS + ) + + def _analyze_image_objects( *, scene: Scene, @@ -386,8 +517,8 @@ def _segment_scene( scene: Scene, vlm_client: OpenAICompatibleVLM, image_segmentation_client: ImageSegmentationClient, -) -> None: - """Add validated table and asset mask paths to a semantic scene.""" +) -> Path: + """Add validated masks and return an asset-only ID overlay image.""" debug_output_root = ( Path(stage_output_root) / "debug" ) # Keeps the mask debug images. @@ -429,6 +560,16 @@ def _segment_scene( vlm_client=vlm_client, image_segmentation_client=image_segmentation_client, ) + asset_masks: list[tuple[str, str]] = [] + for asset in scene.assets: + if asset.mask_path is None: + raise ValueError(f"Asset {asset.id!r} has no validated mask path.") + asset_masks.append((asset.id, asset.mask_path)) + return render_asset_mask_id_overlay( + image_path=image_path, + asset_masks=asset_masks, + output_path=Path(masks_output_root) / "asset_masks_with_ids.png", + ) def _segment_table( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index ccf9af949..29defd4d6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -306,6 +306,67 @@ def render_numbered_mask_candidates( return resolved_output_path +def render_asset_mask_id_overlay( + *, + image_path: str | Path, + asset_masks: list[tuple[str, str | Path]], + output_path: str | Path, +) -> Path: + """Overlay outlined asset masks and stable asset IDs on a scene image. + + The table mask is intentionally omitted so its large contour does not + obscure the asset labels or their visual context in the source image. + """ + asset_ids = [asset_id for asset_id, _ in asset_masks] + if any(not asset_id for asset_id in asset_ids): + raise ValueError("Every asset mask must have a non-empty asset id.") + if len(asset_ids) != len(set(asset_ids)): + raise ValueError("Asset mask ids must be unique.") + + image = Image.open(image_path).convert("RGBA") + overlay = Image.new("RGBA", image.size, (0, 0, 0, 0)) + colors = ( + (239, 83, 80, 255), + (66, 165, 245, 255), + (102, 187, 106, 255), + (255, 202, 40, 255), + (171, 71, 188, 255), + (38, 198, 218, 255), + ) + decoded_masks: list[tuple[str, Image.Image]] = [] + for index, (asset_id, mask_path) in enumerate(asset_masks): + mask = Image.open(mask_path).convert("L") + _require_image_size(mask, image.size) + decoded_masks.append((asset_id, mask)) + color_layer = Image.new("RGBA", image.size, colors[index % len(colors)]) + transparent_layer = Image.new("RGBA", image.size, (0, 0, 0, 0)) + overlay.alpha_composite( + Image.composite( + color_layer, + transparent_layer, + _mask_outer_outline(mask, image.size), + ) + ) + + draw = ImageDraw.Draw(overlay) + font = _load_label_font(image.size) + for asset_id, mask in decoded_masks: + bbox = mask.getbbox() + if bbox is None: + raise ValueError(f"Asset mask {asset_id!r} is empty.") + _draw_number_label( + draw=draw, + label=asset_id, + center=((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2), + font=font, + ) + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + Image.alpha_composite(image, overlay).convert("RGB").save(resolved_output_path) + return resolved_output_path + + def _require_image_size(mask: Image.Image, image_size: tuple[int, int]) -> None: if mask.size != image_size: raise ValueError( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 9ea558a2a..e730b88cb 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -192,12 +192,14 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: "parent_id", "parent_relation", "table_region", + "orientation_state", }: raise ValueError("Scene graph nodes must use the serialized node schema.") object_id = value["object_id"] parent_id = value["parent_id"] parent_relation = value["parent_relation"] table_region = value["table_region"] + orientation_state = value["orientation_state"] if not isinstance(object_id, str) or not isinstance( parent_id, (str, type(None)) ): @@ -206,11 +208,14 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: raise ValueError("Scene graph parent_relation must be 'on' or null.") if table_region is not None and table_region not in TABLE_REGIONS: raise ValueError("Scene graph table_region is invalid.") + if orientation_state not in {None, "standing", "lying"}: + raise ValueError("Scene graph orientation_state is invalid.") return SceneGraphNode( object_id=object_id, parent_id=parent_id, parent_relation=parent_relation, table_region=table_region, + orientation_state=orientation_state, ) @staticmethod diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index b50126751..95658cf16 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -175,12 +175,14 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No "parent_id": None, "parent_relation": None, "table_region": None, + "orientation_state": None, }, { "object_id": "cup", "parent_id": "table", "parent_relation": "on", "table_region": None, + "orientation_state": None, }, ], "relations": [], @@ -195,6 +197,32 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert imported_graph.to_dict() == _scene_graph(scene).to_dict() +def test_scene_graph_importer_restores_node_orientation_state() -> None: + imported_graph = SceneExportImporter._scene_graph_from_data( + { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + { + "object_id": "bottle_001", + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": "standing", + }, + ], + "relations": [], + } + ) + + assert imported_graph.node_by_id()["bottle_001"].orientation_state == "standing" + + def test_scene_export_overwrites_an_existing_scene_export(tmp_path: Path) -> None: table_glb = tmp_path / "table.glb" cup_glb = tmp_path / "cup.glb" diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py index 970a3fc68..d5b6adaf7 100644 --- a/tests/gen_sim/scene_engine/test_scene_graph.py +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -34,6 +34,7 @@ def test_scene_graph_accepts_layered_on_relations() -> None: parent_id="table", parent_relation="on", table_region="center", + orientation_state="standing", ), SceneGraphNode( object_id="cup", @@ -284,6 +285,7 @@ def test_scene_graph_to_dict_serializes_graph_state() -> None: parent_id="table", parent_relation="on", table_region="center", + orientation_state="standing", ), ], ) @@ -297,12 +299,14 @@ def test_scene_graph_to_dict_serializes_graph_state() -> None: "parent_id": None, "parent_relation": None, "table_region": None, + "orientation_state": None, }, { "object_id": "plate", "parent_id": "table", "parent_relation": "on", "table_region": "center", + "orientation_state": "standing", }, ], "relations": [], diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 8b9357f33..5313e75fa 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -20,11 +20,15 @@ import json from pathlib import Path +from PIL import Image, ImageDraw import pytest from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( + render_asset_mask_id_overlay, +) def _response(*, asset_name: str = "cup") -> str: @@ -94,7 +98,44 @@ def complete(self, **_: object) -> str: assert [asset.id for asset in scene.assets] == ["cup_001"] -def test_initial_scene_graph_places_every_asset_on_table() -> None: +def test_asset_mask_id_overlay_excludes_the_table_mask(tmp_path: Path) -> None: + image_path = tmp_path / "scene.png" + table_mask_path = tmp_path / "table_mask.png" + asset_mask_path = tmp_path / "bottle_mask.png" + output_path = tmp_path / "asset_masks_with_ids.png" + image_size = (512, 512) + Image.new("RGB", image_size, "black").save(image_path) + + table_mask = Image.new("L", image_size, 0) + ImageDraw.Draw(table_mask).rectangle((10, 10, 100, 100), fill=255) + table_mask.save(table_mask_path) + asset_mask = Image.new("L", image_size, 0) + ImageDraw.Draw(asset_mask).rectangle((380, 180, 450, 360), fill=255) + asset_mask.save(asset_mask_path) + + rendered_path = render_asset_mask_id_overlay( + image_path=image_path, + asset_masks=[("bottle_001", asset_mask_path)], + output_path=output_path, + ) + + with Image.open(rendered_path) as overlay: + assert overlay.getpixel((10, 10)) == (0, 0, 0) + assert overlay.getpixel((377, 180)) != (0, 0, 0) + + +def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "orientation_states": [ + {"object_id": "table", "orientation_state": None}, + {"object_id": "cup_001", "orientation_state": "lying"}, + ] + } + ) + scene = Scene( objects=[ SceneObject( @@ -114,8 +155,12 @@ def test_initial_scene_graph_places_every_asset_on_table() -> None: ], ) + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( - scene + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] ) assert scene_graph.to_dict() == { @@ -125,13 +170,134 @@ def test_initial_scene_graph_places_every_asset_on_table() -> None: "parent_id": None, "parent_relation": None, "table_region": None, + "orientation_state": None, }, { "object_id": "cup_001", "parent_id": "table", "parent_relation": "on", "table_region": None, + "orientation_state": None, }, ], "relations": [], } + + +def test_scene_graph_initialization_uses_container_orientation_states( + tmp_path: Path, +) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "orientation_states": [ + {"object_id": "table", "orientation_state": None}, + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + {"object_id": "book_001", "orientation_state": "lying"}, + ] + } + ) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + ), + ] + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert scene_graph.node_by_id()["book_001"].orientation_state is None + + +def test_scene_graph_initialization_requires_asset_mask_id_overlay( + tmp_path: Path, +) -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ) + ] + ) + + with pytest.raises(FileNotFoundError, match="Image input not found"): + scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=tmp_path / "missing.png", + vlm_client=object(), # type: ignore[arg-type] + ) + + +def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + center_xy=[0.2, -0.1], + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + center_xy=[-0.1, 0.2], + ), + ] + ) + + simplified_scene_info = ( + scene_understanding._simplify_scene_info_for_graph_initialization( + scene=scene + ) + ) + + assert simplified_scene_info == { + "existing_object_ids": ["table", "bottle_001", "book_001"], + } From 083d614406a771e293b81753cc56128b6bf73ba1 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 18:03:48 +0800 Subject: [PATCH 18/85] fix big-font problem in asset_masks_with_ids; add scene graph based clibration in image-conditioned scene engine pipeline (but only calibrated the upright bottle-like assets currently) --- .../pipeline/generation/scene_generation.py | 119 +++++++++++++++++- .../utils/image_segmentation_utils.py | 60 ++++++++- .../scene_engine/test_scene_generation.py | 95 ++++++++++++++ .../scene_engine/test_scene_understanding.py | 21 ++++ 4 files changed, 287 insertions(+), 8 deletions(-) create mode 100644 tests/gen_sim/scene_engine/test_scene_generation.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 60f91ca81..6103f8cf4 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -22,6 +22,7 @@ import shutil import numpy as np +from scipy.spatial.transform import Rotation import trimesh from shapely.geometry import Polygon @@ -133,6 +134,7 @@ def generate_scene_and_refine( # Layout refinement will start with the table. refined_table_layout, refined_assets_layout = _layout_refinement( scene=scene, # Update this data structure internally. + scene_graph=scene_graph, simready_geometry_output_root=simready_geometry_output_root, # Contains simready assets and their current coarse layout JSON. debug_output_root=debug_output_root, # Keep the table support surface info + optimized layout info (render with matplotlib) for debugging. ) @@ -291,6 +293,7 @@ def _copy_y_up_layout_to_scene_object( def _layout_refinement( *, scene: Scene, + scene_graph: SceneGraph, simready_geometry_output_root: str | Path, debug_output_root: str | Path, ) -> tuple[dict[str, object], list[dict[str, object]]]: @@ -357,7 +360,14 @@ def _layout_refinement( ) ) - # 3. Move all assets as one rigid group so its lowest AABB point is 2cm above + # 3. Correct image-observed standing containers before every geometry-based + # layout stage measures their footprint. + refined_assets_layout = _scene_graph_based_calibration( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + ) + + # 4. Move all assets as one rigid group so its lowest AABB point is 2cm above # the table. This preserves the initial relative poses for the later # gravity simulation, which can settle individual assets physically. @@ -371,7 +381,7 @@ def _layout_refinement( log_info("Scene has no movable assets; skipping support-region clamping.") return refined_table_layout, [] - # 4. Reuse support geometry detected during SimReady processing. + # 5. Reuse support geometry detected during SimReady processing. if ( scene.table is None or scene.table.support_contour_xy is None @@ -395,7 +405,7 @@ def _layout_refinement( ) ) - # 5. Keep the complete clutter rigid in the table plane. A successful + # 6. Keep the complete clutter rigid in the table plane. A successful # result applies one shared z-up XY delta to every AABB, so it preserves # all existing asset-to-asset relations. It is *not* an asset packing # pass: pre-existing overlap is deliberately left to a later optimizer. @@ -421,7 +431,7 @@ def _layout_refinement( ) ) - # 6. Optimize independent asset positions inside the conservative rectangle. + # 7. Optimize independent asset positions inside the conservative rectangle. # The clamp above already used the exact outer contour for the shared shift. overlap_optimizer = AssetsSupportLayoutOptimizer( support_region=table_optimization_rectangle, @@ -437,7 +447,7 @@ def _layout_refinement( refined_assets_layout = overlap_optimizer.optimize() overlap_optimizer.save_overlap_optimization_debug_images() - # 7. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. + # 8. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down # after the simulation. gravity_settler = AssetsGravitySettler( @@ -458,6 +468,105 @@ def _layout_refinement( return refined_table_layout, refined_assets_layout +def _scene_graph_based_calibration( + *, + scene_graph: SceneGraph, + assets_layout: list[dict[str, object]], +) -> list[dict[str, object]]: + """Minimally align graph-marked standing assets with the z-up table frame.""" + # This is the extension point for future image-conditioned scene generation + # calibration. The scene graph may later provide richer image-grounded + # constraints, but the current implementation deliberately consumes only + # ``orientation_state`` to correct standing container axes before layout. + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + nodes_by_id = scene_graph.node_by_id() + calibrated_assets_layout: list[dict[str, object]] = [] + + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset layout must contain a non-empty string id.") + node = nodes_by_id.get(asset_id) + if node is None: + raise ValueError(f"Scene graph does not contain asset {asset_id!r}.") + if node.orientation_state != "standing": + calibrated_assets_layout.append(asset_layout) + continue + + # Conjugate the y-up pose so the SimReady container axis is local z. + z_up_asset_to_table_matrix = ( + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(asset_layout) + @ z_up_to_y_up_matrix + ) + linear_matrix = z_up_asset_to_table_matrix[:3, :3] + # Layout transforms store rotation and per-axis scale in the same matrix. + scale = np.linalg.norm(linear_matrix, axis=0) + if np.any(scale <= 1e-8): + raise ValueError(f"Asset {asset_id!r} has a zero scale axis.") + rotation_matrix = linear_matrix / scale + if not np.allclose(rotation_matrix.T @ rotation_matrix, np.eye(3), atol=1e-6): + raise ValueError(f"Asset {asset_id!r} layout contains shear.") + + local_z_axis_in_table = rotation_matrix[:, 2] + # Treat the long axis as unsigned to avoid an unnecessary 180-degree flip. + target_z_axis = np.array( + [0.0, 0.0, 1.0 if local_z_axis_in_table[2] >= 0.0 else -1.0] + ) + # Left multiplication applies the correction in the table/world frame. + z_up_asset_to_table_matrix[:3, :3] = ( + _minimum_axis_alignment_rotation( + source_axis=local_z_axis_in_table, + target_axis=target_z_axis, + ) + @ rotation_matrix + @ np.diag(scale) + ) + calibrated_assets_layout.append( + transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ z_up_asset_to_table_matrix + @ y_up_to_z_up_matrix, + ) + ) + return calibrated_assets_layout + + +def _minimum_axis_alignment_rotation( + *, + source_axis: np.ndarray, + target_axis: np.ndarray, +) -> np.ndarray: + """Return the smallest proper rotation mapping one nonzero axis to another.""" + source = np.asarray(source_axis, dtype=float) + target = np.asarray(target_axis, dtype=float) + source_norm = np.linalg.norm(source) + target_norm = np.linalg.norm(target) + if source_norm <= 1e-8 or target_norm <= 1e-8: + raise ValueError("Axis alignment requires nonzero axes.") + source /= source_norm + target /= target_norm + + cross_product = np.cross(source, target) + sine = np.linalg.norm(cross_product) + cosine = float(np.clip(np.dot(source, target), -1.0, 1.0)) + if sine <= 1e-8: + if cosine > 0.0: + return np.eye(3) + basis_axis = np.eye(3)[np.argmin(np.abs(source))] + rotation_axis = np.cross(source, basis_axis) + rotation_axis /= np.linalg.norm(rotation_axis) + return Rotation.from_rotvec(np.pi * rotation_axis).as_matrix() + + rotation_axis = cross_product / sine + return Rotation.from_rotvec(np.arctan2(sine, cosine) * rotation_axis).as_matrix() + + def _measure_table_and_assets_in_z_up_world( *, table_layout: dict[str, object], diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index 29defd4d6..302ec2cbf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -349,16 +349,21 @@ def render_asset_mask_id_overlay( ) draw = ImageDraw.Draw(overlay) - font = _load_label_font(image.size) for asset_id, mask in decoded_masks: bbox = mask.getbbox() if bbox is None: raise ValueError(f"Asset mask {asset_id!r} is empty.") + font = _load_asset_id_label_font( + image_size=image.size, + mask_bbox=bbox, + label=asset_id, + ) _draw_number_label( draw=draw, label=asset_id, center=((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2), font=font, + minimum_padding=2, ) resolved_output_path = Path(output_path).expanduser().resolve() @@ -504,6 +509,47 @@ def _union_parent(parents: list[int], first_index: int, second_index: int) -> No def _load_label_font(image_size: tuple[int, int]) -> ImageFont.ImageFont: font_size = max(16, round(min(image_size) / 32)) + return _load_label_font_at_size(font_size) + + +def _load_asset_id_label_font( + *, + image_size: tuple[int, int], + mask_bbox: tuple[int, int, int, int], + label: str, +) -> ImageFont.ImageFont: + """Choose an ID-label font constrained by both image and mask dimensions.""" + # The image sets the readable upper bound; the individual mask then caps it. + image_font_size = min(32, max(8, round(min(image_size) / 48))) + mask_width = mask_bbox[2] - mask_bbox[0] + mask_height = mask_bbox[3] - mask_bbox[1] + maximum_label_width = max(24, round(mask_width * 0.9)) + maximum_label_height = max(16, round(mask_height * 0.75)) + # Measure the complete text-and-background rectangle, not glyphs alone. + probe_draw = ImageDraw.Draw(Image.new("RGBA", image_size)) + smallest_font = _load_label_font_at_size(6) + for font_size in range(image_font_size, 5, -1): + font = _load_label_font_at_size(font_size) + label_bounds = _number_label_bounds( + draw=probe_draw, + label=label, + center=(0.0, 0.0), + font=font, + minimum_padding=2, + ) + if ( + label_bounds[2] - label_bounds[0] <= maximum_label_width + and label_bounds[3] - label_bounds[1] <= maximum_label_height + ): + return font + smallest_font = font + return smallest_font + + +def _load_label_font_at_size(font_size: int) -> ImageFont.ImageFont: + """Load the shared bold label font at one validated pixel size.""" + if font_size < 1: + raise ValueError("Label font size must be positive.") try: return ImageFont.truetype("DejaVuSans-Bold.ttf", font_size) except OSError: @@ -516,10 +562,15 @@ def _draw_number_label( label: str, center: tuple[float, float], font: ImageFont.ImageFont, + minimum_padding: int = 4, ) -> None: """Draw a numbered label with red background and white text at the given center position.""" label_bounds = _number_label_bounds( - draw=draw, label=label, center=center, font=font + draw=draw, + label=label, + center=center, + font=font, + minimum_padding=minimum_padding, ) label_box = draw.textbbox((0, 0), label, font=font) label_width = label_box[2] - label_box[0] @@ -541,12 +592,15 @@ def _number_label_bounds( label: str, center: tuple[float, float], font: ImageFont.ImageFont, + minimum_padding: int = 4, ) -> tuple[int, int, int, int]: """Return the red label rectangle bounds for a label centre.""" label_box = draw.textbbox((0, 0), label, font=font) label_width = label_box[2] - label_box[0] label_height = label_box[3] - label_box[1] - padding = max(4, round(max(label_width, label_height) / 4)) + if minimum_padding < 0: + raise ValueError("Label minimum padding must be non-negative.") + padding = max(minimum_padding, round(max(label_width, label_height) / 4)) x = center[0] - label_width / 2 y = center[1] - label_height / 2 return ( diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py new file mode 100644 index 000000000..6a61a4d1b --- /dev/null +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -0,0 +1,95 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + _scene_graph_based_calibration, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + transform_matrix_to_layout_object, +) + + +def _y_up_layout_from_z_up_rotation( + object_id: str, + rotation_matrix: np.ndarray, +) -> dict[str, object]: + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + z_up_transform = np.eye(4) + z_up_transform[:3, :3] = rotation_matrix + return transform_matrix_to_layout_object( + object_id, + z_up_to_y_up_matrix @ z_up_transform @ y_up_to_z_up_matrix, + ) + + +def _z_up_rotation_from_y_up_layout(layout: dict[str, object]) -> np.ndarray: + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + return ( + y_up_to_z_up_matrix + @ layout_object_to_transform_matrix(layout) + @ np.linalg.inv(y_up_to_z_up_matrix) + )[:3, :3] + + +def test_scene_graph_calibration_makes_standing_asset_vertical() -> None: + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="bottle_001", + parent_id="table", + parent_relation="on", + orientation_state="standing", + ), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + ] + ) + lying_rotation = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + bottle_layout = _y_up_layout_from_z_up_rotation("bottle_001", lying_rotation) + book_layout = _y_up_layout_from_z_up_rotation("book_001", lying_rotation) + + calibrated_layouts = _scene_graph_based_calibration( + scene_graph=scene_graph, + assets_layout=[bottle_layout, book_layout], + ) + + bottle_axis = _z_up_rotation_from_y_up_layout(calibrated_layouts[0])[:, 2] + assert np.isclose(abs(bottle_axis[2]), 1.0) + assert np.allclose( + _z_up_rotation_from_y_up_layout(calibrated_layouts[1]), + lying_rotation, + ) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 5313e75fa..2a9af65b2 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -26,6 +26,7 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation import scene_understanding +from embodichain.gen_sim.scene_engine.pipeline.utils import image_segmentation_utils from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( render_asset_mask_id_overlay, ) @@ -124,6 +125,26 @@ def test_asset_mask_id_overlay_excludes_the_table_mask(tmp_path: Path) -> None: assert overlay.getpixel((377, 180)) != (0, 0, 0) +def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: + image_size = (512, 512) + mask_bbox = (380, 180, 450, 360) + label = "bottle_001" + font = image_segmentation_utils._load_asset_id_label_font( + image_size=image_size, + mask_bbox=mask_bbox, + label=label, + ) + label_bounds = image_segmentation_utils._number_label_bounds( + draw=ImageDraw.Draw(Image.new("RGBA", image_size)), + label=label, + center=(0.0, 0.0), + font=font, + minimum_padding=2, + ) + + assert label_bounds[2] - label_bounds[0] <= round((mask_bbox[2] - mask_bbox[0]) * 0.9) + + def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: class VLM: def complete(self, **_: object) -> str: From c5e9087bffe7f5192d53f19c71bec3d67185ddbf Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 18:04:23 +0800 Subject: [PATCH 19/85] run black --- .../scene_engine/pipeline/generation/scene_generation.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 6103f8cf4..db7ed8ee2 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -529,9 +529,7 @@ def _scene_graph_based_calibration( calibrated_assets_layout.append( transform_matrix_to_layout_object( asset_id, - z_up_to_y_up_matrix - @ z_up_asset_to_table_matrix - @ y_up_to_z_up_matrix, + z_up_to_y_up_matrix @ z_up_asset_to_table_matrix @ y_up_to_z_up_matrix, ) ) return calibrated_assets_layout From 4779f01804f06dc3307e1ddaaafe9c3f45d79390 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Sat, 15 Aug 2026 18:10:49 +0800 Subject: [PATCH 20/85] run black to tests/ --- tests/gen_sim/scene_engine/test_scene_understanding.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 2a9af65b2..cd63bd2b0 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -142,7 +142,9 @@ def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: minimum_padding=2, ) - assert label_bounds[2] - label_bounds[0] <= round((mask_bbox[2] - mask_bbox[0]) * 0.9) + assert label_bounds[2] - label_bounds[0] <= round( + (mask_bbox[2] - mask_bbox[0]) * 0.9 + ) def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: @@ -314,9 +316,7 @@ def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: ) simplified_scene_info = ( - scene_understanding._simplify_scene_info_for_graph_initialization( - scene=scene - ) + scene_understanding._simplify_scene_info_for_graph_initialization(scene=scene) ) assert simplified_scene_info == { From 4a792f01a7384e316e334f4d4252739d6395d92d Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:32:30 +0800 Subject: [PATCH 21/85] replace hard-code bottle-z-up with VLM auto-analyze --- .../pipeline/generation/scene_generation.py | 5 + .../generation/scene_understanding.py | 86 +++++++------- .../pipeline/utils/simready_processor.py | 106 +++++------------- .../scene_engine/test_scene_understanding.py | 83 ++++++++++++-- .../test_simready_processor_utils.py | 22 ++++ 5 files changed, 177 insertions(+), 125 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index db7ed8ee2..7a0e19f35 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -119,6 +119,11 @@ def generate_scene_and_refine( config=SimReadyProcessorConfig( use_vlm_scale=False, use_vlm_rotation=False, + long_axis_object_ids=frozenset( + node.object_id + for node in scene_graph.nodes + if node.orientation_state is not None + ), ), vlm_client=vlm_client, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index c5f340714..9507c1d3a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -150,22 +150,21 @@ text.""" _ORIENTATION_STATE_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. Each visible asset has an outline and an ID label. Determine whether each listed -object is standing, lying, or unknown in the image. +asset is standing, lying, or unknown in the image. -Use "standing" only when an upright container, such as a bottle, can, jar, -flask, or thermos, is resting vertically on its base. Use "lying" only when -such a container rests on its side. Use null for the table, every other object -type, or any uncertain case. +Use a non-null state only for an elongated object with a clear primary long axis. +Use "standing" when its primary axis is approximately vertical to the tabletop. +Use "lying" when its primary axis is approximately parallel to the tabletop. +Use null for every object without a clear primary long axis or when uncertain. -Return JSON only, with exactly this schema. Include every supplied object ID -exactly once and do not add IDs: +Return JSON only, with exactly this schema. Include every supplied asset ID +exactly once. Never include the table or any ID that was not supplied: { "orientation_states": [ {"object_id": "bottle_001", "orientation_state": "standing"}, - {"object_id": "table", "orientation_state": null} + {"object_id": "book_001", "orientation_state": null} ] }""" -_UPRIGHT_CONTAINER_ID_TOKENS = frozenset({"bottle", "can", "jar", "flask", "thermos"}) def understand_scene( @@ -209,6 +208,7 @@ def understand_scene( scene, asset_mask_id_overlay_path=asset_mask_id_overlay_path, vlm_client=vlm_client, + json_max_attempts=json_max_attempts, ) # Write the Updated scene JSON for debugging. @@ -228,6 +228,7 @@ def _initialize_scene_graph_from_segmented_scene( *, asset_mask_id_overlay_path: str | Path, vlm_client: OpenAICompatibleVLM, + json_max_attempts: int = 3, ) -> SceneGraph: """Build the initial graph assuming every segmented asset rests on the table.""" # Get simplified scene info for VLM. @@ -241,20 +242,17 @@ def _initialize_scene_graph_from_segmented_scene( scene_info=scene_info, asset_mask_id_overlay_path=resolved_asset_mask_id_overlay_path, vlm_client=vlm_client, + json_max_attempts=json_max_attempts, ) return SceneGraph( nodes=[ SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), *[ - SceneGraphNode( # semi-hard code. + SceneGraphNode( object_id=asset.id, parent_id=TABLE_OBJECT_ID, - parent_relation="on", - orientation_state=( - orientation_states_by_id[asset.id] - if _is_upright_container_id(asset.id) - else None - ), + parent_relation="on", # semi-hard-code. + orientation_state=orientation_states_by_id[asset.id], ) for asset in scene.assets ], @@ -268,7 +266,7 @@ def _simplify_scene_info_for_graph_initialization( ) -> dict[str, object]: """Return the object metadata needed to initialize an image-based graph.""" return { - "existing_object_ids": [scene_object.id for scene_object in scene.objects], + "asset_ids": [asset.id for asset in scene.assets], } @@ -277,29 +275,42 @@ def _query_orientation_states( scene_info: dict[str, object], asset_mask_id_overlay_path: Path, vlm_client: OpenAICompatibleVLM, + json_max_attempts: int, ) -> dict[str, str | None]: - """Return validated image-observed orientation states keyed by object ID.""" - response_text = vlm_client.complete( - image_path=asset_mask_id_overlay_path, - system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, - user_prompt=json.dumps(scene_info, ensure_ascii=False), - ) - return _parse_orientation_states_response( - response_text=response_text, - existing_object_ids=scene_info["existing_object_ids"], - ) + """Return validated image-observed orientation states keyed by asset ID.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + image_path=asset_mask_id_overlay_path, + system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, + user_prompt=json.dumps(scene_info, ensure_ascii=False), + ) + try: + return _parse_orientation_states_response( + response_text=response_text, + asset_ids=scene_info["asset_ids"], + ) + except ValueError as exc: + last_validation_error = exc + assert last_validation_error is not None + raise ValueError( + "VLM returned invalid orientation-state JSON after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error def _parse_orientation_states_response( *, response_text: str, - existing_object_ids: object, + asset_ids: object, ) -> dict[str, str | None]: - """Parse a complete VLM orientation-state response for known object IDs.""" - if not isinstance(existing_object_ids, list) or not all( - isinstance(object_id, str) for object_id in existing_object_ids + """Parse a complete VLM orientation-state response for known asset IDs.""" + if not isinstance(asset_ids, list) or not all( + isinstance(object_id, str) for object_id in asset_ids ): - raise ValueError("Scene graph initialization requires string object IDs.") + raise ValueError("Scene graph initialization requires string asset IDs.") json_text = _strip_json_code_fence(response_text) try: payload = json.loads(json_text) @@ -335,20 +346,13 @@ def _parse_orientation_states_response( raise ValueError(f"VLM JSON repeats orientation state for {object_id!r}.") orientation_states_by_id[object_id] = orientation_state - if set(orientation_states_by_id) != set(existing_object_ids): + if set(orientation_states_by_id) != set(asset_ids): raise ValueError( - "VLM JSON orientation states must match all existing object IDs." + "VLM JSON orientation states must match all supplied asset IDs." ) return orientation_states_by_id -def _is_upright_container_id(object_id: str) -> bool: - """Return whether an object ID identifies a standardized upright container.""" - return bool( - set(re.findall(r"[a-z0-9]+", object_id.lower())) & _UPRIGHT_CONTAINER_ID_TOKENS - ) - - def _analyze_image_objects( *, scene: Scene, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 55860be7f..7bc041abf 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -19,11 +19,9 @@ from dataclasses import dataclass from pathlib import Path -import re import numpy as np import open3d as o3d -from scipy.spatial import ConvexHull, QhullError from scipy.spatial.transform import Rotation import trimesh @@ -66,14 +64,12 @@ @dataclass(frozen=True) class SimReadyProcessorConfig: - """Object-category policy for SimReady mesh canonicalization.""" + """SceneGraph-conditioned policy for SimReady mesh canonicalization.""" use_vlm_scale: bool = False # Use the VLM-selected asset scale. use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. - upright_container_id_tokens: frozenset[str] = frozenset( - {"bottle", "can", "jar", "flask", "thermos"} - ) # Object-id tokens that enable upright-container standardization. + long_axis_object_ids: frozenset[str] = frozenset() class SimReadyProcessor: @@ -106,8 +102,6 @@ def __init__( self.simready_assets_layout: list[dict[str, object]] | None = None self.config = config if config is not None else SimReadyProcessorConfig() self.vlm_client = vlm_client - if not self.config.upright_container_id_tokens: - raise ValueError("upright_container_id_tokens must not be empty.") if ( self.config.use_vlm_scale or self.config.use_vlm_rotation ) and vlm_client is None: @@ -312,8 +306,6 @@ def _canonicalize_object_mesh( ) if np.any(coarse_scale <= 0): raise ValueError("Coarse object scale values must be positive.") - # We need the object id to determine whether it is a bottle-like object. - # If it does, then we will do a special standardization. (Hard code) if not isinstance(object_id, str) or not object_id: raise ValueError("Scene object id must be a non-empty string.") @@ -324,15 +316,15 @@ def _canonicalize_object_mesh( y_up_to_z_up_transform[:3, :3] = y_up_to_z_up_matrix mesh.apply_transform(y_up_to_z_up_transform) - # Standardize upright containers in temporary z-up coordinates before the - # shared center, scale, and bottom-center preprocessing. - # This is to ensure the action agent can pick up the bottle or can-like objects. - bottle_alignment_matrix = np.eye(3) - if self._is_upright_container_id(object_id): - bottle_alignment_matrix = self._standardize_bottle_z_up(mesh) - bottle_alignment_transform = np.eye(4) - bottle_alignment_transform[:3, :3] = bottle_alignment_matrix - mesh.apply_transform(bottle_alignment_transform) + # Standardize graph-marked elongated assets before shared mesh processing. + # This makes local z their primary axis, so later scene-graph calibration + # can reliably recover the image-observed standing or lying orientation. + long_axis_alignment_matrix = np.eye(3) + if self._requires_long_axis_standardization(object_id): + long_axis_alignment_matrix = self._standardize_long_axis_z_up(mesh) + long_axis_alignment_transform = np.eye(4) + long_axis_alignment_transform[:3, :3] = long_axis_alignment_matrix + mesh.apply_transform(long_axis_alignment_transform) # First make the object's AABB center at the origin. original_aabb_center = mesh.bounds.mean(axis=0) @@ -343,11 +335,11 @@ def _canonicalize_object_mesh( scale_transform[:3, :3] = ( # Actually there's no need to do so, for the scale factor is all equal # in x, y, z axes. - bottle_alignment_matrix + long_axis_alignment_matrix @ y_up_to_z_up_matrix @ np.diag(coarse_scale) @ y_up_to_z_up_matrix.T - @ bottle_alignment_matrix.T + @ long_axis_alignment_matrix.T ) mesh.apply_transform(scale_transform) @@ -367,16 +359,16 @@ def _canonicalize_object_mesh( z_up_to_y_up_transform[:3, :3] = y_up_to_z_up_matrix.T mesh.apply_transform(z_up_to_y_up_transform) - # Compensate the bottle's local rotation so that its coarse world pose does - # not change. - local_bottle_rotation = Rotation.from_matrix( - y_up_to_z_up_matrix.T @ bottle_alignment_matrix @ y_up_to_z_up_matrix + # Compensate the local canonicalization so its coarse world pose does not + # change until layout refinement applies the image-observed correction. + local_long_axis_rotation = Rotation.from_matrix( + y_up_to_z_up_matrix.T @ long_axis_alignment_matrix @ y_up_to_z_up_matrix ) coarse_rotation_matrix = Rotation.from_euler( "xyz", coarse_rot, degrees=True ).as_matrix() rotation = Rotation.from_matrix( - coarse_rotation_matrix @ local_bottle_rotation.inv().as_matrix() + coarse_rotation_matrix @ local_long_axis_rotation.inv().as_matrix() ) # Update the pos. position_offset = y_up_to_z_up_matrix.T @ ( @@ -388,24 +380,18 @@ def _canonicalize_object_mesh( "scale": [1.0, 1.0, 1.0], } - def _is_upright_container_id(self, object_id: str) -> bool: - """Return whether object-id tokens indicate a bottle-like container.""" - # Example: soda_can_0 - # tokens: {"soda", "can", "0"} - # upright_container_id_tokens: {"bottle", "can", "jar"} - # So this returns True because "can" is in the configured token set. - tokens = set(re.findall(r"[a-z0-9]+", object_id.lower())) - return bool(tokens & self.config.upright_container_id_tokens) + def _requires_long_axis_standardization(self, object_id: str) -> bool: + """Return whether graph semantics identified one asset with a long axis.""" + return object_id in self.config.long_axis_object_ids @staticmethod - def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: - """Return a proper rotation that maps a bottle-like mesh's long axis to z-up. - - Thanks to chenjian for this idea! + def _standardize_long_axis_z_up(mesh: trimesh.Trimesh) -> np.ndarray: + """Return a proper rotation that maps a mesh's primary axis to z-up. + Thanks to chanjian's idea. """ if len(mesh.vertices) < 4 or len(mesh.faces) < 4: raise ValueError( - "Bottle standardization requires a non-degenerate triangle mesh." + "Long-axis standardization requires a non-degenerate triangle mesh." ) open3d_mesh = o3d.geometry.TriangleMesh( vertices=o3d.utility.Vector3dVector(mesh.vertices), @@ -419,7 +405,7 @@ def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: # non-finite values. if sampled_points.shape[0] < 4 or not np.all(np.isfinite(sampled_points)): raise ValueError( - "Bottle standardization could not sample valid mesh points." + "Long-axis standardization could not sample valid mesh points." ) centered_points = sampled_points - sampled_points.mean(axis=0) @@ -428,46 +414,12 @@ def _standardize_bottle_z_up(mesh: trimesh.Trimesh) -> np.ndarray: if np.linalg.det(principal_axes) < 0: principal_axes[2, :] *= -1 # in case the SVD returns a reflection. - bottle_rotation = Rotation.from_euler( + long_axis_rotation = Rotation.from_euler( "y", 90.0, degrees=True ).as_matrix() # 3x3 matrix # The first PCA axis is the longest axis; rotate it onto the temporary z axis. - bottle_rotation = bottle_rotation @ principal_axes - standardized_points = (bottle_rotation @ centered_points.T).T - - axis_min = standardized_points[:, 2].min() - axis_max = standardized_points[:, 2].max() - axis_range = axis_max - axis_min - upper_points = standardized_points[ - standardized_points[:, 2] > axis_min + axis_range * 0.8 - ] - lower_points = standardized_points[ - standardized_points[:, 2] < axis_min + axis_range * 0.2 - ] - upper_volume = SimReadyProcessor._convex_hull_volume(upper_points) - lower_volume = SimReadyProcessor._convex_hull_volume(lower_points) - - # Bottles usually have a smaller top (neck) than bottom; flip if necessary. - if upper_volume > lower_volume: - bottle_rotation = ( - Rotation.from_euler("x", 180.0, degrees=True).as_matrix() - @ bottle_rotation - ) - return bottle_rotation - - @staticmethod - def _convex_hull_volume(points: np.ndarray) -> float: - """Return the volume of a non-degenerate point set's convex hull.""" - if points.shape[0] < 4: - raise ValueError( - "Bottle standardization needs at least four points per end." - ) - try: - return float(ConvexHull(points).volume) - except QhullError as exc: - raise ValueError( - "Bottle standardization found a degenerate end volume." - ) from exc + long_axis_rotation = long_axis_rotation @ principal_axes + return long_axis_rotation @staticmethod def _three_floats(value: object, *, field_name: str) -> list[float]: diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index cd63bd2b0..d2b3898da 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -153,8 +153,7 @@ def complete(self, **_: object) -> str: return json.dumps( { "orientation_states": [ - {"object_id": "table", "orientation_state": None}, - {"object_id": "cup_001", "orientation_state": "lying"}, + {"object_id": "cup_001", "orientation_state": None}, ] } ) @@ -207,15 +206,18 @@ def complete(self, **_: object) -> str: } -def test_scene_graph_initialization_uses_container_orientation_states( +def test_scene_graph_initialization_uses_image_orientation_states( tmp_path: Path, ) -> None: class VLM: + def __init__(self) -> None: + self.user_prompt: str | None = None + def complete(self, **_: object) -> str: + self.user_prompt = _["user_prompt"] # type: ignore[assignment,index] return json.dumps( { "orientation_states": [ - {"object_id": "table", "orientation_state": None}, { "object_id": "bottle_001", "orientation_state": "standing", @@ -253,14 +255,81 @@ def complete(self, **_: object) -> str: ] ) + vlm = VLM() + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=vlm, # type: ignore[arg-type] + ) + + assert json.loads(vlm.user_prompt or "{}") == { + "asset_ids": ["bottle_001", "book_001"], + } + assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert scene_graph.node_by_id()["book_001"].orientation_state == "lying" + + +def test_scene_graph_initialization_retries_a_response_containing_table( + tmp_path: Path, +) -> None: + class VLM: + def __init__(self) -> None: + self.responses = [ + json.dumps( + { + "orientation_states": [ + {"object_id": "table", "orientation_state": None}, + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + ] + } + ), + json.dumps( + { + "orientation_states": [ + { + "object_id": "bottle_001", + "orientation_state": "standing", + }, + ] + } + ), + ] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="bottle_001", + kind="asset", + category="bottle", + name="blue bottle", + description="A blue bottle.", + ), + ] + ) + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( scene, asset_mask_id_overlay_path=overlay_path, vlm_client=VLM(), # type: ignore[arg-type] + json_max_attempts=2, ) assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" - assert scene_graph.node_by_id()["book_001"].orientation_state is None def test_scene_graph_initialization_requires_asset_mask_id_overlay( @@ -286,7 +355,7 @@ def test_scene_graph_initialization_requires_asset_mask_id_overlay( ) -def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: +def test_scene_graph_initialization_info_lists_asset_ids() -> None: scene = Scene( objects=[ SceneObject( @@ -320,5 +389,5 @@ def test_scene_graph_initialization_info_lists_existing_object_ids() -> None: ) assert simplified_scene_info == { - "existing_object_ids": ["table", "bottle_001", "book_001"], + "asset_ids": ["bottle_001", "book_001"], } diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py index 51459b56b..13cf73e10 100644 --- a/tests/gen_sim/scene_engine/test_simready_processor_utils.py +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -21,11 +21,33 @@ import pytest import trimesh +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( + SimReadyProcessor, + SimReadyProcessorConfig, +) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( compute_uniform_xy_scale_for_target, ) +def test_simready_long_axis_standardization_uses_graph_selected_ids( + tmp_path: Path, +) -> None: + processor = SimReadyProcessor( + scene=Scene(), + coarse_layout_by_id={}, + coarse_geometry_root=tmp_path / "coarse", + simready_geometry_root=tmp_path / "simready", + config=SimReadyProcessorConfig( + long_axis_object_ids=frozenset({"rolling_pin_001"}), + ), + ) + + assert processor._requires_long_axis_standardization("rolling_pin_001") + assert not processor._requires_long_axis_standardization("bottle_001") + + def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: """Measure y-up GLBs against the VLM's z-up XY target footprint.""" glb_path = tmp_path / "flat_fork.glb" From d3f4365aa13d1ee51dae5a3abb7e8d0e25ce32fc Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:20:28 +0800 Subject: [PATCH 22/85] make assets group layout optimizer more rubust --- .../utils/assets_group_layout_optimizer.py | 58 ++++++++++++++++--- .../scene_engine/test_support_and_layout.py | 35 +++++++++++ 2 files changed, 84 insertions(+), 9 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py index d4021fede..8b3933b09 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_group_layout_optimizer.py @@ -120,24 +120,20 @@ def optimize(self) -> list[dict[str, object]]: base_aabbs = np.stack([aabbs_by_id[asset_id] for asset_id in asset_ids]) offsets = np.zeros((len(asset_ids), 2), dtype=float) if not self._all_contained(safe_support, base_aabbs, offsets): - log_warning( - "AABB overlap optimization requires all input AABBs to be inside " - "the support region." - ) - raise ValueError( - "Overlap optimization requires AABBs already inside support; " - "run AssetsGroupSupportClamp first." - ) + # Independently project each AABB into the rectangular optimization region. + offsets = self._project_aabbs_inside_rectangle(safe_support, base_aabbs) initial_overlaps = self._overlaps(base_aabbs, offsets) + projected_asset_count = int(np.count_nonzero(np.any(offsets != 0.0, axis=1))) log_info( "Support-constrained AABB overlap optimization started: " f"assets={len(asset_ids)}, initial_overlaps={len(initial_overlaps)}, " + f"initial_projections={projected_asset_count}, " f"boundary_margin={self.config.margin_m:.4f} m, " f"aabb_clearance={self.config.aabb_clearance_m:.4f} m, " f"max_rounds={self.config.max_rounds}." ) if not initial_overlaps: # Return directly if there are no overlaps to resolve. - log_info("AABB overlap optimization succeeded without movement.") + log_info("AABB overlap optimization succeeded without pair separation.") self.refined_assets_layout = self._apply_offsets_to_y_up_layouts( asset_ids=asset_ids, offsets=offsets, @@ -213,6 +209,50 @@ def optimize(self) -> list[dict[str, object]]: "inside the detected table support region." ) + @staticmethod + def _project_aabbs_inside_rectangle( + support: Polygon | MultiPolygon, + base_aabbs: np.ndarray, + ) -> np.ndarray: + """Return minimum per-AABB offsets that place AABBs in a rectangle.""" + if not isinstance(support, Polygon) or support.interiors: + raise ValueError( + "Initial AABB projection requires an axis-aligned rectangular " + "support region." + ) + minimum_x, minimum_y, maximum_x, maximum_y = support.bounds + rectangle = Polygon( + [ + (minimum_x, minimum_y), + (maximum_x, minimum_y), + (maximum_x, maximum_y), + (minimum_x, maximum_y), + ] + ) + if not support.equals(rectangle): + raise ValueError( + "Initial AABB projection requires an axis-aligned rectangular " + "support region." + ) + + aabb_minimums, aabb_maximums = base_aabbs.min(axis=1), base_aabbs.max(axis=1) + half_extents = (aabb_maximums - aabb_minimums) / 2.0 + support_minimum = np.array([minimum_x, minimum_y], dtype=float) + support_maximum = np.array([maximum_x, maximum_y], dtype=float) + valid_center_minimums = support_minimum + half_extents + valid_center_maximums = support_maximum - half_extents + if np.any(valid_center_minimums > valid_center_maximums + 1e-9): + raise ValueError( + "An asset AABB is larger than the rectangular support region." + ) + + centers = (aabb_minimums + aabb_maximums) / 2.0 + # A center must stay inset from each boundary by its AABB half extent. + projected_centers = np.clip( + centers, valid_center_minimums, valid_center_maximums + ) + return projected_centers - centers + def _apply_offsets_to_y_up_layouts( self, *, asset_ids: list[str], offsets: np.ndarray ) -> list[dict[str, object]]: diff --git a/tests/gen_sim/scene_engine/test_support_and_layout.py b/tests/gen_sim/scene_engine/test_support_and_layout.py index 98ed30491..fda7f9e08 100644 --- a/tests/gen_sim/scene_engine/test_support_and_layout.py +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -166,6 +166,41 @@ def test_layout_optimizer_resolves_a_simple_pair_overlap() -> None: assert not optimizer._overlaps(base_aabbs, refined_offsets) +def test_layout_optimizer_projects_out_of_bounds_aabb_into_rectangle() -> None: + support = Polygon([(0, 0), (3, 0), (3, 3), (0, 3)]) + layout = _layout("cup", 0.0, 1.5) + aabb = _aabb(-0.5, 1.0, 0.5, 2.0) + optimizer = AssetsSupportLayoutOptimizer( + support_region=support, + assets_aabb_2d_z_up_world_corners_by_id={"cup": aabb}, + assets_layout=[layout], + ) + + refined = optimizer.optimize() + + offset = np.array( + [ + refined[0]["pos"][0] - layout["pos"][0], # type: ignore[index] + layout["pos"][2] - refined[0]["pos"][2], # type: ignore[index] + ] + ) + assert offset[0] == pytest.approx(0.5) + assert optimizer._all_contained(support, np.stack([aabb]), np.stack([offset])) + + +def test_layout_optimizer_rejects_aabb_larger_than_rectangle() -> None: + optimizer = AssetsSupportLayoutOptimizer( + support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), + assets_aabb_2d_z_up_world_corners_by_id={ + "large": _aabb(-0.5, 0.5, 2.5, 1.5) + }, + assets_layout=[_layout("large", 1.0, 1.0)], + ) + + with pytest.raises(ValueError, match="larger than the rectangular support"): + optimizer.optimize() + + def test_layout_optimizer_rejects_unresolvable_overlap() -> None: optimizer = AssetsSupportLayoutOptimizer( support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), From 1d5e3674c39df143fccb7504181b65f64d93b043 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:21:19 +0800 Subject: [PATCH 23/85] delete the hard-coded spatial check in assets' names and descriptionsa --- .../generation/scene_understanding.py | 20 +++-------------- .../scene_engine/test_scene_understanding.py | 22 +++++++++++-------- 2 files changed, 16 insertions(+), 26 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index 9507c1d3a..d93ada019 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -50,12 +50,6 @@ _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} _CATEGORY_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$") -_LOCATION_WORD_PATTERN = re.compile( - r"\b(?:left|right|front|back|center|middle|top|bottom|upper|lower|" - r"foreground|background|near|next|beside|behind|between|on|in|inside|" - r"under|above|below|against)\b", - flags=re.IGNORECASE, -) _SYSTEM_PROMPT = """You inspect one tabletop-scene image. Identify the main table and every visible, physically distinct object that should be segmented and later generated as an independent 3D asset. @@ -82,6 +76,9 @@ 8. For assets, description contains only visible category, material, color, texture, shape, and structural details. Do not mention location, the table, or any relationship to another object. + Structural direction words are allowed when they describe the object itself: + "bottle with a black cap on top" is valid, while "bottle on the left of the + table" is not. Return JSON only: no Markdown, comments, or prose outside this exact schema: { @@ -494,17 +491,6 @@ def _parse_scene_object_fields( f"VLM JSON key {field_name}.category must be a lower-case snake_case " "class name." ) - # Check whether the name and description contain location or relationship words. - if _LOCATION_WORD_PATTERN.search(fields["name"]): - raise ValueError( - f"VLM JSON key {field_name}.name must not contain location or " - "relationship words." - ) - if _LOCATION_WORD_PATTERN.search(fields["description"]): - raise ValueError( - f"VLM JSON key {field_name}.description must not contain location or " - "relationship words." - ) return fields diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index d2b3898da..382e4933d 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -61,19 +61,23 @@ def test_image_object_analysis_parses_code_fence_and_assigns_stable_ids() -> Non assert [asset.id for asset in scene.assets] == ["cup_001"] -def test_image_object_analysis_rejects_location_words_in_object_names() -> None: - with pytest.raises(ValueError, match="must not contain location"): - scene_understanding._parse_image_object_analysis_response( - _response(asset_name="left cup") - ) +def test_image_object_analysis_accepts_name_with_spatial_words() -> None: + scene = scene_understanding._parse_image_object_analysis_response( + _response(asset_name="left cup") + ) + + assert scene.assets[0].name == "left cup" -def test_image_object_analysis_rejects_location_words_in_object_descriptions() -> None: +def test_image_object_analysis_accepts_description_with_structural_words() -> None: response = json.loads(_response()) - response["assets"][0]["description"] = "A small ceramic cup on the table." + response["assets"][0]["description"] = "A small ceramic cup with a lid on top." + + scene = scene_understanding._parse_image_object_analysis_response( + json.dumps(response) + ) - with pytest.raises(ValueError, match="description must not contain location"): - scene_understanding._parse_image_object_analysis_response(json.dumps(response)) + assert scene.assets[0].description == "A small ceramic cup with a lid on top." def test_image_object_analysis_retries_then_updates_scene(tmp_path: Path) -> None: From 9eb4ca103c0cdb727dfab34dad5c914ea5c77b4d Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:44:34 +0800 Subject: [PATCH 24/85] replace jiange's pca heuristic method with VLM operate rotation tools follow the assets' orientation_states --- .../scene_engine/core/scene_edit_plan.py | 20 ++- .../gen_sim/scene_engine/core/scene_graph.py | 4 + .../editing/scene_edit_asset_preparation.py | 10 +- .../editing/scene_edit_understanding.py | 72 +++++++-- .../pipeline/generation/scene_generation.py | 16 +- .../pipeline/utils/simready_processor.py | 140 +++++++----------- .../utils/simready_processor_utils.py | 103 +++++++++---- .../scene_engine/test_scene_edit_plan.py | 32 ++++ .../test_simready_processor_utils.py | 44 +++++- .../scene_engine/test_support_and_layout.py | 4 +- 10 files changed, 300 insertions(+), 145 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 83ad5bd79..0f78b9699 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -21,6 +21,7 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( + OrientationState, SceneConstraintType, SceneGraph, TableRegion, @@ -44,6 +45,7 @@ class SceneEditOperation: category: str | None = None name: str | None = None description: str | None = None + orientation_state: OrientationState | None = None def to_dict(self) -> dict[str, object]: """Serialize one normalized edit operation.""" @@ -56,6 +58,7 @@ def to_dict(self) -> dict[str, object]: "category": self.category, "name": self.name, "description": self.description, + "orientation_state": self.orientation_state, } @@ -84,6 +87,7 @@ def validate(self) -> None: # Edit-plan rules: # - move and delete identify one existing non-table object with object_id. # - add carries generated object_id plus non-empty category, name, and description. + # - add may preserve an explicit standing or lying user placement intent. # - move always supplies target_id and relation; add may omit both. # - table_region is only valid with target_id=table and relation=on. # - target_id and relation are otherwise supplied together or both absent. @@ -164,6 +168,7 @@ def _validate_operation( operation.category, operation.name, operation.description, + operation.orientation_state, ) ): raise ValueError("Delete operations may only specify object_id.") @@ -171,6 +176,13 @@ def _validate_operation( if operation.target_id is None or operation.relation is None: raise ValueError("Move operations must specify target_id and relation.") + existing_orientation_state = self.scene_graph.node_by_id()[ + operation.object_id + ].orientation_state + if operation.orientation_state not in {None, existing_orientation_state}: + raise ValueError( + "Move operations may only preserve the existing orientation_state." + ) self._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, @@ -178,7 +190,11 @@ def _validate_operation( ) if any( value is not None - for value in (operation.category, operation.name, operation.description) + for value in ( + operation.category, + operation.name, + operation.description, + ) ): raise ValueError("Move operations must not declare a new object.") @@ -203,6 +219,8 @@ def _validate_add_operation( for value in (operation.category, operation.name, operation.description) ): raise ValueError("Add operations require category, name, and description.") + if operation.orientation_state not in {None, "standing", "lying"}: + raise ValueError("Add operation orientation_state is invalid.") SceneEditPlan._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index e13f65aba..500980b53 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -193,6 +193,7 @@ def apply_updates( *, deleted_object_ids: set[str], added_object_ids: list[str], + added_orientation_states_by_id: dict[str, OrientationState | None], on_parent_updates: list[tuple[str, str, TableRegion | None]], planar_relation_updates: list[tuple[str, PlanarRelationType, str]], ) -> None: @@ -226,6 +227,8 @@ def apply_updates( raise ValueError( f"Duplicate scene graph nodes: {sorted(duplicate_object_ids)}" ) + if set(added_orientation_states_by_id) != set(added_object_ids): + raise ValueError("Added orientation states must match added node ids.") # New nodes default to the table; later updates replace that parent when needed. self.nodes.extend( @@ -233,6 +236,7 @@ def apply_updates( object_id=object_id, parent_id=TABLE_OBJECT_ID, parent_relation="on", + orientation_state=added_orientation_states_by_id[object_id], ) for object_id in added_object_ids ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index c6ec26fb3..220a91014 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -117,10 +117,18 @@ def prepare_scene_edit_assets( coarse_layout_by_id=_coarse_layouts_by_id(generated_asset_glbs), coarse_geometry_root=stage_output_root / "coarse_geometry", simready_geometry_root=stage_output_root / "simready_geometry", - # Scene editing will later provide the VLM-selected scale and rotation. + # Every added asset uses the VLM's pose and post-pose XY footprint scale. config=SimReadyProcessorConfig( use_vlm_scale=vlm_client is not None, use_vlm_rotation=vlm_client is not None, + # An explicit edit state overrides the default stable tabletop pose. + orientation_states_by_id={ + operation.object_id: operation.orientation_state + for operation in scene_edit_plan.operations + if operation.op == "add" + and operation.object_id is not None + and operation.orientation_state is not None + }, ), vlm_client=vlm_client, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index a648981e5..92720c7ec 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -25,6 +25,7 @@ ) from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( + OrientationState, PlanarRelationType, SceneGraph, SceneGraphNode, @@ -50,7 +51,9 @@ singular snake_case category, name, and description. Multiple add operations may have the same category and name; their final IDs are assigned by the program in operation order. target_id and relation are either both provided - or both null. + or both null. Set orientation_state to standing or lying only when the user + explicitly asks for that placement; otherwise set it to null so the object + uses its natural, physically stable tabletop pose. For every move and every positioned add, target_id must be an Existing object ID and relation must be one of on, left_of, right_of, in_front_of, or behind. @@ -77,11 +80,14 @@ class. name contains only color, material, texture, shape, and object details. description contains only visible category, material, color, texture, shape, and structural details. name and description must not mention position, the -table, or relations to any object. +table, relations to any object, or orientation. orientation_state must be null +unless the user explicitly requests standing/upright/vertical or lying/flat/ +horizontal placement. Follow that explicit user intent even if it is not the +object's natural stable pose. Return JSON only: no Markdown, comments, or prose. Every operation must contain exactly these fields: op, object_id, target_id, relation, table_region, category, -name, and description. Use null for every field that does not apply: +name, description, and orientation_state. Use null for every field that does not apply: { "operations": [ { @@ -92,7 +98,8 @@ "table_region": null, "category": null, "name": null, - "description": null + "description": null, + "orientation_state": null }, { "op": "delete", @@ -102,7 +109,8 @@ "table_region": null, "category": null, "name": null, - "description": null + "description": null, + "orientation_state": null }, { "op": "add", @@ -112,7 +120,8 @@ "table_region": "back_center", "category": "orange", "name": "small orange", - "description": "small round orange with a textured peel" + "description": "small round orange with a textured peel", + "orientation_state": null }, { "op": "add", @@ -122,7 +131,8 @@ "table_region": null, "category": "orange", "name": "small orange", - "description": "small round orange with a textured peel" + "description": "small round orange with a textured peel", + "orientation_state": null }, { "op": "add", @@ -130,14 +140,31 @@ "target_id": null, "relation": null, "table_region": null, - "category": "banana", - "name": "yellow banana", - "description": "curved yellow banana with a green stem" + "category": "bottle", + "name": "blue glass bottle", + "description": "tall transparent blue glass bottle with a narrow neck", + "orientation_state": "standing" + }, + { + "op": "add", + "object_id": null, + "target_id": null, + "relation": null, + "table_region": null, + "category": "fork", + "name": "silver metal fork", + "description": "four-tined silver stainless-steel fork with a plain handle", + "orientation_state": "lying" } ] } -The two orange additions intentionally share category and name. Do not add -fields beyond the required schema.""" +The two orange additions intentionally share category and name. The bottle +example represents an explicit user request to stand it upright, and the fork +example represents an explicit user request to lay it flat. Only add operations +may introduce a new non-null orientation_state. A move may use null or repeat +its existing orientation_state from the supplied scene metadata, but it must not +change that state. Delete operations must use null. Do not add fields beyond the +required schema.""" def understand_scene_edit( @@ -223,6 +250,7 @@ def _apply_scene_edit_plan_to_scene_graph( """Apply the target graph updates implied by add and move operations.""" deleted_object_ids: set[str] = set() added_object_ids: list[str] = [] + added_orientation_states_by_id: dict[str, OrientationState | None] = {} on_parent_updates: list[tuple[str, str, TableRegion | None]] = [] planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] for operation in scene_edit_plan.operations: @@ -234,6 +262,9 @@ def _apply_scene_edit_plan_to_scene_graph( raise ValueError("Add and move operations must have an object_id.") if operation.op == "add": added_object_ids.append(operation.object_id) + added_orientation_states_by_id[operation.object_id] = ( + operation.orientation_state + ) if operation.target_id is None or operation.relation is None: continue if operation.relation == "on": @@ -253,6 +284,7 @@ def _apply_scene_edit_plan_to_scene_graph( scene_graph.apply_updates( deleted_object_ids=deleted_object_ids, added_object_ids=added_object_ids, + added_orientation_states_by_id=added_orientation_states_by_id, on_parent_updates=on_parent_updates, planar_relation_updates=planar_relation_updates, ) @@ -267,6 +299,9 @@ def _simplify_scene_info( table_regions_by_id = { node.object_id: node.table_region for node in scene_graph.nodes } + orientation_states_by_id = { + node.object_id: node.orientation_state for node in scene_graph.nodes + } return { "existing_object_ids": [scene_object.id for scene_object in scene.objects], "objects": [ @@ -277,6 +312,7 @@ def _simplify_scene_info( "description": scene_object.description, "center_xy": scene_object.center_xy, "table_region": table_regions_by_id.get(scene_object.id), + "orientation_state": orientation_states_by_id.get(scene_object.id), } for scene_object in scene.objects ], @@ -351,6 +387,7 @@ def _parse_scene_edit_operations( "category", "name", "description", + "orientation_state", } # Get ids and counts of existing objects to assign new add IDs. assigned_object_ids = {scene_object.id for scene_object in scene.objects} @@ -371,6 +408,7 @@ def _parse_scene_edit_operations( raise ValueError("Scene edit operations must use the required schema.") object_id = _optional_string(value.get("object_id"), field_name="object_id") category = _optional_string(value.get("category"), field_name="category") + orientation_state = _optional_orientation_state(value.get("orientation_state")) if op == "add": if object_id is not None: raise ValueError("VLM add operations must set object_id to null.") @@ -397,6 +435,7 @@ def _parse_scene_edit_operations( description=_optional_string( value.get("description"), field_name="description" ), + orientation_state=orientation_state, ) ) return operations @@ -442,3 +481,12 @@ def _optional_table_region(value: object) -> TableRegion | None: if value not in TABLE_REGIONS: raise ValueError("Scene edit operation table_region is invalid.") return value + + +def _optional_orientation_state(value: object) -> OrientationState | None: + """Validate an optional explicit upright or lying edit intent.""" + if value is None: + return None + if value not in {"standing", "lying"}: + raise ValueError("Scene edit operation orientation_state is invalid.") + return value diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 7a0e19f35..0e1b17f95 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -109,21 +109,24 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } + # Coarse poses already preserve lying and unconstrained assets; only standing + # assets need a VLM semantic-axis correction before later z-up calibration. + standing_orientation_states_by_id = { + node.object_id: node.orientation_state + for node in scene_graph.nodes + if node.orientation_state == "standing" + } simready_processor = SimReadyProcessor( scene=scene, coarse_layout_by_id=coarse_layout_by_id, coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, debug_output_root=debug_output_root, - # Image-to-scene uses the geometry service's coarse scale directly. + # Keep the geometry-server scale and only correct unstable standing poses. config=SimReadyProcessorConfig( use_vlm_scale=False, use_vlm_rotation=False, - long_axis_object_ids=frozenset( - node.object_id - for node in scene_graph.nodes - if node.orientation_state is not None - ), + orientation_states_by_id=standing_orientation_states_by_id, ), vlm_client=vlm_client, ) @@ -498,6 +501,7 @@ def _scene_graph_based_calibration( node = nodes_by_id.get(asset_id) if node is None: raise ValueError(f"Scene graph does not contain asset {asset_id!r}.") + # Only correct the standing assets. if node.orientation_state != "standing": calibrated_assets_layout.append(asset_layout) continue diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 7bc041abf..864e39a59 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -17,15 +17,15 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path import numpy as np -import open3d as o3d from scipy.spatial.transform import Rotation import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import OrientationState from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, @@ -34,8 +34,11 @@ OpenAICompatibleVLM, ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( - query_vlm_object_rotation_and_target_size, + DEFAULT_NEEDED_LAYOUT, + LYING_NEEDED_LAYOUT, + STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, + query_vlm_object_rotation_and_target_size, render_object_front_top_views, rotate_glb_about_x_axis, ) @@ -68,8 +71,8 @@ class SimReadyProcessorConfig: use_vlm_scale: bool = False # Use the VLM-selected asset scale. use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. - - long_axis_object_ids: frozenset[str] = frozenset() + # Explicit graph orientation overrides the default stable tabletop pose. + orientation_states_by_id: dict[str, OrientationState] = field(default_factory=dict) class SimReadyProcessor: @@ -103,7 +106,9 @@ def __init__( self.config = config if config is not None else SimReadyProcessorConfig() self.vlm_client = vlm_client if ( - self.config.use_vlm_scale or self.config.use_vlm_rotation + self.config.use_vlm_scale + or self.config.use_vlm_rotation + or self.config.orientation_states_by_id ) and vlm_client is None: raise ValueError("vlm_client is required when VLM transforms are enabled.") @@ -200,25 +205,34 @@ def _prepare_vlm_rotated_glb( ) -> tuple[Path, list[float] | None]: """Render, query, and optionally bake the VLM-selected x-axis rotation.""" coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" - if not (self.config.use_vlm_scale or self.config.use_vlm_rotation): + orientation_state = self._orientation_state_for_object(scene_object.id) + orientation_pose_required = orientation_state is not None + if not ( + self.config.use_vlm_scale + or self.config.use_vlm_rotation + or orientation_pose_required + ): return coarse_path, None decision = self._vlm_transform_for_object( scene_object, - use_scale=self.config.use_vlm_scale, - use_rotation=self.config.use_vlm_rotation, + needed_layout=self._needed_layout_for_object(scene_object.id), ) rotate_about_x = bool(decision["rotate_about_x"]) - vlm_scale = compute_uniform_xy_scale_for_target( - glb_path=coarse_path, - target_xy_size_cm=decision["target_xy_size_cm"], - rotate_about_x=rotate_about_x, - ) + vlm_scale = None + if self.config.use_vlm_scale: + # The VLM target describes the final, post-rotation z-up XY footprint. + vlm_scale = compute_uniform_xy_scale_for_target( + glb_path=coarse_path, + target_xy_size_cm=decision["target_xy_size_cm"], + rotate_about_x=rotate_about_x, + ) rotated_path = rotate_glb_about_x_axis( input_path=coarse_path, output_path=self.simready_geometry_root / "vlm_rotated" / f"{scene_object.id}.glb", - rotate=rotate_about_x, + rotate=rotate_about_x + and (orientation_pose_required or self.config.use_vlm_rotation), ) # The scale flag controls whether this VLM-derived isotropic scale is used. # Apply the same factor on x, y, and z to preserve the asset's proportions. @@ -227,19 +241,34 @@ def _prepare_vlm_rotated_glb( [vlm_scale, vlm_scale, vlm_scale] if self.config.use_vlm_scale else None, ) + def _orientation_state_for_object(self, object_id: str) -> OrientationState | None: + """Return the explicit graph orientation requested for one object.""" + return self.config.orientation_states_by_id.get(object_id) + + def _needed_layout_for_object(self, object_id: str) -> str: + """Return the VLM layout instruction for one object's graph semantics.""" + return ( + STANDING_NEEDED_LAYOUT + if self._orientation_state_for_object(object_id) == "standing" + else ( + LYING_NEEDED_LAYOUT + if self._orientation_state_for_object(object_id) == "lying" + else DEFAULT_NEEDED_LAYOUT + ) + ) + def _vlm_transform_for_object( self, scene_object: SceneObject, *, - use_scale: bool, - use_rotation: bool, + needed_layout: str, ) -> dict[str, object]: """Render the object and return the validated VLM pose decision.""" - del use_scale, use_rotation assert self.vlm_client is not None coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" - needed_layout = "This asset needs to be place on the table that will not move a lot after simulation." - debug_root = self.simready_geometry_root.parent / "debug" + debug_root = ( + self.debug_output_root or self.simready_geometry_root.parent / "debug" + ) rendered_path = render_object_front_top_views( glb_path=coarse_path, output_path=debug_root / "vlm_views" / f"{scene_object.id}.png", @@ -316,16 +345,6 @@ def _canonicalize_object_mesh( y_up_to_z_up_transform[:3, :3] = y_up_to_z_up_matrix mesh.apply_transform(y_up_to_z_up_transform) - # Standardize graph-marked elongated assets before shared mesh processing. - # This makes local z their primary axis, so later scene-graph calibration - # can reliably recover the image-observed standing or lying orientation. - long_axis_alignment_matrix = np.eye(3) - if self._requires_long_axis_standardization(object_id): - long_axis_alignment_matrix = self._standardize_long_axis_z_up(mesh) - long_axis_alignment_transform = np.eye(4) - long_axis_alignment_transform[:3, :3] = long_axis_alignment_matrix - mesh.apply_transform(long_axis_alignment_transform) - # First make the object's AABB center at the origin. original_aabb_center = mesh.bounds.mean(axis=0) mesh.apply_translation(-original_aabb_center) @@ -333,13 +352,7 @@ def _canonicalize_object_mesh( # Scale the object with the value in the coarse layout. scale_transform = np.eye(4) scale_transform[:3, :3] = ( - # Actually there's no need to do so, for the scale factor is all equal - # in x, y, z axes. - long_axis_alignment_matrix - @ y_up_to_z_up_matrix - @ np.diag(coarse_scale) - @ y_up_to_z_up_matrix.T - @ long_axis_alignment_matrix.T + y_up_to_z_up_matrix @ np.diag(coarse_scale) @ y_up_to_z_up_matrix.T ) mesh.apply_transform(scale_transform) @@ -359,17 +372,7 @@ def _canonicalize_object_mesh( z_up_to_y_up_transform[:3, :3] = y_up_to_z_up_matrix.T mesh.apply_transform(z_up_to_y_up_transform) - # Compensate the local canonicalization so its coarse world pose does not - # change until layout refinement applies the image-observed correction. - local_long_axis_rotation = Rotation.from_matrix( - y_up_to_z_up_matrix.T @ long_axis_alignment_matrix @ y_up_to_z_up_matrix - ) - coarse_rotation_matrix = Rotation.from_euler( - "xyz", coarse_rot, degrees=True - ).as_matrix() - rotation = Rotation.from_matrix( - coarse_rotation_matrix @ local_long_axis_rotation.inv().as_matrix() - ) + rotation = Rotation.from_euler("xyz", coarse_rot, degrees=True) # Update the pos. position_offset = y_up_to_z_up_matrix.T @ ( scale_transform[:3, :3] @ original_aabb_center + scaled_aabb_bottom_center @@ -380,47 +383,6 @@ def _canonicalize_object_mesh( "scale": [1.0, 1.0, 1.0], } - def _requires_long_axis_standardization(self, object_id: str) -> bool: - """Return whether graph semantics identified one asset with a long axis.""" - return object_id in self.config.long_axis_object_ids - - @staticmethod - def _standardize_long_axis_z_up(mesh: trimesh.Trimesh) -> np.ndarray: - """Return a proper rotation that maps a mesh's primary axis to z-up. - Thanks to chanjian's idea. - """ - if len(mesh.vertices) < 4 or len(mesh.faces) < 4: - raise ValueError( - "Long-axis standardization requires a non-degenerate triangle mesh." - ) - open3d_mesh = o3d.geometry.TriangleMesh( - vertices=o3d.utility.Vector3dVector(mesh.vertices), - triangles=o3d.utility.Vector3iVector(mesh.faces), - ) - sampled_points = np.asarray( - open3d_mesh.sample_points_uniformly(number_of_points=10_000).points - ) # (10000, 3) x (x, y, z) - - # Check the number of the points again, and check whether have some - # non-finite values. - if sampled_points.shape[0] < 4 or not np.all(np.isfinite(sampled_points)): - raise ValueError( - "Long-axis standardization could not sample valid mesh points." - ) - - centered_points = sampled_points - sampled_points.mean(axis=0) - # SVD find the longest axis. - _, _, principal_axes = np.linalg.svd(centered_points, full_matrices=False) - if np.linalg.det(principal_axes) < 0: - principal_axes[2, :] *= -1 # in case the SVD returns a reflection. - - long_axis_rotation = Rotation.from_euler( - "y", 90.0, degrees=True - ).as_matrix() # 3x3 matrix - # The first PCA axis is the longest axis; rotate it onto the temporary z axis. - long_axis_rotation = long_axis_rotation @ principal_axes - return long_axis_rotation - @staticmethod def _three_floats(value: object, *, field_name: str) -> list[float]: """Validate and convert a three-value layout field to floats.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py index 539bace70..4523aa3e6 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -70,6 +70,27 @@ rotate_about_x=false and use target_xy_size_cm=[8.0, 8.0]. """ +DEFAULT_NEEDED_LAYOUT = ( + "Place this asset on the table in its natural, physically stable resting " + "orientation. For example, a fork should lie flat on the table rather " + "than stand on an edge." +) +STANDING_NEEDED_LAYOUT = ( + "The scene graph requires this asset to stand vertically on the table, " + "even when its natural stable pose would be lying down. For example, a " + "bottle should stand on its base and a fork should stand upright. If the " + "coarse GLB is lying flat, set rotate_about_x=true so its semantic vertical " + "axis aligns with the z-up world's z axis; if it is already upright, set " + "it to false." +) +LYING_NEEDED_LAYOUT = ( + "The scene graph requires this asset to lie flat on the table, even when " + "its natural stable pose would be standing. For example, a bottle should " + "lie on its side and a fork should lie flat. Choose rotate_about_x so the " + "asset's semantic long axis remains in the tabletop x-y plane rather than " + "along the z-up world's z axis." +) + def render_object_front_top_views( *, @@ -267,17 +288,60 @@ def query_vlm_object_rotation_and_target_size( rendered_views_path: str | Path, vlm_client: OpenAICompatibleVLM, debug_output_path: str | Path | None = None, + json_max_attempts: int = 3, ) -> dict[str, object]: - """Ask the VLM for rotation and post-rotation tabletop footprint.""" - response_text = vlm_client.complete( - system_prompt=_VLM_SYSTEM_PROMPT, - user_prompt=( - f"Object description:\n{scene_object_description}\n\n" - f"Needed layout:\n{needed_layout}\n\n" - "The image contains front view on the left and top view on the right." - ), - image_path=rendered_views_path, - ) + """Ask the VLM for a valid rotation and post-rotation tabletop footprint.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_VLM_SYSTEM_PROMPT, + user_prompt=( + f"Object description:\n{scene_object_description}\n\n" + f"Needed layout:\n{needed_layout}\n\n" + "The image contains front view on the left and top view on the right." + ), + image_path=rendered_views_path, + ) + try: + value = _parse_vlm_rotation_and_target_size_response(response_text) + break + except ValueError as exc: + last_validation_error = exc + else: + assert last_validation_error is not None + raise ValueError( + "VLM transform response is invalid after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + if debug_output_path is not None: + output_path = Path(debug_output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + { + "description": scene_object_description, + "needed_layout": needed_layout, + "rendered_views_path": str( + Path(rendered_views_path).expanduser().resolve() + ), + "vlm_output": value, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return value + + +def _parse_vlm_rotation_and_target_size_response( + response_text: str, +) -> dict[str, object]: + """Validate one VLM rotation-and-scale JSON response.""" try: value = json.loads(_strip_json_code_fence(response_text)) except json.JSONDecodeError as exc: @@ -300,25 +364,6 @@ def query_vlm_object_rotation_and_target_size( or not all(np.isfinite(item) and item > 0 for item in target_size) ): raise ValueError("VLM target_xy_size_cm must contain two positive numbers.") - if debug_output_path is not None: - output_path = Path(debug_output_path).expanduser().resolve() - output_path.parent.mkdir(parents=True, exist_ok=True) - output_path.write_text( - json.dumps( - { - "description": scene_object_description, - "needed_layout": needed_layout, - "rendered_views_path": str( - Path(rendered_views_path).expanduser().resolve() - ), - "vlm_output": value, - }, - indent=2, - ensure_ascii=False, - ) - + "\n", - encoding="utf-8", - ) return value diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index 5579e8d55..a99e04384 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -115,6 +115,7 @@ def test_scene_edit_plan_accepts_add_without_a_position() -> None: "category": "cup", "name": "green cup", "description": "A small green ceramic cup.", + "orientation_state": None, } ] @@ -159,6 +160,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", + "orientation_state": None, }, { "op": "add", @@ -169,6 +171,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", + "orientation_state": "lying", }, ] } @@ -181,6 +184,30 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "orange_002", "orange_003", ] + assert [operation.orientation_state for operation in operations] == [ + None, + "lying", + ] + + +def test_scene_edit_plan_rejects_a_changed_move_orientation_state() -> None: + scene, scene_graph = _scene_and_graph() + scene_graph.node_by_id()["book_001"].orientation_state = "lying" + + with pytest.raises(ValueError, match="may only preserve"): + SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="book_001", + target_id="table", + relation="on", + orientation_state="standing", + ) + ], + ) def test_scene_edit_plan_rejects_targets_outside_the_input_scene() -> None: @@ -426,6 +453,7 @@ def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> No category="cup", name="green cup", description="A small green ceramic cup.", + orientation_state="standing", ) ], ) @@ -438,10 +466,12 @@ def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> No added_node = updated_scene_graph.node_by_id()["cup_001"] assert added_node.parent_id == "table" assert added_node.parent_relation == "on" + assert added_node.orientation_state == "standing" def test_scene_edit_graph_builder_updates_move_on_parent() -> None: scene, scene_graph = _scene_and_graph() + scene_graph.node_by_id()["orange_001"].orientation_state = "lying" plan = SceneEditPlan( scene=scene, scene_graph=scene_graph, @@ -451,6 +481,7 @@ def test_scene_edit_graph_builder_updates_move_on_parent() -> None: object_id="orange_001", target_id="table", relation="on", + orientation_state="lying", ) ], ) @@ -461,6 +492,7 @@ def test_scene_edit_graph_builder_updates_move_on_parent() -> None: ) assert updated_scene_graph.node_by_id()["orange_001"].parent_id == "table" + assert updated_scene_graph.node_by_id()["orange_001"].orientation_state == "lying" def test_scene_edit_graph_builder_adds_planar_relation_with_target_parent() -> None: diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py index 13cf73e10..dd872f913 100644 --- a/tests/gen_sim/scene_engine/test_simready_processor_utils.py +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -27,25 +27,61 @@ SimReadyProcessorConfig, ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + DEFAULT_NEEDED_LAYOUT, + LYING_NEEDED_LAYOUT, + STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, + query_vlm_object_rotation_and_target_size, ) -def test_simready_long_axis_standardization_uses_graph_selected_ids( +def test_simready_pose_layout_uses_graph_orientation_states( tmp_path: Path, ) -> None: + class VLM: + def complete(self, **_: object) -> str: + raise AssertionError("This selection test must not call the VLM.") + processor = SimReadyProcessor( scene=Scene(), coarse_layout_by_id={}, coarse_geometry_root=tmp_path / "coarse", simready_geometry_root=tmp_path / "simready", config=SimReadyProcessorConfig( - long_axis_object_ids=frozenset({"rolling_pin_001"}), + orientation_states_by_id={"bottle_001": "standing", "fork_001": "lying"}, ), + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert processor._orientation_state_for_object("bottle_001") == "standing" + assert processor._orientation_state_for_object("fork_001") == "lying" + assert processor._orientation_state_for_object("knife_001") is None + assert processor._needed_layout_for_object("bottle_001") == STANDING_NEEDED_LAYOUT + assert processor._needed_layout_for_object("fork_001") == LYING_NEEDED_LAYOUT + assert processor._needed_layout_for_object("knife_001") == DEFAULT_NEEDED_LAYOUT + + +def test_vlm_transform_query_retries_an_empty_response(tmp_path: Path) -> None: + class VLM: + def __init__(self) -> None: + self.responses = [ + "", + '{"rotate_about_x": false, "target_xy_size_cm": [8.0, 8.0]}', + ] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + vlm_client = VLM() + decision = query_vlm_object_rotation_and_target_size( + scene_object_description="small blue bottle", + needed_layout=STANDING_NEEDED_LAYOUT, + rendered_views_path=tmp_path / "views.png", + vlm_client=vlm_client, # type: ignore[arg-type] ) - assert processor._requires_long_axis_standardization("rolling_pin_001") - assert not processor._requires_long_axis_standardization("bottle_001") + assert decision == {"rotate_about_x": False, "target_xy_size_cm": [8.0, 8.0]} + assert vlm_client.responses == [] def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: diff --git a/tests/gen_sim/scene_engine/test_support_and_layout.py b/tests/gen_sim/scene_engine/test_support_and_layout.py index fda7f9e08..b6c9df3ff 100644 --- a/tests/gen_sim/scene_engine/test_support_and_layout.py +++ b/tests/gen_sim/scene_engine/test_support_and_layout.py @@ -191,9 +191,7 @@ def test_layout_optimizer_projects_out_of_bounds_aabb_into_rectangle() -> None: def test_layout_optimizer_rejects_aabb_larger_than_rectangle() -> None: optimizer = AssetsSupportLayoutOptimizer( support_region=Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]), - assets_aabb_2d_z_up_world_corners_by_id={ - "large": _aabb(-0.5, 0.5, 2.5, 1.5) - }, + assets_aabb_2d_z_up_world_corners_by_id={"large": _aabb(-0.5, 0.5, 2.5, 1.5)}, assets_layout=[_layout("large", 1.0, 1.0)], ) From ccb4a94b9b33f7bdb69a83ae7837f3b6045e66ce Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:08:42 +0800 Subject: [PATCH 25/85] ignore the fixed 2d aabb overlap when doing the on-table optimization --- .../pipeline/utils/scene_layout_optimizer.py | 16 +++- .../test_scene_layout_optimizer.py | 86 +++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py index ea6387cd2..ed877c734 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py @@ -143,6 +143,7 @@ def optimize_parent_child_xy( root_half_extents_xy=child_half_extents_xy, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, + fixed_root_xy_by_id=fixed_child_xy_by_id, solved_root_xy_by_id=solved_child_xy_by_id, config=self.config, ) @@ -225,6 +226,7 @@ def _optimize_table_root_xy( root_half_extents_xy=root_half_extents_xy, inequality_constraints=inequality_constraints, equality_constraints=equality_constraints, + fixed_root_xy_by_id=fixed_root_xy_by_id, solved_root_xy_by_id=solved_root_xy_by_id, config=config, ) @@ -443,18 +445,26 @@ def _refine_root_collisions( root_half_extents_xy: dict[str, np.ndarray], inequality_constraints: list[tuple[np.ndarray, float]], equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], solved_root_xy_by_id: dict[str, list[float]], config: SceneLayoutOptimizerConfig, ) -> dict[str, list[float]]: - """Add AABB separation constraints until the table roots no longer overlap.""" + """Separate only sibling pairs that include a layout-variable object.""" seen_pairs: set[tuple[str, str]] = set() current_xy_by_id = solved_root_xy_by_id for _ in range(config.max_collision_rounds): - overlaps = _root_aabb_overlaps( + all_overlaps = _root_aabb_overlaps( root_ids=root_ids, root_half_extents_xy=root_half_extents_xy, xy_by_id=current_xy_by_id, ) + # Fixed/fixed siblings are outside this edit and therefore cannot be solved here. + overlaps = [ + overlap + for overlap in all_overlaps + if fixed_root_xy_by_id[overlap[1]] is None + or fixed_root_xy_by_id[overlap[2]] is None + ] if not overlaps: return current_xy_by_id added_constraint_count = 0 @@ -494,6 +504,8 @@ def _refine_root_collisions( root_half_extents_xy=root_half_extents_xy, xy_by_id=current_xy_by_id, ) + if fixed_root_xy_by_id[first_id] is None + or fixed_root_xy_by_id[second_id] is None ] raise ValueError( "Table-root AABB collisions remain after layout refinement: " diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 36b4bfa79..bca20da8e 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -31,9 +31,19 @@ SceneLayoutConstructor, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( + SceneLayoutOptimizer, _table_region_bounds, ) +_TABLE_BOUNDS = [ + [-2.0, -2.0], + [2.0, -2.0], + [2.0, 2.0], + [-2.0, 2.0], +] +_OVERLAPPING_CENTER_XY = [0.0, 0.0] +_ASSET_SIDE_LENGTH_M = 0.2 + def _asset( *, @@ -136,3 +146,79 @@ def test_layout_constructor_places_new_child_on_parent_top( assert placed_cup.center_xy == [0.0, 0.0] # book top is z=0.62 m; cup half-height is 0.1 m and clearance is 0.02 m. assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) + + +def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + first_id, second_id = "first_001", "second_001" + optimizer = SceneLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize_table_root_xy( + assets_by_id={ + first_id: _asset( + object_id=first_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + second_id: _asset( + object_id=second_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + }, + root_ids=[first_id, second_id], + root_seed_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={first_id, second_id}, + fixed_root_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + root_table_regions_by_id={first_id: None, second_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + + assert solved_xy_by_id == { + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + } + + +def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( + tmp_path: Path, +) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + fixed_id, variable_id = "fixed_001", "variable_001" + optimizer = SceneLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize_table_root_xy( + assets_by_id={ + fixed_id: _asset( + object_id=fixed_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + variable_id: _asset( + object_id=variable_id, + glb_path=asset_glb, + ), + }, + root_ids=[fixed_id, variable_id], + root_seed_xy_by_id={ + fixed_id: _OVERLAPPING_CENTER_XY, + variable_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={fixed_id}, + fixed_root_xy_by_id={fixed_id: _OVERLAPPING_CENTER_XY, variable_id: None}, + root_table_regions_by_id={fixed_id: None, variable_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + + assert solved_xy_by_id[fixed_id] == _OVERLAPPING_CENTER_XY + assert np.max(np.abs(solved_xy_by_id[variable_id])) >= _ASSET_SIDE_LENGTH_M - 1e-6 From 031b5f81611ca2f17401081e8b5729629a92c380 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:06:04 +0800 Subject: [PATCH 26/85] finished the table as root layout optimization (but still needs to be tested) --- .../utils/scene_layout_constructor.py | 165 ++-- .../pipeline/utils/scene_layout_optimizer.py | 797 ------------------ .../pipeline/utils/scene_layout_utils.py | 155 ++++ .../utils/table_surface_layout_optimizer.py | 542 ++++++++++++ .../test_scene_layout_optimizer.py | 113 +-- 5 files changed, 815 insertions(+), 957 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py index 38706ac47..9279bd1ae 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -26,9 +26,19 @@ SceneGraph, ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( - SceneLayoutOptimizerConfig, - SceneLayoutOptimizer, +from embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer import ( + ParentSurfaceLayoutOptimizer, + ParentSurfaceLayoutOptimizerConfig, + ParentSurfaceLayoutProblem, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + translate_scene_object_y_up_by_z_up_delta, + update_scene_object_y_up_pose_from_z_up_support, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + TableSurfaceLayoutOptimizer, + TableSurfaceLayoutOptimizerConfig, + TableSurfaceLayoutProblem, ) @@ -66,33 +76,45 @@ def __init__( layout_variable_ids: set[str], generated_scene_objects: list[SceneObject], output_root: str | Path, - config: SceneLayoutOptimizerConfig | None = None, + table_surface_config: TableSurfaceLayoutOptimizerConfig | None = None, + parent_surface_config: ParentSurfaceLayoutOptimizerConfig | None = None, ) -> None: self.formal_scene = formal_scene self.goal_scene_graph = goal_scene_graph self.layout_variable_ids = layout_variable_ids self.generated_scene_objects = generated_scene_objects self.output_root = Path(output_root).expanduser().resolve() - self.layout_optimizer = SceneLayoutOptimizer(config=config) + # Table surface optimizer. + self.table_surface_layout_optimizer = TableSurfaceLayoutOptimizer( + config=table_surface_config + ) + # Parent surface (on) optimizer. + self.parent_surface_layout_optimizer = ParentSurfaceLayoutOptimizer( + config=parent_surface_config + ) self._current_xy_by_id: dict[str, list[float] | None] = {} self._solved_delta_xy_by_id: dict[str, list[float]] = {} self._updated_object_ids: set[str] = set() def construct(self) -> Scene: """Construct table-root layouts before later stacked-group refinement.""" + # Build layout problem. layout_problem = self._build_problem() + # Get current XY centers. self._current_xy_by_id = { object_id: list(initial_xy) if initial_xy is not None else None for object_id, initial_xy in layout_problem.initial_xy_by_id.items() } self._solved_delta_xy_by_id = {} self._updated_object_ids = set() + # Check the group. if ( layout_problem.groups and layout_problem.groups[0].parent_id != TABLE_OBJECT_ID ): raise ValueError("The first layout group must be rooted at the table.") + # Optimize each group in BFS order, propagating solved deltas to descendants. for group in layout_problem.groups: if group.parent_id == TABLE_OBJECT_ID: self._optimize_table_group( @@ -114,56 +136,18 @@ def _optimize_table_group( group: SceneLayoutGroup, ) -> None: """Optimize all direct on-table children before any stacked child groups.""" + table_surface_problem = TableSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, + ) + solved_root_xy_by_id = self.table_surface_layout_optimizer.optimize( + table_surface_problem + ) table = layout_problem.post_edit_scene.table if table is None: raise ValueError("Table group optimization requires a table.") - if table.support_optimization_rect_xy is None: - raise ValueError( - "Table group optimization requires a table support optimization rectangle." - ) - - root_ids = set(group.child_ids) - root_relations = [ - relation - for relation in layout_problem.goal_scene_graph.relations - if relation.source_id in root_ids and relation.target_id in root_ids - ] - root_seed_xy_by_id: dict[str, list[float]] = {} - for root_id in group.child_ids: - inherited_xy = self._current_xy_by_id[root_id] - # New roots start from the table-local origin; imported roots keep their pose. - root_seed_xy_by_id[root_id] = ( - [0.0, 0.0] if inherited_xy is None else list(inherited_xy) - ) - self._current_xy_by_id[root_id] = root_seed_xy_by_id[root_id] - - nodes_by_id = layout_problem.goal_scene_graph.node_by_id() - solved_root_xy_by_id = self.layout_optimizer.optimize_table_root_xy( - assets_by_id={ - asset.id: asset for asset in layout_problem.post_edit_scene.assets - }, - root_ids=group.child_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids={ - root_id - for root_id in group.child_ids - if layout_problem.initial_xy_by_id[root_id] is not None - }, - fixed_root_xy_by_id={ - root_id: ( - None - if root_id in layout_problem.layout_variable_ids - else self._current_xy_by_id[root_id] - ) - for root_id in group.child_ids - }, - root_table_regions_by_id={ - root_id: nodes_by_id[root_id].table_region - for root_id in group.child_ids - }, - table_optimization_rect_xy=table.support_optimization_rect_xy, - root_relations=root_relations, - ) + # Check the table's z. if table.support_surface_z is None and any( root_id in layout_problem.layout_variable_ids for root_id in group.child_ids ): @@ -172,7 +156,7 @@ def _optimize_table_group( asset.id: asset for asset in layout_problem.post_edit_scene.assets } for root_id, solved_xy in solved_root_xy_by_id.items(): - seed_xy = root_seed_xy_by_id[root_id] + seed_xy = table_surface_problem.root_seed_xy_by_id[root_id] delta_xy = [ solved_xy[0] - seed_xy[0], solved_xy[1] - seed_xy[1], @@ -182,10 +166,11 @@ def _optimize_table_group( if root_id in layout_problem.layout_variable_ids: # Direct add/move roots receive a new pose on the table support. assert table.support_surface_z is not None - self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( + update_scene_object_y_up_pose_from_z_up_support( scene_object=assets_by_id[root_id], support_region_z=table.support_surface_z, center_xy=solved_xy, + clearance_m=0.00, # Directly place on the support surface. ) self._updated_object_ids.add(root_id) self._propagate_descendant_delta( @@ -202,6 +187,7 @@ def _propagate_descendant_delta( delta_xy: list[float], ) -> None: """Move every positioned descendant by one solved ancestor XY delta.""" + # A zero root delta cannot change any descendant pose, so skip the subtree walk. if delta_xy == [0.0, 0.0]: return assets_by_id = {asset.id: asset for asset in scene.assets} @@ -219,7 +205,7 @@ def _propagate_descendant_delta( descendant_xy[0] + delta_xy[0], descendant_xy[1] + delta_xy[1], ] - self.layout_optimizer.translate_scene_object_y_up_by_z_up_delta( + translate_scene_object_y_up_by_z_up_delta( scene_object=assets_by_id[descendant_id], delta_xy=delta_xy, ) @@ -233,54 +219,17 @@ def _optimize_parent_group( group: SceneLayoutGroup, ) -> None: """Optimize one settled parent's direct on-children in local XY coordinates.""" - assets_by_id = { - asset.id: asset for asset in layout_problem.post_edit_scene.assets - } - parent = assets_by_id.get(group.parent_id) - if parent is None: - raise ValueError(f"Parent {group.parent_id!r} is not an asset.") - parent_aabb = self.layout_optimizer.scene_object_z_up_world_aabb( - scene_object=parent + parent_surface_problem = ParentSurfaceLayoutProblem.from_layout_problem( + layout_problem=layout_problem, + group=group, + current_xy_by_id=self._current_xy_by_id, ) - parent_aabb_xy = [ - [parent_aabb[0][0], parent_aabb[0][1]], - [parent_aabb[1][0], parent_aabb[1][1]], - ] - parent_center_xy = [ - (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, - (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, - ] - child_seed_xy_by_id: dict[str, list[float]] = {} - for child_id in group.child_ids: - inherited_xy = self._current_xy_by_id[child_id] - # New children start at their parent's current AABB center. - child_seed_xy_by_id[child_id] = ( - parent_center_xy if inherited_xy is None else list(inherited_xy) - ) - self._current_xy_by_id[child_id] = child_seed_xy_by_id[child_id] - - solved_child_xy_by_id = self.layout_optimizer.optimize_parent_child_xy( - assets_by_id=assets_by_id, - child_ids=group.child_ids, - child_seed_xy_by_id=child_seed_xy_by_id, - imported_child_ids={ - child_id - for child_id in group.child_ids - if layout_problem.initial_xy_by_id[child_id] is not None - }, - fixed_child_xy_by_id={ - child_id: ( - None - if child_id in layout_problem.layout_variable_ids - else self._current_xy_by_id[child_id] - ) - for child_id in group.child_ids - }, - parent_aabb_xy=parent_aabb_xy, + # Get results. + solved_child_xy_by_id = self.parent_surface_layout_optimizer.optimize( + parent_surface_problem ) - parent_top_z = parent_aabb[1][2] for child_id, solved_xy in solved_child_xy_by_id.items(): - seed_xy = child_seed_xy_by_id[child_id] + seed_xy = parent_surface_problem.child_seed_xy_by_id[child_id] delta_xy = [ solved_xy[0] - seed_xy[0], solved_xy[1] - seed_xy[1], @@ -289,9 +238,9 @@ def _optimize_parent_group( self._solved_delta_xy_by_id[child_id] = delta_xy if child_id in layout_problem.layout_variable_ids: # Variable children are placed directly above the parent's current top. - self.layout_optimizer.update_scene_object_y_up_pose_from_z_up_support( - scene_object=assets_by_id[child_id], - support_region_z=parent_top_z, + update_scene_object_y_up_pose_from_z_up_support( + scene_object=parent_surface_problem.assets_by_id[child_id], + support_region_z=parent_surface_problem.parent_top_z, center_xy=solved_xy, ) self._updated_object_ids.add(child_id) @@ -303,6 +252,7 @@ def _optimize_parent_group( def _build_problem(self) -> SceneLayoutProblem: """Build post-edit objects and preserve formal-scene centers as seeds.""" + # Validate the graph first. self.goal_scene_graph.validate() graph_object_ids = set(self.goal_scene_graph.node_by_id()) generated_objects_by_id = self._generated_scene_objects_by_id() @@ -326,13 +276,14 @@ def _build_problem(self) -> SceneLayoutProblem: } if post_edit_object_ids != graph_object_ids: raise ValueError("Goal scene graph and post-edit scene have different ids.") + # Get the movable asset ids. if not self.layout_variable_ids.issubset(post_edit_object_ids - {"table"}): raise ValueError( "Only post-edit assets may participate in layout optimization." ) - + # Get initial XY centers. initial_xy_by_id = { - asset.id: self._initial_xy( + asset.id: self._initial_xy( # The assets' center XY should always be updated whenever changes are made. asset, is_generated=asset.id in generated_objects_by_id, ) @@ -343,13 +294,15 @@ def _build_problem(self) -> SceneLayoutProblem: raise ValueError( f"New asset {object_id!r} must participate in layout optimization." ) + # Build the table-rooted BFS groups. + groups = self._build_groups() return SceneLayoutProblem( post_edit_scene=post_edit_scene, goal_scene_graph=self.goal_scene_graph, layout_variable_ids=set(self.layout_variable_ids), initial_xy_by_id=initial_xy_by_id, - groups=self._build_groups(), + groups=groups, ) def _build_groups(self) -> list[SceneLayoutGroup]: diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py deleted file mode 100644 index ed877c734..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_optimizer.py +++ /dev/null @@ -1,797 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - - -from __future__ import annotations - -from dataclasses import dataclass - -import numpy as np -from scipy.optimize import minimize - -from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation -from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( - layout_object_to_transform_matrix, - load_glb_mesh, - transform_matrix_to_layout_object, -) - - -@dataclass(frozen=True) -class SceneLayoutOptimizerConfig: - """Numerical controls shared by each graph-layout solve.""" - - relation_clearance_m: float = 0.03 - collision_margin_m: float = 0.02 - max_slsqp_iterations: int = 500 - slsqp_ftol: float = 1e-6 - max_collision_rounds: int = 8 - max_added_collision_pairs: int = 64 - imported_seed_weight: float = 5.0 - min_center_distance_m: float = 0.01 - min_center_distance_weight: float = 0.05 - - def __post_init__(self) -> None: - """Reject invalid numerical controls before assembling a layout problem.""" - if self.relation_clearance_m < 0.0: - raise ValueError("relation_clearance_m must be non-negative.") - if self.collision_margin_m < 0.0: - raise ValueError("collision_margin_m must be non-negative.") - if self.max_slsqp_iterations <= 0: - raise ValueError("max_slsqp_iterations must be positive.") - if self.slsqp_ftol <= 0.0: - raise ValueError("slsqp_ftol must be positive.") - if self.max_collision_rounds <= 0: - raise ValueError("max_collision_rounds must be positive.") - if self.max_added_collision_pairs <= 0: - raise ValueError("max_added_collision_pairs must be positive.") - - -class SceneLayoutOptimizer: - """Solve graph-constrained XY layouts and apply resulting poses.""" - - def __init__(self, *, config: SceneLayoutOptimizerConfig | None = None) -> None: - self.config = config if config is not None else SceneLayoutOptimizerConfig() - - def optimize_table_root_xy( - self, - *, - assets_by_id: dict[str, SceneObject], - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - fixed_root_xy_by_id: dict[str, list[float] | None], - root_table_regions_by_id: dict[str, str | None], - table_optimization_rect_xy: list[list[float]], - root_relations: list[SceneGraphRelation], - ) -> dict[str, list[float]]: - """Solve direct table-child centers with graph and AABB constraints.""" - return _optimize_table_root_xy( - assets_by_id=assets_by_id, - root_ids=root_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids=imported_root_ids, - fixed_root_xy_by_id=fixed_root_xy_by_id, - root_table_regions_by_id=root_table_regions_by_id, - table_optimization_rect_xy=table_optimization_rect_xy, - root_relations=root_relations, - config=self.config, - ) - - def optimize_parent_child_xy( - self, - *, - assets_by_id: dict[str, SceneObject], - child_ids: list[str], - child_seed_xy_by_id: dict[str, list[float]], - imported_child_ids: set[str], - fixed_child_xy_by_id: dict[str, list[float] | None], - parent_aabb_xy: list[list[float]], - ) -> dict[str, list[float]]: - """Solve direct on-children inside one parent's current XY AABB.""" - child_half_extents_xy = _asset_half_extents_xy( - assets_by_id=assets_by_id, - object_ids=child_ids, - ) - inequality_constraints: list[tuple[np.ndarray, float]] = [] - equality_constraints: list[tuple[np.ndarray, float]] = [] - child_index = {child_id: index for index, child_id in enumerate(child_ids)} - parent_bounds = _bounds_from_points(parent_aabb_xy) - for child_id in child_ids: - _append_aabb_center_bounds( - constraints=inequality_constraints, - root_index=child_index, - root_id=child_id, - bounds=parent_bounds, - half_extents_xy=child_half_extents_xy[child_id], - ) - fixed_xy = fixed_child_xy_by_id[child_id] - if fixed_xy is not None: - _append_fixed_root_constraints( - constraints=equality_constraints, - root_index=child_index, - root_id=child_id, - fixed_xy=fixed_xy, - ) - - solved_child_xy_by_id = _solve_root_xy( - root_ids=child_ids, - root_seed_xy_by_id=child_seed_xy_by_id, - imported_root_ids=imported_child_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=self.config, - ) - return _refine_root_collisions( - root_ids=child_ids, - root_seed_xy_by_id=child_seed_xy_by_id, - imported_root_ids=imported_child_ids, - root_half_extents_xy=child_half_extents_xy, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - fixed_root_xy_by_id=fixed_child_xy_by_id, - solved_root_xy_by_id=solved_child_xy_by_id, - config=self.config, - ) - - @staticmethod - def scene_object_z_up_world_aabb( - *, - scene_object: SceneObject, - ) -> list[list[float]]: - """Return one object's current z-up world AABB as [min, max].""" - return _scene_object_z_up_world_aabb(scene_object=scene_object) - - @staticmethod - def update_scene_object_y_up_pose_from_z_up_support( - *, - scene_object: SceneObject, - support_region_z: float, - center_xy: list[float], - clearance_m: float = 0.02, - ) -> None: - """Place one SimReady asset on a horizontal z-up support region.""" - _update_scene_object_y_up_pose_from_z_up_support( - scene_object=scene_object, - support_region_z=support_region_z, - center_xy=center_xy, - clearance_m=clearance_m, - ) - - @staticmethod - def translate_scene_object_y_up_by_z_up_delta( - *, - scene_object: SceneObject, - delta_xy: list[float], - ) -> None: - """Translate one existing y-up pose by a solved z-up XY delta.""" - _translate_scene_object_y_up_by_z_up_delta( - scene_object=scene_object, - delta_xy=delta_xy, - ) - - -def _optimize_table_root_xy( - *, - assets_by_id: dict[str, SceneObject], - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - fixed_root_xy_by_id: dict[str, list[float] | None], - root_table_regions_by_id: dict[str, str | None], - table_optimization_rect_xy: list[list[float]], - root_relations: list[SceneGraphRelation], - config: SceneLayoutOptimizerConfig, -) -> dict[str, list[float]]: - """Solve direct table-child centers with graph and AABB constraints.""" - root_half_extents_xy = _asset_half_extents_xy( - assets_by_id=assets_by_id, - object_ids=root_ids, - ) - inequality_constraints, equality_constraints = _build_table_root_constraints( - root_ids=root_ids, - root_half_extents_xy=root_half_extents_xy, - root_relations=root_relations, - root_table_regions_by_id=root_table_regions_by_id, - table_optimization_rect_xy=table_optimization_rect_xy, - fixed_root_xy_by_id=fixed_root_xy_by_id, - config=config, - ) - solved_root_xy_by_id = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) - return _refine_root_collisions( - root_ids=root_ids, - root_seed_xy_by_id=root_seed_xy_by_id, - imported_root_ids=imported_root_ids, - root_half_extents_xy=root_half_extents_xy, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - fixed_root_xy_by_id=fixed_root_xy_by_id, - solved_root_xy_by_id=solved_root_xy_by_id, - config=config, - ) - - -def _update_scene_object_y_up_pose_from_z_up_support( - *, - scene_object: SceneObject, - support_region_z: float, - center_xy: list[float], - clearance_m: float = 0.02, -) -> None: - """Place one SimReady asset on a horizontal z-up support region. - - ``SceneObject`` stores poses in y-up before export. The target center and - support height are z-up values because layout optimization uses that frame. - """ - if not np.isfinite(support_region_z): - raise ValueError("support_region_z must be finite.") - if clearance_m < 0.0 or not np.isfinite(clearance_m): - raise ValueError("clearance_m must be finite and non-negative.") - target_xy = _two_floats(center_xy, field_name="center_xy") - rotation_y_up = _three_floats_or_default( - scene_object.rot, - field_name="rot", - default=[0.0, 0.0, 0.0], - ) - mesh = _asset_z_up_mesh_at_zero_translation( - scene_object=scene_object, - rotation_y_up=rotation_y_up, - ) - target_position_z_up = np.array( - [ - target_xy[0] - float(mesh.bounds[:, 0].mean()), - target_xy[1] - float(mesh.bounds[:, 1].mean()), - float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), - ] - ) - z_up_to_y_up = np.linalg.inv(_y_up_to_z_up_matrix()) - # Persist the y-up pose that SceneExporter later converts back to z-up. - scene_object.pos = (z_up_to_y_up[:3, :3] @ target_position_z_up).tolist() - scene_object.rot = rotation_y_up - scene_object.center_xy = target_xy - - -def _translate_scene_object_y_up_by_z_up_delta( - *, - scene_object: SceneObject, - delta_xy: list[float], -) -> None: - """Translate one existing y-up pose by a solved z-up XY delta.""" - dx, dy = _two_floats(delta_xy, field_name="delta_xy") - current_pos = _three_floats_or_default( - scene_object.pos, - field_name="pos", - default=None, - ) - # z-up x maps to y-up x, while z-up y maps to negative y-up z. - scene_object.pos = [ - current_pos[0] + dx, - current_pos[1], - current_pos[2] - dy, - ] - if scene_object.center_xy is not None: - scene_object.center_xy = [ - scene_object.center_xy[0] + dx, - scene_object.center_xy[1] + dy, - ] - - -def _scene_object_z_up_world_aabb( - *, - scene_object: SceneObject, -) -> list[list[float]]: - """Measure one current SceneObject pose in z-up world coordinates.""" - position_y_up = _three_floats_or_default( - scene_object.pos, - field_name="pos", - default=None, - ) - mesh = _asset_z_up_mesh_at_zero_translation(scene_object=scene_object) - position_z_up = _y_up_to_z_up_matrix()[:3, :3] @ np.asarray( - position_y_up, - dtype=float, - ) - mesh.apply_translation(position_z_up) - return mesh.bounds.tolist() - - -def _build_table_root_constraints( - *, - root_ids: list[str], - root_half_extents_xy: dict[str, np.ndarray], - root_relations: list[SceneGraphRelation], - root_table_regions_by_id: dict[str, str | None], - table_optimization_rect_xy: list[list[float]], - fixed_root_xy_by_id: dict[str, list[float] | None], - config: SceneLayoutOptimizerConfig, -) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: - """Build hard table, region, planar, and fixed-root constraints.""" - root_index = {root_id: index for index, root_id in enumerate(root_ids)} - table_bounds = _bounds_from_points(table_optimization_rect_xy) - inequality_constraints: list[tuple[np.ndarray, float]] = [] - equality_constraints: list[tuple[np.ndarray, float]] = [] - - for root_id in root_ids: - region_bounds = _table_region_bounds( - table_bounds=table_bounds, - table_region=root_table_regions_by_id[root_id], - ) - _append_aabb_center_bounds( - constraints=inequality_constraints, - root_index=root_index, - root_id=root_id, - bounds=region_bounds, - half_extents_xy=root_half_extents_xy[root_id], - ) - fixed_xy = fixed_root_xy_by_id[root_id] - if fixed_xy is not None: - _append_fixed_root_constraints( - constraints=equality_constraints, - root_index=root_index, - root_id=root_id, - fixed_xy=fixed_xy, - ) - - for relation in root_relations: - _append_planar_relation_constraint( - constraints=inequality_constraints, - root_index=root_index, - source_id=relation.source_id, - relation=relation.relation, - target_id=relation.target_id, - source_half_extents_xy=root_half_extents_xy[relation.source_id], - target_half_extents_xy=root_half_extents_xy[relation.target_id], - relation_clearance_m=config.relation_clearance_m, - ) - return inequality_constraints, equality_constraints - - -def _solve_root_xy( - *, - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - inequality_constraints: list[tuple[np.ndarray, float]], - equality_constraints: list[tuple[np.ndarray, float]], - config: SceneLayoutOptimizerConfig, -) -> dict[str, list[float]]: - """Solve one root-group center model with the legacy SLSQP settings.""" - root_index = {root_id: index for index, root_id in enumerate(root_ids)} - initial_xy = np.asarray( - [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float - ) - x0 = initial_xy.reshape(-1) - - def unpack(values: np.ndarray) -> dict[str, list[float]]: - return { - root_id: [float(values[2 * index]), float(values[2 * index + 1])] - for root_id, index in root_index.items() - } - - def objective(values: np.ndarray) -> float: - coordinates = values.reshape(-1, 2) - loss = 0.0 - for root_id, index in root_index.items(): - if root_id in imported_root_ids: - delta = coordinates[index] - initial_xy[index] - loss += config.imported_seed_weight * float(delta @ delta) - for first_index in range(len(root_ids)): - for second_index in range(first_index + 1, len(root_ids)): - distance = float( - np.linalg.norm(coordinates[first_index] - coordinates[second_index]) - ) - shortfall = max(0.0, config.min_center_distance_m - distance) - loss += config.min_center_distance_weight * shortfall**2 - return loss - - constraints: list[dict[str, object]] = [] - for row, bound in inequality_constraints: - constraints.append( - { - "type": "ineq", - "fun": lambda values, row=row, bound=bound: bound - float(row @ values), - } - ) - for row, bound in equality_constraints: - constraints.append( - { - "type": "eq", - "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, - } - ) - - result = minimize( - objective, - x0, - method="SLSQP", - constraints=constraints, - options={ - "maxiter": config.max_slsqp_iterations, - "ftol": config.slsqp_ftol, - "disp": False, - }, - ) - if not result.success: - raise ValueError(f"Table layout optimization failed: {result.message}") - return unpack(np.asarray(result.x, dtype=float)) - - -def _refine_root_collisions( - *, - root_ids: list[str], - root_seed_xy_by_id: dict[str, list[float]], - imported_root_ids: set[str], - root_half_extents_xy: dict[str, np.ndarray], - inequality_constraints: list[tuple[np.ndarray, float]], - equality_constraints: list[tuple[np.ndarray, float]], - fixed_root_xy_by_id: dict[str, list[float] | None], - solved_root_xy_by_id: dict[str, list[float]], - config: SceneLayoutOptimizerConfig, -) -> dict[str, list[float]]: - """Separate only sibling pairs that include a layout-variable object.""" - seen_pairs: set[tuple[str, str]] = set() - current_xy_by_id = solved_root_xy_by_id - for _ in range(config.max_collision_rounds): - all_overlaps = _root_aabb_overlaps( - root_ids=root_ids, - root_half_extents_xy=root_half_extents_xy, - xy_by_id=current_xy_by_id, - ) - # Fixed/fixed siblings are outside this edit and therefore cannot be solved here. - overlaps = [ - overlap - for overlap in all_overlaps - if fixed_root_xy_by_id[overlap[1]] is None - or fixed_root_xy_by_id[overlap[2]] is None - ] - if not overlaps: - return current_xy_by_id - added_constraint_count = 0 - for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: - pair_key = tuple(sorted((first_id, second_id))) - if pair_key in seen_pairs: - continue - inequality_constraints.append( - _aabb_separation_constraint( - root_ids=root_ids, - first_id=first_id, - second_id=second_id, - first_half_extents_xy=root_half_extents_xy[first_id], - second_half_extents_xy=root_half_extents_xy[second_id], - first_xy=current_xy_by_id[first_id], - second_xy=current_xy_by_id[second_id], - collision_margin_m=config.collision_margin_m, - ) - ) - seen_pairs.add(pair_key) - added_constraint_count += 1 - if added_constraint_count == 0: - break - current_xy_by_id = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=current_xy_by_id, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) - - remaining_pairs = [ - f"{first_id}/{second_id}" - for _, first_id, second_id in _root_aabb_overlaps( - root_ids=root_ids, - root_half_extents_xy=root_half_extents_xy, - xy_by_id=current_xy_by_id, - ) - if fixed_root_xy_by_id[first_id] is None - or fixed_root_xy_by_id[second_id] is None - ] - raise ValueError( - "Table-root AABB collisions remain after layout refinement: " - f"{remaining_pairs}." - ) - - -def _asset_half_extents_xy( - *, - assets_by_id: dict[str, SceneObject], - object_ids: list[str], -) -> dict[str, np.ndarray]: - """Measure each asset's oriented z-up footprint around its XY center.""" - half_extents_xy: dict[str, np.ndarray] = {} - for object_id in object_ids: - asset = assets_by_id.get(object_id) - if asset is None: - raise ValueError(f"Table root {object_id!r} is not an asset.") - half_extents_xy[object_id] = _asset_half_extent_xy(asset) - return half_extents_xy - - -def _asset_half_extent_xy(asset: SceneObject) -> np.ndarray: - """Measure one SimReady GLB with its current orientation and scale.""" - mesh = _asset_z_up_mesh_at_zero_translation(scene_object=asset) - return (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 - - -def _asset_z_up_mesh_at_zero_translation( - *, - scene_object: SceneObject, - rotation_y_up: list[float] | None = None, -): - """Load one SimReady GLB in z-up with orientation and scale but no position.""" - asset = scene_object - if asset.simready_glb_path is None: - raise ValueError(f"Asset {asset.id!r} has no SimReady GLB path.") - y_up_layout = { - "id": asset.id, - "rot": ( - rotation_y_up - if rotation_y_up is not None - else _three_floats_or_default( - asset.rot, - field_name="rot", - default=[0.0, 0.0, 0.0], - ) - ), - "pos": [0.0, 0.0, 0.0], - "scale": _three_floats_or_default( - asset.scale, - field_name="scale", - default=[1.0, 1.0, 1.0], - ), - } - y_up_to_z_up = _y_up_to_z_up_matrix() - z_up_layout = transform_matrix_to_layout_object( - asset.id, - y_up_to_z_up - @ layout_object_to_transform_matrix(y_up_layout) - @ np.linalg.inv(y_up_to_z_up), - ) - mesh = load_glb_mesh(asset.simready_glb_path) - mesh.apply_transform(y_up_to_z_up) - mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) - return mesh - - -def _y_up_to_z_up_matrix() -> np.ndarray: - """Return the coordinate transform used by SceneExporter and layout stages.""" - matrix = np.eye(4) - matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) - return matrix - - -def _two_floats(value: object, *, field_name: str) -> list[float]: - """Validate one finite two-value vector.""" - if not isinstance(value, (list, tuple)) or len(value) != 2: - raise ValueError(f"{field_name} must contain two values.") - vector = [float(component) for component in value] - if not np.all(np.isfinite(vector)): - raise ValueError(f"{field_name} must contain finite values.") - return vector - - -def _three_floats_or_default( - value: object, - *, - field_name: str, - default: list[float] | None, -) -> list[float]: - """Return a finite three-value vector or the canonical SimReady default.""" - if value is None: - if default is None: - raise ValueError(f"{field_name} must contain three values.") - return list(default) - if not isinstance(value, (list, tuple)) or len(value) != 3: - raise ValueError(f"{field_name} must contain three values.") - vector = [float(component) for component in value] - if not np.all(np.isfinite(vector)): - raise ValueError(f"{field_name} must contain finite values.") - return vector - - -def _bounds_from_points(points: list[list[float]]) -> np.ndarray: - """Return [[min_x, min_y], [max_x, max_y]] from finite XY points.""" - coordinates = np.asarray(points, dtype=float) - if coordinates.ndim != 2 or coordinates.shape[1] != 2 or len(coordinates) < 2: - raise ValueError("XY bounds must contain at least two points.") - if not np.all(np.isfinite(coordinates)): - raise ValueError("XY bounds must contain finite values.") - return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) - - -def _table_region_bounds( - *, - table_bounds: np.ndarray, - table_region: str | None, -) -> np.ndarray: - """Return the requested 3x3 table region, with y increasing toward front.""" - if table_region is None: - return table_bounds.copy() - column_by_region = { - "left_back": 0, - "left_center": 0, - "left_front": 0, - "back_center": 1, - "center": 1, - "front_center": 1, - "right_back": 2, - "right_center": 2, - "right_front": 2, - } - row_by_region = { - "left_back": 0, - "back_center": 0, - "right_back": 0, - "left_center": 1, - "center": 1, - "right_center": 1, - "left_front": 2, - "front_center": 2, - "right_front": 2, - } - if table_region not in column_by_region: - raise ValueError(f"Unsupported table region {table_region!r}.") - minimum, maximum = table_bounds - cell_size = (maximum - minimum) / 3.0 - region_minimum = minimum + cell_size * np.array( - [column_by_region[table_region], row_by_region[table_region]] - ) - return np.stack([region_minimum, region_minimum + cell_size]) - - -def _append_aabb_center_bounds( - *, - constraints: list[tuple[np.ndarray, float]], - root_index: dict[str, int], - root_id: str, - bounds: np.ndarray, - half_extents_xy: np.ndarray, -) -> None: - """Keep one root's complete AABB inside the given rectangular bounds.""" - minimum = bounds[0] + half_extents_xy - maximum = bounds[1] - half_extents_xy - if np.any(minimum > maximum): - raise ValueError( - f"Asset {root_id!r} cannot fit inside its assigned table region." - ) - variable_count = 2 * len(root_index) - root_offset = 2 * root_index[root_id] - for axis in range(2): - upper_row = np.zeros(variable_count) - upper_row[root_offset + axis] = 1.0 - constraints.append((upper_row, float(maximum[axis]))) - lower_row = np.zeros(variable_count) - lower_row[root_offset + axis] = -1.0 - constraints.append((lower_row, -float(minimum[axis]))) - - -def _append_fixed_root_constraints( - *, - constraints: list[tuple[np.ndarray, float]], - root_index: dict[str, int], - root_id: str, - fixed_xy: list[float], -) -> None: - """Use equality constraints so unchanged formal objects remain fixed.""" - variable_count = 2 * len(root_index) - root_offset = 2 * root_index[root_id] - for axis, coordinate in enumerate(fixed_xy): - row = np.zeros(variable_count) - row[root_offset + axis] = 1.0 - constraints.append((row, float(coordinate))) - - -def _append_planar_relation_constraint( - *, - constraints: list[tuple[np.ndarray, float]], - root_index: dict[str, int], - source_id: str, - relation: str, - target_id: str, - source_half_extents_xy: np.ndarray, - target_half_extents_xy: np.ndarray, - relation_clearance_m: float, -) -> None: - """Require directional relations to clear both sibling AABB footprints.""" - if source_id not in root_index or target_id not in root_index: - raise ValueError("Table-root planar relations must reference table roots.") - axis, source_sign = { - "left_of": (0, 1.0), - "right_of": (0, -1.0), - "behind": (1, 1.0), - "in_front_of": (1, -1.0), - }.get(relation, (None, None)) - if axis is None or source_sign is None: - raise ValueError(f"Unsupported planar relation {relation!r}.") - row = np.zeros(2 * len(root_index)) - row[2 * root_index[source_id] + axis] = source_sign - row[2 * root_index[target_id] + axis] = -source_sign - required_distance = ( - source_half_extents_xy[axis] - + target_half_extents_xy[axis] - + relation_clearance_m - ) - constraints.append((row, -float(required_distance))) - - -def _root_aabb_overlaps( - *, - root_ids: list[str], - root_half_extents_xy: dict[str, np.ndarray], - xy_by_id: dict[str, list[float]], -) -> list[tuple[float, str, str]]: - """Return root pairs whose current XY AABBs overlap without a margin.""" - overlaps: list[tuple[float, str, str]] = [] - for first_index, first_id in enumerate(root_ids): - first_xy = np.asarray(xy_by_id[first_id], dtype=float) - first_half_extents = root_half_extents_xy[first_id] - for second_id in root_ids[first_index + 1 :]: - second_xy = np.asarray(xy_by_id[second_id], dtype=float) - second_half_extents = root_half_extents_xy[second_id] - overlap_xy = np.minimum( - first_xy + first_half_extents, - second_xy + second_half_extents, - ) - np.maximum( - first_xy - first_half_extents, - second_xy - second_half_extents, - ) - if np.all(overlap_xy > 1e-9): - overlaps.append((float(np.min(overlap_xy)), first_id, second_id)) - return sorted(overlaps, reverse=True) - - -def _aabb_separation_constraint( - *, - root_ids: list[str], - first_id: str, - second_id: str, - first_half_extents_xy: np.ndarray, - second_half_extents_xy: np.ndarray, - first_xy: list[float], - second_xy: list[float], - collision_margin_m: float, -) -> tuple[np.ndarray, float]: - """Separate one overlapping pair along its shallowest penetration axis.""" - root_index = {root_id: index for index, root_id in enumerate(root_ids)} - first_xy_array = np.asarray(first_xy, dtype=float) - second_xy_array = np.asarray(second_xy, dtype=float) - overlap_xy = np.minimum( - first_xy_array + first_half_extents_xy, - second_xy_array + second_half_extents_xy, - ) - np.maximum( - first_xy_array - first_half_extents_xy, - second_xy_array - second_half_extents_xy, - ) - axis = int(np.argmin(overlap_xy)) - first_is_lower = first_xy_array[axis] < second_xy_array[axis] or ( - first_xy_array[axis] == second_xy_array[axis] and first_id < second_id - ) - row = np.zeros(2 * len(root_ids)) - first_coefficient = 1.0 if first_is_lower else -1.0 - row[2 * root_index[first_id] + axis] = first_coefficient - row[2 * root_index[second_id] + axis] = -first_coefficient - required_distance = ( - first_half_extents_xy[axis] + second_half_extents_xy[axis] + collision_margin_m - ) - return row, -float(required_distance) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py new file mode 100644 index 000000000..100f5b269 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py @@ -0,0 +1,155 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import numpy as np + +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + load_glb_mesh, + transform_matrix_to_layout_object, +) + + +def update_scene_object_y_up_pose_from_z_up_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float = 0.02, +) -> None: + """Place a SimReady asset on a horizontal z-up support region.""" + if ( + not np.isfinite(support_region_z) + or clearance_m < 0.0 + or not np.isfinite(clearance_m) + ): + raise ValueError("support_region_z and clearance_m must be finite and valid.") + target_xy = two_floats(center_xy, field_name="center_xy") + rotation_y_up = three_floats_or_default( + scene_object.rot, field_name="rot", default=[0.0, 0.0, 0.0] + ) + mesh = load_scene_object_z_up_mesh( + scene_object=scene_object, rotation_y_up=rotation_y_up + ) + target_position_z_up = np.array( + [ + target_xy[0] - float(mesh.bounds[:, 0].mean()), + target_xy[1] - float(mesh.bounds[:, 1].mean()), + float(support_region_z) + clearance_m - float(mesh.bounds[0, 2]), + ] + ) + scene_object.pos = ( + np.linalg.inv(y_up_to_z_up_matrix())[:3, :3] @ target_position_z_up + ).tolist() + scene_object.rot = rotation_y_up + scene_object.center_xy = target_xy + + +def translate_scene_object_y_up_by_z_up_delta( + *, scene_object: SceneObject, delta_xy: list[float] +) -> None: + """Translate an existing y-up pose by a solved z-up XY delta.""" + dx, dy = two_floats(delta_xy, field_name="delta_xy") + position = three_floats_or_default(scene_object.pos, field_name="pos", default=None) + scene_object.pos = [position[0] + dx, position[1], position[2] - dy] + if scene_object.center_xy is not None: + scene_object.center_xy = [ + scene_object.center_xy[0] + dx, + scene_object.center_xy[1] + dy, + ] + + +def measure_scene_object_z_up_world_aabb( + *, scene_object: SceneObject +) -> list[list[float]]: + """Measure one current SceneObject pose in z-up world coordinates.""" + position_y_up = three_floats_or_default( + scene_object.pos, field_name="pos", default=None + ) + mesh = load_scene_object_z_up_mesh(scene_object=scene_object) + mesh.apply_translation( + y_up_to_z_up_matrix()[:3, :3] @ np.asarray(position_y_up, dtype=float) + ) + return mesh.bounds.tolist() + + +def load_scene_object_z_up_mesh( + *, scene_object: SceneObject, rotation_y_up: list[float] | None = None +): + """Load a SimReady mesh in z-up with orientation and scale but no translation.""" + if scene_object.simready_glb_path is None: + raise ValueError(f"Asset {scene_object.id!r} has no SimReady GLB path.") + y_up_layout = { + "id": scene_object.id, + "rot": ( + rotation_y_up + if rotation_y_up is not None + else three_floats_or_default( + scene_object.rot, field_name="rot", default=[0.0, 0.0, 0.0] + ) + ), + "pos": [0.0, 0.0, 0.0], + "scale": three_floats_or_default( + scene_object.scale, field_name="scale", default=[1.0, 1.0, 1.0] + ), + } + y_up_to_z_up = y_up_to_z_up_matrix() + z_up_layout = transform_matrix_to_layout_object( + scene_object.id, + y_up_to_z_up + @ layout_object_to_transform_matrix(y_up_layout) + @ np.linalg.inv(y_up_to_z_up), + ) + mesh = load_glb_mesh(scene_object.simready_glb_path) + mesh.apply_transform(y_up_to_z_up) + mesh.apply_transform(layout_object_to_transform_matrix(z_up_layout)) + return mesh + + +def y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate transform used by layout and export stages.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix + + +def two_floats(value: object, *, field_name: str) -> list[float]: + """Validate and return one finite two-value vector.""" + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError(f"{field_name} must contain two values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"{field_name} must contain finite values.") + return result + + +def three_floats_or_default( + value: object, *, field_name: str, default: list[float] | None +) -> list[float]: + """Validate three finite values, or return a canonical default.""" + if value is None: + if default is None: + raise ValueError(f"{field_name} must contain three values.") + return list(default) + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"{field_name} must contain three values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"{field_name} must contain finite values.") + return result diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py new file mode 100644 index 000000000..38fe7e6c4 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -0,0 +1,542 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + load_scene_object_z_up_mesh, +) + +if TYPE_CHECKING: + from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutGroup, + SceneLayoutProblem, + ) + + +@dataclass +class TableSurfaceLayoutProblem: + """All scene-graph and geometry inputs for one table-surface solve.""" + + assets_by_id: dict[str, SceneObject] + root_ids: list[str] + root_seed_xy_by_id: dict[str, list[float]] + imported_root_ids: set[str] + fixed_root_xy_by_id: dict[str, list[float] | None] + root_table_regions_by_id: dict[str, str | None] + table_optimization_rect_xy: list[list[float]] + root_relations: list[SceneGraphRelation] + + @classmethod + def from_layout_problem( + cls, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + current_xy_by_id: dict[str, list[float] | None], + ) -> TableSurfaceLayoutProblem: + """Build one table-surface problem without mutating layout state.""" + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table group optimization requires a table.") + if table.support_optimization_rect_xy is None: + raise ValueError( + "Table group optimization requires a table support optimization rectangle." + ) + root_ids = set(group.child_ids) + nodes_by_id = layout_problem.goal_scene_graph.node_by_id() + root_seed_xy_by_id = {} + for root_id in group.child_ids: + inherited_xy = current_xy_by_id[root_id] + # New roots begin from the table origin; imported roots retain their seed. + root_seed_xy_by_id[root_id] = ( + [0.0, 0.0] if inherited_xy is None else list(inherited_xy) + ) + return cls( + assets_by_id={ + asset.id: asset for asset in layout_problem.post_edit_scene.assets + }, + root_ids=group.child_ids, + root_seed_xy_by_id=root_seed_xy_by_id, + imported_root_ids={ + root_id + for root_id in group.child_ids + if layout_problem.initial_xy_by_id[root_id] is not None + }, + fixed_root_xy_by_id={ + root_id: ( + None + if root_id in layout_problem.layout_variable_ids + else current_xy_by_id[root_id] + ) + for root_id in group.child_ids + }, + root_table_regions_by_id={ + root_id: nodes_by_id[root_id].table_region + for root_id in group.child_ids + }, + table_optimization_rect_xy=table.support_optimization_rect_xy, + root_relations=[ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in root_ids and relation.target_id in root_ids + ], + ) + + +@dataclass(frozen=True) +class TableSurfaceLayoutOptimizerConfig: + """Numerical controls for one direct-table sibling layout solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid controls before assembling table-surface constraints.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class TableSurfaceLayoutOptimizer: + """Solve direct table children with table, relation, and collision constraints.""" + + def __init__( + self, + *, + config: TableSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.config = ( + config if config is not None else TableSurfaceLayoutOptimizerConfig() + ) + + def optimize( + self, + problem: TableSurfaceLayoutProblem, + ) -> dict[str, list[float]]: + """Return the table-frame XY centers satisfying this atomic problem.""" + # Measure only this sibling group from the complete scene-asset index. + root_half_extents_xy = _asset_half_extents_xy( + assets_by_id=problem.assets_by_id, + object_ids=problem.root_ids, + ) + # Equality constraints for fixed roots, and inequality constraints for table-region and planar-relation bounds. + inequality_constraints, equality_constraints = _build_constraints( + problem=problem, + root_half_extents_xy=root_half_extents_xy, + config=self.config, + ) + # Solve with the SLSQP optimizer. + solved_root_xy_by_id = _solve_root_xy( + root_ids=problem.root_ids, + root_seed_xy_by_id=problem.root_seed_xy_by_id, + imported_root_ids=problem.imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=problem.root_ids, + root_seed_xy_by_id=problem.root_seed_xy_by_id, + imported_root_ids=problem.imported_root_ids, + root_half_extents_xy=root_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + fixed_root_xy_by_id=problem.fixed_root_xy_by_id, + solved_root_xy_by_id=solved_root_xy_by_id, + config=self.config, + ) + + +def _build_constraints( + *, + problem: TableSurfaceLayoutProblem, + root_half_extents_xy: dict[str, np.ndarray], + config: TableSurfaceLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard table-region, planar-relation, and fixed-root constraints.""" + # Objects which need to be optimized. + root_index = {root_id: index for index, root_id in enumerate(problem.root_ids)} + table_bounds = _bounds_from_points(problem.table_optimization_rect_xy) + # Initi constraints. + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + for root_id in problem.root_ids: + # Get the table region bound for this root asset. + region_bounds = _table_region_bounds( + table_bounds=table_bounds, + table_region=problem.root_table_regions_by_id[root_id], + ) + # Add AABB constraints for each root's center inside the table region. + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=root_id, + bounds=region_bounds, + half_extents_xy=root_half_extents_xy[root_id], + ) + fixed_xy = problem.fixed_root_xy_by_id[root_id] + if fixed_xy is not None: + # Add fixed-root constraints for each root with a fixed XY center. + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=root_id, + fixed_xy=fixed_xy, + ) + for relation in problem.root_relations: + # Add planar-relation constraints for each sibling relation in this group. + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=root_half_extents_xy[relation.source_id], + target_half_extents_xy=root_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _table_region_bounds( + *, + table_bounds: np.ndarray, + table_region: str | None, +) -> np.ndarray: + """Return the requested 3x3 table region, with y increasing toward front.""" + if table_region is None: + return table_bounds.copy() + column_by_region = { + "left_back": 0, + "left_center": 0, + "left_front": 0, + "back_center": 1, + "center": 1, + "front_center": 1, + "right_back": 2, + "right_center": 2, + "right_front": 2, + } + row_by_region = { + "left_back": 0, + "back_center": 0, + "right_back": 0, + "left_center": 1, + "center": 1, + "right_center": 1, + "left_front": 2, + "front_center": 2, + "right_front": 2, + } + if table_region not in column_by_region: + raise ValueError(f"Unsupported table region {table_region!r}.") + minimum, maximum = table_bounds + # 9-grid. + cell_size = (maximum - minimum) / 3.0 + region_minimum = minimum + cell_size * np.array( + [column_by_region[table_region], row_by_region[table_region]] + ) + return np.stack([region_minimum, region_minimum + cell_size]) + + +def _asset_half_extents_xy( + *, assets_by_id: dict[str, SceneObject], object_ids: list[str] +) -> dict[str, np.ndarray]: + """Measure each optimized asset's oriented z-up XY half-extents.""" + result = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Table root {object_id!r} is not an asset.") + mesh = load_scene_object_z_up_mesh(scene_object=asset) + result[object_id] = (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + return result + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + coordinates = np.asarray(points, dtype=float) + if ( + coordinates.ndim != 2 + or coordinates.shape[1] != 2 + or len(coordinates) < 2 + or not np.all(np.isfinite(coordinates)) + ): + raise ValueError("XY bounds must contain at least two finite points.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + minimum, maximum = bounds[0] + half_extents_xy, bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError( + f"Asset {root_id!r} cannot fit inside its assigned table region." + ) + # root_id is the sibling whose center is constrained in this AABB bound. + offset, count = 2 * root_index[root_id], 2 * len(root_index) + # offset selects this root's XY pair; count is the full flattened XY vector size. + for axis in range(2): + upper, lower = np.zeros(count), np.zeros(count) + upper[offset + axis], lower[offset + axis] = 1.0, -1.0 + constraints.extend( + [(upper, float(maximum[axis])), (lower, -float(minimum[axis]))] + ) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + offset, count = 2 * root_index[root_id], 2 * len(root_index) + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(count) + row[offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + axis, sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_id not in root_index or target_id not in root_index: + raise ValueError(f"Unsupported table-root planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = sign + row[2 * root_index[target_id] + axis] = -sign + constraints.append( + ( + row, + -float( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ), + ) + ) + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: TableSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + # Init with XY-seeds. + initial = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + + def objective(values: np.ndarray) -> float: + xy = values.reshape(-1, 2) + loss = 0.0 + for index, root_id in enumerate(root_ids): + if root_id in imported_root_ids: + delta = xy[index] - initial[index] + loss += config.imported_seed_weight * float(delta @ delta) + return loss + + constraints = [ + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + for row, bound in inequality_constraints + ] + [ + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + for row, bound in equality_constraints + ] + result = minimize( + objective, + initial.reshape(-1), + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise ValueError(f"Table layout optimization failed: {result.message}") + return { + root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] + for index, root_id in enumerate(root_ids) + } + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + solved_root_xy_by_id: dict[str, list[float]], + config: TableSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + # Get current SLSQP solution. + current = solved_root_xy_by_id + seen: set[tuple[str, str]] = set() + for _ in range(config.max_collision_rounds): + # Fine overlaps. + overlaps = [ + pair + for pair in _root_aabb_overlaps( + root_ids=root_ids, half_extents=root_half_extents_xy, xy_by_id=current + ) + if fixed_root_xy_by_id[pair[1]] is None + or fixed_root_xy_by_id[pair[2]] is None + ] + if not overlaps: + return current + added = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + key = tuple(sorted((first_id, second_id))) + if key in seen: + continue + # Add a new SLSQP constraint to separate this overlapping pair. + inequality_constraints.append( + _aabb_separation_constraint( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ) + ) + seen.add(key) + added += 1 + if not added: + break + current = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + raise ValueError("Table-root AABB collisions remain after layout refinement.") + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return all overlapping root pairs with their minimum XY overlap distance.""" + result = [] + for index, first_id in enumerate(root_ids): + for second_id in root_ids[index + 1 :]: + overlap = np.minimum( + np.asarray(xy_by_id[first_id]) + half_extents[first_id], + np.asarray(xy_by_id[second_id]) + half_extents[second_id], + ) - np.maximum( + np.asarray(xy_by_id[first_id]) - half_extents[first_id], + np.asarray(xy_by_id[second_id]) - half_extents[second_id], + ) + if np.all(overlap > 1e-9): + result.append((float(np.min(overlap)), first_id, second_id)) + return sorted(result, reverse=True) + + +def _aabb_separation_constraint( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], + margin: float, +) -> tuple[np.ndarray, float]: + first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) + # Positive overlap on both axes means these two center-based AABBs intersect. + overlap = np.minimum( + first + half_extents[first_id], second + half_extents[second_id] + ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) + # Separate along the least-penetrating axis to require the smallest local shift. + axis = int(np.argmin(overlap)) + # Preserve the current order on that axis; object IDs break an exact tie deterministically. + lower = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + index = {root_id: i for i, root_id in enumerate(root_ids)} + # One row addresses the x/y variable pair of each root in the flattened solver vector. + row = np.zeros(2 * len(root_ids)) + sign = 1.0 if lower else -1.0 + row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign + # row @ values <= bound keeps the selected AABB faces apart by the requested margin. + return row, -float( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index bca20da8e..9b650469c 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -30,8 +30,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( SceneLayoutConstructor, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_optimizer import ( - SceneLayoutOptimizer, +from embodichain.gen_sim.scene_engine.pipeline.utils.table_surface_layout_optimizer import ( + TableSurfaceLayoutOptimizer, + TableSurfaceLayoutProblem, _table_region_bounds, ) @@ -152,34 +153,36 @@ def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: asset_glb = tmp_path / "asset.glb" trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) first_id, second_id = "first_001", "second_001" - optimizer = SceneLayoutOptimizer() - - solved_xy_by_id = optimizer.optimize_table_root_xy( - assets_by_id={ - first_id: _asset( - object_id=first_id, - glb_path=asset_glb, - center_xy=_OVERLAPPING_CENTER_XY, - ), - second_id: _asset( - object_id=second_id, - glb_path=asset_glb, - center_xy=_OVERLAPPING_CENTER_XY, - ), - }, - root_ids=[first_id, second_id], - root_seed_xy_by_id={ - first_id: _OVERLAPPING_CENTER_XY, - second_id: _OVERLAPPING_CENTER_XY, - }, - imported_root_ids={first_id, second_id}, - fixed_root_xy_by_id={ - first_id: _OVERLAPPING_CENTER_XY, - second_id: _OVERLAPPING_CENTER_XY, - }, - root_table_regions_by_id={first_id: None, second_id: None}, - table_optimization_rect_xy=_TABLE_BOUNDS, - root_relations=[], + optimizer = TableSurfaceLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + first_id: _asset( + object_id=first_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + second_id: _asset( + object_id=second_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + }, + root_ids=[first_id, second_id], + root_seed_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={first_id, second_id}, + fixed_root_xy_by_id={ + first_id: _OVERLAPPING_CENTER_XY, + second_id: _OVERLAPPING_CENTER_XY, + }, + root_table_regions_by_id={first_id: None, second_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) ) assert solved_xy_by_id == { @@ -194,30 +197,32 @@ def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( asset_glb = tmp_path / "asset.glb" trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) fixed_id, variable_id = "fixed_001", "variable_001" - optimizer = SceneLayoutOptimizer() - - solved_xy_by_id = optimizer.optimize_table_root_xy( - assets_by_id={ - fixed_id: _asset( - object_id=fixed_id, - glb_path=asset_glb, - center_xy=_OVERLAPPING_CENTER_XY, - ), - variable_id: _asset( - object_id=variable_id, - glb_path=asset_glb, - ), - }, - root_ids=[fixed_id, variable_id], - root_seed_xy_by_id={ - fixed_id: _OVERLAPPING_CENTER_XY, - variable_id: _OVERLAPPING_CENTER_XY, - }, - imported_root_ids={fixed_id}, - fixed_root_xy_by_id={fixed_id: _OVERLAPPING_CENTER_XY, variable_id: None}, - root_table_regions_by_id={fixed_id: None, variable_id: None}, - table_optimization_rect_xy=_TABLE_BOUNDS, - root_relations=[], + optimizer = TableSurfaceLayoutOptimizer() + + solved_xy_by_id = optimizer.optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + fixed_id: _asset( + object_id=fixed_id, + glb_path=asset_glb, + center_xy=_OVERLAPPING_CENTER_XY, + ), + variable_id: _asset( + object_id=variable_id, + glb_path=asset_glb, + ), + }, + root_ids=[fixed_id, variable_id], + root_seed_xy_by_id={ + fixed_id: _OVERLAPPING_CENTER_XY, + variable_id: _OVERLAPPING_CENTER_XY, + }, + imported_root_ids={fixed_id}, + fixed_root_xy_by_id={fixed_id: _OVERLAPPING_CENTER_XY, variable_id: None}, + root_table_regions_by_id={fixed_id: None, variable_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) ) assert solved_xy_by_id[fixed_id] == _OVERLAPPING_CENTER_XY From ae4095e315239ad53e7e0ac278d69c6887ee9859 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:15:55 +0800 Subject: [PATCH 27/85] finish a simple on-relationship optimizer --- .../utils/parent_surface_layout_optimizer.py | 489 ++++++++++++++++++ .../utils/scene_layout_constructor.py | 3 +- .../test_scene_layout_optimizer.py | 46 +- 3 files changed, 535 insertions(+), 3 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py new file mode 100644 index 000000000..0032f1d66 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -0,0 +1,489 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import numpy as np +from scipy.optimize import minimize + +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + load_scene_object_z_up_mesh, + measure_scene_object_z_up_world_aabb, +) + +if TYPE_CHECKING: + from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( + SceneLayoutGroup, + SceneLayoutProblem, + ) + + +@dataclass +class ParentSurfaceLayoutProblem: + """All geometry and layout-state inputs for one parent-surface solve.""" + + assets_by_id: dict[str, SceneObject] + child_ids: list[str] + child_seed_xy_by_id: dict[str, list[float]] + imported_child_ids: set[str] + fixed_child_xy_by_id: dict[str, list[float] | None] + parent_aabb_xy: list[list[float]] + parent_top_z: float + child_relations: list[SceneGraphRelation] + + @classmethod + def from_layout_problem( + cls, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + current_xy_by_id: dict[str, list[float] | None], + ) -> ParentSurfaceLayoutProblem: + """Build one parent-surface problem without mutating layout state.""" + assets_by_id = { + asset.id: asset for asset in layout_problem.post_edit_scene.assets + } + child_ids = set(group.child_ids) + parent = assets_by_id.get(group.parent_id) + if parent is None: + raise ValueError(f"Parent {group.parent_id!r} is not an asset.") + parent_aabb = measure_scene_object_z_up_world_aabb(scene_object=parent) + parent_aabb_xy = [ + [parent_aabb[0][0], parent_aabb[0][1]], + [parent_aabb[1][0], parent_aabb[1][1]], + ] + parent_center_xy = [ + (parent_aabb[0][0] + parent_aabb[1][0]) / 2.0, + (parent_aabb[0][1] + parent_aabb[1][1]) / 2.0, + ] + child_seed_xy_by_id = {} + for child_id in group.child_ids: + inherited_xy = current_xy_by_id[child_id] + # New children begin from the solved parent's AABB center. + child_seed_xy_by_id[child_id] = ( + parent_center_xy if inherited_xy is None else list(inherited_xy) + ) + return cls( + assets_by_id=assets_by_id, + child_ids=group.child_ids, + child_seed_xy_by_id=child_seed_xy_by_id, + imported_child_ids={ + child_id + for child_id in group.child_ids + if layout_problem.initial_xy_by_id[child_id] is not None + }, + fixed_child_xy_by_id={ + child_id: ( + None + if child_id in layout_problem.layout_variable_ids + else current_xy_by_id[child_id] + ) + for child_id in group.child_ids + }, + parent_aabb_xy=parent_aabb_xy, + parent_top_z=parent_aabb[1][2], + child_relations=[ + relation + for relation in layout_problem.goal_scene_graph.relations + if relation.source_id in child_ids and relation.target_id in child_ids + ], + ) + + +@dataclass(frozen=True) +class ParentSurfaceLayoutOptimizerConfig: + """Numerical controls for one non-table parent-surface sibling solve.""" + + relation_clearance_m: float = 0.03 + collision_margin_m: float = 0.02 + max_slsqp_iterations: int = 500 + slsqp_ftol: float = 1e-6 + max_collision_rounds: int = 8 + max_added_collision_pairs: int = 64 + imported_seed_weight: float = 5.0 + min_center_distance_m: float = 0.01 + min_center_distance_weight: float = 0.05 + + def __post_init__(self) -> None: + """Reject invalid controls before assembling parent-surface constraints.""" + if self.relation_clearance_m < 0.0: + raise ValueError("relation_clearance_m must be non-negative.") + if self.collision_margin_m < 0.0: + raise ValueError("collision_margin_m must be non-negative.") + if self.max_slsqp_iterations <= 0: + raise ValueError("max_slsqp_iterations must be positive.") + if self.slsqp_ftol <= 0.0: + raise ValueError("slsqp_ftol must be positive.") + if self.max_collision_rounds <= 0: + raise ValueError("max_collision_rounds must be positive.") + if self.max_added_collision_pairs <= 0: + raise ValueError("max_added_collision_pairs must be positive.") + + +class ParentSurfaceLayoutOptimizer: + """Solve direct ``on`` children inside one parent's current XY footprint.""" + + def __init__( + self, + *, + config: ParentSurfaceLayoutOptimizerConfig | None = None, + ) -> None: + self.config = ( + config if config is not None else ParentSurfaceLayoutOptimizerConfig() + ) + + def optimize( + self, + problem: ParentSurfaceLayoutProblem, + ) -> dict[str, list[float]]: + """Return sibling XY centers inside the parent AABB without overlap.""" + child_half_extents_xy = _asset_half_extents_xy( + assets_by_id=problem.assets_by_id, + object_ids=problem.child_ids, + ) + inequality_constraints, equality_constraints = _build_constraints( + problem=problem, + child_half_extents_xy=child_half_extents_xy, + config=self.config, + ) + solved_child_xy_by_id = _solve_root_xy( + root_ids=problem.child_ids, + root_seed_xy_by_id=problem.child_seed_xy_by_id, + imported_root_ids=problem.imported_child_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=self.config, + ) + return _refine_root_collisions( + root_ids=problem.child_ids, + root_seed_xy_by_id=problem.child_seed_xy_by_id, + imported_root_ids=problem.imported_child_ids, + root_half_extents_xy=child_half_extents_xy, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + fixed_root_xy_by_id=problem.fixed_child_xy_by_id, + solved_root_xy_by_id=solved_child_xy_by_id, + config=self.config, + ) + + +def _build_constraints( + *, + problem: ParentSurfaceLayoutProblem, + child_half_extents_xy: dict[str, np.ndarray], + config: ParentSurfaceLayoutOptimizerConfig, +) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: + """Build hard parent-AABB, planar-relation, and fixed-child constraints.""" + root_index = {child_id: index for index, child_id in enumerate(problem.child_ids)} + parent_bounds = _bounds_from_points(problem.parent_aabb_xy) + inequality_constraints: list[tuple[np.ndarray, float]] = [] + equality_constraints: list[tuple[np.ndarray, float]] = [] + for child_id in problem.child_ids: + # Keep each child's 2D AABB inside the parent AABB support proxy. + _append_aabb_center_bounds( + constraints=inequality_constraints, + root_index=root_index, + root_id=child_id, + bounds=parent_bounds, + half_extents_xy=child_half_extents_xy[child_id], + ) + fixed_xy = problem.fixed_child_xy_by_id[child_id] + if fixed_xy is not None: + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=child_id, + fixed_xy=fixed_xy, + ) + for relation in problem.child_relations: + # Apply planar relations between direct on-children of this parent. + _append_planar_relation_constraint( + constraints=inequality_constraints, + root_index=root_index, + source_id=relation.source_id, + relation=relation.relation, + target_id=relation.target_id, + source_half_extents_xy=child_half_extents_xy[relation.source_id], + target_half_extents_xy=child_half_extents_xy[relation.target_id], + relation_clearance_m=config.relation_clearance_m, + ) + return inequality_constraints, equality_constraints + + +def _asset_half_extents_xy( + *, assets_by_id: dict[str, SceneObject], object_ids: list[str] +) -> dict[str, np.ndarray]: + """Measure each optimized child asset's z-up XY half-extents.""" + result = {} + for object_id in object_ids: + asset = assets_by_id.get(object_id) + if asset is None: + raise ValueError(f"Parent child {object_id!r} is not an asset.") + mesh = load_scene_object_z_up_mesh(scene_object=asset) + result[object_id] = (mesh.bounds[1, :2] - mesh.bounds[0, :2]) / 2.0 + return result + + +def _bounds_from_points(points: list[list[float]]) -> np.ndarray: + """Return finite XY minimum and maximum bounds from polygon points.""" + coordinates = np.asarray(points, dtype=float) + if ( + coordinates.ndim != 2 + or coordinates.shape[1] != 2 + or len(coordinates) < 2 + or not np.all(np.isfinite(coordinates)) + ): + raise ValueError("XY bounds must contain at least two finite points.") + return np.stack([coordinates.min(axis=0), coordinates.max(axis=0)]) + + +def _append_aabb_center_bounds( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + bounds: np.ndarray, + half_extents_xy: np.ndarray, +) -> None: + """Constrain one child AABB center to lie completely inside XY bounds.""" + minimum, maximum = bounds[0] + half_extents_xy, bounds[1] - half_extents_xy + if np.any(minimum > maximum): + raise ValueError(f"Asset {root_id!r} cannot fit inside its parent AABB.") + # root_id is the child whose center is constrained in this AABB bound. + offset, count = 2 * root_index[root_id], 2 * len(root_index) + # offset selects this child's XY pair; count is the full flattened XY vector size. + for axis in range(2): + upper, lower = np.zeros(count), np.zeros(count) + upper[offset + axis], lower[offset + axis] = 1.0, -1.0 + constraints.extend( + [(upper, float(maximum[axis])), (lower, -float(minimum[axis]))] + ) + + +def _append_fixed_root_constraints( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + root_id: str, + fixed_xy: list[float], +) -> None: + """Lock one fixed child's center to its imported XY coordinates.""" + offset, count = 2 * root_index[root_id], 2 * len(root_index) + for axis, coordinate in enumerate(fixed_xy): + row = np.zeros(count) + row[offset + axis] = 1.0 + constraints.append((row, float(coordinate))) + + +def _append_planar_relation_constraint( + *, + constraints: list[tuple[np.ndarray, float]], + root_index: dict[str, int], + source_id: str, + relation: str, + target_id: str, + source_half_extents_xy: np.ndarray, + target_half_extents_xy: np.ndarray, + relation_clearance_m: float, +) -> None: + """Append one world-XY separation constraint for sibling planar semantics.""" + axis, sign = { + "left_of": (0, 1.0), + "right_of": (0, -1.0), + "behind": (1, 1.0), + "in_front_of": (1, -1.0), + }.get(relation, (None, None)) + if axis is None or source_id not in root_index or target_id not in root_index: + raise ValueError(f"Unsupported parent-child planar relation {relation!r}.") + row = np.zeros(2 * len(root_index)) + row[2 * root_index[source_id] + axis] = sign + row[2 * root_index[target_id] + axis] = -sign + constraints.append( + ( + row, + -float( + source_half_extents_xy[axis] + + target_half_extents_xy[axis] + + relation_clearance_m + ), + ) + ) + + +def _solve_root_xy( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + config: ParentSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Solve one parent child-group's XY positions with SLSQP.""" + initial = np.asarray( + [root_seed_xy_by_id[root_id] for root_id in root_ids], dtype=float + ) + + def objective(values: np.ndarray) -> float: + xy = values.reshape(-1, 2) + loss = 0.0 + for index, root_id in enumerate(root_ids): + if root_id in imported_root_ids: + delta = xy[index] - initial[index] + loss += config.imported_seed_weight * float(delta @ delta) + return loss + + constraints = [ + { + "type": "ineq", + "fun": lambda values, row=row, bound=bound: bound - float(row @ values), + } + for row, bound in inequality_constraints + ] + [ + { + "type": "eq", + "fun": lambda values, row=row, bound=bound: float(row @ values) - bound, + } + for row, bound in equality_constraints + ] + result = minimize( + objective, + initial.reshape(-1), + method="SLSQP", + constraints=constraints, + options={ + "maxiter": config.max_slsqp_iterations, + "ftol": config.slsqp_ftol, + "disp": False, + }, + ) + if not result.success: + raise ValueError(f"Parent layout optimization failed: {result.message}") + return { + root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] + for index, root_id in enumerate(root_ids) + } + + +def _refine_root_collisions( + *, + root_ids: list[str], + root_seed_xy_by_id: dict[str, list[float]], + imported_root_ids: set[str], + root_half_extents_xy: dict[str, np.ndarray], + inequality_constraints: list[tuple[np.ndarray, float]], + equality_constraints: list[tuple[np.ndarray, float]], + fixed_root_xy_by_id: dict[str, list[float] | None], + solved_root_xy_by_id: dict[str, list[float]], + config: ParentSurfaceLayoutOptimizerConfig, +) -> dict[str, list[float]]: + """Iteratively add separation constraints for overlapping child AABBs.""" + current = solved_root_xy_by_id + seen: set[tuple[str, str]] = set() + for _ in range(config.max_collision_rounds): + overlaps = [ + pair + for pair in _root_aabb_overlaps( + root_ids=root_ids, half_extents=root_half_extents_xy, xy_by_id=current + ) + if fixed_root_xy_by_id[pair[1]] is None + or fixed_root_xy_by_id[pair[2]] is None + ] + if not overlaps: + return current + added = 0 + for _, first_id, second_id in overlaps[: config.max_added_collision_pairs]: + key = tuple(sorted((first_id, second_id))) + if key in seen: + continue + inequality_constraints.append( + _aabb_separation_constraint( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ) + ) + seen.add(key) + added += 1 + if not added: + break + current = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=inequality_constraints, + equality_constraints=equality_constraints, + config=config, + ) + raise ValueError("Parent-child AABB collisions remain after layout refinement.") + + +def _root_aabb_overlaps( + *, + root_ids: list[str], + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], +) -> list[tuple[float, str, str]]: + """Return overlapping child pairs with their minimum XY overlap distance.""" + result = [] + for index, first_id in enumerate(root_ids): + for second_id in root_ids[index + 1 :]: + overlap = np.minimum( + np.asarray(xy_by_id[first_id]) + half_extents[first_id], + np.asarray(xy_by_id[second_id]) + half_extents[second_id], + ) - np.maximum( + np.asarray(xy_by_id[first_id]) - half_extents[first_id], + np.asarray(xy_by_id[second_id]) - half_extents[second_id], + ) + if np.all(overlap > 1e-9): + result.append((float(np.min(overlap)), first_id, second_id)) + return sorted(result, reverse=True) + + +def _aabb_separation_constraint( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + xy_by_id: dict[str, list[float]], + margin: float, +) -> tuple[np.ndarray, float]: + """Return one least-penetration AABB separation inequality.""" + first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) + overlap = np.minimum( + first + half_extents[first_id], second + half_extents[second_id] + ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) + axis = int(np.argmin(overlap)) + lower = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + index = {root_id: i for i, root_id in enumerate(root_ids)} + row = np.zeros(2 * len(root_ids)) + sign = 1.0 if lower else -1.0 + row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign + return row, -float( + half_extents[first_id][axis] + half_extents[second_id][axis] + margin + ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py index 9279bd1ae..95c77d712 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -170,7 +170,7 @@ def _optimize_table_group( scene_object=assets_by_id[root_id], support_region_z=table.support_surface_z, center_xy=solved_xy, - clearance_m=0.00, # Directly place on the support surface. + clearance_m=0.00, # Directly place on the support surface. ) self._updated_object_ids.add(root_id) self._propagate_descendant_delta( @@ -242,6 +242,7 @@ def _optimize_parent_group( scene_object=parent_surface_problem.assets_by_id[child_id], support_region_z=parent_surface_problem.parent_top_z, center_xy=solved_xy, + clearance_m=0.00, # Directly place on the parent's top surface. ) self._updated_object_ids.add(child_id) self._propagate_descendant_delta( diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 9b650469c..39a44ecc7 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -25,8 +25,13 @@ from embodichain.gen_sim.scene_engine.core.scene_graph import ( SceneGraph, SceneGraphNode, + SceneGraphRelation, ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer import ( + ParentSurfaceLayoutOptimizer, + ParentSurfaceLayoutProblem, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( SceneLayoutConstructor, ) @@ -44,6 +49,7 @@ ] _OVERLAPPING_CENTER_XY = [0.0, 0.0] _ASSET_SIDE_LENGTH_M = 0.2 +_RELATION_CLEARANCE_M = 0.03 def _asset( @@ -145,8 +151,8 @@ def test_layout_constructor_places_new_child_on_parent_top( asset for asset in post_edit_scene.assets if asset.id == "cup_001" ) assert placed_cup.center_xy == [0.0, 0.0] - # book top is z=0.62 m; cup half-height is 0.1 m and clearance is 0.02 m. - assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) + # Book top is z=0.62 m; cup half-height is 0.1 m with zero support clearance. + assert np.allclose(placed_cup.pos, [0.0, 0.72, 0.0]) def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: @@ -227,3 +233,39 @@ def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( assert solved_xy_by_id[fixed_id] == _OVERLAPPING_CENTER_XY assert np.max(np.abs(solved_xy_by_id[variable_id])) >= _ASSET_SIDE_LENGTH_M - 1e-6 + + +def test_parent_optimizer_applies_sibling_planar_relation(tmp_path: Path) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + left_id, right_id = "left_001", "right_001" + + solved_xy_by_id = ParentSurfaceLayoutOptimizer().optimize( + ParentSurfaceLayoutProblem( + assets_by_id={ + left_id: _asset(object_id=left_id, glb_path=asset_glb), + right_id: _asset(object_id=right_id, glb_path=asset_glb), + }, + child_ids=[left_id, right_id], + child_seed_xy_by_id={ + left_id: _OVERLAPPING_CENTER_XY, + right_id: _OVERLAPPING_CENTER_XY, + }, + imported_child_ids=set(), + fixed_child_xy_by_id={left_id: None, right_id: None}, + parent_aabb_xy=_TABLE_BOUNDS, + parent_top_z=0.0, + child_relations=[ + SceneGraphRelation( + source_id=left_id, + relation="left_of", + target_id=right_id, + ) + ], + ) + ) + + assert ( + solved_xy_by_id[right_id][0] - solved_xy_by_id[left_id][0] + >= _ASSET_SIDE_LENGTH_M + _RELATION_CLEARANCE_M - 1e-6 + ) From 73d3c7faa04f1532a0ea2f7c448e96f237d932bf Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 17:11:04 +0800 Subject: [PATCH 28/85] add the heuristic attaches in collision optimization --- .../utils/parent_surface_layout_optimizer.py | 107 +++++++++++++---- .../utils/table_surface_layout_optimizer.py | 110 +++++++++++++----- .../test_scene_layout_optimizer.py | 71 +++++++++++ 3 files changed, 236 insertions(+), 52 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py index 0032f1d66..6d2a3b0a5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -186,6 +186,10 @@ def optimize( ) +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + def _build_constraints( *, problem: ParentSurfaceLayoutProblem, @@ -377,7 +381,9 @@ def objective(values: np.ndarray) -> float: }, ) if not result.success: - raise ValueError(f"Parent layout optimization failed: {result.message}") + raise _LayoutInfeasibleError( + f"Parent layout optimization failed: {result.message}" + ) return { root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] for index, root_id in enumerate(root_ids) @@ -415,28 +421,51 @@ def _refine_root_collisions( key = tuple(sorted((first_id, second_id))) if key in seen: continue - inequality_constraints.append( - _aabb_separation_constraint( + # Earlier pair updates may already have separated this stale overlap. + if key not in { + tuple(sorted((first, second))) + for _, first, second in _root_aabb_overlaps( root_ids=root_ids, - first_id=first_id, - second_id=second_id, half_extents=root_half_extents_xy, xy_by_id=current, - margin=config.collision_margin_m, ) - ) - seen.add(key) - added += 1 + }: + continue + for separation_constraint in _aabb_separation_constraints( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ): + # Keep a candidate only when it is compatible with all hard constraints. + try: + solved_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=[ + *inequality_constraints, + separation_constraint, + ], + equality_constraints=equality_constraints, + config=config, + ) + except _LayoutInfeasibleError: + continue + inequality_constraints.append(separation_constraint) + current = solved_xy_by_id + seen.add(key) + added += 1 + break + else: + raise ValueError( + "Parent-child AABB pair has no feasible separation direction: " + f"{first_id!r}, {second_id!r}." + ) if not added: break - current = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=current, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) raise ValueError("Parent-child AABB collisions remain after layout refinement.") @@ -462,7 +491,7 @@ def _root_aabb_overlaps( return sorted(result, reverse=True) -def _aabb_separation_constraint( +def _aabb_separation_constraints( *, root_ids: list[str], first_id: str, @@ -470,19 +499,47 @@ def _aabb_separation_constraint( half_extents: dict[str, np.ndarray], xy_by_id: dict[str, list[float]], margin: float, -) -> tuple[np.ndarray, float]: - """Return one least-penetration AABB separation inequality.""" +) -> list[tuple[np.ndarray, float]]: + """Return ordered feasible-direction candidates for one overlapping AABB pair.""" first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) overlap = np.minimum( first + half_extents[first_id], second + half_extents[second_id] ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) - axis = int(np.argmin(overlap)) - lower = first[axis] < second[axis] or ( - first[axis] == second[axis] and first_id < second_id - ) + axes = np.argsort(overlap) + constraints = [] + for axis in axes: + current_order = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + for first_is_lower in (current_order, not current_order): + constraints.append( + _aabb_separation_constraint_for_direction( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=half_extents, + axis=int(axis), + first_is_lower=first_is_lower, + margin=margin, + ) + ) + return constraints + + +def _aabb_separation_constraint_for_direction( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + axis: int, + first_is_lower: bool, + margin: float, +) -> tuple[np.ndarray, float]: + """Return one directed AABB separation inequality on a selected axis.""" index = {root_id: i for i, root_id in enumerate(root_ids)} row = np.zeros(2 * len(root_ids)) - sign = 1.0 if lower else -1.0 + sign = 1.0 if first_is_lower else -1.0 row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign return row, -float( half_extents[first_id][axis] + half_extents[second_id][axis] + margin diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py index 38fe7e6c4..9d874f239 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -186,6 +186,10 @@ def optimize( ) +class _LayoutInfeasibleError(ValueError): + """Internal marker for an SLSQP failure while testing one collision direction.""" + + def _build_constraints( *, problem: TableSurfaceLayoutProblem, @@ -424,7 +428,9 @@ def objective(values: np.ndarray) -> float: }, ) if not result.success: - raise ValueError(f"Table layout optimization failed: {result.message}") + raise _LayoutInfeasibleError( + f"Table layout optimization failed: {result.message}" + ) return { root_id: [float(result.x[2 * index]), float(result.x[2 * index + 1])] for index, root_id in enumerate(root_ids) @@ -463,29 +469,51 @@ def _refine_root_collisions( key = tuple(sorted((first_id, second_id))) if key in seen: continue - # Add a new SLSQP constraint to separate this overlapping pair. - inequality_constraints.append( - _aabb_separation_constraint( + # Earlier pair updates may already have separated this stale overlap. + if key not in { + tuple(sorted((first, second))) + for _, first, second in _root_aabb_overlaps( root_ids=root_ids, - first_id=first_id, - second_id=second_id, half_extents=root_half_extents_xy, xy_by_id=current, - margin=config.collision_margin_m, ) - ) - seen.add(key) - added += 1 + }: + continue + for separation_constraint in _aabb_separation_constraints( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=root_half_extents_xy, + xy_by_id=current, + margin=config.collision_margin_m, + ): + # Keep a candidate only when it is compatible with all hard constraints. + try: + solved_xy_by_id = _solve_root_xy( + root_ids=root_ids, + root_seed_xy_by_id=current, + imported_root_ids=imported_root_ids, + inequality_constraints=[ + *inequality_constraints, + separation_constraint, + ], + equality_constraints=equality_constraints, + config=config, + ) + except _LayoutInfeasibleError: + continue + inequality_constraints.append(separation_constraint) + current = solved_xy_by_id + seen.add(key) + added += 1 + break + else: + raise ValueError( + "Table-root AABB pair has no feasible separation direction: " + f"{first_id!r}, {second_id!r}." + ) if not added: break - current = _solve_root_xy( - root_ids=root_ids, - root_seed_xy_by_id=current, - imported_root_ids=imported_root_ids, - inequality_constraints=inequality_constraints, - equality_constraints=equality_constraints, - config=config, - ) raise ValueError("Table-root AABB collisions remain after layout refinement.") @@ -511,7 +539,7 @@ def _root_aabb_overlaps( return sorted(result, reverse=True) -def _aabb_separation_constraint( +def _aabb_separation_constraints( *, root_ids: list[str], first_id: str, @@ -519,22 +547,50 @@ def _aabb_separation_constraint( half_extents: dict[str, np.ndarray], xy_by_id: dict[str, list[float]], margin: float, -) -> tuple[np.ndarray, float]: +) -> list[tuple[np.ndarray, float]]: + """Return ordered feasible-direction candidates for one overlapping AABB pair.""" first, second = np.asarray(xy_by_id[first_id]), np.asarray(xy_by_id[second_id]) # Positive overlap on both axes means these two center-based AABBs intersect. overlap = np.minimum( first + half_extents[first_id], second + half_extents[second_id] ) - np.maximum(first - half_extents[first_id], second - half_extents[second_id]) - # Separate along the least-penetrating axis to require the smallest local shift. - axis = int(np.argmin(overlap)) - # Preserve the current order on that axis; object IDs break an exact tie deterministically. - lower = first[axis] < second[axis] or ( - first[axis] == second[axis] and first_id < second_id - ) + # Try the least-penetrating axis first, but permit order reversal if required. + axes = np.argsort(overlap) + constraints = [] + for axis in axes: + current_order = first[axis] < second[axis] or ( + first[axis] == second[axis] and first_id < second_id + ) + for first_is_lower in (current_order, not current_order): + constraints.append( + _aabb_separation_constraint_for_direction( + root_ids=root_ids, + first_id=first_id, + second_id=second_id, + half_extents=half_extents, + axis=int(axis), + first_is_lower=first_is_lower, + margin=margin, + ) + ) + return constraints + + +def _aabb_separation_constraint_for_direction( + *, + root_ids: list[str], + first_id: str, + second_id: str, + half_extents: dict[str, np.ndarray], + axis: int, + first_is_lower: bool, + margin: float, +) -> tuple[np.ndarray, float]: + """Return one directed AABB separation inequality on a selected axis.""" index = {root_id: i for i, root_id in enumerate(root_ids)} # One row addresses the x/y variable pair of each root in the flattened solver vector. row = np.zeros(2 * len(root_ids)) - sign = 1.0 if lower else -1.0 + sign = 1.0 if first_is_lower else -1.0 row[2 * index[first_id] + axis], row[2 * index[second_id] + axis] = sign, -sign # row @ values <= bound keeps the selected AABB faces apart by the requested margin. return row, -float( diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 39a44ecc7..6698d5c83 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -50,6 +50,10 @@ _OVERLAPPING_CENTER_XY = [0.0, 0.0] _ASSET_SIDE_LENGTH_M = 0.2 _RELATION_CLEARANCE_M = 0.03 +_COLLISION_MARGIN_M = 0.02 +_BOARD_XY_SIZE_M = 0.6 +_CAN_XY_SIZE_M = 0.1 +_PENCIL_XY_SIZE_M = [0.04, 0.2] def _asset( @@ -235,6 +239,73 @@ def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( assert np.max(np.abs(solved_xy_by_id[variable_id])) >= _ASSET_SIDE_LENGTH_M - 1e-6 +def test_table_optimizer_can_reverse_a_collision_order(tmp_path: Path) -> None: + board_glb = tmp_path / "board.glb" + can_glb = tmp_path / "can.glb" + pencil_glb = tmp_path / "pencil.glb" + # SimReady GLBs are y-up, so z-up XY uses the source XZ extents. + trimesh.creation.box( + extents=[_BOARD_XY_SIZE_M, _ASSET_SIDE_LENGTH_M, _BOARD_XY_SIZE_M] + ).export(board_glb) + trimesh.creation.box( + extents=[_CAN_XY_SIZE_M, _ASSET_SIDE_LENGTH_M, _CAN_XY_SIZE_M] + ).export(can_glb) + trimesh.creation.box( + extents=[ + _PENCIL_XY_SIZE_M[0], + _ASSET_SIDE_LENGTH_M, + _PENCIL_XY_SIZE_M[1], + ] + ).export(pencil_glb) + board_id, pencil_id, can_id = "board_001", "pencil_001", "can_001" + board_xy, pencil_xy, can_xy = [0.0, 0.0], [-0.2, 0.0], [-0.4, 0.0] + + solved_xy_by_id = TableSurfaceLayoutOptimizer().optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + board_id: _asset( + object_id=board_id, + glb_path=board_glb, + center_xy=board_xy, + ), + pencil_id: _asset(object_id=pencil_id, glb_path=pencil_glb), + can_id: _asset( + object_id=can_id, + glb_path=can_glb, + center_xy=can_xy, + ), + }, + root_ids=[board_id, pencil_id, can_id], + root_seed_xy_by_id={ + board_id: board_xy, + pencil_id: pencil_xy, + can_id: can_xy, + }, + imported_root_ids={board_id, can_id}, + fixed_root_xy_by_id={ + board_id: board_xy, + pencil_id: None, + can_id: can_xy, + }, + root_table_regions_by_id={ + board_id: None, + pencil_id: None, + can_id: None, + }, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + ) + + expected_pencil_x_upper_bound = ( + can_xy[0] + - _CAN_XY_SIZE_M / 2.0 + - _PENCIL_XY_SIZE_M[0] / 2.0 + - _COLLISION_MARGIN_M + ) + assert solved_xy_by_id[pencil_id][0] <= expected_pencil_x_upper_bound + 1e-6 + + def test_parent_optimizer_applies_sibling_planar_relation(tmp_path: Path) -> None: asset_glb = tmp_path / "asset.glb" trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) From a03cda1a87b32d36a819bc7b463d1cbff6cb13a5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Wed, 19 Aug 2026 18:13:15 +0800 Subject: [PATCH 29/85] replace the assets gravity settler with gravity settler --- .../pipeline/generation/scene_generation.py | 40 +- .../pipeline/utils/assets_gravity_settler.py | 349 ----------------- .../pipeline/utils/gravity_settler.py | 356 ++++++++++++++++++ .../scene_engine/test_gravity_settler.py | 81 ++++ 4 files changed, 465 insertions(+), 361 deletions(-) delete mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py create mode 100644 tests/gen_sim/scene_engine/test_gravity_settler.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 0e1b17f95..7010e5f14 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -44,8 +44,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( AssetsSupportLayoutOptimizer, ) -from embodichain.gen_sim.scene_engine.pipeline.utils.assets_gravity_settler import ( - AssetsGravitySettler, +from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( + GravitySettleBody, + GravitySettler, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, @@ -455,16 +456,31 @@ def _layout_refinement( refined_assets_layout = overlap_optimizer.optimize() overlap_optimizer.save_overlap_optimization_debug_images() - # 8. Gravity simulation, to let all the assets to be stable and placed well on the table's support surface. - # Notice that: we do not consider the assets like a bottle, which should be standing on the table but laid down - # after the simulation. - gravity_settler = AssetsGravitySettler( - scene=scene, - table_layout=refined_table_layout, - assets_layout=refined_assets_layout, - geometry_root=simready_geometry_output_root, - ) - refined_assets_layout = gravity_settler.settle() + # 8. The initial image graph has one on-table level, so every asset settles + # dynamically against the table in this first generic gravity pass. + assets_by_id = {asset.id: asset for asset in scene.assets} + # All the assets are dynamic; the table is static. + settled_pose_by_id = GravitySettler( + table_body=GravitySettleBody( + scene_object=scene.table, + y_up_layout=refined_table_layout, + ), + participant_bodies=[ + GravitySettleBody( + scene_object=assets_by_id[str(asset_layout["id"])], + y_up_layout=asset_layout, + ) + for asset_layout in refined_assets_layout + ], + dynamic_asset_ids=set(assets_by_id), + static_asset_ids=set(), + ).settle() + # Update. + for asset_layout in refined_assets_layout: + asset_id = str(asset_layout["id"]) + settled_pose = settled_pose_by_id[asset_id] + asset_layout["pos"] = settled_pose["pos"] + asset_layout["rot"] = settled_pose["rot"] # Update the scene data structure with the final layout and spatial metadata. _update_scene_final_y_up_layout_and_z_up_centers( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py deleted file mode 100644 index 31e9e8443..000000000 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/assets_gravity_settler.py +++ /dev/null @@ -1,349 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Sequence - -import numpy as np -from scipy.spatial.transform import Rotation -import trimesh - -from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_object import ( - ObjectPhysics, - SceneObject, -) -from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( - layout_object_to_transform_matrix, - load_glb_mesh, - transform_matrix_to_layout_object, -) -from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg -from embodichain.lab.sim.shapes import MeshCfg -from embodichain.utils.logger import log_info - - -@dataclass(frozen=True) -class AssetsGravitySettlerConfig: - """Physics controls for table-top asset settling.""" - - clearance_m: float = 0.02 # Initial gap between each asset and the table top. - settle_steps: int = 300 # Fixed number of simulator steps to execute. - physics_dt: float = 1.0 / 100.0 # Physics timestep in seconds. - sim_device: str = "cpu" # Simulation device requested from EmbodiChain Lab. - - -class AssetsGravitySettler: - """Settle all assets together on one kinematic table in a z-up simulation.""" - - def __init__( - self, - *, - scene: Scene, - table_layout: dict[str, object], - assets_layout: list[dict[str, object]], - geometry_root: str | Path, - config: AssetsGravitySettlerConfig | None = None, - ) -> None: - self.scene = scene - self.table_layout = table_layout - self.assets_layout = assets_layout - self.geometry_root = Path(geometry_root).expanduser().resolve() - self.settled_assets_layout: list[dict[str, object]] | None = None - self.config = config if config is not None else AssetsGravitySettlerConfig() - # Check. - if self.config.clearance_m < 0.0: - raise ValueError("Gravity-settle clearance_m must be non-negative.") - if self.config.settle_steps <= 0: - raise ValueError("Gravity-settle settle_steps must be positive.") - if self.config.physics_dt <= 0.0: - raise ValueError("Gravity-settle physics_dt must be positive.") - - def settle(self) -> list[dict[str, object]]: - """Run gravity settling and return the resulting y-up asset layouts.""" - self.settled_assets_layout = None - if not self.assets_layout: - self.settled_assets_layout = [] - log_info("Scene has no movable assets; skipping gravity settling.") - return self.settled_assets_layout - - table_id = self._require_layout_id(self.table_layout, name="Table") - table_object = self._require_scene_object(table_id, kind="table") - asset_ids: set[str] = set() - asset_objects_by_id: dict[str, SceneObject] = {} - for asset_layout in self.assets_layout: - asset_id = self._require_layout_id(asset_layout, name="Asset") - if asset_id in asset_ids: - raise ValueError(f"Asset layouts contain duplicate id {asset_id!r}.") - asset_ids.add(asset_id) - asset_objects_by_id[asset_id] = self._require_scene_object( - asset_id, kind="asset" - ) - expected_asset_ids = {asset.id for asset in self.scene.assets} - if asset_ids != expected_asset_ids: - raise ValueError( - "Gravity-settle layouts must contain exactly the scene asset ids." - ) - - y_up_to_z_up_matrix = np.eye(4) - y_up_to_z_up_matrix[:3, :3] = np.array( - [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] - ) - z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) - table_info = self._prepare_sim_body( - layout_object=self.table_layout, - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - table_world_mesh = self._mesh_to_z_up_world_for_aabb( - y_up_mesh=table_info["mesh"], - z_up_rigid_layout=table_info["rigid_layout"], - z_up_scale=table_info["z_up_scale"], - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - table_top_z = float(table_world_mesh.bounds[1, 2]) - - prepared_assets: dict[str, dict[str, object]] = {} - for asset_layout in self.assets_layout: - asset_id = str(asset_layout["id"]) - asset_info = self._prepare_sim_body( - layout_object=asset_layout, - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - asset_world_mesh = self._mesh_to_z_up_world_for_aabb( - y_up_mesh=asset_info["mesh"], - z_up_rigid_layout=asset_info["rigid_layout"], - z_up_scale=asset_info["z_up_scale"], - y_up_to_z_up_matrix=y_up_to_z_up_matrix, - ) - asset_bottom_z = float(asset_world_mesh.bounds[0, 2]) - asset_info["rigid_layout"]["pos"][2] += ( - table_top_z + self.config.clearance_m - asset_bottom_z - ) - prepared_assets[asset_id] = asset_info - - log_info( - "Gravity settling started: " - f"assets={len(prepared_assets)}, steps={self.config.settle_steps}, " - f"physics_dt={self.config.physics_dt:.4f} s." - ) - sim = SimulationManager( - SimulationManagerCfg( - headless=True, - physics_dt=self.config.physics_dt, - sim_device=self.config.sim_device, - ) - ) - try: - # Add table. - sim.add_rigid_object( - RigidObjectCfg( - uid=table_id, - shape=MeshCfg(fpath=str(table_info["mesh_path"])), - init_pos=tuple(table_info["rigid_layout"]["pos"]), - init_rot=tuple( - self._simulation_euler_xyz_degrees(table_info["rigid_layout"]) - ), - body_scale=tuple(table_info["y_up_scale"]), - attrs=self._rigid_body_attrs(table_object.physics), - body_type=table_object.physics.body_type, - max_convex_hull_num=table_object.physics.max_convex_hull_num, - acd_method="vhacd", - ) - ) - # Add assets. - simulated_assets: dict[str, object] = {} - for asset_id, asset_info in prepared_assets.items(): - rigid_layout = asset_info["rigid_layout"] - simulated_assets[asset_id] = sim.add_rigid_object( - RigidObjectCfg( - uid=asset_id, - shape=MeshCfg(fpath=str(asset_info["mesh_path"])), - init_pos=tuple(rigid_layout["pos"]), - init_rot=tuple( - self._simulation_euler_xyz_degrees(rigid_layout) - ), - body_scale=tuple(asset_info["y_up_scale"]), - attrs=self._rigid_body_attrs( - asset_objects_by_id[asset_id].physics - ), - body_type=asset_objects_by_id[asset_id].physics.body_type, - max_convex_hull_num=( - asset_objects_by_id[asset_id].physics.max_convex_hull_num - ), - acd_method="vhacd", - ) - ) - # Run simulation to settle all assets. - sim.update(step=self.config.settle_steps) - - # Update the final layouts. - settled_layout_by_id: dict[str, dict[str, object]] = {} - for asset_id, simulated_asset in simulated_assets.items(): - final_rigid_pose_z_up = np.asarray( - simulated_asset.get_local_pose(to_matrix=True)[0] - .detach() - .cpu() - .numpy(), - dtype=float, - ) - scale_matrix = np.eye(4) - scale_matrix[:3, :3] = np.diag(prepared_assets[asset_id]["z_up_scale"]) - final_z_up_layout_matrix = final_rigid_pose_z_up @ scale_matrix - settled_layout_by_id[asset_id] = transform_matrix_to_layout_object( - asset_id, - z_up_to_y_up_matrix - @ final_z_up_layout_matrix - @ y_up_to_z_up_matrix, - ) - finally: - # Release resources. - sim.destroy(exit_process=False) - SimulationManager.flush_cleanup_queue() - - self.settled_assets_layout = [ - settled_layout_by_id[str(asset_layout["id"])] - for asset_layout in self.assets_layout - ] - log_info("Gravity settling completed for all assets.") - return self.settled_assets_layout - - def _prepare_sim_body( - self, - *, - layout_object: dict[str, object], - y_up_to_z_up_matrix: np.ndarray, - ) -> dict[str, object]: - """Load one y-up GLB and prepare its z-up simulation pose.""" - object_id = self._require_layout_id(layout_object, name="Layout object") - source_mesh_path = self.geometry_root / f"{object_id}.glb" - source_mesh = load_glb_mesh(source_mesh_path) - z_up_layout = self._convert_layout_coordinate_system( - layout_object, - source_to_target_matrix=y_up_to_z_up_matrix, - ) - return { - "mesh_path": source_mesh_path, - "mesh": source_mesh, - "rigid_layout": { - "id": object_id, - "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), - "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), - "scale": [1.0, 1.0, 1.0], - }, - "y_up_scale": self._three_floats( - layout_object.get("scale"), field_name="scale" - ), - "z_up_scale": self._three_floats( - z_up_layout.get("scale"), field_name="scale" - ), - } - - def _require_scene_object(self, object_id: str, *, kind: str) -> SceneObject: - """Return one physics-ready scene object with the expected semantic kind.""" - matching_objects = [ - scene_object - for scene_object in self.scene.objects - if scene_object.id == object_id - ] - if len(matching_objects) != 1: - raise ValueError( - f"Gravity settling requires exactly one scene object {object_id!r}." - ) - scene_object = matching_objects[0] - if scene_object.kind != kind: - raise ValueError( - f"Scene object {object_id!r} must have kind {kind!r} before " - "gravity settling." - ) - if scene_object.physics is None: - raise ValueError( - f"Scene object {object_id!r} has no SimReady physics settings." - ) - return scene_object - - @staticmethod - def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: - """Convert persisted SceneObject physics attributes into Lab config.""" - if physics is None: - raise ValueError("Gravity settling requires SimReady physics settings.") - return RigidBodyAttributesCfg(**physics.attrs) - - @staticmethod - def _mesh_to_z_up_world_for_aabb( - *, - y_up_mesh: trimesh.Trimesh, - z_up_rigid_layout: dict[str, object], - z_up_scale: Sequence[float], - y_up_to_z_up_matrix: np.ndarray, - ) -> trimesh.Trimesh: - """Transform a y-up mesh into its z-up world pose for AABB measurement.""" - mesh = y_up_mesh.copy() - mesh.apply_transform(y_up_to_z_up_matrix) - scale_matrix = np.eye(4) - scale_matrix[:3, :3] = np.diag(z_up_scale) - mesh.apply_transform(scale_matrix) - mesh.apply_transform(layout_object_to_transform_matrix(z_up_rigid_layout)) - return mesh - - @staticmethod - def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: - """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" - layout_rotation = Rotation.from_euler( - "xyz", - AssetsGravitySettler._three_floats( - layout_object.get("rot"), field_name="rot" - ), - degrees=True, - ) - return layout_rotation.as_euler("XYZ", degrees=True).tolist() - - @staticmethod - def _convert_layout_coordinate_system( - layout_object: dict[str, object], - *, - source_to_target_matrix: np.ndarray, - ) -> dict[str, object]: - """Convert one layout object between coordinate frames through its matrix.""" - return transform_matrix_to_layout_object( - str(layout_object["id"]), - source_to_target_matrix - @ layout_object_to_transform_matrix(layout_object) - @ np.linalg.inv(source_to_target_matrix), - ) - - @staticmethod - def _require_layout_id(layout_object: dict[str, object], *, name: str) -> str: - """Check id.""" - object_id = layout_object.get("id") - if not isinstance(object_id, str) or not object_id: - raise ValueError(f"{name} layout must contain a non-empty string id.") - return object_id - - @staticmethod - def _three_floats(value: object, *, field_name: str) -> list[float]: - """Check three values.""" - if not isinstance(value, list) or len(value) != 3: - raise ValueError(f"Layout field {field_name} must contain three values.") - try: - return [float(item) for item in value] - except (TypeError, ValueError) as exc: - raise ValueError( - f"Layout field {field_name} must contain numeric values." - ) from exc diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py new file mode 100644 index 000000000..b5078c75f --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py @@ -0,0 +1,356 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( + layout_object_to_transform_matrix, + transform_matrix_to_layout_object, +) +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg +from embodichain.utils.logger import log_info + + +@dataclass(frozen=True) +class GravitySettlerConfig: + """Physics controls for one caller-defined gravity-settlement pass.""" + + settle_steps: int = 300 + physics_dt: float = 1.0 / 100.0 + sim_device: str = "cpu" + + def __post_init__(self) -> None: + """Reject invalid numerical controls before starting a simulation.""" + if self.settle_steps <= 0: + raise ValueError("Gravity-settle settle_steps must be positive.") + if self.physics_dt <= 0.0: + raise ValueError("Gravity-settle physics_dt must be positive.") + + +@dataclass(frozen=True) +class GravitySettleBody: + """One scene object and its latest complete y-up pipeline layout.""" + + scene_object: SceneObject + y_up_layout: dict[str, object] + + +class GravitySettler: + """Settle caller-selected dynamic assets against a mandatory table body. + + All supplied layouts use Scene Engine's y-up pipeline convention. The + settler converts them to z-up only at the simulation boundary. The caller + explicitly classifies every participant as dynamic or static; static + participants and the table are kinematic collision bodies. + """ + + def __init__( + self, + *, + table_body: GravitySettleBody, + participant_bodies: list[GravitySettleBody], + dynamic_asset_ids: set[str], + static_asset_ids: set[str], + config: GravitySettlerConfig | None = None, + ) -> None: + self.table_body = table_body + self.participant_bodies = participant_bodies + self.dynamic_asset_ids = set(dynamic_asset_ids) + self.static_asset_ids = set(static_asset_ids) + self.config = config if config is not None else GravitySettlerConfig() + + def settle(self) -> dict[str, dict[str, list[float]]]: + """Return final y-up poses for dynamic participants only. + + Input layouts are used as-is. Placement clearance and support-surface + alignment remain the responsibility of the calling layout optimizer. + Static participants and every object's scale are unchanged, so they are + deliberately omitted from the result. + """ + # Check table. + table = self.table_body.scene_object + if table.kind != "table": + raise ValueError("Gravity settling requires a table body.") + table_id = self._require_body_layout_id(self.table_body, name="Table") + + participant_bodies_by_id: dict[str, GravitySettleBody] = {} + for participant_body in self.participant_bodies: + asset_id = self._require_body_layout_id( + participant_body, name="Participant asset" + ) + if asset_id == table_id: + raise ValueError( + "Gravity-settle participants cannot include the table." + ) + if participant_body.scene_object.kind != "asset": + raise ValueError( + f"Gravity-settle participant {asset_id!r} must be an asset body." + ) + if asset_id in participant_bodies_by_id: + raise ValueError( + f"Gravity-settle participant assets repeat id {asset_id!r}." + ) + participant_bodies_by_id[asset_id] = participant_body + + participant_ids = set(participant_bodies_by_id) + classified_ids = self.dynamic_asset_ids | self.static_asset_ids + if self.dynamic_asset_ids & self.static_asset_ids: + raise ValueError( + "Gravity-settle dynamic and static asset IDs must not overlap." + ) + if classified_ids != participant_ids: + raise ValueError( + "Gravity-settle dynamic and static asset IDs must exactly match " + f"participants; participants={sorted(participant_ids)}, " + f"classified={sorted(classified_ids)}." + ) + if not self.dynamic_asset_ids: + log_info("Gravity settle has no dynamic participants; skipping simulation.") + return {} + + y_up_to_z_up_matrix = self._y_up_to_z_up_matrix() + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + table_info = self._prepare_sim_body( + body=self.table_body, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + participant_infos_by_id = { + asset_id: self._prepare_sim_body( + body=participant_body, + y_up_to_z_up_matrix=y_up_to_z_up_matrix, + ) + for asset_id, participant_body in participant_bodies_by_id.items() + } + + log_info( + "Gravity settling started: " + f"dynamic_assets={len(self.dynamic_asset_ids)}, " + f"kinematic_assets={len(participant_infos_by_id) - len(self.dynamic_asset_ids)}, " + f"steps={self.config.settle_steps}, " + f"physics_dt={self.config.physics_dt:.4f} s." + ) + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + physics_dt=self.config.physics_dt, + sim_device=self.config.sim_device, + ) + ) + try: + self._add_sim_body( + sim=sim, + object_id=table_id, + body_info=table_info, + physics=table.physics, + body_type="kinematic", + ) + simulated_assets: dict[str, object] = {} + for asset_id, asset_info in participant_infos_by_id.items(): + simulated_assets[asset_id] = self._add_sim_body( + sim=sim, + object_id=asset_id, + body_info=asset_info, + physics=participant_bodies_by_id[asset_id].scene_object.physics, + body_type=( + "dynamic" if asset_id in self.dynamic_asset_ids else "kinematic" + ), + ) + sim.update(step=self.config.settle_steps) + + settled_pose_by_id: dict[str, dict[str, list[float]]] = {} + for asset_id in self.dynamic_asset_ids: + simulated_asset = simulated_assets[asset_id] + final_rigid_pose_z_up = np.asarray( + simulated_asset.get_local_pose(to_matrix=True)[0] + .detach() + .cpu() + .numpy(), + dtype=float, + ) + scale_matrix = np.eye(4) + scale_matrix[:3, :3] = np.diag( + participant_infos_by_id[asset_id]["z_up_scale"] + ) + final_y_up_layout = transform_matrix_to_layout_object( + asset_id, + z_up_to_y_up_matrix + @ final_rigid_pose_z_up + @ scale_matrix + @ y_up_to_z_up_matrix, + ) + settled_pose_by_id[asset_id] = { + "pos": self._three_floats( + final_y_up_layout.get("pos"), field_name="pos" + ), + "rot": self._three_floats( + final_y_up_layout.get("rot"), field_name="rot" + ), + } + finally: + sim.destroy(exit_process=False) + SimulationManager.flush_cleanup_queue() + + log_info("Gravity settling completed for participating assets.") + return settled_pose_by_id + + def _prepare_sim_body( + self, + *, + body: GravitySettleBody, + y_up_to_z_up_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one supplied y-up layout into a simulator body pose.""" + scene_object = body.scene_object + if scene_object.simready_glb_path is None: + raise ValueError( + f"Gravity-settle object {scene_object.id!r} has no SimReady GLB path." + ) + mesh_path = Path(scene_object.simready_glb_path).expanduser().resolve() + if not mesh_path.is_file(): + raise FileNotFoundError( + f"Gravity-settle GLB for {scene_object.id!r} not found: {mesh_path}" + ) + y_up_layout = body.y_up_layout + z_up_layout = self._convert_layout_coordinate_system( + y_up_layout, + source_to_target_matrix=y_up_to_z_up_matrix, + ) + return { + "mesh_path": mesh_path, + "rigid_layout": { + "id": scene_object.id, + "rot": self._three_floats(z_up_layout.get("rot"), field_name="rot"), + "pos": self._three_floats(z_up_layout.get("pos"), field_name="pos"), + "scale": [1.0, 1.0, 1.0], + }, + "y_up_scale": self._three_floats( + y_up_layout.get("scale"), field_name="scale" + ), + "z_up_scale": self._three_floats( + z_up_layout.get("scale"), field_name="scale" + ), + } + + def _add_sim_body( + self, + *, + sim: SimulationManager, + object_id: str, + body_info: dict[str, object], + physics: ObjectPhysics | None, + body_type: str, + ) -> object: + """Add one supplied body with a pass-specific dynamic or kinematic type.""" + rigid_layout = body_info["rigid_layout"] + if not isinstance(rigid_layout, dict): + raise ValueError("Gravity-settle body has invalid rigid layout.") + return sim.add_rigid_object( + RigidObjectCfg( + uid=object_id, + shape=MeshCfg(fpath=str(body_info["mesh_path"])), + init_pos=tuple( + self._three_floats(rigid_layout.get("pos"), field_name="pos") + ), + init_rot=tuple(self._simulation_euler_xyz_degrees(rigid_layout)), + body_scale=tuple( + self._three_floats(body_info["y_up_scale"], field_name="scale") + ), + attrs=self._rigid_body_attrs(physics), + body_type=body_type, + max_convex_hull_num=self._max_convex_hull_num(physics), + acd_method="vhacd", + ) + ) + + @staticmethod + def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: + """Convert persisted collision material data into one Lab config.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return RigidBodyAttributesCfg(**physics.attrs) + + @staticmethod + def _max_convex_hull_num(physics: ObjectPhysics | None) -> int: + """Read the persisted collision-hull budget after validating physics.""" + if physics is None: + raise ValueError("Gravity settling requires SimReady physics settings.") + return physics.max_convex_hull_num + + @staticmethod + def _require_body_layout_id(body: GravitySettleBody, *, name: str) -> str: + """Validate that a body layout belongs to its scene object.""" + object_id = body.y_up_layout.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{name} layout requires a non-empty string id.") + if object_id != body.scene_object.id: + raise ValueError( + f"{name} layout id {object_id!r} does not match its scene object." + ) + return object_id + + @staticmethod + def _three_floats(value: object, *, field_name: str) -> list[float]: + """Return three finite layout values as Python floats.""" + if not isinstance(value, (list, tuple)) or len(value) != 3: + raise ValueError(f"Gravity-settle {field_name} must contain three values.") + result = [float(component) for component in value] + if not np.all(np.isfinite(result)): + raise ValueError(f"Gravity-settle {field_name} must contain finite values.") + return result + + @staticmethod + def _simulation_euler_xyz_degrees(layout_object: dict[str, object]) -> list[float]: + """Convert lowercase-xyz layout rotation to SimulationManager's XYZ order.""" + layout_rotation = Rotation.from_euler( + "xyz", + GravitySettler._three_floats(layout_object.get("rot"), field_name="rot"), + degrees=True, + ) + return layout_rotation.as_euler("XYZ", degrees=True).tolist() + + @staticmethod + def _convert_layout_coordinate_system( + layout_object: dict[str, object], + *, + source_to_target_matrix: np.ndarray, + ) -> dict[str, object]: + """Convert one complete layout through the y-up/z-up basis change.""" + return transform_matrix_to_layout_object( + str(layout_object["id"]), + source_to_target_matrix + @ layout_object_to_transform_matrix(layout_object) + @ np.linalg.inv(source_to_target_matrix), + ) + + @staticmethod + def _y_up_to_z_up_matrix() -> np.ndarray: + """Return the coordinate conversion used by Scene Engine layouts.""" + matrix = np.eye(4) + matrix[:3, :3] = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]]) + return matrix diff --git a/tests/gen_sim/scene_engine/test_gravity_settler.py b/tests/gen_sim/scene_engine/test_gravity_settler.py new file mode 100644 index 000000000..5829fa39a --- /dev/null +++ b/tests/gen_sim/scene_engine/test_gravity_settler.py @@ -0,0 +1,81 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( + GravitySettleBody, + GravitySettler, +) + +_TABLE_ID = "table" +_ASSET_ID = "cube_001" +_IDENTITY_LAYOUT = { + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], +} + + +def _table_body() -> GravitySettleBody: + return GravitySettleBody( + scene_object=SceneObject( + id=_TABLE_ID, + kind="table", + category="table", + name="table", + description="table", + ), + y_up_layout={"id": _TABLE_ID, **_IDENTITY_LAYOUT}, + ) + + +def _asset_body() -> GravitySettleBody: + return GravitySettleBody( + scene_object=SceneObject( + id=_ASSET_ID, + kind="asset", + category="cube", + name="cube", + description="cube", + ), + y_up_layout={"id": _ASSET_ID, **_IDENTITY_LAYOUT}, + ) + + +def test_gravity_settler_returns_no_poses_without_dynamic_assets() -> None: + settled_pose_by_id = GravitySettler( + table_body=_table_body(), + participant_bodies=[_asset_body()], + dynamic_asset_ids=set(), + static_asset_ids={_ASSET_ID}, + ).settle() + + assert settled_pose_by_id == {} + + +def test_gravity_settler_rejects_dynamic_assets_outside_participants() -> None: + with pytest.raises(ValueError, match="exactly match participants"): + GravitySettler( + table_body=_table_body(), + participant_bodies=[], + dynamic_asset_ids={_ASSET_ID}, + static_asset_ids=set(), + ).settle() From 6f8e26ea8b9e58c2e91e4b50383b0f4070336bf0 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:06:00 +0800 Subject: [PATCH 30/85] use gravity settler in parent surface layout --- .../utils/parent_surface_layout_optimizer.py | 45 ++++++++ .../utils/scene_layout_constructor.py | 101 ++++++++++++++++- .../pipeline/utils/scene_layout_utils.py | 10 ++ .../test_scene_layout_optimizer.py | 104 +++++++++++++++++- 4 files changed, 256 insertions(+), 4 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py index 6d2a3b0a5..de8514dee 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/parent_surface_layout_optimizer.py @@ -25,9 +25,14 @@ from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraphRelation from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( + GravitySettleBody, + GravitySettler, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( load_scene_object_z_up_mesh, measure_scene_object_z_up_world_aabb, + scene_object_y_up_layout, ) if TYPE_CHECKING: @@ -185,6 +190,46 @@ def optimize( config=self.config, ) + def settle_dynamic_children( + self, + *, + table: SceneObject, + parent: SceneObject, + problem: ParentSurfaceLayoutProblem, + dynamic_child_ids: set[str], + ) -> dict[str, dict[str, list[float]]]: + """Settle variable children against their parent and fixed siblings. + + Children must already have been placed 2cm above ``parent`` by the + caller. The physical table remains mandatory, while the parent and + unedited siblings become kinematic collision bodies for this pass. + """ + child_ids = set(problem.child_ids) + if not dynamic_child_ids.issubset(child_ids): + raise ValueError("Only parent-surface children may be dynamic.") + if not dynamic_child_ids: + return {} + participant_ids = {parent.id, *child_ids} + if parent.id in child_ids: + raise ValueError("A parent-surface group cannot contain its parent.") + return GravitySettler( + table_body=GravitySettleBody( + scene_object=table, + y_up_layout=scene_object_y_up_layout(table), + ), + participant_bodies=[ + GravitySettleBody( + scene_object=problem.assets_by_id[object_id], + y_up_layout=scene_object_y_up_layout( + problem.assets_by_id[object_id] + ), + ) + for object_id in participant_ids + ], + dynamic_asset_ids=dynamic_child_ids, + static_asset_ids=participant_ids - dynamic_child_ids, + ).settle() + class _LayoutInfeasibleError(ValueError): """Internal marker for an SLSQP failure while testing one collision direction.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py index 95c77d712..b22a84c63 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_constructor.py @@ -31,7 +31,13 @@ ParentSurfaceLayoutOptimizerConfig, ParentSurfaceLayoutProblem, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( + GravitySettleBody, + GravitySettler, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + measure_scene_object_z_up_world_aabb, + scene_object_y_up_layout, translate_scene_object_y_up_by_z_up_delta, update_scene_object_y_up_pose_from_z_up_support, ) @@ -170,7 +176,7 @@ def _optimize_table_group( scene_object=assets_by_id[root_id], support_region_z=table.support_surface_z, center_xy=solved_xy, - clearance_m=0.00, # Directly place on the support surface. + clearance_m=0.02, # Lift dynamic roots before gravity settling. ) self._updated_object_ids.add(root_id) self._propagate_descendant_delta( @@ -178,6 +184,47 @@ def _optimize_table_group( root_id=root_id, delta_xy=delta_xy, ) + self._settle_table_group_dynamic_roots( + layout_problem=layout_problem, + group=group, + assets_by_id=assets_by_id, + ) + + def _settle_table_group_dynamic_roots( + self, + *, + layout_problem: SceneLayoutProblem, + group: SceneLayoutGroup, + assets_by_id: dict[str, SceneObject], + ) -> None: + """Settle edited table roots while their unedited siblings stay fixed.""" + dynamic_root_ids = set(group.child_ids) & layout_problem.layout_variable_ids + if not dynamic_root_ids: + return + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Table gravity settling requires a table.") + + settled_pose_by_id = GravitySettler( + table_body=GravitySettleBody( + scene_object=table, + y_up_layout=scene_object_y_up_layout(table), + ), + participant_bodies=[ + GravitySettleBody( + scene_object=assets_by_id[child_id], + y_up_layout=scene_object_y_up_layout(assets_by_id[child_id]), + ) + for child_id in group.child_ids + ], + dynamic_asset_ids=dynamic_root_ids, + static_asset_ids=set(group.child_ids) - dynamic_root_ids, + ).settle() + self._apply_settled_dynamic_poses( + scene=layout_problem.post_edit_scene, + assets_by_id=assets_by_id, + settled_pose_by_id=settled_pose_by_id, + ) def _propagate_descendant_delta( self, @@ -242,7 +289,7 @@ def _optimize_parent_group( scene_object=parent_surface_problem.assets_by_id[child_id], support_region_z=parent_surface_problem.parent_top_z, center_xy=solved_xy, - clearance_m=0.00, # Directly place on the parent's top surface. + clearance_m=0.02, # Leave the standard settling clearance. ) self._updated_object_ids.add(child_id) self._propagate_descendant_delta( @@ -250,6 +297,56 @@ def _optimize_parent_group( root_id=child_id, delta_xy=delta_xy, ) + table = layout_problem.post_edit_scene.table + if table is None: + raise ValueError("Parent gravity settling requires a table.") + dynamic_child_ids = set(group.child_ids) & layout_problem.layout_variable_ids + settled_pose_by_id = ( + self.parent_surface_layout_optimizer.settle_dynamic_children( + table=table, + parent=parent_surface_problem.assets_by_id[group.parent_id], + problem=parent_surface_problem, + dynamic_child_ids=dynamic_child_ids, + ) + ) + self._apply_settled_dynamic_poses( + scene=layout_problem.post_edit_scene, + assets_by_id=parent_surface_problem.assets_by_id, + settled_pose_by_id=settled_pose_by_id, + ) + + def _apply_settled_dynamic_poses( + self, + *, + scene: Scene, + assets_by_id: dict[str, SceneObject], + settled_pose_by_id: dict[str, dict[str, list[float]]], + ) -> None: + """Write dynamic y-up poses and preserve their descendants' XY frame.""" + for object_id, settled_pose in settled_pose_by_id.items(): + previous_center_xy = self._current_xy_by_id[object_id] + asset = assets_by_id[object_id] + asset.pos = settled_pose["pos"] + asset.rot = settled_pose["rot"] + # Later BFS groups must start from the settled z-up AABB center. + settled_aabb = measure_scene_object_z_up_world_aabb(scene_object=asset) + settled_center_xy = [ + (settled_aabb[0][0] + settled_aabb[1][0]) / 2.0, + (settled_aabb[0][1] + settled_aabb[1][1]) / 2.0, + ] + asset.center_xy = settled_center_xy + self._current_xy_by_id[object_id] = settled_center_xy + self._updated_object_ids.add(object_id) + if previous_center_xy is not None: + # Existing descendants retain their relative XY placement after settling. + self._propagate_descendant_delta( + scene=scene, + root_id=object_id, + delta_xy=[ + settled_center_xy[0] - previous_center_xy[0], + settled_center_xy[1] - previous_center_xy[1], + ], + ) def _build_problem(self) -> SceneLayoutProblem: """Build post-edit objects and preserve formal-scene centers as seeds.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py index 100f5b269..855342920 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py @@ -129,6 +129,16 @@ def y_up_to_z_up_matrix() -> np.ndarray: return matrix +def scene_object_y_up_layout(scene_object: SceneObject) -> dict[str, object]: + """Adapt one persisted SceneObject pose to a complete y-up layout object.""" + return { + "id": scene_object.id, + "rot": scene_object.rot, + "pos": scene_object.pos, + "scale": scene_object.scale, + } + + def two_floats(value: object, *, field_name: str) -> list[float]: """Validate and return one finite two-value vector.""" if not isinstance(value, (list, tuple)) or len(value) != 2: diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index 6698d5c83..bd5239c70 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -32,6 +32,8 @@ ParentSurfaceLayoutOptimizer, ParentSurfaceLayoutProblem, ) +import embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer as parent_surface_layout_optimizer_module +import embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor as scene_layout_constructor_module from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_constructor import ( SceneLayoutConstructor, ) @@ -54,6 +56,8 @@ _BOARD_XY_SIZE_M = 0.6 _CAN_XY_SIZE_M = 0.1 _PENCIL_XY_SIZE_M = [0.04, 0.2] +_SETTLED_DYNAMIC_POS_Y_UP = [0.25, 0.5, -0.4] +_SETTLED_DYNAMIC_ROT_Y_UP = [0.0, 0.0, 0.0] def _asset( @@ -98,6 +102,7 @@ def test_table_regions_put_front_at_larger_y() -> None: def test_layout_constructor_places_new_child_on_parent_top( tmp_path: Path, + monkeypatch, ) -> None: book_glb = tmp_path / "book.glb" cup_glb = tmp_path / "cup.glb" @@ -142,6 +147,25 @@ def test_layout_constructor_places_new_child_on_parent_top( ), ] ) + settler_inputs: dict[str, object] = {} + + class _FakeGravitySettler: + def __init__(self, **kwargs: object) -> None: + settler_inputs.update(kwargs) + + def settle(self) -> dict[str, dict[str, list[float]]]: + return { + "cup_001": { + "pos": [0.0, 0.74, 0.0], + "rot": [0.0, 0.0, 0.0], + } + } + + monkeypatch.setattr( + parent_surface_layout_optimizer_module, + "GravitySettler", + _FakeGravitySettler, + ) post_edit_scene = SceneLayoutConstructor( formal_scene=Scene(objects=[table, book]), @@ -155,8 +179,84 @@ def test_layout_constructor_places_new_child_on_parent_top( asset for asset in post_edit_scene.assets if asset.id == "cup_001" ) assert placed_cup.center_xy == [0.0, 0.0] - # Book top is z=0.62 m; cup half-height is 0.1 m with zero support clearance. - assert np.allclose(placed_cup.pos, [0.0, 0.72, 0.0]) + # Book top is z=0.62 m; cup half-height plus support clearance is 0.12 m. + assert np.allclose(placed_cup.pos, [0.0, 0.74, 0.0]) + assert settler_inputs["dynamic_asset_ids"] == {"cup_001"} + assert settler_inputs["static_asset_ids"] == {"book_001"} + + +def test_layout_constructor_settles_dynamic_table_roots( + tmp_path: Path, + monkeypatch, +) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + table = SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + rot=[0.0, 0.0, 0.0], + pos=[0.0, 0.0, 0.0], + scale=[1.0, 1.0, 1.0], + support_surface_z=0.0, + support_optimization_rect_xy=_TABLE_BOUNDS, + ) + fixed_asset = _asset( + object_id="fixed_001", + glb_path=asset_glb, + center_xy=[-0.5, 0.0], + ) + dynamic_asset = _asset(object_id="dynamic_001", glb_path=asset_glb) + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="fixed_001", parent_id="table", parent_relation="on" + ), + SceneGraphNode( + object_id="dynamic_001", parent_id="table", parent_relation="on" + ), + ] + ) + settler_inputs: dict[str, object] = {} + + class _FakeGravitySettler: + def __init__(self, **kwargs: object) -> None: + settler_inputs.update(kwargs) + + def settle(self) -> dict[str, dict[str, list[float]]]: + return { + "dynamic_001": { + "pos": _SETTLED_DYNAMIC_POS_Y_UP, + "rot": _SETTLED_DYNAMIC_ROT_Y_UP, + } + } + + monkeypatch.setattr( + scene_layout_constructor_module, + "GravitySettler", + _FakeGravitySettler, + ) + + post_edit_scene = SceneLayoutConstructor( + formal_scene=Scene(objects=[table, fixed_asset]), + goal_scene_graph=graph, + layout_variable_ids={"dynamic_001"}, + generated_scene_objects=[dynamic_asset], + output_root=tmp_path, + ).construct() + + settled_dynamic_asset = next( + asset for asset in post_edit_scene.assets if asset.id == "dynamic_001" + ) + assert settler_inputs["dynamic_asset_ids"] == {"dynamic_001"} + assert settler_inputs["static_asset_ids"] == {"fixed_001"} + assert settled_dynamic_asset.pos == _SETTLED_DYNAMIC_POS_Y_UP + assert settled_dynamic_asset.rot == _SETTLED_DYNAMIC_ROT_Y_UP + # y-up [x, y, z] maps to z-up tabletop XY [x, -z]. + assert settled_dynamic_asset.center_xy == [0.25, 0.4] def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: From a84be90cacf3e06afceae639e80235ac9a8aaf24 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:18:31 +0800 Subject: [PATCH 31/85] Updated docs --- docs/source/features/generative_sim/scene_engine.md | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index e4255163e..f321f29a5 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -44,6 +44,16 @@ overwrites its `scene_config.json`, `scene_graph.json`, `scene.json`, and final `mesh_assets`; intermediate generation and edit artifacts remain available for debugging. +### Edit Flow + +- **Import and understanding**: reads and validates the existing export, then + resolves the instruction into `add`, `move`, and `delete` operations and an + updated scene graph. +- **Asset preparation**: only `add` operations generate an image, segmentation, + geometry, and SimReady asset. Move-only and delete-only edits skip this stage. +- **Layout and export**: refines the edited layout using the updated scene graph + and writes the resulting scene back to `scene_export`. + ## Configuration Scene Engine reads the LLM, segmentation, image-generation, and @@ -89,7 +99,7 @@ The important final outputs are: scene_output/ |-- scene_understanding/ # Object analysis, masks, and stage JSON |-- scene_generation/ # Generated, SimReady, and layout-debug artifacts -|-- scene_editing/ # Present after edits; generated asset/debug artifacts +|-- scene_editing/ # Present after edits; asset-preparation and layout-optimization artifacts `-- scene_export/ |-- mesh_assets/ # Final GLBs |-- scene_config.json # Exported z-up scene description From 654165c8bce4cb7cd658c6c9cc092bb216176a60 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 20 Aug 2026 11:36:49 +0800 Subject: [PATCH 32/85] fix(scene-engine): validate edit categories and preserve batched planar relations --- .../gen_sim/scene_engine/core/scene_graph.py | 5 ++- .../editing/scene_edit_understanding.py | 17 +++++-- .../scene_engine/test_scene_edit_plan.py | 29 ++++++++++++ .../gen_sim/scene_engine/test_scene_graph.py | 45 +++++++++++++++++++ 4 files changed, 92 insertions(+), 4 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index 500980b53..b43c7353a 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -251,8 +251,11 @@ def apply_updates( # Resolve chained planar parent inheritance before adding final relations. self._resolve_planar_parent_updates(planar_relation_updates) - for source_id, relation, target_id in planar_relation_updates: + # Clear stale relations before appending this batch so chained updates + # cannot remove a relation that an earlier update just requested. + for source_id in {source_id for source_id, _, _ in planar_relation_updates}: self._clear_incident_planar_relations(source_id) + for source_id, relation, target_id in planar_relation_updates: self.relations.append( SceneGraphRelation( source_id=source_id, diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 92720c7ec..3205a7a48 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -18,6 +18,7 @@ from __future__ import annotations import json +import re from embodichain.gen_sim.scene_engine.core.scene_edit_plan import ( SceneEditOperation, @@ -37,6 +38,8 @@ OpenAICompatibleVLM, ) +_ADD_CATEGORY_PATTERN = re.compile(r"[a-z][a-z0-9_]*") + _EDIT_SYSTEM_PROMPT = """You convert one user instruction into edits for an existing tabletop scene. Use an existing object ID only when it appears in the supplied Existing object @@ -412,10 +415,9 @@ def _parse_scene_edit_operations( if op == "add": if object_id is not None: raise ValueError("VLM add operations must set object_id to null.") - if category is None: - raise ValueError("VLM add operations must provide a category.") + category = _validated_add_category(category) # Add operation should generate new id here. - # Never believe the LLM could always generate a valid id. + # The validated category keeps generated IDs safe for asset output paths. object_id = _next_add_object_id( category=category, category_counts=category_counts, @@ -467,6 +469,15 @@ def _optional_string(value: object, *, field_name: str) -> str | None: return value.strip() +def _validated_add_category(category: str | None) -> str: + """Return an add category that is safe to embed in a generated object ID.""" + if category is None or _ADD_CATEGORY_PATTERN.fullmatch(category) is None: + raise ValueError( + "VLM add operation category must be lower-case singular snake_case." + ) + return category + + def _optional_relation(value: object) -> str | None: if value is None: return None diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index a99e04384..a68293668 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -190,6 +190,35 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: ] +@pytest.mark.parametrize( + "unsafe_category", + ["../outside", "cup/../outside", r"cup\\outside"], +) +def test_scene_edit_parser_rejects_path_traversal_add_categories( + unsafe_category: str, +) -> None: + """Reject unsafe categories before generated IDs reach asset output paths.""" + scene, _ = _scene_and_graph() + draft = { + "operations": [ + { + "op": "add", + "object_id": None, + "target_id": None, + "relation": None, + "table_region": None, + "category": unsafe_category, + "name": "unsafe cup", + "description": "A small cup.", + "orientation_state": None, + } + ] + } + + with pytest.raises(ValueError, match="category must be lower-case"): + _parse_scene_edit_operations(draft, scene=scene) + + def test_scene_edit_plan_rejects_a_changed_move_orientation_state() -> None: scene, scene_graph = _scene_and_graph() scene_graph.node_by_id()["book_001"].orientation_state = "lying" diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py index d5b6adaf7..c6cbd3243 100644 --- a/tests/gen_sim/scene_engine/test_scene_graph.py +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -276,6 +276,51 @@ def test_scene_graph_materializes_inverse_planar_relations() -> None: ] +def test_scene_graph_batch_planar_updates_preserve_chained_constraints() -> None: + """Keep both requested relations when one update targets another source.""" + graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="plate", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="cup", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="spoon", + parent_id="table", + parent_relation="on", + ), + ] + ) + + graph.apply_updates( + deleted_object_ids=set(), + added_object_ids=[], + added_orientation_states_by_id={}, + on_parent_updates=[], + planar_relation_updates=[ + ("plate", "left_of", "cup"), + ("cup", "in_front_of", "spoon"), + ], + ) + + assert { + (relation.source_id, relation.relation, relation.target_id) + for relation in graph.relations + } == { + ("plate", "left_of", "cup"), + ("cup", "right_of", "plate"), + ("cup", "in_front_of", "spoon"), + ("spoon", "behind", "cup"), + } + + def test_scene_graph_to_dict_serializes_graph_state() -> None: graph = SceneGraph( nodes=[ From e92e978a2dbe43229d048ed3495d6c546d804b30 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 20 Aug 2026 20:11:32 +0800 Subject: [PATCH 33/85] Outside-table-2d-region checker now does not check those who did not participate in the optimization --- .../utils/table_surface_layout_optimizer.py | 24 +++++++------- .../test_scene_layout_optimizer.py | 32 +++++++++++++++++++ 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py index 9d874f239..be6f0062d 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/table_surface_layout_optimizer.py @@ -196,7 +196,7 @@ def _build_constraints( root_half_extents_xy: dict[str, np.ndarray], config: TableSurfaceLayoutOptimizerConfig, ) -> tuple[list[tuple[np.ndarray, float]], list[tuple[np.ndarray, float]]]: - """Build hard table-region, planar-relation, and fixed-root constraints.""" + """Build variable table-region, planar-relation, and fixed-root constraints.""" # Objects which need to be optimized. root_index = {root_id: index for index, root_id in enumerate(problem.root_ids)} table_bounds = _bounds_from_points(problem.table_optimization_rect_xy) @@ -204,12 +204,21 @@ def _build_constraints( inequality_constraints: list[tuple[np.ndarray, float]] = [] equality_constraints: list[tuple[np.ndarray, float]] = [] for root_id in problem.root_ids: - # Get the table region bound for this root asset. + fixed_xy = problem.fixed_root_xy_by_id[root_id] + if fixed_xy is not None: + # Imported, unedited roots retain their pose even when partly off-table. + _append_fixed_root_constraints( + constraints=equality_constraints, + root_index=root_index, + root_id=root_id, + fixed_xy=fixed_xy, + ) + continue + # Only layout variables must keep their complete AABB inside the table region. region_bounds = _table_region_bounds( table_bounds=table_bounds, table_region=problem.root_table_regions_by_id[root_id], ) - # Add AABB constraints for each root's center inside the table region. _append_aabb_center_bounds( constraints=inequality_constraints, root_index=root_index, @@ -217,15 +226,6 @@ def _build_constraints( bounds=region_bounds, half_extents_xy=root_half_extents_xy[root_id], ) - fixed_xy = problem.fixed_root_xy_by_id[root_id] - if fixed_xy is not None: - # Add fixed-root constraints for each root with a fixed XY center. - _append_fixed_root_constraints( - constraints=equality_constraints, - root_index=root_index, - root_id=root_id, - fixed_xy=fixed_xy, - ) for relation in problem.root_relations: # Add planar-relation constraints for each sibling relation in this group. _append_planar_relation_constraint( diff --git a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py index bd5239c70..7dab3640b 100644 --- a/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py +++ b/tests/gen_sim/scene_engine/test_scene_layout_optimizer.py @@ -301,6 +301,38 @@ def test_table_optimizer_ignores_fixed_sibling_overlap(tmp_path: Path) -> None: } +def test_table_optimizer_keeps_fixed_sibling_partly_outside_table( + tmp_path: Path, +) -> None: + asset_glb = tmp_path / "asset.glb" + trimesh.creation.box(extents=[_ASSET_SIDE_LENGTH_M] * 3).export(asset_glb) + fixed_id, variable_id = "fixed_001", "variable_001" + fixed_xy = [1.95, 0.0] # Its 0.2 m AABB extends beyond the x=2 m table edge. + + solved_xy_by_id = TableSurfaceLayoutOptimizer().optimize( + TableSurfaceLayoutProblem( + assets_by_id={ + fixed_id: _asset( + object_id=fixed_id, + glb_path=asset_glb, + center_xy=fixed_xy, + ), + variable_id: _asset(object_id=variable_id, glb_path=asset_glb), + }, + root_ids=[fixed_id, variable_id], + root_seed_xy_by_id={fixed_id: fixed_xy, variable_id: [0.0, 0.0]}, + imported_root_ids={fixed_id}, + fixed_root_xy_by_id={fixed_id: fixed_xy, variable_id: None}, + root_table_regions_by_id={fixed_id: None, variable_id: None}, + table_optimization_rect_xy=_TABLE_BOUNDS, + root_relations=[], + ) + ) + + assert solved_xy_by_id[fixed_id] == fixed_xy + assert np.all(np.abs(solved_xy_by_id[variable_id]) <= 1.9 + 1e-6) + + def test_table_optimizer_separates_variable_sibling_from_fixed_sibling( tmp_path: Path, ) -> None: From ad2306398b442b96d85e8e1d3d47ea10a5669ced Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:07:12 +0800 Subject: [PATCH 34/85] fix(scene-engine): delete the function export --- embodichain/gen_sim/scene_engine/core/scene_edit_plan.py | 2 -- .../pipeline/editing/scene_edit_asset_preparation.py | 2 -- 2 files changed, 4 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 0f78b9699..902f6eda7 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -28,8 +28,6 @@ TABLE_OBJECT_ID, ) -__all__ = ["SceneEditOperation", "SceneEditPlan"] - SceneEditOperationType = Literal["add", "move", "delete"] diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index 220a91014..aa19e97e1 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -50,8 +50,6 @@ SimReadyProcessorConfig, ) -__all__ = ["prepare_scene_edit_assets"] - @dataclass(frozen=True) class _AddedAssetInfo: From 301d7eb8f13145e25b3b4513a31f559e29fa1fb3 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Mon, 24 Aug 2026 17:24:26 +0800 Subject: [PATCH 35/85] finished scene-graph based on realtion addtion in image as input pipeline(not finished, but can be ran) --- .../pipeline/generation/scene_generation.py | 684 +++++++++++++++++- .../generation/scene_understanding.py | 166 +++-- .../scene_engine/test_scene_generation.py | 216 ++++++ .../scene_engine/test_scene_understanding.py | 218 +++++- 4 files changed, 1169 insertions(+), 115 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 7010e5f14..ef14e4d79 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -17,15 +17,22 @@ from __future__ import annotations +from collections import deque import json from pathlib import Path import shutil +import matplotlib import numpy as np from scipy.spatial.transform import Rotation import trimesh from shapely.geometry import Polygon +matplotlib.use("Agg") + +import matplotlib.pyplot as plt +from matplotlib.patches import Rectangle + from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, ) @@ -48,12 +55,21 @@ GravitySettleBody, GravitySettler, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.parent_surface_layout_optimizer import ( + ParentSurfaceLayoutOptimizer, + ParentSurfaceLayoutProblem, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, load_glb_mesh, quaternion_wxyz_to_euler_xyz_degrees, transform_matrix_to_layout_object, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + measure_scene_object_z_up_world_aabb, + scene_object_y_up_layout, + update_scene_object_y_up_pose_from_z_up_support, +) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( SimReadyProcessor, SimReadyProcessorConfig, @@ -216,10 +232,57 @@ def _generate_coarse_results_from_masks( json.dumps(coarse_layout, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) + _write_baked_coarse_layout_debug_glbs( + coarse_layout=coarse_layout, + coarse_geometry_root=coarse_geometry_output_root, + debug_output_root=debug_output_root, + ) # Nothing to be returned. return None +def _write_baked_coarse_layout_debug_glbs( + *, + coarse_layout: list[dict[str, object]], + coarse_geometry_root: str | Path, + debug_output_root: str | Path, +) -> Path: + """Write per-object and combined GLBs with each coarse pose baked in.""" + resolved_geometry_root = Path(coarse_geometry_root).expanduser().resolve() + baked_root = ( + Path(debug_output_root).expanduser().resolve() / "coarse_layout_baked_glbs" + ) + baked_root.mkdir(parents=True, exist_ok=True) + + baked_scene = trimesh.Scene() + baked_object_ids: set[str] = set() + for layout_object in coarse_layout: + object_id = layout_object.get("id") + if not isinstance(object_id, str) or not object_id: + raise ValueError( + "Each coarse layout object must contain a non-empty string id." + ) + if object_id in baked_object_ids: + raise ValueError(f"Coarse layout contains duplicate id {object_id!r}.") + baked_object_ids.add(object_id) + + # The geometry server declares these poses in the GLB y-up convention. + baked_mesh = load_glb_mesh(resolved_geometry_root / f"{object_id}.glb") + baked_mesh.apply_transform(layout_object_to_transform_matrix(layout_object)) + baked_mesh.export(baked_root / f"{object_id}.glb", file_type="glb") + baked_scene.add_geometry( + baked_mesh, + node_name=object_id, + geom_name=object_id, + ) + + if not baked_object_ids: + raise ValueError("Cannot write baked coarse debug GLBs for an empty layout.") + baked_scene.export(baked_root / "scene.glb", file_type="glb") + log_info(f"Wrote baked coarse-pose debug GLBs: {baked_root}") + return baked_root + + def _update_scene_final_y_up_layout_and_z_up_centers( *, scene: Scene, @@ -376,19 +439,39 @@ def _layout_refinement( assets_layout=refined_assets_layout, ) - # 4. Move all assets as one rigid group so its lowest AABB point is 2cm above - # the table. This preserves the initial relative poses for the later - # gravity simulation, which can settle individual assets physically. + # 4. Only direct on-table children participate in the table-level layout + # stages. Their descendants follow each solved root transform until their + # own parent-surface optimization is introduced in a later BFS pass. + table_root_ids = _table_on_asset_ids(scene_graph=scene_graph, table_id=table_id) + table_root_layouts = _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids=table_root_ids, + ) + if not table_root_layouts: + log_info("Scene has no on-table assets; skipping table layout refinement.") + return refined_table_layout, refined_assets_layout + # Move the direct on-table assets as one rigid group so their lowest AABB + # point is 2cm above the table. Descendants retain their relative transforms. + table_root_matrices_before_align = _layout_matrices_by_id(table_root_layouts) group_table_aligner = AssetsGroupTableAligner( table_layout=refined_table_layout, - assets_layout=refined_assets_layout, + assets_layout=table_root_layouts, geometry_root=simready_geometry_output_root, ) - refined_table_layout, refined_assets_layout = group_table_aligner.align() - if not refined_assets_layout: - log_info("Scene has no movable assets; skipping support-region clamping.") - return refined_table_layout, [] + # Align. + refined_table_layout, aligned_table_root_layouts = group_table_aligner.align() + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=aligned_table_root_layouts, + root_matrices_before_update=table_root_matrices_before_align, + ) + # After update the layouts, re-select those who are direct on-table children. + table_root_layouts = _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids=table_root_ids, + ) # 5. Reuse support geometry detected during SimReady processing. if ( @@ -406,60 +489,91 @@ def _layout_refinement( or table_optimization_rectangle.is_empty ): raise ValueError("Scene table optimization rectangle is not valid.") - _, assets_aabb_2d_z_up_world_corners_by_id = ( + _, table_root_aabb_2d_z_up_world_corners_by_id = ( _measure_table_and_assets_in_z_up_world( table_layout=refined_table_layout, - assets_layout=refined_assets_layout, + assets_layout=table_root_layouts, geometry_root=simready_geometry_output_root, ) ) - # 6. Keep the complete clutter rigid in the table plane. A successful - # result applies one shared z-up XY delta to every AABB, so it preserves - # all existing asset-to-asset relations. It is *not* an asset packing - # pass: pre-existing overlap is deliberately left to a later optimizer. + # 6. Keep the direct on-table clutter rigid in the table plane. A successful + # result applies one shared z-up XY delta to every root AABB, and the same + # transform is propagated to each root's descendants. It is *not* an asset + # packing pass: pre-existing overlap is deliberately left to a later optimizer. + table_root_matrices_before_clamp = _layout_matrices_by_id(table_root_layouts) group_clamp = AssetsGroupSupportClamp( support_region=table_support_polygon, assets_aabb_2d_z_up_world_corners_by_id=( - assets_aabb_2d_z_up_world_corners_by_id + table_root_aabb_2d_z_up_world_corners_by_id ), - assets_layout=refined_assets_layout, + assets_layout=table_root_layouts, debug_output_root=debug_output_root, ) - refined_assets_layout = group_clamp.clamp() + # Clamp. + clamped_table_root_layouts = group_clamp.clamp() + # Save debug images. group_clamp.save_group_clamp_debug_images() + # Update to descendants. + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=clamped_table_root_layouts, + root_matrices_before_update=table_root_matrices_before_clamp, + ) + # Re-select. + table_root_layouts = _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids=table_root_ids, + ) # The clamp returns y-up layouts; measure their resulting z-up AABBs again # so the following independent optimizer consumes the same world-frame # geometry as every other stage. - _, clamped_assets_aabb_2d_z_up_world_corners_by_id = ( + _, clamped_table_root_aabb_2d_z_up_world_corners_by_id = ( _measure_table_and_assets_in_z_up_world( table_layout=refined_table_layout, - assets_layout=refined_assets_layout, + assets_layout=table_root_layouts, geometry_root=simready_geometry_output_root, ) ) - # 7. Optimize independent asset positions inside the conservative rectangle. + # 7. Optimize independent on-table root positions inside the conservative + # rectangle. # The clamp above already used the exact outer contour for the shared shift. + table_root_matrices_before_optimization = _layout_matrices_by_id(table_root_layouts) overlap_optimizer = AssetsSupportLayoutOptimizer( support_region=table_optimization_rectangle, assets_aabb_2d_z_up_world_corners_by_id=( - clamped_assets_aabb_2d_z_up_world_corners_by_id + clamped_table_root_aabb_2d_z_up_world_corners_by_id ), - assets_layout=refined_assets_layout, + assets_layout=table_root_layouts, debug_output_root=debug_output_root, ) # Render this stage separately from the rigid group clamp. The latter # intentionally preserves pre-existing overlaps, while this figure shows # whether independent AABB separation actually resolved them. - refined_assets_layout = overlap_optimizer.optimize() + optimized_table_root_layouts = overlap_optimizer.optimize() + # Save debug image. overlap_optimizer.save_overlap_optimization_debug_images() + # Update. + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=optimized_table_root_layouts, + root_matrices_before_update=table_root_matrices_before_optimization, + ) + # Re-select. + table_root_layouts = _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids=table_root_ids, + ) - # 8. The initial image graph has one on-table level, so every asset settles - # dynamically against the table in this first generic gravity pass. + # 8. Settle only the direct on-table roots. Their unoptimized descendants + # stay outside this simulation and inherit each root's final pose delta. assets_by_id = {asset.id: asset for asset in scene.assets} - # All the assets are dynamic; the table is static. + table_root_matrices_before_settle = _layout_matrices_by_id(table_root_layouts) + # Settle. settled_pose_by_id = GravitySettler( table_body=GravitySettleBody( scene_object=scene.table, @@ -470,17 +584,40 @@ def _layout_refinement( scene_object=assets_by_id[str(asset_layout["id"])], y_up_layout=asset_layout, ) - for asset_layout in refined_assets_layout + for asset_layout in table_root_layouts ], - dynamic_asset_ids=set(assets_by_id), + dynamic_asset_ids=set(table_root_ids), static_asset_ids=set(), ).settle() - # Update. - for asset_layout in refined_assets_layout: + settled_table_root_layouts: list[dict[str, object]] = [] + for asset_layout in table_root_layouts: asset_id = str(asset_layout["id"]) settled_pose = settled_pose_by_id[asset_id] - asset_layout["pos"] = settled_pose["pos"] - asset_layout["rot"] = settled_pose["rot"] + settled_table_root_layouts.append( + { + **asset_layout, + "pos": settled_pose["pos"], + "rot": settled_pose["rot"], + } + ) + # Update the descendants. + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=settled_table_root_layouts, + root_matrices_before_update=table_root_matrices_before_settle, + ) + + # 9. The table roots are now stable, so refine each non-table on-parent + # group in BFS order. Every child begins only after its parent is current. + refined_assets_layout = _refine_on_children_bfs( + scene=scene, + scene_graph=scene_graph, + table_layout=refined_table_layout, + assets_layout=refined_assets_layout, + table_root_ids=table_root_ids, + debug_output_root=debug_output_root, + ) # Update the scene data structure with the final layout and spatial metadata. _update_scene_final_y_up_layout_and_z_up_centers( @@ -492,6 +629,489 @@ def _layout_refinement( return refined_table_layout, refined_assets_layout +def _refine_on_children_bfs( + *, + scene: Scene, + scene_graph: SceneGraph, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + table_root_ids: set[str], + debug_output_root: str | Path, +) -> list[dict[str, object]]: + """Refine every non-table ``on`` group after its parent has settled.""" + if scene.table is None: + raise ValueError("Cannot refine parent-surface groups without a table.") + _sync_scene_y_up_poses_from_layouts( + scene=scene, + table_layout=table_layout, + assets_layout=assets_layout, + ) + assets_by_id = {asset.id: asset for asset in scene.assets} + children_by_parent: dict[str, list[str]] = {} + for node in scene_graph.nodes: + if node.parent_id is not None and node.parent_relation == "on": + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + parent_surface_optimizer = ParentSurfaceLayoutOptimizer() + pending_parent_ids = deque(table_root_ids) + refined_assets_layout = assets_layout + while pending_parent_ids: + parent_id = pending_parent_ids.popleft() + child_ids = children_by_parent.get(parent_id, []) + if not child_ids: + continue + parent = assets_by_id.get(parent_id) + if parent is None: + raise ValueError(f"Non-table parent {parent_id!r} is not an asset.") + + parent_aabb = measure_scene_object_z_up_world_aabb(scene_object=parent) + parent_aabb_xy = [ + [float(parent_aabb[0][0]), float(parent_aabb[0][1])], + [float(parent_aabb[1][0]), float(parent_aabb[1][1])], + ] + child_aabbs_xy_by_id = { + child_id: _scene_object_z_up_aabb_xy(scene_object=assets_by_id[child_id]) + for child_id in child_ids + } + child_seed_xy_by_id, projected_child_aabbs_xy_by_id = ( + _project_child_aabb_centers_into_parent_aabb( + parent_aabb_xy=parent_aabb_xy, + child_aabbs_xy_by_id=child_aabbs_xy_by_id, + ) + ) + projected_child_ids = [ + child_id + for child_id in child_ids + if not np.allclose( + child_aabbs_xy_by_id[child_id], + projected_child_aabbs_xy_by_id[child_id], + ) + ] + if projected_child_ids: + log_info( + "Projected " + f"{len(projected_child_ids)} child AABBs into parent {parent_id!r}: " + f"{projected_child_ids}." + ) + else: + log_info( + f"All direct children are already inside parent {parent_id!r}'s AABB." + ) + _render_parent_child_aabb_transition( + parent_id=parent_id, + parent_aabb_xy=parent_aabb_xy, + before_child_aabbs_xy_by_id=child_aabbs_xy_by_id, + after_child_aabbs_xy_by_id=projected_child_aabbs_xy_by_id, + before_title="Before parent-AABB projection", + after_title="After parent-AABB projection", + output_path=( + Path(debug_output_root) + / f"parent_{parent_id}_child_aabb_projection_2d.png" + ), + ) + child_id_set = set(child_ids) + # All image-observed children are movable and start from their nearest + # parent-AABB-valid image seed. + parent_surface_problem = ParentSurfaceLayoutProblem( + assets_by_id=assets_by_id, + child_ids=child_ids, + child_seed_xy_by_id=child_seed_xy_by_id, + imported_child_ids=child_id_set, + fixed_child_xy_by_id={child_id: None for child_id in child_ids}, + parent_aabb_xy=parent_aabb_xy, + parent_top_z=float(parent_aabb[1][2]), + child_relations=[ + relation + for relation in scene_graph.relations + if relation.source_id in child_id_set + and relation.target_id in child_id_set + ], + ) + solved_child_xy_by_id = parent_surface_optimizer.optimize( + parent_surface_problem + ) + optimized_child_aabbs_xy_by_id = { + child_id: _translated_aabb_xy( + aabb_xy=projected_child_aabbs_xy_by_id[child_id], + delta_xy=( + np.asarray(solved_child_xy_by_id[child_id], dtype=float) + - np.asarray(child_seed_xy_by_id[child_id], dtype=float) + ), + ) + for child_id in child_ids + } + _render_parent_child_aabb_transition( + parent_id=parent_id, + parent_aabb_xy=parent_aabb_xy, + before_child_aabbs_xy_by_id=projected_child_aabbs_xy_by_id, + after_child_aabbs_xy_by_id=optimized_child_aabbs_xy_by_id, + before_title="Before parent-child AABB optimization", + after_title="After parent-child AABB optimization", + output_path=( + Path(debug_output_root) + / f"parent_{parent_id}_child_aabb_optimization_2d.png" + ), + ) + + child_matrices_before_placement = _layout_matrices_by_id( + _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids=child_id_set, + ) + ) + for child_id, solved_xy in solved_child_xy_by_id.items(): + # Place each child above the parent's current top before gravity settles it. + update_scene_object_y_up_pose_from_z_up_support( + scene_object=assets_by_id[child_id], + support_region_z=parent_surface_problem.parent_top_z, + center_xy=solved_xy, + clearance_m=0.02, + ) + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=[ + scene_object_y_up_layout(assets_by_id[child_id]) + for child_id in child_ids + ], + root_matrices_before_update=child_matrices_before_placement, + ) + _sync_scene_y_up_poses_from_layouts( + scene=scene, + table_layout=table_layout, + assets_layout=refined_assets_layout, + ) + + child_matrices_before_settle = _layout_matrices_by_id( + _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids=child_id_set, + ) + ) + settled_pose_by_id = parent_surface_optimizer.settle_dynamic_children( + table=scene.table, + parent=parent, + problem=parent_surface_problem, + dynamic_child_ids=child_id_set, + ) + for child_id, settled_pose in settled_pose_by_id.items(): + assets_by_id[child_id].pos = settled_pose["pos"] + assets_by_id[child_id].rot = settled_pose["rot"] + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=[ + scene_object_y_up_layout(assets_by_id[child_id]) + for child_id in child_ids + ], + root_matrices_before_update=child_matrices_before_settle, + ) + _sync_scene_y_up_poses_from_layouts( + scene=scene, + table_layout=table_layout, + assets_layout=refined_assets_layout, + ) + pending_parent_ids.extend(child_ids) + + return refined_assets_layout + + +def _sync_scene_y_up_poses_from_layouts( + *, + scene: Scene, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], +) -> None: + """Synchronize current y-up poses without rewriting final spatial metadata.""" + if scene.table is None: + raise ValueError("Cannot synchronize layouts without a table.") + _copy_y_up_layout_to_scene_object(scene.table, table_layout) + assets_by_id = {asset.id: asset for asset in scene.assets} + synced_asset_ids: set[str] = set() + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + asset = assets_by_id.get(asset_id) + if asset is None: + raise ValueError(f"Current layout contains unknown asset {asset_id!r}.") + if asset_id in synced_asset_ids: + raise ValueError(f"Current layout contains duplicate asset {asset_id!r}.") + _copy_y_up_layout_to_scene_object(asset, asset_layout) + synced_asset_ids.add(asset_id) + if synced_asset_ids != set(assets_by_id): + raise ValueError("Current layouts do not cover every scene asset.") + + +def _scene_object_z_up_aabb_xy(*, scene_object: SceneObject) -> np.ndarray: + """Measure one current scene object's z-up XY AABB bounds.""" + aabb = measure_scene_object_z_up_world_aabb(scene_object=scene_object) + return np.asarray( + [ + [float(aabb[0][0]), float(aabb[0][1])], + [float(aabb[1][0]), float(aabb[1][1])], + ] + ) + + +def _project_child_aabb_centers_into_parent_aabb( + *, + parent_aabb_xy: list[list[float]], + child_aabbs_xy_by_id: dict[str, np.ndarray], +) -> tuple[dict[str, list[float]], dict[str, np.ndarray]]: + """Project child AABB centers to their nearest parent-AABB-valid positions.""" + parent_bounds = np.asarray(parent_aabb_xy, dtype=float) + if parent_bounds.shape != (2, 2) or not np.all(np.isfinite(parent_bounds)): + raise ValueError("Parent AABB must contain two finite XY corners.") + parent_minimum, parent_maximum = parent_bounds + projected_center_xy_by_id: dict[str, list[float]] = {} + projected_aabbs_xy_by_id: dict[str, np.ndarray] = {} + for child_id, child_aabb_xy in child_aabbs_xy_by_id.items(): + child_bounds = np.asarray(child_aabb_xy, dtype=float) + if child_bounds.shape != (2, 2) or not np.all(np.isfinite(child_bounds)): + raise ValueError(f"Child {child_id!r} has an invalid XY AABB.") + child_minimum, child_maximum = child_bounds + half_extents_xy = (child_maximum - child_minimum) / 2.0 + center_xy = (child_minimum + child_maximum) / 2.0 + legal_minimum = parent_minimum + half_extents_xy + legal_maximum = parent_maximum - half_extents_xy + if np.any(legal_minimum > legal_maximum): + raise ValueError(f"Asset {child_id!r} cannot fit inside its parent AABB.") + projected_center_xy = np.clip(center_xy, legal_minimum, legal_maximum) + projected_center_xy_by_id[child_id] = projected_center_xy.tolist() + projected_aabbs_xy_by_id[child_id] = _translated_aabb_xy( + aabb_xy=child_bounds, + delta_xy=projected_center_xy - center_xy, + ) + return projected_center_xy_by_id, projected_aabbs_xy_by_id + + +def _translated_aabb_xy(*, aabb_xy: np.ndarray, delta_xy: np.ndarray) -> np.ndarray: + """Translate one validated XY AABB by a finite two-dimensional delta.""" + bounds = np.asarray(aabb_xy, dtype=float) + delta = np.asarray(delta_xy, dtype=float) + if bounds.shape != (2, 2) or delta.shape != (2,): + raise ValueError( + "AABB translation requires two XY corners and a two-value delta." + ) + if not np.all(np.isfinite(bounds)) or not np.all(np.isfinite(delta)): + raise ValueError("AABB translation requires finite values.") + return bounds + delta + + +def _render_parent_child_aabb_transition( + *, + parent_id: str, + parent_aabb_xy: list[list[float]], + before_child_aabbs_xy_by_id: dict[str, np.ndarray], + after_child_aabbs_xy_by_id: dict[str, np.ndarray], + before_title: str, + after_title: str, + output_path: str | Path, +) -> None: + """Render one parent AABB and its direct-child AABBs before and after a stage.""" + parent_bounds = np.asarray(parent_aabb_xy, dtype=float) + if parent_bounds.shape != (2, 2) or not np.all(np.isfinite(parent_bounds)): + raise ValueError("Parent AABB must contain two finite XY corners.") + if set(before_child_aabbs_xy_by_id) != set(after_child_aabbs_xy_by_id): + raise ValueError("Parent-child debug AABB states must contain identical IDs.") + + output = Path(output_path).expanduser().resolve() + output.parent.mkdir(parents=True, exist_ok=True) + figure, axes = plt.subplots(1, 2, figsize=(14, 7), dpi=160, constrained_layout=True) + for axis, child_aabbs_xy_by_id, title, color in ( + (axes[0], before_child_aabbs_xy_by_id, before_title, "tab:blue"), + (axes[1], after_child_aabbs_xy_by_id, after_title, "tab:green"), + ): + _draw_parent_child_aabb_state( + axis=axis, + parent_id=parent_id, + parent_bounds=parent_bounds, + child_aabbs_xy_by_id=child_aabbs_xy_by_id, + title=title, + child_color=color, + ) + figure.savefig(output, bbox_inches="tight") + plt.close(figure) + + +def _draw_parent_child_aabb_state( + *, + axis: object, + parent_id: str, + parent_bounds: np.ndarray, + child_aabbs_xy_by_id: dict[str, np.ndarray], + title: str, + child_color: str, +) -> None: + """Draw one parent AABB and direct children in a single z-up XY panel.""" + parent_minimum, parent_maximum = parent_bounds + axis.add_patch( + Rectangle( + parent_minimum, + *(parent_maximum - parent_minimum), + facecolor="tab:orange", + edgecolor="saddlebrown", + alpha=0.25, + linewidth=2.0, + ) + ) + axis.text( + *parent_bounds.mean(axis=0), + parent_id, + ha="center", + va="center", + fontsize=10, + bbox={"facecolor": "white", "alpha": 0.8, "edgecolor": "none"}, + ) + all_bounds = [parent_bounds] + for child_id, child_aabb_xy in child_aabbs_xy_by_id.items(): + child_bounds = np.asarray(child_aabb_xy, dtype=float) + if child_bounds.shape != (2, 2) or not np.all(np.isfinite(child_bounds)): + raise ValueError(f"Child {child_id!r} has an invalid debug XY AABB.") + child_minimum, child_maximum = child_bounds + axis.add_patch( + Rectangle( + child_minimum, + *(child_maximum - child_minimum), + facecolor=child_color, + edgecolor=child_color, + alpha=0.3, + linewidth=1.5, + ) + ) + axis.text( + *child_bounds.mean(axis=0), + child_id, + ha="center", + va="center", + fontsize=9, + bbox={"facecolor": "white", "alpha": 0.8, "edgecolor": "none"}, + ) + all_bounds.append(child_bounds) + combined_bounds = np.vstack(all_bounds) + padding = max(float(np.ptp(combined_bounds, axis=0).max()) * 0.1, 0.02) + axis.set_xlim( + combined_bounds[:, 0].min() - padding, combined_bounds[:, 0].max() + padding + ) + axis.set_ylim( + combined_bounds[:, 1].min() - padding, combined_bounds[:, 1].max() + padding + ) + axis.set_aspect("equal", adjustable="box") + axis.set_xlabel("x (z-up world)") + axis.set_ylabel("y (z-up world)") + axis.set_title(title) + axis.grid(True, alpha=0.25) + + +def _table_on_asset_ids(*, scene_graph: SceneGraph, table_id: str) -> set[str]: + """Return the IDs of assets that directly rest on the table.""" + return { + node.object_id + for node in scene_graph.nodes + if node.parent_id == table_id and node.parent_relation == "on" + } + + +def _select_asset_layouts( + *, + assets_layout: list[dict[str, object]], + asset_ids: set[str], +) -> list[dict[str, object]]: + """Return the requested asset layouts without discarding other scene assets.""" + selected_assets_layout = [ + asset_layout + for asset_layout in assets_layout + if asset_layout.get("id") in asset_ids + ] + selected_ids = { + asset_id + for asset_layout in selected_assets_layout + if isinstance(asset_id := asset_layout.get("id"), str) + } + if selected_ids != asset_ids: + raise ValueError( + "Scene graph on-table assets and refined asset layouts do not match." + ) + return selected_assets_layout + + +def _layout_matrices_by_id( + assets_layout: list[dict[str, object]], +) -> dict[str, np.ndarray]: + """Return complete layout transforms keyed by their validated asset IDs.""" + matrices_by_id: dict[str, np.ndarray] = {} + for asset_layout in assets_layout: + asset_id = asset_layout.get("id") + if not isinstance(asset_id, str) or not asset_id: + raise ValueError("Each asset layout must contain a non-empty string id.") + if asset_id in matrices_by_id: + raise ValueError(f"Asset layouts repeat id {asset_id!r}.") + matrices_by_id[asset_id] = layout_object_to_transform_matrix(asset_layout) + return matrices_by_id + + +def _apply_root_layout_updates_to_descendant_subtrees( + *, + scene_graph: SceneGraph, + assets_layout: list[dict[str, object]], + updated_root_layouts: list[dict[str, object]], + root_matrices_before_update: dict[str, np.ndarray], +) -> list[dict[str, object]]: + """Replace solved root layouts and propagate each complete pose delta downward.""" + assets_layout_by_id = { + asset_id: asset_layout + for asset_layout in assets_layout + if isinstance(asset_id := asset_layout.get("id"), str) and asset_id + } + if len(assets_layout_by_id) != len(assets_layout): + raise ValueError( + "Each asset layout must contain one unique non-empty string id." + ) + updated_root_layouts_by_id = { + asset_id: asset_layout + for asset_layout in updated_root_layouts + if isinstance(asset_id := asset_layout.get("id"), str) and asset_id + } + if len(updated_root_layouts_by_id) != len(updated_root_layouts): + raise ValueError("Updated root layouts must have unique non-empty string ids.") + if set(updated_root_layouts_by_id) != set(root_matrices_before_update): + raise ValueError("Updated root layouts do not match the saved root poses.") + + children_by_parent: dict[str, list[str]] = {} + for node in scene_graph.nodes: + if node.parent_id is not None: + children_by_parent.setdefault(node.parent_id, []).append(node.object_id) + + for root_id, root_layout in updated_root_layouts_by_id.items(): + if root_id not in assets_layout_by_id: + raise ValueError(f"Updated root layout {root_id!r} is not an asset layout.") + # Compute the world-frame delta from the previous root pose to the updated root pose. + root_matrix_after_update = layout_object_to_transform_matrix(root_layout) + root_delta = root_matrix_after_update @ np.linalg.inv( + root_matrices_before_update[root_id] + ) + assets_layout_by_id[root_id] = root_layout + + # A parent pose update moves every descendant in the same world frame. + pending_descendant_ids = list(children_by_parent.get(root_id, [])) + while pending_descendant_ids: + descendant_id = pending_descendant_ids.pop(0) + descendant_layout = assets_layout_by_id.get(descendant_id) + if descendant_layout is None: + raise ValueError( + f"Scene graph descendant {descendant_id!r} has no asset layout." + ) + assets_layout_by_id[descendant_id] = transform_matrix_to_layout_object( + descendant_id, + root_delta @ layout_object_to_transform_matrix(descendant_layout), + ) + # Add the next generation of descendants to the pending list. + pending_descendant_ids.extend(children_by_parent.get(descendant_id, [])) + + return [ + assets_layout_by_id[str(asset_layout["id"])] for asset_layout in assets_layout + ] + + def _scene_graph_based_calibration( *, scene_graph: SceneGraph, diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index d93ada019..4f7839916 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -145,23 +145,34 @@ Return JSON only, with exactly one key: assignments. It must be null or an array of asset_id and mask_index objects. Do not include Markdown or any other text.""" -_ORIENTATION_STATE_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. -Each visible asset has an outline and an ID label. Determine whether each listed -asset is standing, lying, or unknown in the image. - -Use a non-null state only for an elongated object with a clear primary long axis. -Use "standing" when its primary axis is approximately vertical to the tabletop. -Use "lying" when its primary axis is approximately parallel to the tabletop. +_INITIAL_SCENE_GRAPH_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. +Each visible asset has an outline and an ID label. Build a support graph for the +listed assets and determine each asset's image-observed orientation state. + +Every asset must have exactly one direct support parent with relation "on". +Use "table" when the asset directly rests on the table. Use another supplied +asset ID only when the image clearly shows the asset resting directly on that +asset's top surface. Do not infer an on relationship from 2D overlap alone. +When the direct support parent is uncertain, use "table". The table is a fixed +support ID, not an output node: never include it in nodes. + +Use a non-null orientation_state only for an elongated object with a clear +primary long axis. Use "standing" when that axis is approximately vertical to +the tabletop, and "lying" when it is approximately parallel to the tabletop. Use null for every object without a clear primary long axis or when uncertain. +The orientation state describes the asset itself and is independent of its +support parent. -Return JSON only, with exactly this schema. Include every supplied asset ID -exactly once. Never include the table or any ID that was not supplied: -{ - "orientation_states": [ - {"object_id": "bottle_001", "orientation_state": "standing"}, - {"object_id": "book_001", "orientation_state": null} - ] -}""" +Examples: +- A bottle directly on the table is upright: + {"nodes": [{"object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", "orientation_state": "standing"}]} +- A pen lies flat on a book, and the book is on the table: + {"nodes": [{"object_id": "book_001", "parent_id": "table", "parent_relation": "on", "orientation_state": null}, {"object_id": "pen_001", "parent_id": "book_001", "parent_relation": "on", "orientation_state": "lying"}]} +- A round cup directly on the table has no reliable long axis: + {"nodes": [{"object_id": "cup_001", "parent_id": "table", "parent_relation": "on", "orientation_state": null}]} + +Return JSON only, with exactly one key: nodes. Include every supplied asset ID +exactly once and no unknown IDs. Do not include Markdown or any other text.""" def understand_scene( @@ -227,7 +238,7 @@ def _initialize_scene_graph_from_segmented_scene( vlm_client: OpenAICompatibleVLM, json_max_attempts: int = 3, ) -> SceneGraph: - """Build the initial graph assuming every segmented asset rests on the table.""" + """Build the initial image-observed support graph for segmented assets.""" # Get simplified scene info for VLM. scene_info = _simplify_scene_info_for_graph_initialization(scene=scene) resolved_asset_mask_id_overlay_path = _validate_image_path( @@ -235,26 +246,13 @@ def _initialize_scene_graph_from_segmented_scene( ) if scene.table is None: raise ValueError("Cannot initialize a scene graph without a table.") - orientation_states_by_id = _query_orientation_states( + # Return a validated scene graph. + return _query_initial_scene_graph( scene_info=scene_info, asset_mask_id_overlay_path=resolved_asset_mask_id_overlay_path, vlm_client=vlm_client, json_max_attempts=json_max_attempts, ) - return SceneGraph( - nodes=[ - SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), - *[ - SceneGraphNode( - object_id=asset.id, - parent_id=TABLE_OBJECT_ID, - parent_relation="on", # semi-hard-code. - orientation_state=orientation_states_by_id[asset.id], - ) - for asset in scene.assets - ], - ], - ) def _simplify_scene_info_for_graph_initialization( @@ -263,91 +261,117 @@ def _simplify_scene_info_for_graph_initialization( ) -> dict[str, object]: """Return the object metadata needed to initialize an image-based graph.""" return { - "asset_ids": [asset.id for asset in scene.assets], + "assets": [ + { + "id": asset.id, + "category": asset.category, + "name": asset.name, + "description": asset.description, + } + for asset in scene.assets + ], } -def _query_orientation_states( +def _query_initial_scene_graph( *, scene_info: dict[str, object], asset_mask_id_overlay_path: Path, vlm_client: OpenAICompatibleVLM, json_max_attempts: int, -) -> dict[str, str | None]: - """Return validated image-observed orientation states keyed by asset ID.""" +) -> SceneGraph: + """Return one validated image-observed support graph.""" if json_max_attempts < 1: raise ValueError("json_max_attempts must be at least 1.") last_validation_error: ValueError | None = None for _ in range(json_max_attempts): + # Get response. response_text = vlm_client.complete( image_path=asset_mask_id_overlay_path, - system_prompt=_ORIENTATION_STATE_SYSTEM_PROMPT, + system_prompt=_INITIAL_SCENE_GRAPH_SYSTEM_PROMPT, user_prompt=json.dumps(scene_info, ensure_ascii=False), ) try: - return _parse_orientation_states_response( + # Validate. + return _parse_initial_scene_graph_response( response_text=response_text, - asset_ids=scene_info["asset_ids"], + assets=scene_info["assets"], ) except ValueError as exc: last_validation_error = exc assert last_validation_error is not None raise ValueError( - "VLM returned invalid orientation-state JSON after " + "VLM returned invalid initial scene-graph JSON after " f"{json_max_attempts} attempts: {last_validation_error}" ) from last_validation_error -def _parse_orientation_states_response( +def _parse_initial_scene_graph_response( *, response_text: str, - asset_ids: object, -) -> dict[str, str | None]: - """Parse a complete VLM orientation-state response for known asset IDs.""" - if not isinstance(asset_ids, list) or not all( - isinstance(object_id, str) for object_id in asset_ids + assets: object, +) -> SceneGraph: + """Parse a complete VLM support graph response for known scene assets.""" + if not isinstance(assets, list) or not all( + isinstance(asset, dict) and isinstance(asset.get("id"), str) for asset in assets ): - raise ValueError("Scene graph initialization requires string asset IDs.") + raise ValueError("Scene graph initialization requires assets with string ids.") + asset_ids = [asset["id"] for asset in assets] json_text = _strip_json_code_fence(response_text) try: payload = json.loads(json_text) except json.JSONDecodeError as exc: raise ValueError(f"VLM response is not valid JSON: {exc.msg}") from exc - if not isinstance(payload, dict) or set(payload) != {"orientation_states"}: - raise ValueError("VLM JSON must contain exactly the key: orientation_states.") - states_value = payload["orientation_states"] - if not isinstance(states_value, list): - raise ValueError("VLM JSON key orientation_states must be an array.") - - orientation_states_by_id: dict[str, str | None] = {} - for index, state_value in enumerate(states_value): - if not isinstance(state_value, dict) or set(state_value) != { + if not isinstance(payload, dict) or set(payload) != {"nodes"}: + raise ValueError("VLM JSON must contain exactly the key: nodes.") + nodes_value = payload["nodes"] + if not isinstance(nodes_value, list): + raise ValueError("VLM JSON key nodes must be an array.") + + nodes_by_id: dict[str, SceneGraphNode] = {} + for index, node_value in enumerate(nodes_value): + if not isinstance(node_value, dict) or set(node_value) != { "object_id", + "parent_id", + "parent_relation", "orientation_state", }: raise ValueError( - "VLM JSON orientation_states[" - f"{index}] must contain exactly object_id and orientation_state." + "VLM JSON nodes[" + f"{index}] must contain exactly object_id, parent_id, " + "parent_relation, and orientation_state." ) - object_id = state_value["object_id"] - orientation_state = state_value["orientation_state"] + object_id = node_value["object_id"] + parent_id = node_value["parent_id"] + parent_relation = node_value["parent_relation"] + orientation_state = node_value["orientation_state"] if not isinstance(object_id, str) or not object_id: - raise ValueError( - f"VLM JSON orientation_states[{index}].object_id is invalid." - ) + raise ValueError(f"VLM JSON nodes[{index}].object_id is invalid.") + if not isinstance(parent_id, str) or not parent_id: + raise ValueError(f"VLM JSON nodes[{index}].parent_id is invalid.") + if parent_relation != "on": + raise ValueError(f"VLM JSON nodes[{index}].parent_relation is invalid.") if orientation_state not in {None, "standing", "lying"}: - raise ValueError( - f"VLM JSON orientation_states[{index}].orientation_state is invalid." - ) - if object_id in orientation_states_by_id: - raise ValueError(f"VLM JSON repeats orientation state for {object_id!r}.") - orientation_states_by_id[object_id] = orientation_state + raise ValueError(f"VLM JSON nodes[{index}].orientation_state is invalid.") + if object_id in nodes_by_id: + raise ValueError(f"VLM JSON repeats scene graph node for {object_id!r}.") + nodes_by_id[object_id] = SceneGraphNode( + object_id=object_id, + parent_id=parent_id, + parent_relation=parent_relation, + orientation_state=orientation_state, + ) - if set(orientation_states_by_id) != set(asset_ids): + if set(nodes_by_id) != set(asset_ids): raise ValueError( - "VLM JSON orientation states must match all supplied asset IDs." + "VLM JSON scene graph nodes must match all supplied asset IDs." ) - return orientation_states_by_id + return SceneGraph( + nodes=[ + SceneGraphNode(object_id=TABLE_OBJECT_ID, parent_id=None), + *(nodes_by_id[asset_id] for asset_id in asset_ids), + ] + ) def _analyze_image_objects( diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py index 6a61a4d1b..a43db730c 100644 --- a/tests/gen_sim/scene_engine/test_scene_generation.py +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -17,14 +17,20 @@ from __future__ import annotations import numpy as np +import pytest from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene_graph import ( SceneGraph, SceneGraphNode, ) +from embodichain.gen_sim.scene_engine.core.scene import Scene, SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + _apply_root_layout_updates_to_descendant_subtrees, + _project_child_aabb_centers_into_parent_aabb, + _refine_on_children_bfs, _scene_graph_based_calibration, + _table_on_asset_ids, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, @@ -93,3 +99,213 @@ def test_scene_graph_calibration_makes_standing_asset_vertical() -> None: _z_up_rotation_from_y_up_layout(calibrated_layouts[1]), lying_rotation, ) + + +def test_table_root_update_propagates_its_pose_delta_to_descendants() -> None: + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", + parent_id="table", + parent_relation="on", + ), + SceneGraphNode( + object_id="pen_001", + parent_id="book_001", + parent_relation="on", + ), + ] + ) + book_layout = { + "id": "book_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.1, 0.2, 0.3], + "scale": [1.0, 1.0, 1.0], + } + pen_layout = { + "id": "pen_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.2, 0.25, 0.35], + "scale": [1.0, 1.0, 1.0], + } + updated_book_layout = { + **book_layout, + "pos": [0.6, -0.1, 0.4], + } + + refined_layouts = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=[book_layout, pen_layout], + updated_root_layouts=[updated_book_layout], + root_matrices_before_update={ + "book_001": layout_object_to_transform_matrix(book_layout) + }, + ) + + assert _table_on_asset_ids(scene_graph=scene_graph, table_id="table") == { + "book_001" + } + assert np.allclose( + layout_object_to_transform_matrix(refined_layouts[1])[:3, 3], + [0.7, -0.05, 0.45], + ) + + +def test_on_children_bfs_refines_every_non_table_parent( + monkeypatch, + tmp_path, +) -> None: + class FakeParentSurfaceLayoutOptimizer: + refined_child_ids_by_call: list[list[str]] = [] + + def optimize(self, problem): + self.refined_child_ids_by_call.append(problem.child_ids) + return problem.child_seed_xy_by_id + + def settle_dynamic_children(self, **_: object): + return {} + + def fake_measure_scene_object_z_up_world_aabb(*, scene_object: SceneObject): + assert scene_object.pos is not None + x, y_up, z_up_negative_y = scene_object.pos + return [ + [x - 0.05, -z_up_negative_y - 0.05, y_up - 0.05], + [x + 0.05, -z_up_negative_y + 0.05, y_up + 0.05], + ] + + def fake_place_on_support( + *, + scene_object: SceneObject, + support_region_z: float, + center_xy: list[float], + clearance_m: float, + ) -> None: + scene_object.pos = [ + center_xy[0], + support_region_z + clearance_m, + -center_xy[1], + ] + + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation.ParentSurfaceLayoutOptimizer", + FakeParentSurfaceLayoutOptimizer, + ) + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation.measure_scene_object_z_up_world_aabb", + fake_measure_scene_object_z_up_world_aabb, + ) + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation.update_scene_object_y_up_pose_from_z_up_support", + fake_place_on_support, + ) + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="book", + description="book", + ), + SceneObject( + id="pen_001", + kind="asset", + category="pen", + name="pen", + description="pen", + ), + SceneObject( + id="eraser_001", + kind="asset", + category="eraser", + name="eraser", + description="eraser", + ), + ] + ) + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="book_001", parent_id="table", parent_relation="on" + ), + SceneGraphNode( + object_id="pen_001", parent_id="book_001", parent_relation="on" + ), + SceneGraphNode( + object_id="eraser_001", parent_id="pen_001", parent_relation="on" + ), + ] + ) + table_layout = { + "id": "table", + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + assets_layout = [ + { + "id": "book_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.1, 0.0], + "scale": [1.0, 1.0, 1.0], + }, + { + "id": "pen_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.01, 0.2, 0.0], + "scale": [1.0, 1.0, 1.0], + }, + { + "id": "eraser_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.02, 0.3, 0.0], + "scale": [1.0, 1.0, 1.0], + }, + ] + + _refine_on_children_bfs( + scene=scene, + scene_graph=scene_graph, + table_layout=table_layout, + assets_layout=assets_layout, + table_root_ids={"book_001"}, + debug_output_root=tmp_path, + ) + + assert FakeParentSurfaceLayoutOptimizer.refined_child_ids_by_call == [ + ["pen_001"], + ["eraser_001"], + ] + assert (tmp_path / "parent_book_001_child_aabb_projection_2d.png").is_file() + assert (tmp_path / "parent_book_001_child_aabb_optimization_2d.png").is_file() + + +def test_parent_aabb_projection_uses_the_nearest_valid_child_center() -> None: + projected_centers, projected_aabbs = _project_child_aabb_centers_into_parent_aabb( + parent_aabb_xy=[[0.0, 0.0], [1.0, 1.0]], + child_aabbs_xy_by_id={ + "pen_001": np.array([[-0.4, 0.3], [0.0, 0.7]]), + }, + ) + + assert projected_centers == {"pen_001": [0.2, 0.5]} + assert np.allclose(projected_aabbs["pen_001"], [[0.0, 0.3], [0.4, 0.7]]) + + +def test_parent_aabb_projection_rejects_a_child_that_cannot_fit() -> None: + with pytest.raises(ValueError, match="cannot fit inside its parent AABB"): + _project_child_aabb_centers_into_parent_aabb( + parent_aabb_xy=[[0.0, 0.0], [1.0, 1.0]], + child_aabbs_xy_by_id={ + "book_001": np.array([[0.0, 0.0], [1.1, 0.2]]), + }, + ) diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 382e4933d..16611e866 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -151,13 +151,18 @@ def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: ) -def test_initial_scene_graph_places_every_asset_on_table(tmp_path: Path) -> None: +def test_initial_scene_graph_places_an_asset_on_the_table(tmp_path: Path) -> None: class VLM: def complete(self, **_: object) -> str: return json.dumps( { - "orientation_states": [ - {"object_id": "cup_001", "orientation_state": None}, + "nodes": [ + { + "object_id": "cup_001", + "parent_id": "table", + "parent_relation": "on", + "orientation_state": None, + }, ] } ) @@ -210,7 +215,7 @@ def complete(self, **_: object) -> str: } -def test_scene_graph_initialization_uses_image_orientation_states( +def test_scene_graph_initialization_uses_image_support_and_orientation( tmp_path: Path, ) -> None: class VLM: @@ -221,12 +226,19 @@ def complete(self, **_: object) -> str: self.user_prompt = _["user_prompt"] # type: ignore[assignment,index] return json.dumps( { - "orientation_states": [ + "nodes": [ { "object_id": "bottle_001", + "parent_id": "table", + "parent_relation": "on", "orientation_state": "standing", }, - {"object_id": "book_001", "orientation_state": "lying"}, + { + "object_id": "book_001", + "parent_id": "table", + "parent_relation": "on", + "orientation_state": "lying", + }, ] } ) @@ -267,7 +279,20 @@ def complete(self, **_: object) -> str: ) assert json.loads(vlm.user_prompt or "{}") == { - "asset_ids": ["bottle_001", "book_001"], + "assets": [ + { + "id": "bottle_001", + "category": "bottle", + "name": "blue bottle", + "description": "A blue bottle.", + }, + { + "id": "book_001", + "category": "book", + "name": "red book", + "description": "A red book.", + }, + ], } assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" assert scene_graph.node_by_id()["book_001"].orientation_state == "lying" @@ -281,10 +306,17 @@ def __init__(self) -> None: self.responses = [ json.dumps( { - "orientation_states": [ - {"object_id": "table", "orientation_state": None}, + "nodes": [ + { + "object_id": "table", + "parent_id": "table", + "parent_relation": "on", + "orientation_state": None, + }, { "object_id": "bottle_001", + "parent_id": "table", + "parent_relation": "on", "orientation_state": "standing", }, ] @@ -292,9 +324,11 @@ def __init__(self) -> None: ), json.dumps( { - "orientation_states": [ + "nodes": [ { "object_id": "bottle_001", + "parent_id": "table", + "parent_relation": "on", "orientation_state": "standing", }, ] @@ -336,6 +370,91 @@ def complete(self, **_: object) -> str: assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" +def test_scene_graph_initialization_retries_a_response_with_a_parent_cycle( + tmp_path: Path, +) -> None: + class VLM: + def __init__(self) -> None: + self.responses = [ + json.dumps( + { + "nodes": [ + { + "object_id": "book_001", + "parent_id": "pen_001", + "parent_relation": "on", + "orientation_state": None, + }, + { + "object_id": "pen_001", + "parent_id": "book_001", + "parent_relation": "on", + "orientation_state": "lying", + }, + ] + } + ), + json.dumps( + { + "nodes": [ + { + "object_id": "book_001", + "parent_id": "table", + "parent_relation": "on", + "orientation_state": None, + }, + { + "object_id": "pen_001", + "parent_id": "book_001", + "parent_relation": "on", + "orientation_state": "lying", + }, + ] + } + ), + ] + + def complete(self, **_: object) -> str: + return self.responses.pop(0) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + ), + SceneObject( + id="pen_001", + kind="asset", + category="pen", + name="black pen", + description="A black pen.", + ), + ] + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + json_max_attempts=2, + ) + + assert scene_graph.node_by_id()["pen_001"].parent_id == "book_001" + + def test_scene_graph_initialization_requires_asset_mask_id_overlay( tmp_path: Path, ) -> None: @@ -359,7 +478,7 @@ def test_scene_graph_initialization_requires_asset_mask_id_overlay( ) -def test_scene_graph_initialization_info_lists_asset_ids() -> None: +def test_scene_graph_initialization_info_lists_asset_metadata() -> None: scene = Scene( objects=[ SceneObject( @@ -393,5 +512,80 @@ def test_scene_graph_initialization_info_lists_asset_ids() -> None: ) assert simplified_scene_info == { - "asset_ids": ["bottle_001", "book_001"], + "assets": [ + { + "id": "bottle_001", + "category": "bottle", + "name": "blue bottle", + "description": "A blue bottle.", + }, + { + "id": "book_001", + "category": "book", + "name": "red book", + "description": "A red book.", + }, + ], } + + +def test_initial_scene_graph_places_a_lying_asset_on_an_asset(tmp_path: Path) -> None: + class VLM: + def complete(self, **_: object) -> str: + return json.dumps( + { + "nodes": [ + { + "object_id": "book_001", + "parent_id": "table", + "parent_relation": "on", + "orientation_state": None, + }, + { + "object_id": "pen_001", + "parent_id": "book_001", + "parent_relation": "on", + "orientation_state": "lying", + }, + ] + } + ) + + overlay_path = tmp_path / "asset_masks_with_ids.png" + overlay_path.write_bytes(b"png") + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="wooden table", + description="A wooden table.", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="red book", + description="A red book.", + ), + SceneObject( + id="pen_001", + kind="asset", + category="pen", + name="black pen", + description="A black pen.", + ), + ] + ) + + scene_graph = scene_understanding._initialize_scene_graph_from_segmented_scene( + scene, + asset_mask_id_overlay_path=overlay_path, + vlm_client=VLM(), # type: ignore[arg-type] + ) + + pen = scene_graph.node_by_id()["pen_001"] + assert pen.parent_id == "book_001" + assert pen.parent_relation == "on" + assert pen.orientation_state == "lying" From 86ef235111d26feda15432b18614840225af5b9b Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:28:45 +0800 Subject: [PATCH 36/85] Delete the debug code for coarse layout from server output --- .../pipeline/generation/scene_generation.py | 47 ------------------- 1 file changed, 47 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index ef14e4d79..338f2786f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -232,57 +232,10 @@ def _generate_coarse_results_from_masks( json.dumps(coarse_layout, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) - _write_baked_coarse_layout_debug_glbs( - coarse_layout=coarse_layout, - coarse_geometry_root=coarse_geometry_output_root, - debug_output_root=debug_output_root, - ) # Nothing to be returned. return None -def _write_baked_coarse_layout_debug_glbs( - *, - coarse_layout: list[dict[str, object]], - coarse_geometry_root: str | Path, - debug_output_root: str | Path, -) -> Path: - """Write per-object and combined GLBs with each coarse pose baked in.""" - resolved_geometry_root = Path(coarse_geometry_root).expanduser().resolve() - baked_root = ( - Path(debug_output_root).expanduser().resolve() / "coarse_layout_baked_glbs" - ) - baked_root.mkdir(parents=True, exist_ok=True) - - baked_scene = trimesh.Scene() - baked_object_ids: set[str] = set() - for layout_object in coarse_layout: - object_id = layout_object.get("id") - if not isinstance(object_id, str) or not object_id: - raise ValueError( - "Each coarse layout object must contain a non-empty string id." - ) - if object_id in baked_object_ids: - raise ValueError(f"Coarse layout contains duplicate id {object_id!r}.") - baked_object_ids.add(object_id) - - # The geometry server declares these poses in the GLB y-up convention. - baked_mesh = load_glb_mesh(resolved_geometry_root / f"{object_id}.glb") - baked_mesh.apply_transform(layout_object_to_transform_matrix(layout_object)) - baked_mesh.export(baked_root / f"{object_id}.glb", file_type="glb") - baked_scene.add_geometry( - baked_mesh, - node_name=object_id, - geom_name=object_id, - ) - - if not baked_object_ids: - raise ValueError("Cannot write baked coarse debug GLBs for an empty layout.") - baked_scene.export(baked_root / "scene.glb", file_type="glb") - log_info(f"Wrote baked coarse-pose debug GLBs: {baked_root}") - return baked_root - - def _update_scene_final_y_up_layout_and_z_up_centers( *, scene: Scene, From 0865b90579052fa8f9764cf402c1b93758d9c587 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:46:28 +0800 Subject: [PATCH 37/85] support lying object VLM-based simready --- .../pipeline/generation/scene_generation.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 338f2786f..e601836c5 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -126,12 +126,12 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - # Coarse poses already preserve lying and unconstrained assets; only standing - # assets need a VLM semantic-axis correction before later z-up calibration. - standing_orientation_states_by_id = { + # Explicit graph orientations use VLM front/top views before SimReady + # canonicalization; null leaves the geometry-server pose unchanged. + orientation_states_by_id = { node.object_id: node.orientation_state for node in scene_graph.nodes - if node.orientation_state == "standing" + if node.orientation_state is not None } simready_processor = SimReadyProcessor( scene=scene, @@ -139,11 +139,11 @@ def generate_scene_and_refine( coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, debug_output_root=debug_output_root, - # Keep the geometry-server scale and only correct unstable standing poses. + # Keep geometry-server scale while applying explicit orientation semantics. config=SimReadyProcessorConfig( use_vlm_scale=False, use_vlm_rotation=False, - orientation_states_by_id=standing_orientation_states_by_id, + orientation_states_by_id=orientation_states_by_id, ), vlm_client=vlm_client, ) From 3f8cba3dc235a7cc88a7f68c15759336f3ef137d Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Tue, 25 Aug 2026 20:18:44 +0800 Subject: [PATCH 38/85] add: 1. articulated or not 2. segmented RGBA image for each asset --- .../gen_sim/scene_engine/core/scene_object.py | 4 + .../generation/scene_understanding.py | 69 +++++++++++++-- .../utils/image_segmentation_utils.py | 61 +++++++++++++ .../pipeline/utils/scene_exporter.py | 1 + .../pipeline/utils/scene_importer.py | 4 + .../test_scene_core_and_export.py | 3 + tests/gen_sim/scene_engine/test_scene_edit.py | 2 + .../scene_engine/test_scene_understanding.py | 87 +++++++++++++++++++ 8 files changed, 224 insertions(+), 7 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index 2b868e3c3..251dd0762 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -61,7 +61,9 @@ class SceneObject: category: str # Semantic category identified by scene understanding. name: str # Human-readable visual name. description: str # Detailed semantic and spatial description. + is_articulated: bool = False # Whether this object has movable links or joints. mask_path: str | None = None # Absolute path to the validated binary image mask. + visible_rgba_path: str | None = None # None for future unsegmented objects. simready_glb_path: str | None = None # Absolute path to the canonical SimReady GLB. rot: list[float] | None = None # Final y-up Euler XYZ rotation in degrees. pos: list[float] | None = None # Final y-up world position in metres. @@ -80,7 +82,9 @@ def to_dict(self) -> dict[str, object]: "category": self.category, "name": self.name, "description": self.description, + "is_articulated": self.is_articulated, "mask_path": self.mask_path, + "visible_rgba_path": self.visible_rgba_path, "simready_glb_path": self.simready_glb_path, "rot": self.rot, "pos": self.pos, diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index 4f7839916..f8fb48d50 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -45,11 +45,13 @@ render_image_without_masks, render_numbered_mask_candidates, save_binary_mask, + save_visible_rgba_crop, union_overlapping_mask_candidates, ) _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} _CATEGORY_PATTERN = re.compile(r"^[a-z][a-z0-9_]*$") +_VISIBLE_RGBA_IMAGE_SIZE = (512, 512) _SYSTEM_PROMPT = """You inspect one tabletop-scene image. Identify the main table and every visible, physically distinct object that should be segmented and later generated as an independent 3D asset. @@ -79,26 +81,39 @@ Structural direction words are allowed when they describe the object itself: "bottle with a black cap on top" is valid, while "bottle on the left of the table" is not. +9. is_articulated is true only for articulated objects with functional movable + parts that matter in simulation. Typical true examples are a microwave with + a door, cabinet, button, or drawer. Treat every other object as false unless + its independently movable links or joints are clearly visible; in particular, + rigid objects such as bottles, mugs, bowls, books, utensils, and boxes are false. Return JSON only: no Markdown, comments, or prose outside this exact schema: { "table": { "category": "coffee_table", "name": "light wood coffee table", - "description": "low rectangular light wood coffee table with a smooth wood surface" + "description": "low rectangular light wood coffee table with a smooth wood surface", + "is_articulated": false }, "assets": [ { "category": "mug", "name": "blue ceramic mug", - "description": "small blue ceramic mug with a curved handle" + "description": "small blue ceramic mug with a curved handle", + "is_articulated": false + }, + { + "category": "drawer", + "name": "white storage drawer", + "description": "white rectangular storage drawer with a horizontal pull handle", + "is_articulated": true } ] } For two identical blue mugs, output two asset entries with the same category, name, and description. Do not infer objects that are not visible. Use an empty assets array when no objects are visible. Every field must be a non-empty -string.""" +string. is_articulated must be a boolean.""" _USER_PROMPT = "Analyze the provided image and return only the required JSON object." @@ -490,18 +505,19 @@ def _parse_scene_object_fields( value: object, *, field_name: str, -) -> dict[str, str]: +) -> dict[str, str | bool]: if not isinstance(value, dict) or set(value) != { "category", "name", "description", + "is_articulated", }: raise ValueError( f"VLM JSON key {field_name} must contain exactly category, name, and " - "description." + "description, and is_articulated." ) - fields = {} + fields: dict[str, str | bool] = {} for key in ("category", "name", "description"): raw_value = value[key] if not isinstance(raw_value, str) or not raw_value.strip(): @@ -510,7 +526,14 @@ def _parse_scene_object_fields( ) fields[key] = raw_value.strip() - if not _CATEGORY_PATTERN.fullmatch(fields["category"]): + is_articulated = value["is_articulated"] + if not isinstance(is_articulated, bool): + raise ValueError(f"VLM JSON key {field_name}.is_articulated must be a boolean.") + fields["is_articulated"] = is_articulated + + category = fields["category"] + assert isinstance(category, str) + if not _CATEGORY_PATTERN.fullmatch(category): raise ValueError( f"VLM JSON key {field_name}.category must be a lower-case snake_case " "class name." @@ -539,8 +562,12 @@ def _segment_scene( masks_output_root = ( Path(stage_output_root) / "masks" ) # Keeps the validated masked images of each assets (include the table) + object_images_output_root = ( + Path(stage_output_root) / "object_images" + ) # Keeps fixed-size RGBA visual observations outside the debug directory. debug_output_root.mkdir() masks_output_root.mkdir() + object_images_output_root.mkdir() # Segment the table and assets with VLM validation separately. _segment_assets( @@ -574,6 +601,12 @@ def _segment_scene( vlm_client=vlm_client, image_segmentation_client=image_segmentation_client, ) + # Save the visible RGBA observations. + _save_visible_rgba_observations( + image_path=image_path, + output_root=object_images_output_root, + scene=scene, + ) asset_masks: list[tuple[str, str]] = [] for asset in scene.assets: if asset.mask_path is None: @@ -586,6 +619,28 @@ def _segment_scene( ) +def _save_visible_rgba_observations( + *, + image_path: str | Path, + output_root: str | Path, + scene: Scene, +) -> None: + """Save visual evidence only for objects with a validated binary mask.""" + for scene_object in scene.objects: + # No mask means no trustworthy image observation; later tools can use semantics instead. + scene_object.visible_rgba_path = None + if scene_object.mask_path is None: + continue + scene_object.visible_rgba_path = str( + save_visible_rgba_crop( + image_path=image_path, + mask_path=scene_object.mask_path, + output_path=Path(output_root) / f"{scene_object.id}_rgba.png", + output_size=_VISIBLE_RGBA_IMAGE_SIZE, + ) + ) + + def _segment_table( image_path: str | Path, validation_image_path: str | Path, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py index 302ec2cbf..a5c9bff72 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/image_segmentation_utils.py @@ -18,6 +18,7 @@ from __future__ import annotations from dataclasses import dataclass +import math from pathlib import Path from typing import Any @@ -185,6 +186,66 @@ def save_binary_mask( return resolved_output_path +def save_visible_rgba_crop( + *, + image_path: str | Path, + mask_path: str | Path, + output_path: str | Path, + output_size: tuple[int, int] = (512, 512), + padding_ratio: float = 0.1, +) -> Path: + """Save a fixed-size RGBA crop of one object's visible source-image pixels.""" + if len(output_size) != 2 or not all( + isinstance(value, int) and value > 0 for value in output_size + ): + raise ValueError("RGBA crop output_size must contain two positive integers.") + if padding_ratio < 0: + raise ValueError("RGBA crop padding_ratio must be non-negative.") + + with Image.open(image_path) as loaded_image: + image = loaded_image.convert("RGB") + with Image.open(mask_path) as loaded_mask: + mask = loaded_mask.convert("L") + _require_image_size(mask, image.size) + bbox = mask.getbbox() + if bbox is None: + raise ValueError("Cannot save an RGBA crop for an empty binary mask.") + + left, top, right, bottom = bbox + padding = math.ceil(max(right - left, bottom - top) * padding_ratio) + crop_bounds = ( + max(0, left - padding), + max(0, top - padding), + min(image.width, right + padding), + min(image.height, bottom + padding), + ) + image_crop = image.crop(crop_bounds) + mask_crop = mask.crop(crop_bounds) + scale = min( + output_size[0] / image_crop.width, + output_size[1] / image_crop.height, + ) + resized_size = ( + max(1, round(image_crop.width * scale)), + max(1, round(image_crop.height * scale)), + ) + rgba_crop = image_crop.resize(resized_size, Image.Resampling.LANCZOS).convert( + "RGBA" + ) + rgba_crop.putalpha(mask_crop.resize(resized_size, Image.Resampling.NEAREST)) + + rgba_canvas = Image.new("RGBA", output_size, (0, 0, 0, 0)) + paste_xy = ( + (output_size[0] - resized_size[0]) // 2, + (output_size[1] - resized_size[1]) // 2, + ) + rgba_canvas.alpha_composite(rgba_crop, dest=paste_xy) + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + rgba_canvas.save(resolved_output_path) + return resolved_output_path + + def render_image_without_masks( *, image_path: str | Path, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index 7451296b4..c8f9a4591 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -213,6 +213,7 @@ def _scene_object_config( "category": scene_object.category, "name": scene_object.name, "description": scene_object.description, + "is_articulated": scene_object.is_articulated, "shape": { "shape_type": "Mesh", "fpath": asset_relative_path, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index e730b88cb..bf5d9caea 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -280,6 +280,9 @@ def _scene_object_from_export_entry( entry.get("support_optimization_rect_xy"), field_name=f"{uid}.support_optimization_rect_xy", ) + is_articulated = entry.get("is_articulated") + if not isinstance(is_articulated, bool): + raise ValueError(f"Scene object {uid!r} is_articulated must be a boolean.") pos_y_up = _Z_UP_TO_Y_UP_ROTATION @ np.asarray(pos_z_up, dtype=float) rotation_z_up = Rotation.from_euler("XYZ", rot_z_up, degrees=True).as_matrix() @@ -302,6 +305,7 @@ def _scene_object_from_export_entry( default=uid, ), description=str(entry.get("description") or uid), + is_articulated=is_articulated, simready_glb_path=str(glb_path), rot=rot_y_up.tolist(), pos=pos_y_up.tolist(), diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 95658cf16..19df695ee 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -146,6 +146,7 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No physics=_physics("dynamic"), ) asset.center_xy = [0.25, -0.5] + asset.is_articulated = True scene = Scene(objects=[table, asset]) export_path = SceneExporter( @@ -163,6 +164,7 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert entry["uid"] == "cup" assert entry["category"] == "asset" assert entry["name"] == "cup" + assert entry["is_articulated"] is True assert entry["body_type"] == "dynamic" assert entry["init_pos"] == [1.0, -3.0, 2.0] assert entry["body_scale"] == [1.0, 2.0, 3.0] @@ -194,6 +196,7 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert [asset.id for asset in imported_scene.assets] == ["cup"] assert imported_scene.assets[0].category == "asset" assert imported_scene.assets[0].name == "cup" + assert imported_scene.assets[0].is_articulated is True assert imported_graph.to_dict() == _scene_graph(scene).to_dict() diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py index fd1a9e933..72ee061a3 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit.py +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -55,6 +55,7 @@ def _write_scene_export( "init_pos": [0.0, 0.0, 0.0], "init_rot": [0.0, 0.0, 0.0], "body_scale": [1.0, 1.0, 1.0], + "is_articulated": False, "max_convex_hull_num": 16, } ) @@ -72,6 +73,7 @@ def _write_scene_export( "init_pos": [1.0, -3.0, 2.0], "init_rot": [0.0, 0.0, 0.0], "body_scale": [1.0, 2.0, 3.0], + "is_articulated": False, "center_xy": [1.0, -3.0], "max_convex_hull_num": 32, } diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 16611e866..24c8b6261 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -29,6 +29,7 @@ from embodichain.gen_sim.scene_engine.pipeline.utils import image_segmentation_utils from embodichain.gen_sim.scene_engine.pipeline.utils.image_segmentation_utils import ( render_asset_mask_id_overlay, + save_visible_rgba_crop, ) @@ -39,12 +40,14 @@ def _response(*, asset_name: str = "cup") -> str: "category": "dining_table", "name": "wooden table", "description": "A rectangular wooden table.", + "is_articulated": False, }, "assets": [ { "category": "cup", "name": asset_name, "description": "A small ceramic cup.", + "is_articulated": False, } ], } @@ -59,6 +62,19 @@ def test_image_object_analysis_parses_code_fence_and_assigns_stable_ids() -> Non assert scene.table is not None assert scene.table.id == "table" assert [asset.id for asset in scene.assets] == ["cup_001"] + assert scene.assets[0].is_articulated is False + + +def test_image_object_analysis_preserves_articulated_object_metadata() -> None: + response = json.loads(_response()) + response["assets"][0]["category"] = "drawer" + response["assets"][0]["is_articulated"] = True + + scene = scene_understanding._parse_image_object_analysis_response( + json.dumps(response) + ) + + assert scene.assets[0].is_articulated is True def test_image_object_analysis_accepts_name_with_spatial_words() -> None: @@ -129,6 +145,77 @@ def test_asset_mask_id_overlay_excludes_the_table_mask(tmp_path: Path) -> None: assert overlay.getpixel((377, 180)) != (0, 0, 0) +def test_visible_rgba_crop_preserves_object_pixels_on_a_fixed_canvas( + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + mask_path = tmp_path / "cup_mask.png" + output_path = tmp_path / "object_images" / "cup_001_rgba.png" + image_size = (100, 60) + Image.new("RGB", image_size, "black").save(image_path) + with Image.open(image_path) as image: + image = image.copy() + ImageDraw.Draw(image).rectangle((30, 20, 70, 40), fill="red") + image.save(image_path) + + mask = Image.new("L", image_size, 0) + ImageDraw.Draw(mask).rectangle((30, 20, 70, 40), fill=255) + mask.save(mask_path) + + rendered_path = save_visible_rgba_crop( + image_path=image_path, + mask_path=mask_path, + output_path=output_path, + ) + + with Image.open(rendered_path) as rgba: + assert rgba.mode == "RGBA" + assert rgba.size == (512, 512) + assert rgba.getpixel((256, 256))[:3] == (255, 0, 0) + assert rgba.getpixel((0, 0))[3] == 0 + + +def test_visible_rgba_observations_leave_unsegmented_objects_empty( + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + mask_path = tmp_path / "cup_mask.png" + image_size = (40, 40) + Image.new("RGB", image_size, "red").save(image_path) + mask = Image.new("L", image_size, 0) + ImageDraw.Draw(mask).rectangle((10, 10, 30, 30), fill=255) + mask.save(mask_path) + scene = Scene( + objects=[ + SceneObject( + id="cup_001", + kind="asset", + category="cup", + name="red cup", + description="A red cup.", + mask_path=str(mask_path), + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="blue book", + description="A blue book.", + ), + ] + ) + + scene_understanding._save_visible_rgba_observations( + image_path=image_path, + output_root=tmp_path / "object_images", + scene=scene, + ) + + assert scene.assets[0].visible_rgba_path is not None + assert Path(scene.assets[0].visible_rgba_path).is_file() + assert scene.assets[1].visible_rgba_path is None + + def test_asset_mask_id_label_font_fits_the_mask_bbox() -> None: image_size = (512, 512) mask_bbox = (380, 180, 450, 360) From a9afa6620e728fd3e6be12879ada8b84c57c0ae7 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 27 Aug 2026 15:54:38 +0800 Subject: [PATCH 39/85] 1. VLM judge lying or standing 2. VLM give a yaw 3. updated render tool: now add a small slice of floor 4. need to be test --- .../pipeline/generation/scene_generation.py | 169 ++++- .../pipeline/utils/simready_processor.py | 60 +- .../utils/simready_processor_utils.py | 600 +++++++++++++++--- .../pipeline/utils/visual_yaw_optimizer.py | 494 ++++++++++++++ .../scene_engine/test_scene_generation.py | 262 ++++++++ .../test_simready_processor_utils.py | 158 ++++- 6 files changed, 1604 insertions(+), 139 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/visual_yaw_optimizer.py diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index e601836c5..2a7e526dd 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -37,7 +37,10 @@ GeometryGenerationClient, ) from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + TABLE_OBJECT_ID, + SceneGraph, +) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( OpenAICompatibleVLM, @@ -74,6 +77,9 @@ SimReadyProcessor, SimReadyProcessorConfig, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.visual_yaw_optimizer import ( + VisualYawOptimizer, +) from embodichain.utils.logger import log_info _SUPPORTED_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png"} @@ -126,12 +132,12 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - # Explicit graph orientations use VLM front/top views before SimReady - # canonicalization; null leaves the geometry-server pose unchanged. + # Every graph asset, including null, receives a VLM check before SimReady + # canonicalization; null requests its natural stable tabletop pose. orientation_states_by_id = { node.object_id: node.orientation_state for node in scene_graph.nodes - if node.orientation_state is not None + if node.object_id != TABLE_OBJECT_ID } simready_processor = SimReadyProcessor( scene=scene, @@ -148,6 +154,18 @@ def generate_scene_and_refine( vlm_client=vlm_client, ) simready_assets_layout = simready_processor.process_assets() + # Replace unreliable coarse rotations with VLM-observed canonical z-up yaw. + visual_yaws_by_id = _optimize_simready_asset_visual_yaws( + scene=scene, + simready_assets_layout=simready_assets_layout, + coarse_layout_by_id=coarse_layout_by_id, + vlm_client=vlm_client, + debug_output_root=debug_output_root / "visual_yaw", + ) + simready_assets_layout = _apply_visual_yaws_to_simready_asset_layouts( + simready_assets_layout=simready_assets_layout, + z_up_yaws_degrees_by_id=visual_yaws_by_id, + ) simready_table_layout = simready_processor.process_table() # Concat then save the table info and the assets info in one JSON file. simready_layout = [simready_table_layout, *simready_assets_layout] @@ -172,6 +190,86 @@ def generate_scene_and_refine( return scene +def _optimize_simready_asset_visual_yaws( + *, + scene: Scene, + simready_assets_layout: list[dict[str, object]], + coarse_layout_by_id: dict[str, dict[str, object]], + vlm_client: OpenAICompatibleVLM, + debug_output_root: str | Path, +) -> dict[str, float]: + """Query one absolute canonical z-up yaw for every observed SimReady asset.""" + assets_by_id = {asset.id: asset for asset in scene.assets} + layout_ids = [layout.get("id") for layout in simready_assets_layout] + if not all(isinstance(layout_id, str) and layout_id for layout_id in layout_ids): + raise ValueError("Every SimReady asset layout must contain a non-empty id.") + if len(layout_ids) != len(set(layout_ids)): + raise ValueError("SimReady asset layouts must have unique ids.") + if set(layout_ids) != set(assets_by_id): + raise ValueError("SimReady asset layouts must match the scene asset ids.") + + yaws_degrees_by_id: dict[str, float] = {} + for asset_layout in simready_assets_layout: + layout_id = asset_layout["id"] + assert isinstance(layout_id, str) + coarse_layout = coarse_layout_by_id.get(layout_id) + if coarse_layout is None: + raise ValueError(f"Coarse layout does not contain asset {layout_id!r}.") + coarse_scale = coarse_layout.get("scale") + if not isinstance(coarse_scale, list): + raise ValueError( + f"Coarse layout scale for asset {layout_id!r} must be a list." + ) + yaws_degrees_by_id[layout_id] = VisualYawOptimizer( + scene_object=assets_by_id[layout_id], + baked_scale_y_up=coarse_scale, + vlm_client=vlm_client, + debug_output_root=debug_output_root, + ).optimize_z_up_yaw_degrees() + return yaws_degrees_by_id + + +def _apply_visual_yaws_to_simready_asset_layouts( + *, + simready_assets_layout: list[dict[str, object]], + z_up_yaws_degrees_by_id: dict[str, float], +) -> list[dict[str, object]]: + """Keep SimReady positions but replace each coarse rotation with canonical yaw.""" + layout_ids = [layout.get("id") for layout in simready_assets_layout] + if not all(isinstance(layout_id, str) and layout_id for layout_id in layout_ids): + raise ValueError("Every SimReady asset layout must contain a non-empty id.") + typed_layout_ids = [str(layout_id) for layout_id in layout_ids] + if len(typed_layout_ids) != len(set(typed_layout_ids)): + raise ValueError("SimReady asset layouts must have unique ids.") + if set(z_up_yaws_degrees_by_id) != set(typed_layout_ids): + raise ValueError("Visual yaws must match the SimReady asset layout ids.") + + y_up_to_z_up_matrix = np.eye(4) + y_up_to_z_up_matrix[:3, :3] = Rotation.from_euler( + "x", 90.0, degrees=True + ).as_matrix() + z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) + yawed_layouts: list[dict[str, object]] = [] + for asset_layout, asset_id in zip(simready_assets_layout, typed_layout_ids): + original_matrix = layout_object_to_transform_matrix(asset_layout) + z_up_yaw_matrix = np.eye(4) + z_up_yaw_matrix[:3, :3] = Rotation.from_euler( + "z", z_up_yaws_degrees_by_id[asset_id], degrees=True + ).as_matrix() + # The canonical SimReady pose replaces the coarse layout rotation. + canonical_y_up_matrix = ( + z_up_to_y_up_matrix @ z_up_yaw_matrix @ y_up_to_z_up_matrix + ) + canonical_y_up_matrix[:3, 3] = original_matrix[:3, 3] + canonical_y_up_matrix[:3, :3] = canonical_y_up_matrix[:3, :3] @ np.diag( + np.linalg.norm(original_matrix[:3, :3], axis=0) + ) + yawed_layouts.append( + transform_matrix_to_layout_object(asset_id, canonical_y_up_matrix) + ) + return yawed_layouts + + def _generate_coarse_results_from_masks( image_path: str | Path, debug_output_root: str | Path, @@ -404,23 +502,15 @@ def _layout_refinement( log_info("Scene has no on-table assets; skipping table layout refinement.") return refined_table_layout, refined_assets_layout - # Move the direct on-table assets as one rigid group so their lowest AABB - # point is 2cm above the table. Descendants retain their relative transforms. - table_root_matrices_before_align = _layout_matrices_by_id(table_root_layouts) - group_table_aligner = AssetsGroupTableAligner( - table_layout=refined_table_layout, - assets_layout=table_root_layouts, - geometry_root=simready_geometry_output_root, - ) - # Align. - refined_table_layout, aligned_table_root_layouts = group_table_aligner.align() - refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + # Each on-table root needs its own support height; its descendants follow it. + refined_table_layout, refined_assets_layout = _align_table_roots_individually( scene_graph=scene_graph, + table_layout=refined_table_layout, assets_layout=refined_assets_layout, - updated_root_layouts=aligned_table_root_layouts, - root_matrices_before_update=table_root_matrices_before_align, + table_root_ids=table_root_ids, + geometry_root=simready_geometry_output_root, ) - # After update the layouts, re-select those who are direct on-table children. + # Re-select direct table children after their independent vertical placement. table_root_layouts = _select_asset_layouts( assets_layout=refined_assets_layout, asset_ids=table_root_ids, @@ -582,6 +672,49 @@ def _layout_refinement( return refined_table_layout, refined_assets_layout +def _align_table_roots_individually( + *, + scene_graph: SceneGraph, + table_layout: dict[str, object], + assets_layout: list[dict[str, object]], + table_root_ids: set[str], + geometry_root: str | Path, +) -> tuple[dict[str, object], list[dict[str, object]]]: + """Place each direct on-table root above the table and move its subtree.""" + refined_table_layout = table_layout + refined_assets_layout = assets_layout + # Preserve layout order while each root receives an independent z correction. + initial_table_root_layouts = _select_asset_layouts( + assets_layout=assets_layout, + asset_ids=table_root_ids, + ) + for initial_root_layout in initial_table_root_layouts: + root_id = initial_root_layout.get("id") + if not isinstance(root_id, str) or not root_id: + raise ValueError( + "Every direct on-table layout must contain a non-empty id." + ) + current_root_layouts = _select_asset_layouts( + assets_layout=refined_assets_layout, + asset_ids={root_id}, + ) + root_matrices_before_align = _layout_matrices_by_id(current_root_layouts) + aligned_table_layout, aligned_root_layouts = AssetsGroupTableAligner( + table_layout=refined_table_layout, + assets_layout=current_root_layouts, + geometry_root=geometry_root, + ).align() + refined_table_layout = aligned_table_layout + # Propagating the complete root delta keeps descendants attached to it. + refined_assets_layout = _apply_root_layout_updates_to_descendant_subtrees( + scene_graph=scene_graph, + assets_layout=refined_assets_layout, + updated_root_layouts=aligned_root_layouts, + root_matrices_before_update=root_matrices_before_align, + ) + return refined_table_layout, refined_assets_layout + + def _refine_on_children_bfs( *, scene: Scene, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 864e39a59..fb8d4c2d9 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -38,8 +38,10 @@ LYING_NEEDED_LAYOUT, STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, - query_vlm_object_rotation_and_target_size, + query_vlm_pose_switch_candidate, + query_vlm_object_pose_and_target_size, render_object_front_top_views, + render_object_pose_switch_candidates, rotate_glb_about_x_axis, ) from embodichain.gen_sim.scene_engine.pipeline.utils.table_support_surface import ( @@ -72,7 +74,9 @@ class SimReadyProcessorConfig: use_vlm_scale: bool = False # Use the VLM-selected asset scale. use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. # Explicit graph orientation overrides the default stable tabletop pose. - orientation_states_by_id: dict[str, OrientationState] = field(default_factory=dict) + orientation_states_by_id: dict[str, OrientationState | None] = field( + default_factory=dict + ) class SimReadyProcessor: @@ -205,8 +209,10 @@ def _prepare_vlm_rotated_glb( ) -> tuple[Path, list[float] | None]: """Render, query, and optionally bake the VLM-selected x-axis rotation.""" coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" - orientation_state = self._orientation_state_for_object(scene_object.id) - orientation_pose_required = orientation_state is not None + # A graph entry, including null, requests a VLM check of the desired pose. + orientation_pose_required = ( + scene_object.id in self.config.orientation_states_by_id + ) if not ( self.config.use_vlm_scale or self.config.use_vlm_rotation @@ -217,23 +223,49 @@ def _prepare_vlm_rotated_glb( scene_object, needed_layout=self._needed_layout_for_object(scene_object.id), ) - rotate_about_x = bool(decision["rotate_about_x"]) - vlm_scale = None - if self.config.use_vlm_scale: - # The VLM target describes the final, post-rotation z-up XY footprint. - vlm_scale = compute_uniform_xy_scale_for_target( + pose_action = decision["pose_action"] + if pose_action not in {"keep_current", "rotate_to_required_pose"}: + raise ValueError("VLM pose_action is not a supported semantic action.") + selected_x_rotation_degrees = 0.0 + if pose_action == "rotate_to_required_pose": + # The temporary A/B renders resolve the otherwise ambiguous flip direction. + assert self.vlm_client is not None + candidate_views_path = render_object_pose_switch_candidates( glb_path=coarse_path, - target_xy_size_cm=decision["target_xy_size_cm"], - rotate_about_x=rotate_about_x, + output_path=( + self.debug_output_root + or self.simready_geometry_root.parent / "debug" + ) + / "vlm_pose_candidates" + / f"{scene_object.id}.png", + ) + selected_x_rotation_degrees, _ = query_vlm_pose_switch_candidate( + scene_object_description=scene_object.description, + needed_layout=self._needed_layout_for_object(scene_object.id), + rendered_candidates_path=candidate_views_path, + vlm_client=self.vlm_client, + debug_output_path=( + self.debug_output_root + or self.simready_geometry_root.parent / "debug" + ) + / "vlm_pose_candidates" + / f"{scene_object.id}.json", ) rotated_path = rotate_glb_about_x_axis( input_path=coarse_path, output_path=self.simready_geometry_root / "vlm_rotated" / f"{scene_object.id}.glb", - rotate=rotate_about_x - and (orientation_pose_required or self.config.use_vlm_rotation), + rotation_degrees=selected_x_rotation_degrees, ) + vlm_scale = None + if self.config.use_vlm_scale: + # Measure the selected pose because the VLM target is its final XY footprint. + vlm_scale = compute_uniform_xy_scale_for_target( + glb_path=rotated_path, + target_xy_size_cm=decision["target_xy_size_cm"], + rotate_about_x=False, + ) # The scale flag controls whether this VLM-derived isotropic scale is used. # Apply the same factor on x, y, and z to preserve the asset's proportions. return ( @@ -274,7 +306,7 @@ def _vlm_transform_for_object( output_path=debug_root / "vlm_views" / f"{scene_object.id}.png", ) # Both semantic questions are always answered in one multimodal call. - return query_vlm_object_rotation_and_target_size( + return query_vlm_object_pose_and_target_size( scene_object_description=scene_object.description, needed_layout=needed_layout, rendered_views_path=rendered_path, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py index 4523aa3e6..3b57463ca 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -31,45 +31,91 @@ OpenAICompatibleVLM, ) -_VLM_SYSTEM_PROMPT = """You inspect one isolated 3D object from front and top views. -Use the object description and the rendered views together. +_VLM_SYSTEM_PROMPT = """You inspect one isolated 3D object from elevated oblique and top views. +The object is temporarily visual-normalized and placed on a small neutral +support patch. The patch conveys only local contact and up/down context; do +not infer real-world scale from it. Use the object description, needed layout, +and rendered views together. -Use the rendered views and the needed layout to decide whether the object should -be rotated around its own center by +90 degrees around the z-up world's x axis. +Choose whether the current asset pose already satisfies the needed layout or +whether it needs one pose switch. Do not reason about coordinate-axis rotations. The z-up world is right-handed: x is left-right, y is front-back, and z is up. -In the composed image, FRONT VIEW is the left panel: x is horizontal and z is -vertical; the upper-right marker shows the positive z direction. TOP VIEW is -the right panel: x is horizontal and y is vertical; the upper-right markers -show the positive x and y directions. -Do not confuse the top view with looking at the object from above in the image -description: it is a projection along the z axis onto the x-y plane. -After deciding and applying that rotation, estimate the object's desired AABB -footprint on the x-y plane in real-world centimetres. The first value is the x -size and the second value is the y size. +The OBLIQUE VIEW has x horizontal with visible z height and floor contact; TOP +VIEW is a projection along z onto the x-y plane. For a lying object, its long +dimension is visible in TOP VIEW while the OBLIQUE VIEW shows only a small +vertical thickness. For a standing object, the OBLIQUE VIEW shows its main +height along z while TOP VIEW has a compact footprint. Do not request a pose +switch merely because an object is rotated within the x-y plane; visual yaw is +resolved later. + +Treat TOP VIEW as decisive for a clearly visible long, thin silhouette. Its +screen vertical direction is world +y, not world +z: a knife or whisk that +spans the TOP VIEW from end to end is lying in the x-y table plane even if it +looks vertical on screen or foreshortened in the OBLIQUE VIEW. Return +rotate_to_required_pose only when both views clearly support a required pose +switch. When the needed layout asks for a natural stable pose rather than an +explicit standing or lying state, default to keep_current unless an unstable +tip-, edge-, or tiny-contact placement is unambiguous in both views. + +After deciding, estimate the object's desired AABB footprint on the x-y plane +in real-world centimetres. The first value is the x size and the second value +is the y size. Return JSON only with exactly this schema: { - "rotate_about_x": false, + "pose_action": "keep_current", + "reason": "brief visual justification", "target_xy_size_cm": [12.0, 5.0] } +pose_action must be exactly one of: keep_current, rotate_to_required_pose. +Use keep_current when the current views already satisfy the needed pose. Use +rotate_to_required_pose only when the front/top evidence clearly shows that +the current asset is standing but must lie, or lying but must stand. + Examples: -- Fork lying flat on a table: in FRONT VIEW the fork is mostly a thin - horizontal line; in TOP VIEW its length is visible. Keep it flat with - rotate_about_x=false, and use the tabletop footprint, for example - target_xy_size_cm=[15.0, 3.0]. -- Fork placed in a pen holder: the desired fork is upright, so its long axis is - approximately z. If the input coarse fork is lying in the x-y plane, set - rotate_about_x=true; if the input coarse fork is already upright, set it to - false. The target is the footprint inside the holder, not the fork's full - length, for example target_xy_size_cm=[3.0, 3.0]. -- Fork requested to lie flat on a table even when the input coarse fork is - upright: set rotate_about_x=true and estimate the final flat footprint, for - example target_xy_size_cm=[15.0, 3.0]. -- Bottle already standing on its flat base: keep it upright with - rotate_about_x=false and use target_xy_size_cm=[8.0, 8.0]. +- Fork requested to lie flat: if its length is already visible in TOP VIEW and + the OBLIQUE VIEW is thin, use keep_current; otherwise use rotate_to_required_pose. + Use target_xy_size_cm=[15.0, 3.0]. +- Fork requested upright in a holder: if the OBLIQUE VIEW already shows the fork + height and TOP VIEW is compact, use keep_current; otherwise use + rotate_to_required_pose. Use target_xy_size_cm=[3.0, 3.0]. +- Bottle requested standing on its base: use keep_current only when it already + appears upright in the OBLIQUE VIEW; otherwise use rotate_to_required_pose and use + target_xy_size_cm=[8.0, 8.0]. +""" + +_VLM_POSE_CANDIDATE_SYSTEM_PROMPT = """You select the physically and semantically +correct pose for one isolated 3D object. The image shows two temporary +rendered candidates of the same asset, each visual-normalized and placed on a +small neutral support patch in the same fixed, slightly elevated +robot-manipulation view. The LEFT panel is CANDIDATE A and the RIGHT panel is +CANDIDATE B. The candidates differ only in which of two opposite 90-degree +pose switches was used. + +Use the object description and needed layout to choose the candidate that +satisfies the requested lying, standing, or natural stable resting pose. Pay +attention to obvious up/down semantics: a pan or bowl opening should face up, +a bottle should rest on its base when standing, and a long tool should not be +upside down when the distinction is visible. Do not infer an exact coordinate +axis direction; simply choose the visually correct candidate. + +Return JSON only with exactly this schema: +{ + "selected_candidate": "a", + "reason": "brief visual justification" +} + +selected_candidate must be exactly one of: a, b. Do not return any other keys +or prose outside the JSON object. """ +_POSE_CANDIDATE_X_ROTATIONS_DEGREES = {"a": 90.0, "b": -90.0} +_VLM_ORTHOGRAPHIC_SCALE = 1.75 # Leave floor context around normalized assets. +_VISUAL_MAX_OBJECT_EXTENT = 0.75 # Shared temporary extent for VLM pose views. +_VISUAL_FLOOR_SCALE = 1.3 # Preserve a small border around the asset footprint. +_VISUAL_MIN_FLOOR_SIZE = 0.45 # Keep a contact cue for compact upright objects. + DEFAULT_NEEDED_LAYOUT = ( "Place this asset on the table in its natural, physically stable resting " "orientation. For example, a fork should lie flat on the table rather " @@ -78,17 +124,12 @@ STANDING_NEEDED_LAYOUT = ( "The scene graph requires this asset to stand vertically on the table, " "even when its natural stable pose would be lying down. For example, a " - "bottle should stand on its base and a fork should stand upright. If the " - "coarse GLB is lying flat, set rotate_about_x=true so its semantic vertical " - "axis aligns with the z-up world's z axis; if it is already upright, set " - "it to false." + "bottle should stand on its base and a fork should stand upright." ) LYING_NEEDED_LAYOUT = ( "The scene graph requires this asset to lie flat on the table, even when " "its natural stable pose would be standing. For example, a bottle should " - "lie on its side and a fork should lie flat. Choose rotate_about_x so the " - "asset's semantic long axis remains in the tabletop x-y plane rather than " - "along the z-up world's z axis." + "lie on its side and a fork should lie flat." ) @@ -98,7 +139,7 @@ def render_object_front_top_views( output_path: str | Path, resolution: int = 512, ) -> Path: - """Render fixed z-up front/top views and compose them horizontally.""" + """Render grounded fixed z-up oblique/top views and compose them horizontally.""" if resolution <= 0: raise ValueError("resolution must be positive.") source_path = Path(glb_path).expanduser().resolve() @@ -122,7 +163,8 @@ def render_object_front_top_views( _run_blender_operation_silently( lambda: bpy.ops.import_scene.gltf(filepath=str(source_path)) ) - if not any(obj.type == "MESH" for obj in bpy.context.scene.objects): + mesh_objects = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"] + if not mesh_objects: raise ValueError(f"GLB contains no mesh objects: {source_path}") scene = bpy.context.scene # Eevee renders imported GLB materials and textures instead of Workbench previews. @@ -159,7 +201,13 @@ def render_object_front_top_views( scene.collection.objects.link(camera) scene.camera = camera camera.data.type = "ORTHO" - camera.data.ortho_scale = 1.25 + # A wider framing preserves floor and orientation context for VLM inspection. + camera.data.ortho_scale = _VLM_ORTHOGRAPHIC_SCALE + + _place_meshes_on_visual_floor( + bpy=bpy, + mesh_objects=mesh_objects, + ) def render_view(path: Path, location: tuple[float, float, float]) -> None: camera.location = location @@ -171,58 +219,285 @@ def render_view(path: Path, location: tuple[float, float, float]) -> None: scene.render.filepath = str(path) _run_blender_operation_silently(lambda: bpy.ops.render.render(write_still=True)) - # Blender uses a right-handed z-up world; front is viewed along +y. - render_view(front_path, (0.0, -3.0, 0.0)) - render_view(top_path, (0.0, 0.0, 3.0)) - with Image.open(front_path) as front, Image.open(top_path) as top: - composed = Image.new("RGB", (resolution * 2, resolution), "white") - composed.paste(front.convert("RGB"), (0, 0)) - composed.paste(top.convert("RGB"), (resolution, 0)) - draw = ImageDraw.Draw(composed) - # Use a readable scaled font for the panel labels when available. - try: - font = ImageFont.truetype( - "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", - max(24, resolution // 16), + # A diagonal oblique view avoids looking straight down a long planar tool. + render_view(front_path, (3.0, -4.0, 3.0)) + render_view(top_path, (0.0, 0.0, 4.0)) + try: + with Image.open(front_path) as front, Image.open(top_path) as top: + composed = Image.new("RGB", (resolution * 2, resolution), "white") + composed.paste(front.convert("RGB"), (0, 0)) + composed.paste(top.convert("RGB"), (resolution, 0)) + draw = ImageDraw.Draw(composed) + # Use a readable scaled font for the panel labels when available. + try: + font = ImageFont.truetype( + "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", + max(24, resolution // 16), + ) + except OSError: + font = ImageFont.load_default() + # Label each panel so the VLM and manual debugging can distinguish views. + for label, origin in ( + ("OBLIQUE VIEW", (0, 0)), + ("TOP VIEW", (resolution, 0)), + ): + x, y = origin + text_box = draw.textbbox((x + 16, y + 16), label, font=font) + draw.rectangle( + ( + text_box[0] - 8, + text_box[1] - 6, + text_box[2] + 8, + text_box[3] + 6, + ), + fill="white", + ) + draw.text((x + 16, y + 16), label, fill="black", font=font) + # Mark the positive axes used by each projection for VLM interpretation. + _draw_arrow( + draw, + (resolution - 62, 62), + (resolution - 62, 20), + "+Z", + font, + color="blue", ) - except OSError: - font = ImageFont.load_default() - # Label each panel so the VLM and manual debugging can distinguish views. - for label, origin in (("FRONT VIEW", (0, 0)), ("TOP VIEW", (resolution, 0))): - x, y = origin - text_box = draw.textbbox((x + 16, y + 16), label, font=font) - draw.rectangle( - (text_box[0] - 8, text_box[1] - 6, text_box[2] + 8, text_box[3] + 6), - fill="white", + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 42, 62), + "+X", + font, + color="red", ) - draw.text((x + 16, y + 16), label, fill="black", font=font) - # Mark the positive axes used by each projection for VLM interpretation. - _draw_arrow( - draw, - (resolution - 62, 62), - (resolution - 62, 20), - "+Z", - font, - color="blue", + _draw_arrow( + draw, + (2 * resolution - 92, 62), + (2 * resolution - 92, 20), + "+Y", + font, + color="green", + ) + composed.save(output_path) + finally: + front_path.unlink(missing_ok=True) + top_path.unlink(missing_ok=True) + return output_path + + +def render_object_pose_switch_candidates( + *, + glb_path: str | Path, + output_path: str | Path, + resolution: int = 512, +) -> Path: + """Render opposite temporary x-rotation candidates without writing GLBs.""" + if resolution <= 0: + raise ValueError("resolution must be positive.") + source_path = Path(glb_path).expanduser().resolve() + if not source_path.is_file(): + raise FileNotFoundError(f"GLB for candidate rendering not found: {source_path}") + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + candidate_paths = { + candidate_id: output_path.with_name( + f".{output_path.stem}_candidate_{candidate_id}.png" ) - _draw_arrow( - draw, - (2 * resolution - 92, 62), - (2 * resolution - 42, 62), - "+X", - font, - color="red", + for candidate_id in _POSE_CANDIDATE_X_ROTATIONS_DEGREES + } + try: + for ( + candidate_id, + rotation_degrees, + ) in _POSE_CANDIDATE_X_ROTATIONS_DEGREES.items(): + _render_grounded_pose_candidate( + glb_path=source_path, + x_rotation_degrees=rotation_degrees, + output_path=candidate_paths[candidate_id], + resolution=resolution, + ) + with ( + Image.open(candidate_paths["a"]) as candidate_a, + Image.open(candidate_paths["b"]) as candidate_b, + ): + panel_size = candidate_a.size + composed = Image.new("RGB", (panel_size[0] * 2, panel_size[1]), "white") + composed.paste(candidate_a.convert("RGB"), (0, 0)) + composed.paste(candidate_b.convert("RGB"), (panel_size[0], 0)) + draw = ImageDraw.Draw(composed) + font = _label_font(panel_size[1]) + _draw_panel_label(draw, "CANDIDATE A", (0, 0), font) + _draw_panel_label(draw, "CANDIDATE B", (panel_size[0], 0), font) + composed.save(output_path) + finally: + for candidate_path in candidate_paths.values(): + candidate_path.unlink(missing_ok=True) + return output_path + + +def _render_grounded_pose_candidate( + *, + glb_path: Path, + x_rotation_degrees: float, + output_path: Path, + resolution: int, +) -> None: + """Render one temporary x-rotated candidate in the shared visual frame.""" + try: + import bpy + from mathutils import Matrix, Vector + except ImportError as exc: + raise RuntimeError( + "Blender's bpy is required for SimReady candidate rendering." + ) from exc + _run_blender_operation_silently( + lambda: bpy.ops.wm.read_factory_settings(use_empty=True) + ) + _run_blender_operation_silently( + lambda: bpy.ops.import_scene.gltf(filepath=str(glb_path)) + ) + mesh_objects = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"] + if not mesh_objects: + raise ValueError(f"GLB contains no mesh objects: {glb_path}") + scene = bpy.context.scene + _configure_vlm_render_scene( + scene=scene, resolution=resolution, bpy=bpy, Vector=Vector + ) + root = bpy.data.objects.new("PoseCandidateRoot", None) + scene.collection.objects.link(root) + for mesh_object in mesh_objects: + original_world_matrix = mesh_object.matrix_world.copy() + mesh_object.parent = root + mesh_object.matrix_parent_inverse = root.matrix_world.inverted() + mesh_object.matrix_world = original_world_matrix + center = _mesh_world_bounds(mesh_objects)[0] + x_rotation = Rotation.from_euler("x", x_rotation_degrees, degrees=True).as_matrix() + rotation_transform = np.eye(4) + rotation_transform[:3, :3] = x_rotation + root.matrix_world = Matrix( + _translation_matrix(center) @ rotation_transform @ _translation_matrix(-center) + ) + _place_meshes_on_visual_floor(bpy=bpy, mesh_objects=mesh_objects) + _add_oblique_camera(scene=scene, bpy=bpy, Vector=Vector) + scene.render.filepath = str(output_path) + _run_blender_operation_silently(lambda: bpy.ops.render.render(write_still=True)) + + +def _configure_vlm_render_scene( + *, scene: object, resolution: int, bpy: object, Vector: object +) -> None: + """Configure the shared illuminated, opaque Blender render scene.""" + try: + scene.render.engine = "BLENDER_EEVEE_NEXT" + except TypeError: + scene.render.engine = "BLENDER_EEVEE" + scene.render.resolution_x = resolution + scene.render.resolution_y = resolution + scene.render.resolution_percentage = 100 + scene.render.image_settings.file_format = "PNG" + scene.render.film_transparent = False + if scene.world is None: + scene.world = bpy.data.worlds.new("VLM_World") + scene.world.color = (0.08, 0.08, 0.08) + for name, location, energy in ( + ("VLM_Key", (2.0, -2.0, 3.0), 700.0), + ("VLM_Fill", (-2.0, 1.0, 2.0), 400.0), + ): + light_data = bpy.data.lights.new(name, type="AREA") + light_data.energy = energy + light_data.shape = "DISK" + light_data.size = 4.0 + light = bpy.data.objects.new(name, light_data) + light.location = location + light.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - light.location) + .to_track_quat("-Z", "Y") + .to_euler() + ) + scene.collection.objects.link(light) + + +def _add_oblique_camera(*, scene: object, bpy: object, Vector: object) -> None: + """Add the fixed, wider z-up oblique camera used for pose candidates.""" + camera_data = bpy.data.cameras.new("VLM_CandidateCamera") + camera = bpy.data.objects.new("VLM_CandidateCamera", camera_data) + scene.collection.objects.link(camera) + scene.camera = camera + camera.data.type = "ORTHO" + camera.data.ortho_scale = _VLM_ORTHOGRAPHIC_SCALE + camera.location = (0.0, -4.0, 5.0) + camera.rotation_euler = ( + (Vector((0.0, 0.0, 0.3)) - camera.location).to_track_quat("-Z", "Y").to_euler() + ) + + +def _place_meshes_on_visual_floor(*, bpy: object, mesh_objects: list[object]) -> None: + """Normalize one temporary object and place it on a local support patch.""" + from mathutils import Matrix, Vector + + _run_blender_operation_silently(lambda: bpy.context.view_layer.update()) + minimum, maximum = _mesh_world_aabb(mesh_objects) + extent = maximum - minimum + largest_extent = float(np.max(extent)) + if not np.isfinite(largest_extent) or largest_extent <= 0.0: + raise ValueError("VLM pose rendering requires a mesh with positive extent.") + center = (minimum + maximum) * 0.5 + visual_scale = _VISUAL_MAX_OBJECT_EXTENT / largest_extent + normalization_transform = np.eye(4) + normalization_transform[:3, :3] *= visual_scale + normalization_transform[:3, 3] = center - visual_scale * center + for mesh_object in mesh_objects: + mesh_object.matrix_world = ( + Matrix(normalization_transform) @ mesh_object.matrix_world ) - _draw_arrow( - draw, - (2 * resolution - 92, 62), - (2 * resolution - 92, 20), - "+Y", - font, - color="green", + _run_blender_operation_silently(lambda: bpy.context.view_layer.update()) + + minimum, maximum = _mesh_world_aabb(mesh_objects) + center = (minimum + maximum) * 0.5 + for mesh_object in mesh_objects: + # Assign in world space so this remains correct under a rotated parent root. + mesh_object.matrix_world.translation -= Vector( + (center[0], center[1], minimum[2]) ) - composed.save(output_path) - return output_path + _run_blender_operation_silently(lambda: bpy.context.view_layer.update()) + floor_size = max( + _VISUAL_MIN_FLOOR_SIZE, + _VISUAL_FLOOR_SCALE * float(np.max(maximum[:2] - minimum[:2])), + ) + bpy.ops.mesh.primitive_plane_add(size=floor_size, location=(0.0, 0.0, 0.0)) + floor = bpy.context.object + floor.name = "VLM_VisualFloor" + material = bpy.data.materials.new("VLM_VisualFloorMaterial") + material.diffuse_color = (0.22, 0.22, 0.22, 1.0) + floor.data.materials.append(material) + + +def _mesh_world_bounds(mesh_objects: list[object]) -> tuple[np.ndarray, float]: + """Return a world-space AABB centre and its minimum z coordinate.""" + minimum, maximum = _mesh_world_aabb(mesh_objects) + return (minimum + maximum) * 0.5, float(minimum[2]) + + +def _mesh_world_aabb(mesh_objects: list[object]) -> tuple[np.ndarray, np.ndarray]: + """Return world-space AABB minimum and maximum for one temporary object.""" + from mathutils import Vector + + points = np.asarray( + [ + tuple(mesh_object.matrix_world @ Vector(corner)) + for mesh_object in mesh_objects + for corner in mesh_object.bound_box + ], + dtype=float, + ) + return points.min(axis=0), points.max(axis=0) + + +def _translation_matrix(translation: np.ndarray) -> np.ndarray: + """Return one homogeneous translation transform.""" + transform = np.eye(4) + transform[:3, 3] = translation + return transform def _run_blender_operation_silently(operation: Callable[[], object]) -> object: @@ -281,7 +556,34 @@ def _draw_arrow( draw.text((int(tip_x + 8), int(tip_y - 8)), label, fill=color, font=font) -def query_vlm_object_rotation_and_target_size( +def _label_font(image_height: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + """Return a readable label font without requiring a system font.""" + try: + return ImageFont.truetype( + "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", + max(20, image_height // 20), + ) + except OSError: + return ImageFont.load_default() + + +def _draw_panel_label( + draw: ImageDraw.ImageDraw, + label: str, + origin: tuple[int, int], + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, +) -> None: + """Draw one panel label on an opaque backing rectangle.""" + x, y = origin + text_box = draw.textbbox((x + 16, y + 16), label, font=font) + draw.rectangle( + (text_box[0] - 8, text_box[1] - 6, text_box[2] + 8, text_box[3] + 6), + fill="white", + ) + draw.text((x + 16, y + 16), label, fill="black", font=font) + + +def query_vlm_object_pose_and_target_size( *, scene_object_description: str, needed_layout: str, @@ -290,7 +592,7 @@ def query_vlm_object_rotation_and_target_size( debug_output_path: str | Path | None = None, json_max_attempts: int = 3, ) -> dict[str, object]: - """Ask the VLM for a valid rotation and post-rotation tabletop footprint.""" + """Ask the VLM whether to preserve or switch the semantic pose.""" if json_max_attempts < 1: raise ValueError("json_max_attempts must be at least 1.") last_validation_error: ValueError | None = None @@ -300,12 +602,12 @@ def query_vlm_object_rotation_and_target_size( user_prompt=( f"Object description:\n{scene_object_description}\n\n" f"Needed layout:\n{needed_layout}\n\n" - "The image contains front view on the left and top view on the right." + "The image contains an OBLIQUE VIEW on the left and TOP VIEW on the right." ), image_path=rendered_views_path, ) try: - value = _parse_vlm_rotation_and_target_size_response(response_text) + value = _parse_vlm_pose_and_target_size_response(response_text) break except ValueError as exc: last_validation_error = exc @@ -338,24 +640,96 @@ def query_vlm_object_rotation_and_target_size( return value -def _parse_vlm_rotation_and_target_size_response( +def query_vlm_pose_switch_candidate( + *, + scene_object_description: str, + needed_layout: str, + rendered_candidates_path: str | Path, + vlm_client: OpenAICompatibleVLM, + debug_output_path: str | Path | None = None, + json_max_attempts: int = 3, +) -> tuple[float, str]: + """Ask the VLM to choose one opposite temporary semantic-pose candidate.""" + if json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + last_validation_error: ValueError | None = None + for _ in range(json_max_attempts): + response_text = vlm_client.complete( + system_prompt=_VLM_POSE_CANDIDATE_SYSTEM_PROMPT, + user_prompt=( + f"Object description:\n{scene_object_description}\n\n" + f"Needed layout:\n{needed_layout}\n\n" + "Choose the correct grounded candidate from the paired image." + ), + image_path=rendered_candidates_path, + ) + try: + selected_candidate, reason = _parse_vlm_pose_candidate_response( + response_text + ) + break + except ValueError as exc: + last_validation_error = exc + else: + assert last_validation_error is not None + raise ValueError( + "VLM pose-candidate response is invalid after " + f"{json_max_attempts} attempts: {last_validation_error}" + ) from last_validation_error + + selected_rotation_degrees = _POSE_CANDIDATE_X_ROTATIONS_DEGREES[selected_candidate] + if debug_output_path is not None: + output_path = Path(debug_output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps( + { + "description": scene_object_description, + "needed_layout": needed_layout, + "rendered_candidates_path": str( + Path(rendered_candidates_path).expanduser().resolve() + ), + "selected_candidate": selected_candidate, + "x_rotation_degrees": selected_rotation_degrees, + "reason": reason, + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return selected_rotation_degrees, reason + + +def _parse_vlm_pose_and_target_size_response( response_text: str, ) -> dict[str, object]: - """Validate one VLM rotation-and-scale JSON response.""" + """Validate one VLM semantic-pose and scale JSON response.""" try: value = json.loads(_strip_json_code_fence(response_text)) except json.JSONDecodeError as exc: raise ValueError(f"VLM transform response is not valid JSON: {exc}") from exc if not isinstance(value, dict) or set(value) != { - "rotate_about_x", + "pose_action", + "reason", "target_xy_size_cm", }: raise ValueError( - "VLM transform response must contain exactly rotate_about_x and " + "VLM transform response must contain exactly pose_action, reason, and " "target_xy_size_cm." ) - if not isinstance(value["rotate_about_x"], bool): - raise ValueError("VLM rotate_about_x must be boolean.") + pose_action = value["pose_action"] + if not isinstance(pose_action, str) or pose_action not in { + "keep_current", + "rotate_to_required_pose", + }: + raise ValueError( + "VLM pose_action must be keep_current or rotate_to_required_pose." + ) + reason = value["reason"] + if not isinstance(reason, str) or not reason.strip(): + raise ValueError("VLM pose reason must be a non-empty string.") target_size = value["target_xy_size_cm"] if ( not isinstance(target_size, list) @@ -367,6 +741,28 @@ def _parse_vlm_rotation_and_target_size_response( return value +def _parse_vlm_pose_candidate_response(response_text: str) -> tuple[str, str]: + """Validate one VLM binary semantic-pose candidate selection.""" + try: + value = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError( + f"VLM pose-candidate response is not valid JSON: {exc}" + ) from exc + if not isinstance(value, dict) or set(value) != {"selected_candidate", "reason"}: + raise ValueError( + "VLM pose-candidate response must contain exactly selected_candidate and " + "reason." + ) + selected_candidate = value["selected_candidate"] + reason = value["reason"] + if selected_candidate not in _POSE_CANDIDATE_X_ROTATIONS_DEGREES: + raise ValueError("VLM selected_candidate must be a or b.") + if not isinstance(reason, str) or not reason.strip(): + raise ValueError("VLM pose-candidate reason must be a non-empty string.") + return selected_candidate, reason.strip() + + def compute_uniform_xy_scale_for_target( *, glb_path: str | Path, @@ -407,11 +803,9 @@ def rotate_glb_about_x_axis( *, input_path: str | Path, output_path: str | Path, - rotate: bool, + rotation_degrees: float, ) -> Path: - """Bake an optional +90-degree x-axis rotation around the mesh centre.""" - # Current coarse layouts are either flat on xy with possible random z rotation, - # or upright with almost no random y rotation, so this x-axis toggle is enough. + """Bake one x-axis rotation around the mesh AABB centre.""" source_path = Path(input_path).expanduser().resolve() destination_path = Path(output_path).expanduser().resolve() destination_path.parent.mkdir(parents=True, exist_ok=True) @@ -421,11 +815,15 @@ def rotate_glb_about_x_axis( ) if not isinstance(mesh, trimesh.Trimesh): raise ValueError(f"GLB is not a mesh: {source_path}") - if rotate: + if not np.isfinite(rotation_degrees): + raise ValueError("rotation_degrees must be finite.") + if rotation_degrees != 0.0: center = mesh.bounds.mean(axis=0) mesh.apply_translation(-center) transform = np.eye(4) - transform[:3, :3] = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + transform[:3, :3] = Rotation.from_euler( + "x", rotation_degrees, degrees=True + ).as_matrix() mesh.apply_transform(transform) mesh.apply_translation(center) mesh.export(destination_path, file_type="glb") diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/visual_yaw_optimizer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/visual_yaw_optimizer.py new file mode 100644 index 000000000..558364217 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/visual_yaw_optimizer.py @@ -0,0 +1,494 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""VLM-guided visual yaw selection for canonical SimReady assets.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +from PIL import Image, ImageDraw, ImageFont +from scipy.spatial.transform import Rotation + +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( + _run_blender_operation_silently, +) + +_RENDER_RESOLUTION = 512 +_ALLOWED_CLOCKWISE_YAWS_DEGREES = frozenset(range(0, 360, 15)) +_VISUAL_FLOOR_SCALE = 1.3 # Preserve a small border around the canonical footprint. +_VISUAL_MIN_FLOOR_SIZE = 0.45 # Keep a contact cue for compact upright assets. +_VLM_SYSTEM_PROMPT = """You align one isolated canonical 3D asset to an image observation. +The image has two panels. The LEFT panel is a fixed, slightly distant +orthographic oblique view from above of the canonical SimReady asset in a +right-handed z-up world, resting on a small neutral support patch. It has no +coarse-layout position or rotation applied. The blue arrow marks world +z and +the red arrow marks world +x. The RIGHT panel is the asset's visible RGBA crop +from the source image, which is also a robot-manipulation view. + +Choose the absolute clockwise yaw about world z that makes the LEFT asset best +match the visible direction of the RIGHT observation. The output yaw is +measured from the current canonical LEFT view, not from any coarse-layout +rotation. The two views can have different camera poses, perspective, +occlusion, and segmentation error, so do not attempt pixel-level alignment. +Keep 0 unless there is a severe, clearly supported directional mismatch, such +as a handle, spout, blade, long axis, label, or asymmetric silhouette facing +the wrong direction. Clockwise means a physical rotation when looking down +from world +z toward the table, not a screen-plane rotation in either panel. +Use this yaw compass: 0 keeps the canonical asset unchanged; 90 clockwise +turns its world +x direction toward world -y; 180 reverses it; 270 clockwise +turns world +x toward world +y. Choose the nearest 15-degree value only after +first deciding whether a nonzero correction is clearly necessary. + +Return JSON only with exactly this schema: +{ + "clockwise_yaw_degrees": 0, + "reason": "brief visual justification" +} + +clockwise_yaw_degrees must be one of 0, 15, 30, ..., 345. Do not return any +other keys or prose outside the JSON object.""" + + +class VisualYawOptimizer: + """Ask a VLM for one absolute simulator-world z-up yaw for an asset. + + The SimReady GLB has already resolved standing, lying, or stable semantic + pose and baked the geometry-server scale. Rendering applies that scale's + inverse only to a temporary Blender root, so every VLM query sees the + same normalized canonical asset without modifying the saved GLB. + """ + + def __init__( + self, + *, + scene_object: SceneObject, + baked_scale_y_up: list[float], + vlm_client: OpenAICompatibleVLM, + debug_output_root: str | Path, + json_max_attempts: int = 3, + ) -> None: + self._scene_object = scene_object + self._baked_scale_y_up = baked_scale_y_up + self._vlm_client = vlm_client + self._debug_output_root = Path(debug_output_root).expanduser().resolve() + self._json_max_attempts = json_max_attempts + + def optimize_z_up_yaw_degrees(self) -> float: + """Return the absolute canonical yaw in z-up world degrees. + + Positive returned angles are counterclockwise in the z-up world. The + VLM reports clockwise image angles, so this method converts signs at + the boundary before scene generation applies the result to a layout. + """ + self._validate_inputs() + if self._scene_object.visible_rgba_path is None: + return 0.0 + + # Only the two composed comparisons are debug artifacts; raw renders are temporary. + before_path = ( + self._debug_output_root / f".{self._scene_object.id}_canonical.png" + ) + after_path = self._debug_output_root / f".{self._scene_object.id}_yawed.png" + source_path = Path(self._scene_object.visible_rgba_path).expanduser().resolve() + vlm_input_path = ( + self._debug_output_root / f"{self._scene_object.id}_vlm_input.png" + ) + try: + _render_canonical_oblique_view( + glb_path=self._scene_object.simready_glb_path, + baked_scale_y_up=self._baked_scale_y_up, + z_up_yaw_degrees=0.0, + output_path=before_path, + resolution=_RENDER_RESOLUTION, + ) + _compose_render_and_source( + rendered_path=before_path, + source_rgba_path=source_path, + rendered_label="CANONICAL OBLIQUE VIEW", + source_label="SOURCE RGBA", + output_path=vlm_input_path, + ) + clockwise_yaw_degrees, reason = self._query_clockwise_yaw(vlm_input_path) + z_up_yaw_degrees = -float(clockwise_yaw_degrees) + _render_canonical_oblique_view( + glb_path=self._scene_object.simready_glb_path, + baked_scale_y_up=self._baked_scale_y_up, + z_up_yaw_degrees=z_up_yaw_degrees, + output_path=after_path, + resolution=_RENDER_RESOLUTION, + ) + _compose_render_and_source( + rendered_path=after_path, + source_rgba_path=source_path, + rendered_label="YAWED OBLIQUE VIEW", + source_label="SOURCE RGBA", + output_path=self._debug_output_root + / f"{self._scene_object.id}_yaw_result.png", + ) + finally: + before_path.unlink(missing_ok=True) + after_path.unlink(missing_ok=True) + (self._debug_output_root / f"{self._scene_object.id}.json").write_text( + json.dumps( + { + "object_id": self._scene_object.id, + "clockwise_yaw_degrees": clockwise_yaw_degrees, + "z_up_yaw_degrees": z_up_yaw_degrees, + "reason": reason, + "vlm_input_path": str(vlm_input_path), + }, + indent=2, + ensure_ascii=False, + ) + + "\n", + encoding="utf-8", + ) + return z_up_yaw_degrees + + def _query_clockwise_yaw(self, vlm_input_path: Path) -> tuple[int, str]: + """Query and validate the VLM's discrete clockwise yaw selection.""" + last_error: ValueError | None = None + for _ in range(self._json_max_attempts): + response_text = self._vlm_client.complete( + system_prompt=_VLM_SYSTEM_PROMPT, + user_prompt=( + f"Object description: {self._scene_object.description}\n" + "Select the canonical asset's yaw from the paired image." + ), + image_path=vlm_input_path, + ) + try: + return _parse_clockwise_yaw_response(response_text) + except ValueError as exc: + last_error = exc + assert last_error is not None + raise ValueError( + "VLM visual yaw response is invalid after " + f"{self._json_max_attempts} attempts: {last_error}" + ) from last_error + + def _validate_inputs(self) -> None: + if self._json_max_attempts < 1: + raise ValueError("json_max_attempts must be at least 1.") + if self._scene_object.simready_glb_path is None: + raise ValueError( + "Visual yaw optimization requires a SimReady GLB for " + f"{self._scene_object.id!r}." + ) + if not Path(self._scene_object.simready_glb_path).is_file(): + raise FileNotFoundError( + "Visual yaw optimization SimReady GLB not found for " + f"{self._scene_object.id!r}: {self._scene_object.simready_glb_path}" + ) + if len(self._baked_scale_y_up) != 3 or any( + not np.isfinite(value) or value <= 0.0 for value in self._baked_scale_y_up + ): + raise ValueError("Visual yaw optimization requires three positive scales.") + if ( + self._scene_object.visible_rgba_path is not None + and not Path(self._scene_object.visible_rgba_path).is_file() + ): + raise FileNotFoundError( + "Visual yaw optimization RGBA observation not found for " + f"{self._scene_object.id!r}: " + f"{self._scene_object.visible_rgba_path}" + ) + + +def _render_canonical_oblique_view( + *, + glb_path: str | Path, + baked_scale_y_up: list[float], + z_up_yaw_degrees: float, + output_path: str | Path, + resolution: int, +) -> Path: + """Render a temporary unscaled GLB from a fixed z-up oblique view.""" + source_path = Path(glb_path).expanduser().resolve() + destination_path = Path(output_path).expanduser().resolve() + destination_path.parent.mkdir(parents=True, exist_ok=True) + try: + import bpy + from mathutils import Matrix, Vector + except ImportError as exc: + raise RuntimeError( + "Blender's bpy is required for visual-yaw rendering." + ) from exc + + _run_blender_operation_silently( + lambda: bpy.ops.wm.read_factory_settings(use_empty=True) + ) + _run_blender_operation_silently( + lambda: bpy.ops.import_scene.gltf(filepath=str(source_path)) + ) + mesh_objects = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"] + if not mesh_objects: + raise ValueError(f"GLB contains no mesh objects: {source_path}") + + scene = bpy.context.scene + try: + scene.render.engine = "BLENDER_EEVEE_NEXT" + except TypeError: + scene.render.engine = "BLENDER_EEVEE" + scene.render.resolution_x = resolution + scene.render.resolution_y = resolution + scene.render.resolution_percentage = 100 + scene.render.image_settings.file_format = "PNG" + scene.render.film_transparent = False + if scene.world is None: + scene.world = bpy.data.worlds.new("VisualYawWorld") + scene.world.color = (0.08, 0.08, 0.08) + + # Blender imports the y-up GLB into z-up. Undo only baked y-up scale here. + y_up_to_z_up = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() + inverse_scale_z_up = ( + y_up_to_z_up + @ np.diag(1.0 / np.asarray(baked_scale_y_up, dtype=float)) + @ y_up_to_z_up.T + ) + yaw_z_up = Rotation.from_euler("z", z_up_yaw_degrees, degrees=True).as_matrix() + root = bpy.data.objects.new("VisualYawRoot", None) + scene.collection.objects.link(root) + for mesh_object in mesh_objects: + original_world_matrix = mesh_object.matrix_world.copy() + mesh_object.parent = root + mesh_object.matrix_parent_inverse = root.matrix_world.inverted() + mesh_object.matrix_world = original_world_matrix + transform = np.eye(4) + transform[:3, :3] = yaw_z_up @ inverse_scale_z_up + root.matrix_world = Matrix(transform) + visual_floor_size = _center_root_on_visual_floor( + bpy=bpy, + mesh_objects=mesh_objects, + root=root, + ) + _add_visual_floor(bpy=bpy, size=visual_floor_size) + + light_data = bpy.data.lights.new("VisualYawKey", type="AREA") + light_data.energy = 900.0 + light_data.shape = "DISK" + light_data.size = 4.0 + light = bpy.data.objects.new("VisualYawKey", light_data) + light.location = (2.0, -2.0, 4.0) + light.rotation_euler = ( + (Vector((0.0, 0.0, 0.0)) - light.location).to_track_quat("-Z", "Y").to_euler() + ) + scene.collection.objects.link(light) + camera_data = bpy.data.cameras.new("VisualYawCamera") + camera = bpy.data.objects.new("VisualYawCamera", camera_data) + scene.collection.objects.link(camera) + scene.camera = camera + camera.data.type = "ORTHO" + # Keep a wider fixed robot view with floor context, without coarse-layout scale. + camera.data.ortho_scale = 1.75 + camera.location = (0.0, -4.0, 5.0) + camera.rotation_euler = ( + (Vector((0.0, 0.0, 0.3)) - camera.location).to_track_quat("-Z", "Y").to_euler() + ) + scene.render.filepath = str(destination_path) + _run_blender_operation_silently(lambda: bpy.ops.render.render(write_still=True)) + _draw_oblique_view_axes(destination_path) + return destination_path + + +def _center_root_on_visual_floor( + *, bpy: object, mesh_objects: list[object], root: object +) -> float: + """Translate a temporary transformed root so its mesh rests at the origin.""" + from mathutils import Matrix, Vector + + _run_blender_operation_silently(lambda: bpy.context.view_layer.update()) + points = np.asarray( + [ + tuple(mesh_object.matrix_world @ Vector(corner)) + for mesh_object in mesh_objects + for corner in mesh_object.bound_box + ], + dtype=float, + ) + minimum = points.min(axis=0) + maximum = points.max(axis=0) + translation = np.eye(4) + translation[:3, 3] = [ + -(minimum[0] + maximum[0]) * 0.5, + -(minimum[1] + maximum[1]) * 0.5, + -minimum[2], + ] + root.matrix_world = Matrix(translation) @ root.matrix_world + _run_blender_operation_silently(lambda: bpy.context.view_layer.update()) + return max( + _VISUAL_MIN_FLOOR_SIZE, + _VISUAL_FLOOR_SCALE * float(np.max(maximum[:2] - minimum[:2])), + ) + + +def _add_visual_floor(*, bpy: object, size: float) -> None: + """Add a footprint-scaled support patch for local contact context.""" + bpy.ops.mesh.primitive_plane_add(size=size, location=(0.0, 0.0, 0.0)) + floor = bpy.context.object + floor.name = "VisualYawFloor" + material = bpy.data.materials.new("VisualYawFloorMaterial") + material.diffuse_color = (0.22, 0.22, 0.22, 1.0) + floor.data.materials.append(material) + + +def _compose_render_and_source( + *, + rendered_path: str | Path, + source_rgba_path: str | Path, + rendered_label: str, + source_label: str, + output_path: str | Path, +) -> Path: + """Compose a rendered oblique view and an alpha-aware source observation.""" + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + with Image.open(rendered_path) as rendered, Image.open(source_rgba_path) as source: + panel_size = rendered.size + source = source.convert("RGBA").resize(panel_size, Image.Resampling.LANCZOS) + source_panel = Image.new("RGB", panel_size, (36, 36, 36)) + source_panel.paste(source, mask=source.getchannel("A")) + composed = Image.new("RGB", (panel_size[0] * 2, panel_size[1]), "white") + composed.paste(rendered.convert("RGB"), (0, 0)) + composed.paste(source_panel, (panel_size[0], 0)) + draw = ImageDraw.Draw(composed) + font = _label_font(panel_size[1]) + _draw_panel_label(draw, rendered_label, (0, 0), font) + _draw_panel_label(draw, source_label, (panel_size[0], 0), font) + composed.save(output_path) + return output_path + + +def _draw_oblique_view_axes(rendered_path: Path) -> None: + """Mark the fixed z-up axes used by the VLM yaw convention.""" + with Image.open(rendered_path) as rendered: + image = rendered.convert("RGB") + draw = ImageDraw.Draw(image) + font = _label_font(image.height) + origin = (image.width - 84, 84) + _draw_arrow(draw, origin, (image.width - 28, 84), "+X", font, "red") + _draw_arrow(draw, origin, (image.width - 84, 28), "+Z", font, "blue") + draw.text((18, image.height - 42), "FIXED OBLIQUE VIEW", fill="white", font=font) + image.save(rendered_path) + + +def _parse_clockwise_yaw_response(response_text: str) -> tuple[int, str]: + """Parse the exact discrete-yaw response contract from the VLM.""" + try: + value = json.loads(_strip_json_code_fence(response_text)) + except json.JSONDecodeError as exc: + raise ValueError(f"VLM visual yaw response is not valid JSON: {exc}") from exc + if not isinstance(value, dict) or set(value) != {"clockwise_yaw_degrees", "reason"}: + raise ValueError( + "VLM visual yaw response must contain exactly clockwise_yaw_degrees and reason." + ) + yaw = value["clockwise_yaw_degrees"] + reason = value["reason"] + if ( + isinstance(yaw, bool) + or not isinstance(yaw, int) + or yaw not in _ALLOWED_CLOCKWISE_YAWS_DEGREES + ): + raise ValueError( + "VLM clockwise_yaw_degrees must be a 15-degree value in [0, 345]." + ) + if not isinstance(reason, str) or not reason.strip(): + raise ValueError("VLM visual yaw reason must be a non-empty string.") + return yaw, reason.strip() + + +def _strip_json_code_fence(response_text: str) -> str: + """Remove one optional Markdown JSON fence from a VLM response.""" + stripped = response_text.strip() + if not stripped.startswith("```"): + return stripped + lines = stripped.splitlines() + if lines and lines[0].startswith("```"): + lines = lines[1:] + if lines and lines[-1].strip() == "```": + lines = lines[:-1] + return "\n".join(lines).strip() + + +def _label_font(image_height: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: + """Return a readable panel-label font without requiring a system font.""" + try: + return ImageFont.truetype( + "/usr/share/fonts/truetype/ubuntu/Ubuntu-B.ttf", + max(20, image_height // 20), + ) + except OSError: + return ImageFont.load_default() + + +def _draw_panel_label( + draw: ImageDraw.ImageDraw, + label: str, + origin: tuple[int, int], + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, +) -> None: + """Draw one panel label on an opaque backing rectangle.""" + x, y = origin + text_box = draw.textbbox((x + 16, y + 16), label, font=font) + draw.rectangle( + (text_box[0] - 8, text_box[1] - 6, text_box[2] + 8, text_box[3] + 6), + fill="white", + ) + draw.text((x + 16, y + 16), label, fill="black", font=font) + + +def _draw_arrow( + draw: ImageDraw.ImageDraw, + start: tuple[int, int], + end: tuple[int, int], + label: str, + font: ImageFont.FreeTypeFont | ImageFont.ImageFont, + color: str, +) -> None: + """Draw one labeled axis arrow.""" + dx, dy = end[0] - start[0], end[1] - start[1] + length = max(abs(dx), abs(dy)) + if length == 0: + raise ValueError("Axis arrow start and end must differ.") + unit_x, unit_y = dx / length, dy / length + perpendicular_x, perpendicular_y = -unit_y, unit_x + head_length = 14.0 + head_width = 8.0 + tip_x, tip_y = end + base_x = tip_x - unit_x * head_length + base_y = tip_y - unit_y * head_length + draw.line((*start, *end), fill=color, width=4) + draw.polygon( + ( + (tip_x, tip_y), + ( + base_x + perpendicular_x * head_width, + base_y + perpendicular_y * head_width, + ), + ( + base_x - perpendicular_x * head_width, + base_y - perpendicular_y * head_width, + ), + ), + fill=color, + ) + draw.text((int(tip_x + 8), int(tip_y - 8)), label, fill=color, font=font) diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py index a43db730c..a1a76f113 100644 --- a/tests/gen_sim/scene_engine/test_scene_generation.py +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -16,8 +16,11 @@ from __future__ import annotations +from pathlib import Path + import numpy as np import pytest +from PIL import Image from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene_graph import ( @@ -26,12 +29,18 @@ ) from embodichain.gen_sim.scene_engine.core.scene import Scene, SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + _align_table_roots_individually, + _apply_visual_yaws_to_simready_asset_layouts, _apply_root_layout_updates_to_descendant_subtrees, + _optimize_simready_asset_visual_yaws, _project_child_aabb_centers_into_parent_aabb, _refine_on_children_bfs, _scene_graph_based_calibration, _table_on_asset_ids, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.visual_yaw_optimizer import ( + VisualYawOptimizer, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, transform_matrix_to_layout_object, @@ -101,6 +110,166 @@ def test_scene_graph_calibration_makes_standing_asset_vertical() -> None: ) +def test_visual_yaw_optimizer_returns_zero_for_unobserved_asset( + tmp_path, +) -> None: + glb_path = tmp_path / "book_001.glb" + glb_path.write_bytes(b"glTF") + scene_object = SceneObject( + id="book_001", + kind="asset", + category="book", + name="book", + description="book", + simready_glb_path=str(glb_path), + ) + + yaw_delta_degrees = VisualYawOptimizer( + scene_object=scene_object, + baked_scale_y_up=[1.0, 1.0, 1.0], + vlm_client=object(), + debug_output_root=tmp_path / "visual_yaw", + ).optimize_z_up_yaw_degrees() + + assert yaw_delta_degrees == 0.0 + + +def test_visual_yaw_optimizer_queries_vlm_and_saves_yawed_debug_image( + monkeypatch, + tmp_path, +) -> None: + glb_path = tmp_path / "book_001.glb" + glb_path.write_bytes(b"glTF") + rgba_path = tmp_path / "book_001_rgba.png" + Image.new("RGBA", (64, 64), (255, 0, 0, 255)).save(rgba_path) + + rendered_yaws: list[float] = [] + + def fake_render(*, z_up_yaw_degrees, output_path, **_) -> None: + rendered_yaws.append(z_up_yaw_degrees) + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + Image.new("RGB", (512, 512), "black").save(output_path) + + class FakeVLM: + image_paths: list[Path] = [] + + def complete(self, *, image_path, **_) -> str: + self.image_paths.append(Path(image_path)) + return '{"clockwise_yaw_degrees": 90, "reason": "long axis"}' + + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.utils.visual_yaw_optimizer._render_canonical_oblique_view", + fake_render, + ) + fake_vlm = FakeVLM() + yaw_degrees = VisualYawOptimizer( + scene_object=SceneObject( + id="book_001", + kind="asset", + category="book", + name="book", + description="book", + simready_glb_path=str(glb_path), + visible_rgba_path=str(rgba_path), + ), + baked_scale_y_up=[1.0, 1.0, 1.0], + vlm_client=fake_vlm, + debug_output_root=tmp_path / "visual_yaw", + ).optimize_z_up_yaw_degrees() + + assert yaw_degrees == -90.0 + assert rendered_yaws == [0.0, -90.0] + assert fake_vlm.image_paths == [tmp_path / "visual_yaw" / "book_001_vlm_input.png"] + assert (tmp_path / "visual_yaw" / "book_001_yaw_result.png").is_file() + + +def test_simready_visual_yaw_queries_every_asset(monkeypatch, tmp_path) -> None: + queried_asset_ids: list[str] = [] + + class FakeVisualYawOptimizer: + def __init__( + self, + *, + scene_object, + baked_scale_y_up, + vlm_client, + debug_output_root, + ) -> None: + assert baked_scale_y_up == [1.0, 2.0, 3.0] + assert vlm_client is fake_vlm_client + assert debug_output_root == tmp_path / "visual_yaw" + queried_asset_ids.append(scene_object.id) + + def optimize_z_up_yaw_degrees(self) -> float: + return 15.0 + + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation.VisualYawOptimizer", + FakeVisualYawOptimizer, + ) + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="book", + description="book", + simready_glb_path=str(tmp_path / "book_001.glb"), + ), + SceneObject( + id="cup_001", + kind="asset", + category="cup", + name="cup", + description="cup", + simready_glb_path=str(tmp_path / "cup_001.glb"), + ), + ] + ) + fake_vlm_client = object() + + yaw_deltas_by_id = _optimize_simready_asset_visual_yaws( + scene=scene, + simready_assets_layout=[{"id": "book_001"}, {"id": "cup_001"}], + coarse_layout_by_id={ + "book_001": {"scale": [1.0, 2.0, 3.0]}, + "cup_001": {"scale": [1.0, 2.0, 3.0]}, + }, + vlm_client=fake_vlm_client, + debug_output_root=tmp_path / "visual_yaw", + ) + + assert queried_asset_ids == ["book_001", "cup_001"] + assert yaw_deltas_by_id == {"book_001": 15.0, "cup_001": 15.0} + + +def test_visual_yaws_replace_coarse_rotations_but_preserve_positions() -> None: + yawed_layout = _apply_visual_yaws_to_simready_asset_layouts( + simready_assets_layout=[ + { + "id": "book_001", + "rot": [20.0, -15.0, 40.0], + "pos": [0.1, 0.2, 0.3], + "scale": [1.0, 1.0, 1.0], + } + ], + z_up_yaws_degrees_by_id={"book_001": 45.0}, + )[0] + + expected_z_up_yaw = Rotation.from_euler("z", 45.0, degrees=True).as_matrix() + assert np.allclose(_z_up_rotation_from_y_up_layout(yawed_layout), expected_z_up_yaw) + assert np.allclose(yawed_layout["pos"], [0.1, 0.2, 0.3]) + + def test_table_root_update_propagates_its_pose_delta_to_descendants() -> None: scene_graph = SceneGraph( nodes=[ @@ -152,6 +321,99 @@ def test_table_root_update_propagates_its_pose_delta_to_descendants() -> None: ) +def test_table_roots_align_independently_and_move_only_their_descendants( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + class FakeAssetsGroupTableAligner: + aligned_root_ids: list[str] = [] + + def __init__(self, *, table_layout, assets_layout, geometry_root) -> None: + assert geometry_root == tmp_path / "geometry" + assert len(assets_layout) == 1 + self.table_layout = table_layout + self.root_layout = assets_layout[0] + + def align(self): + root_id = self.root_layout["id"] + assert isinstance(root_id, str) + self.aligned_root_ids.append(root_id) + vertical_delta_by_id = {"board_001": 0.4, "bottle_001": -0.2} + return self.table_layout, [ + { + **self.root_layout, + "pos": [ + *self.root_layout["pos"][:2], + self.root_layout["pos"][2] + vertical_delta_by_id[root_id], + ], + } + ] + + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation.AssetsGroupTableAligner", + FakeAssetsGroupTableAligner, + ) + scene_graph = SceneGraph( + nodes=[ + SceneGraphNode(object_id="table", parent_id=None), + SceneGraphNode( + object_id="board_001", parent_id="table", parent_relation="on" + ), + SceneGraphNode( + object_id="knife_001", parent_id="board_001", parent_relation="on" + ), + SceneGraphNode( + object_id="bottle_001", parent_id="table", parent_relation="on" + ), + ] + ) + table_layout = { + "id": "table", + "rot": [0.0, 0.0, 0.0], + "pos": [0.0, 0.0, 0.0], + "scale": [1.0, 1.0, 1.0], + } + assets_layout = [ + { + "id": "board_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.1, 0.2, 0.3], + "scale": [1.0, 1.0, 1.0], + }, + { + "id": "knife_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.15, 0.25, 0.35], + "scale": [1.0, 1.0, 1.0], + }, + { + "id": "bottle_001", + "rot": [0.0, 0.0, 0.0], + "pos": [0.5, 0.6, 0.8], + "scale": [1.0, 1.0, 1.0], + }, + ] + + _, aligned_assets_layout = _align_table_roots_individually( + scene_graph=scene_graph, + table_layout=table_layout, + assets_layout=assets_layout, + table_root_ids={"board_001", "bottle_001"}, + geometry_root=tmp_path / "geometry", + ) + + aligned_layouts_by_id = { + str(asset_layout["id"]): asset_layout for asset_layout in aligned_assets_layout + } + assert FakeAssetsGroupTableAligner.aligned_root_ids == [ + "board_001", + "bottle_001", + ] + assert np.allclose(aligned_layouts_by_id["board_001"]["pos"], [0.1, 0.2, 0.7]) + assert np.allclose(aligned_layouts_by_id["knife_001"]["pos"], [0.15, 0.25, 0.75]) + assert np.allclose(aligned_layouts_by_id["bottle_001"]["pos"], [0.5, 0.6, 0.6]) + + def test_on_children_bfs_refines_every_non_table_parent( monkeypatch, tmp_path, diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py index dd872f913..959412af1 100644 --- a/tests/gen_sim/scene_engine/test_simready_processor_utils.py +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -21,7 +21,7 @@ import pytest import trimesh -from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene import Scene, SceneObject from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor import ( SimReadyProcessor, SimReadyProcessorConfig, @@ -31,7 +31,8 @@ LYING_NEEDED_LAYOUT, STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, - query_vlm_object_rotation_and_target_size, + query_vlm_pose_switch_candidate, + query_vlm_object_pose_and_target_size, ) @@ -48,7 +49,11 @@ def complete(self, **_: object) -> str: coarse_geometry_root=tmp_path / "coarse", simready_geometry_root=tmp_path / "simready", config=SimReadyProcessorConfig( - orientation_states_by_id={"bottle_001": "standing", "fork_001": "lying"}, + orientation_states_by_id={ + "bottle_001": "standing", + "fork_001": "lying", + "knife_001": None, + }, ), vlm_client=VLM(), # type: ignore[arg-type] ) @@ -66,24 +71,165 @@ class VLM: def __init__(self) -> None: self.responses = [ "", - '{"rotate_about_x": false, "target_xy_size_cm": [8.0, 8.0]}', + '{"pose_action": "keep_current", "reason": "already upright", ' + '"target_xy_size_cm": [8.0, 8.0]}', ] def complete(self, **_: object) -> str: return self.responses.pop(0) vlm_client = VLM() - decision = query_vlm_object_rotation_and_target_size( + decision = query_vlm_object_pose_and_target_size( scene_object_description="small blue bottle", needed_layout=STANDING_NEEDED_LAYOUT, rendered_views_path=tmp_path / "views.png", vlm_client=vlm_client, # type: ignore[arg-type] ) - assert decision == {"rotate_about_x": False, "target_xy_size_cm": [8.0, 8.0]} + assert decision == { + "pose_action": "keep_current", + "reason": "already upright", + "target_xy_size_cm": [8.0, 8.0], + } assert vlm_client.responses == [] +def test_vlm_pose_candidate_query_returns_selected_rotation(tmp_path: Path) -> None: + class VLM: + def complete(self, **_: object) -> str: + return '{"selected_candidate": "b", "reason": "pan opening faces up"}' + + selected_rotation_degrees, reason = query_vlm_pose_switch_candidate( + scene_object_description="small saucepan", + needed_layout=LYING_NEEDED_LAYOUT, + rendered_candidates_path=tmp_path / "candidates.png", + vlm_client=VLM(), # type: ignore[arg-type] + ) + + assert selected_rotation_degrees == -90.0 + assert reason == "pan opening faces up" + + +def test_simready_pose_switch_uses_the_vlm_selected_candidate( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + scene_object = SceneObject( + id="pan_001", + kind="asset", + category="pan", + name="pan", + description="small saucepan", + ) + processor = SimReadyProcessor( + scene=Scene(objects=[scene_object]), + coarse_layout_by_id={"pan_001": {}}, + coarse_geometry_root=tmp_path / "coarse", + simready_geometry_root=tmp_path / "simready", + config=SimReadyProcessorConfig(orientation_states_by_id={"pan_001": "lying"}), + vlm_client=object(), # type: ignore[arg-type] + ) + calls: dict[str, object] = {} + + def fake_render_candidates(**kwargs: object) -> Path: + calls["candidate_render"] = kwargs + return tmp_path / "candidates.png" + + def fake_query_candidate(**kwargs: object) -> tuple[float, str]: + calls["candidate_query"] = kwargs + return -90.0, "opening up" + + def fake_rotate(**kwargs: object) -> Path: + calls["rotate"] = kwargs + return tmp_path / "rotated.glb" + + monkeypatch.setattr( + processor, + "_vlm_transform_for_object", + lambda *_args, **_kwargs: { + "pose_action": "rotate_to_required_pose", + "reason": "current pan is upright", + "target_xy_size_cm": [20.0, 20.0], + }, + ) + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor.render_object_pose_switch_candidates", + fake_render_candidates, + ) + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor.query_vlm_pose_switch_candidate", + fake_query_candidate, + ) + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor.rotate_glb_about_x_axis", + fake_rotate, + ) + + rotated_path, vlm_scale = processor._prepare_vlm_rotated_glb(scene_object) + + assert rotated_path == tmp_path / "rotated.glb" + assert vlm_scale is None + assert calls["candidate_render"] == { + "glb_path": tmp_path / "coarse" / "pan_001.glb", + "output_path": tmp_path / "debug" / "vlm_pose_candidates" / "pan_001.png", + } + assert calls["candidate_query"] == { + "scene_object_description": "small saucepan", + "needed_layout": LYING_NEEDED_LAYOUT, + "rendered_candidates_path": tmp_path / "candidates.png", + "vlm_client": processor.vlm_client, + "debug_output_path": tmp_path + / "debug" + / "vlm_pose_candidates" + / "pan_001.json", + } + assert calls["rotate"] == { + "input_path": tmp_path / "coarse" / "pan_001.glb", + "output_path": tmp_path / "simready" / "vlm_rotated" / "pan_001.glb", + "rotation_degrees": -90.0, + } + + +def test_simready_null_orientation_checks_the_default_stable_pose( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + scene_object = SceneObject( + id="book_001", + kind="asset", + category="book", + name="book", + description="small book", + ) + processor = SimReadyProcessor( + scene=Scene(objects=[scene_object]), + coarse_layout_by_id={"book_001": {}}, + coarse_geometry_root=tmp_path / "coarse", + simready_geometry_root=tmp_path / "simready", + config=SimReadyProcessorConfig(orientation_states_by_id={"book_001": None}), + vlm_client=object(), # type: ignore[arg-type] + ) + requested_layouts: list[str] = [] + + def fake_transform(*_args: object, needed_layout: str) -> dict[str, object]: + requested_layouts.append(needed_layout) + return { + "pose_action": "keep_current", + "reason": "book is already flat", + "target_xy_size_cm": [20.0, 15.0], + } + + monkeypatch.setattr(processor, "_vlm_transform_for_object", fake_transform) + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor.rotate_glb_about_x_axis", + lambda **_kwargs: tmp_path / "rotated.glb", + ) + + processor._prepare_vlm_rotated_glb(scene_object) + + assert requested_layouts == [DEFAULT_NEEDED_LAYOUT] + + def test_uniform_scale_uses_the_z_up_tabletop_footprint(tmp_path: Path) -> None: """Measure y-up GLBs against the VLM's z-up XY target footprint.""" glb_path = tmp_path / "flat_fork.glb" From 3aeede400edcb275873bd4ca6ad331ee015d1c9e Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Thu, 27 Aug 2026 17:14:10 +0800 Subject: [PATCH 40/85] add dongge's client, but failed --- .../features/generative_sim/scene_engine.md | 10 +- .../clients/articulated_generation.py | 310 ++++++++++++++++++ tests/gen_sim/scene_engine/test_clients.py | 69 +++- 3 files changed, 386 insertions(+), 3 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/clients/articulated_generation.py diff --git a/docs/source/features/generative_sim/scene_engine.md b/docs/source/features/generative_sim/scene_engine.md index f321f29a5..4fb75a781 100644 --- a/docs/source/features/generative_sim/scene_engine.md +++ b/docs/source/features/generative_sim/scene_engine.md @@ -57,7 +57,9 @@ debugging. ## Configuration Scene Engine reads the LLM, segmentation, image-generation, and -geometry-generation settings from `embodichain/gen_sim/.env`: +geometry-generation settings from `embodichain/gen_sim/.env`. The same file +also contains the optional articulation-server connection used by the +articulated-asset generation client: ```bash OPENAI_API_KEY="your-api-key" @@ -83,6 +85,12 @@ SCENE_ENGINE_GEOMETRY_GENERATION_TIMEOUT_S=600 SCENE_ENGINE_GEOMETRY_GENERATION_MAX_ATTEMPTS=3 SCENE_ENGINE_GEOMETRY_GENERATION_HEALTH_PATH="/health" SCENE_ENGINE_GEOMETRY_GENERATION_OBJECTS_PATH="/generate_multiple_objects" + +SCENE_ENGINE_ARTICULATED_GENERATION_BASE_URL="http://host:port" +SCENE_ENGINE_ARTICULATED_GENERATION_TIMEOUT_S=7200 +SCENE_ENGINE_ARTICULATED_GENERATION_MAX_ATTEMPTS=3 +SCENE_ENGINE_ARTICULATED_GENERATION_HEALTH_PATH="/health" +SCENE_ENGINE_ARTICULATED_GENERATION_GENERATE_PATH="/generate_articulation" ``` ## Processing Flow diff --git a/embodichain/gen_sim/scene_engine/clients/articulated_generation.py b/embodichain/gen_sim/scene_engine/clients/articulated_generation.py new file mode 100644 index 000000000..d1c824c1b --- /dev/null +++ b/embodichain/gen_sim/scene_engine/clients/articulated_generation.py @@ -0,0 +1,310 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +import json +import mimetypes +from pathlib import Path +import time +from typing import Any +from urllib.parse import urljoin, urlsplit + +import requests + +from embodichain.gen_sim.scene_engine.configs.environment import ( + read_scene_engine_env_values, +) + + +class ArticulatedGenerationClient: + """Request one articulated asset from articulation-server.""" + + def __init__( + self, + *, + base_url: str, + timeout_s: int, + max_attempts: int, + health_path: str, + generate_path: str, + session: requests.Session | None = None, + ) -> None: + self._base_url = _validate_base_url(base_url) + self._timeout_s = timeout_s + self._max_attempts = max_attempts + self._health_path = _validate_relative_path(health_path, "health_path") + self._generate_path = _validate_relative_path(generate_path, "generate_path") + self._session = session or requests.Session() + + @classmethod + def from_dotenv(cls) -> "ArticulatedGenerationClient": + """Create a client from its required ``gen_sim/.env`` settings.""" + return cls(**_load_dotenv_config()) + + def check_health(self) -> None: + """Raise when articulation-server does not report ``ok=true``.""" + response_data = self._request_json("get", self._health_path) + if response_data.get("ok") is not True: + raise RuntimeError( + "Articulated Generation Server health response does not contain ok=true." + ) + + def close(self) -> None: + """Close the HTTP session owned by this client.""" + self._session.close() + + def generate_articulated_object( + self, + *, + prompt: str, + image_path: str | Path | None = None, + ) -> dict[str, Any]: + """Request one articulated asset and return the server JSON response. + + Args: + prompt: Text description supplied to articulation-server. + image_path: Optional image observation sent with the prompt. + + Returns: + JSON object returned by ``/generate_articulation``. + + Raises: + FileNotFoundError: If ``image_path`` does not identify a file. + RuntimeError: If the server request or response is invalid. + ValueError: If the prompt is invalid. + """ + prompt = prompt.strip() + if not prompt: + raise ValueError("Articulated generation prompt must not be empty.") + + resolved_image_path = _resolve_optional_image_path(image_path) + return self._request_generation( + prompt=prompt, + image_path=resolved_image_path, + ) + + def generate_articulated_usdc( + self, + *, + prompt: str, + image_path: str | Path, + output_path: str | Path, + ) -> Path: + """Generate one articulated USDC asset and write it to ``output_path``.""" + response_data = self.generate_articulated_object( + prompt=prompt, + image_path=image_path, + ) + request_id = response_data.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise RuntimeError( + "Articulated Generation Server response must contain a request_id." + ) + + deadline = time.monotonic() + self._timeout_s + while True: + response_data = self._request_json("get", f"/tasks/{request_id}") + status = response_data.get("status") + if status == "succeeded": + break + if status in {"failed", "cancelled"}: + raise RuntimeError( + "Articulated Generation Server generation failed: " + f"{response_data.get('error', 'unknown error')}" + ) + if status not in {"queued", "running", "waiting"}: + raise RuntimeError( + "Articulated Generation Server returned unknown generation status: " + f"{status!r}." + ) + if time.monotonic() >= deadline: + raise RuntimeError( + "Articulated Generation Server did not finish within " + f"{self._timeout_s} seconds." + ) + time.sleep(1.0) + + result = response_data.get("result") + artifacts = result.get("artifacts") if isinstance(result, dict) else None + usdc_path = artifacts.get("usdc") if isinstance(artifacts, dict) else None + if not isinstance(usdc_path, str) or not usdc_path: + raise RuntimeError( + "Articulated Generation Server completed without a usdc." + ) + _validate_relative_path(usdc_path, "usdc artifact path") + + resolved_output_path = Path(output_path).expanduser().resolve() + resolved_output_path.parent.mkdir(parents=True, exist_ok=True) + temporary_output_path = resolved_output_path.with_name( + f".{resolved_output_path.name}.part" + ) + try: + temporary_output_path.write_bytes(self._request_content("get", usdc_path)) + temporary_output_path.replace(resolved_output_path) + except BaseException: + temporary_output_path.unlink(missing_ok=True) + raise + return resolved_output_path + + def _request_generation( + self, + *, + prompt: str, + image_path: Path | None, + ) -> dict[str, Any]: + if image_path is None: + return self._request_json( + "post", + self._generate_path, + json={"prompt": prompt}, + ) + + with image_path.open("rb") as image_file: + return self._request_json( + "post", + self._generate_path, + data={"prompt": prompt}, + files={ + "image": ( + image_path.name, + image_file, + mimetypes.guess_type(image_path.name)[0] + or "application/octet-stream", + ) + }, + ) + + def _request_json( + self, + method: str, + path: str, + **request_kwargs: object, + ) -> dict[str, Any]: + content = self._request_content(method, path, **request_kwargs) + try: + response_data = json.loads(content.decode("utf-8")) + except (UnicodeDecodeError, ValueError) as exc: + raise RuntimeError( + "Articulated Generation Server response is not valid JSON." + ) from exc + if not isinstance(response_data, dict): + raise RuntimeError( + "Articulated Generation Server response must be a JSON object." + ) + return response_data + + def _request_content( + self, + method: str, + path: str, + **request_kwargs: object, + ) -> bytes: + last_error: Exception | None = None + for _ in range(self._max_attempts): + try: + response = getattr(self._session, method)( + self._url(path), + timeout=self._timeout_s, + **request_kwargs, + ) + response.raise_for_status() + return response.content + except requests.RequestException as exc: + last_error = exc + + assert last_error is not None + raise RuntimeError( + "Articulated Generation Server request failed after " + f"{self._max_attempts} attempts." + ) from last_error + + def _url(self, path: str) -> str: + return urljoin(self._base_url, path.lstrip("/")) + + +def _load_dotenv_config() -> dict[str, Any]: + values = read_scene_engine_env_values( + "SCENE_ENGINE_ARTICULATED_GENERATION_BASE_URL", + "SCENE_ENGINE_ARTICULATED_GENERATION_TIMEOUT_S", + "SCENE_ENGINE_ARTICULATED_GENERATION_MAX_ATTEMPTS", + "SCENE_ENGINE_ARTICULATED_GENERATION_HEALTH_PATH", + "SCENE_ENGINE_ARTICULATED_GENERATION_GENERATE_PATH", + ) + timeout_s = _positive_int( + values["SCENE_ENGINE_ARTICULATED_GENERATION_TIMEOUT_S"], + "SCENE_ENGINE_ARTICULATED_GENERATION_TIMEOUT_S", + ) + max_attempts = _positive_int( + values["SCENE_ENGINE_ARTICULATED_GENERATION_MAX_ATTEMPTS"], + "SCENE_ENGINE_ARTICULATED_GENERATION_MAX_ATTEMPTS", + ) + return { + "base_url": values["SCENE_ENGINE_ARTICULATED_GENERATION_BASE_URL"].strip(), + "timeout_s": timeout_s, + "max_attempts": max_attempts, + "health_path": values[ + "SCENE_ENGINE_ARTICULATED_GENERATION_HEALTH_PATH" + ].strip(), + "generate_path": values[ + "SCENE_ENGINE_ARTICULATED_GENERATION_GENERATE_PATH" + ].strip(), + } + + +def _positive_int(value: str, key: str) -> int: + try: + parsed_value = int(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{key} must be an integer.") from exc + if parsed_value < 1: + raise ValueError(f"{key} must be at least 1.") + return parsed_value + + +def _validate_base_url(base_url: str) -> str: + parsed_url = urlsplit(base_url) + if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc: + raise ValueError("base_url must be an absolute HTTP(S) URL.") + if parsed_url.query or parsed_url.fragment: + raise ValueError("base_url must not contain a query or fragment.") + return base_url.rstrip("/") + "/" + + +def _validate_relative_path(path: str, field_name: str) -> str: + if not path.strip(): + raise ValueError(f"{field_name} must be a non-empty relative URL path.") + parsed_url = urlsplit(path) + if ( + parsed_url.scheme + or parsed_url.netloc + or parsed_url.query + or parsed_url.fragment + ): + raise ValueError(f"{field_name} must be a relative URL path.") + return path + + +def _resolve_optional_image_path(image_path: str | Path | None) -> Path | None: + if image_path is None: + return None + resolved_image_path = Path(image_path).expanduser().resolve() + if not resolved_image_path.is_file(): + raise FileNotFoundError( + f"Articulated generation input not found: {resolved_image_path}" + ) + return resolved_image_path diff --git a/tests/gen_sim/scene_engine/test_clients.py b/tests/gen_sim/scene_engine/test_clients.py index 948db9c73..e3b10e02b 100644 --- a/tests/gen_sim/scene_engine/test_clients.py +++ b/tests/gen_sim/scene_engine/test_clients.py @@ -17,11 +17,13 @@ from __future__ import annotations +import json from pathlib import Path from typing import Any import pytest +from embodichain.gen_sim.scene_engine.clients import articulated_generation from embodichain.gen_sim.scene_engine.clients import geometry_generation from embodichain.gen_sim.scene_engine.clients import image_generation from embodichain.gen_sim.scene_engine.clients import image_segmentation @@ -39,7 +41,7 @@ def __init__( headers: dict[str, str] | None = None, ) -> None: self._payload = payload - self.content = content + self.content = content or json.dumps(payload).encode("utf-8") self.headers = headers or {} def raise_for_status(self) -> None: @@ -60,7 +62,7 @@ def __init__( self.get_calls: list[tuple[str, int]] = [] self.post_call: dict[str, object] | None = None - def get(self, url: str, *, timeout: int) -> _Response: + def get(self, url: str, *, timeout: int, **_: object) -> _Response: self.get_calls.append((url, timeout)) return _Response(self.get_payload) @@ -96,6 +98,13 @@ def test_clients_load_their_required_dotenv_values( "SCENE_ENGINE_IMAGE_GENERATION_HEALTH_PATH": "/health", "SCENE_ENGINE_IMAGE_GENERATION_BY_PROMPT_PATH": "/generate_image_by_prompt", } + articulated_generation_values = { + "SCENE_ENGINE_ARTICULATED_GENERATION_BASE_URL": "http://articulation/", + "SCENE_ENGINE_ARTICULATED_GENERATION_TIMEOUT_S": "7200", + "SCENE_ENGINE_ARTICULATED_GENERATION_MAX_ATTEMPTS": "2", + "SCENE_ENGINE_ARTICULATED_GENERATION_HEALTH_PATH": "/health", + "SCENE_ENGINE_ARTICULATED_GENERATION_GENERATE_PATH": "/generate_articulation", + } llm_values = { "OPENAI_API_KEY": "test-key", "OPENAI_MODEL": "test-model", @@ -116,6 +125,11 @@ def test_clients_load_their_required_dotenv_values( "read_scene_engine_env_values", lambda *_: image_generation_values, ) + monkeypatch.setattr( + articulated_generation, + "read_scene_engine_env_values", + lambda *_: articulated_generation_values, + ) monkeypatch.setattr( load_config, "read_scene_engine_env_values", lambda *_: llm_values ) @@ -123,6 +137,9 @@ def test_clients_load_their_required_dotenv_values( geometry_client = geometry_generation.GeometryGenerationClient.from_dotenv() segmentation_client = image_segmentation.ImageSegmentationClient.from_dotenv() image_generation_client = image_generation.ImageGenerationClient.from_dotenv() + articulated_generation_client = ( + articulated_generation.ArticulatedGenerationClient.from_dotenv() + ) llm_client_config = load_config.load_llm_config() assert geometry_client._base_url == "http://geometry" @@ -134,6 +151,7 @@ def test_clients_load_their_required_dotenv_values( image_generation_client._generate_image_by_prompt_path == "/generate_image_by_prompt" ) + assert articulated_generation_client._base_url == "http://articulation/" assert llm_client_config.default_query == {"api-version": "1"} assert llm_client_config.base_url == "http://llm/v1" @@ -184,16 +202,63 @@ def test_service_health_checks_use_the_configured_health_path() -> None: generate_image_by_prompt_path="/generate_image_by_prompt", session=image_generation_session, ) + articulated_generation_session = _Session(get_payload={"ok": True}) + articulated_generation_client = articulated_generation.ArticulatedGenerationClient( + base_url="http://articulation", + timeout_s=30, + max_attempts=1, + health_path="/health", + generate_path="/generate_articulation", + session=articulated_generation_session, + ) geometry_client.check_health() segmentation_client.check_health() image_generation_client.check_health() + articulated_generation_client.check_health() assert geometry_session.get_calls == [("http://geometry/health", 10)] assert segmentation_session.get_calls == [("http://segment/health", 30)] assert image_generation_session.get_calls == [ ("http://image-generation/health", 10) ] + assert articulated_generation_session.get_calls == [ + ("http://articulation/health", 30) + ] + + +def test_articulated_generation_client_posts_image_and_returns_server_json( + tmp_path: Path, +) -> None: + class ArticulatedGenerationSession(_Session): + def __init__(self) -> None: + super().__init__(get_payload={}) + + def post(self, url: str, **kwargs: object) -> _Response: + self.post_call = {"url": url, **kwargs} + return _Response({"status": "accepted"}, content=b'{"status":"accepted"}') + + image_path = tmp_path / "reference.png" + image_path.write_bytes(b"png") + session = ArticulatedGenerationSession() + client = articulated_generation.ArticulatedGenerationClient( + base_url="http://articulation", + timeout_s=30, + max_attempts=1, + health_path="/health", + generate_path="/generate_articulation", + session=session, + ) + + response_data = client.generate_articulated_object( + prompt="a cabinet with one opening door", + image_path=image_path, + ) + + assert response_data == {"status": "accepted"} + assert session.post_call is not None + assert session.post_call["url"] == "http://articulation/generate_articulation" + assert session.post_call["data"] == {"prompt": "a cabinet with one opening door"} def test_image_generation_client_posts_prompt_and_writes_png( From acd43875c53b45cb4595a14627a1ebfce93035e1 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 28 Aug 2026 09:47:08 +0800 Subject: [PATCH 41/85] VLM-auto gives a yaw, but still have a small bug in coarse layout: sometimes the table's coarse layout is weird. --- .../scene_engine/core/scene_edit_plan.py | 31 ++--- .../gen_sim/scene_engine/core/scene_graph.py | 38 +++--- .../editing/scene_edit_asset_preparation.py | 14 ++- .../editing/scene_edit_understanding.py | 76 ++++++----- .../pipeline/generation/scene_generation.py | 118 +----------------- .../generation/scene_understanding.py | 39 +++--- .../pipeline/utils/scene_importer.py | 14 ++- .../pipeline/utils/simready_processor.py | 31 ++--- .../utils/simready_processor_utils.py | 17 +-- .../test_scene_core_and_export.py | 33 ++++- .../scene_engine/test_scene_edit_plan.py | 58 +++++---- .../scene_engine/test_scene_generation.py | 52 -------- .../gen_sim/scene_engine/test_scene_graph.py | 11 +- .../scene_engine/test_scene_understanding.py | 45 ++++--- .../test_simready_processor_utils.py | 39 +++--- 15 files changed, 248 insertions(+), 368 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 902f6eda7..a459eb523 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -21,7 +21,6 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( - OrientationState, SceneConstraintType, SceneGraph, TableRegion, @@ -43,7 +42,7 @@ class SceneEditOperation: category: str | None = None name: str | None = None description: str | None = None - orientation_state: OrientationState | None = None + pose_description: str | None = None def to_dict(self) -> dict[str, object]: """Serialize one normalized edit operation.""" @@ -56,7 +55,7 @@ def to_dict(self) -> dict[str, object]: "category": self.category, "name": self.name, "description": self.description, - "orientation_state": self.orientation_state, + "pose_description": self.pose_description, } @@ -85,7 +84,7 @@ def validate(self) -> None: # Edit-plan rules: # - move and delete identify one existing non-table object with object_id. # - add carries generated object_id plus non-empty category, name, and description. - # - add may preserve an explicit standing or lying user placement intent. + # - add and move may carry an explicit self-pose description. # - move always supplies target_id and relation; add may omit both. # - table_region is only valid with target_id=table and relation=on. # - target_id and relation are otherwise supplied together or both absent. @@ -166,7 +165,7 @@ def _validate_operation( operation.category, operation.name, operation.description, - operation.orientation_state, + operation.pose_description, ) ): raise ValueError("Delete operations may only specify object_id.") @@ -174,13 +173,8 @@ def _validate_operation( if operation.target_id is None or operation.relation is None: raise ValueError("Move operations must specify target_id and relation.") - existing_orientation_state = self.scene_graph.node_by_id()[ - operation.object_id - ].orientation_state - if operation.orientation_state not in {None, existing_orientation_state}: - raise ValueError( - "Move operations may only preserve the existing orientation_state." - ) + if operation.pose_description is not None: + self._validate_pose_description(operation.pose_description) self._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, @@ -217,14 +211,23 @@ def _validate_add_operation( for value in (operation.category, operation.name, operation.description) ): raise ValueError("Add operations require category, name, and description.") - if operation.orientation_state not in {None, "standing", "lying"}: - raise ValueError("Add operation orientation_state is invalid.") + if operation.pose_description is not None: + SceneEditPlan._validate_pose_description(operation.pose_description) SceneEditPlan._validate_position_reference( operation=operation, existing_object_ids=existing_object_ids, deleted_object_ids=deleted_object_ids, ) + @staticmethod + def _validate_pose_description(pose_description: str) -> None: + if ( + not isinstance(pose_description, str) + or not pose_description.strip() + or len(pose_description) > 240 + ): + raise ValueError("pose_description must be a short non-empty string.") + @staticmethod def _validate_position_reference( *, diff --git a/embodichain/gen_sim/scene_engine/core/scene_graph.py b/embodichain/gen_sim/scene_engine/core/scene_graph.py index b43c7353a..2efde829f 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_graph.py +++ b/embodichain/gen_sim/scene_engine/core/scene_graph.py @@ -54,23 +54,22 @@ # A PlanarRelation with B, then A and B must have the same parent node. PlanarRelationType = Literal["left_of", "right_of", "in_front_of", "behind"] SceneConstraintType = SupportRelationType | PlanarRelationType -OrientationState = Literal["standing", "lying"] @dataclass class SceneGraphNode: """One object node in the edit-time scene hierarchy. - ``orientation_state`` is an image-derived placement semantic, rather than - an edge to the node itself or an exact three-dimensional transform. + ``pose_description`` records the object's desired pose relative to its + direct support, rather than an exact three-dimensional transform. """ object_id: str parent_id: str | None parent_relation: SupportRelationType | None = None table_region: TableRegion | None = None - # Preserves image-observed placement semantics for later pose refinement. - orientation_state: OrientationState | None = None + # Preserves image-observed or user-requested pose semantics for refinement. + pose_description: str | None = None def __post_init__(self) -> None: """Validate local node fields before graph-level checks.""" @@ -78,16 +77,20 @@ def __post_init__(self) -> None: raise ValueError("object_id must be non-empty.") if self.table_region not in {None, *TABLE_REGIONS}: raise ValueError("table_region is invalid.") - if self.orientation_state not in {None, "standing", "lying"}: - raise ValueError("orientation_state is invalid.") + if self.pose_description is not None and ( + not isinstance(self.pose_description, str) + or not self.pose_description.strip() + or len(self.pose_description) > 240 + ): + raise ValueError("pose_description is invalid.") # If the node is the table. if self.object_id == TABLE_OBJECT_ID: if self.parent_id is not None: raise ValueError("table must not have a parent.") if self.parent_relation is not None: raise ValueError("table must not have a parent relation.") - if self.orientation_state is not None: - raise ValueError("table must not have an orientation state.") + if self.pose_description is not None: + raise ValueError("table must not have a pose description.") # If the node is not the table. elif self.parent_id is None: raise ValueError("non-table nodes must have a parent.") @@ -101,7 +104,7 @@ def to_dict(self) -> dict[str, object]: "parent_id": self.parent_id, "parent_relation": self.parent_relation, "table_region": self.table_region, - "orientation_state": self.orientation_state, + "pose_description": self.pose_description, } @@ -193,7 +196,8 @@ def apply_updates( *, deleted_object_ids: set[str], added_object_ids: list[str], - added_orientation_states_by_id: dict[str, OrientationState | None], + added_pose_descriptions_by_id: dict[str, str | None], + pose_description_updates_by_id: dict[str, str], on_parent_updates: list[tuple[str, str, TableRegion | None]], planar_relation_updates: list[tuple[str, PlanarRelationType, str]], ) -> None: @@ -227,8 +231,8 @@ def apply_updates( raise ValueError( f"Duplicate scene graph nodes: {sorted(duplicate_object_ids)}" ) - if set(added_orientation_states_by_id) != set(added_object_ids): - raise ValueError("Added orientation states must match added node ids.") + if set(added_pose_descriptions_by_id) != set(added_object_ids): + raise ValueError("Added pose descriptions must match added node ids.") # New nodes default to the table; later updates replace that parent when needed. self.nodes.extend( @@ -236,11 +240,17 @@ def apply_updates( object_id=object_id, parent_id=TABLE_OBJECT_ID, parent_relation="on", - orientation_state=added_orientation_states_by_id[object_id], + pose_description=added_pose_descriptions_by_id[object_id], ) for object_id in added_object_ids ) + for object_id, pose_description in pose_description_updates_by_id.items(): + node = self.node_by_id().get(object_id) + if node is None or object_id == TABLE_OBJECT_ID: + raise ValueError("Pose descriptions may update only existing assets.") + node.pose_description = pose_description + # Apply support-parent changes before planar updates need the final parent. for object_id, parent_id, table_region in on_parent_updates: self._set_on_parent( diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index aa19e97e1..670436e39 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -119,13 +119,15 @@ def prepare_scene_edit_assets( config=SimReadyProcessorConfig( use_vlm_scale=vlm_client is not None, use_vlm_rotation=vlm_client is not None, - # An explicit edit state overrides the default stable tabletop pose. - orientation_states_by_id={ - operation.object_id: operation.orientation_state + # Explicit edit pose descriptions override the default stable pose. + pose_descriptions_by_id={ + operation.object_id: operation.pose_description for operation in scene_edit_plan.operations - if operation.op == "add" - and operation.object_id is not None - and operation.orientation_state is not None + if ( + operation.op == "add" + and operation.object_id is not None + and operation.pose_description is not None + ) }, ), vlm_client=vlm_client, diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py index 3205a7a48..be54d41e1 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_understanding.py @@ -26,7 +26,6 @@ ) from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( - OrientationState, PlanarRelationType, SceneGraph, SceneGraphNode, @@ -54,9 +53,8 @@ singular snake_case category, name, and description. Multiple add operations may have the same category and name; their final IDs are assigned by the program in operation order. target_id and relation are either both provided - or both null. Set orientation_state to standing or lying only when the user - explicitly asks for that placement; otherwise set it to null so the object - uses its natural, physically stable tabletop pose. + or both null. Set pose_description only when the user explicitly requests a + self-pose; otherwise set it to null. For every move and every positioned add, target_id must be an Existing object ID and relation must be one of on, left_of, right_of, in_front_of, or behind. @@ -83,14 +81,16 @@ class. name contains only color, material, texture, shape, and object details. description contains only visible category, material, color, texture, shape, and structural details. name and description must not mention position, the -table, relations to any object, or orientation. orientation_state must be null -unless the user explicitly requests standing/upright/vertical or lying/flat/ -horizontal placement. Follow that explicit user intent even if it is not the -object's natural stable pose. +table, relations to any object, or orientation. pose_description is a short +English sentence describing only the asset relative to its direct support, for +example "Stand upright on its base.", "Lie flat on the support surface.", or +"Rest stably with the cooking surface facing upward." It must be null unless +the user explicitly requests a self-pose. Follow that explicit user intent even +if it is not the object's natural stable pose. Return JSON only: no Markdown, comments, or prose. Every operation must contain exactly these fields: op, object_id, target_id, relation, table_region, category, -name, description, and orientation_state. Use null for every field that does not apply: +name, description, and pose_description. Use null for every field that does not apply: { "operations": [ { @@ -102,7 +102,7 @@ "category": null, "name": null, "description": null, - "orientation_state": null + "pose_description": null }, { "op": "delete", @@ -113,7 +113,7 @@ "category": null, "name": null, "description": null, - "orientation_state": null + "pose_description": null }, { "op": "add", @@ -124,7 +124,7 @@ "category": "orange", "name": "small orange", "description": "small round orange with a textured peel", - "orientation_state": null + "pose_description": null }, { "op": "add", @@ -135,7 +135,7 @@ "category": "orange", "name": "small orange", "description": "small round orange with a textured peel", - "orientation_state": null + "pose_description": null }, { "op": "add", @@ -146,7 +146,7 @@ "category": "bottle", "name": "blue glass bottle", "description": "tall transparent blue glass bottle with a narrow neck", - "orientation_state": "standing" + "pose_description": "Stand upright on its base." }, { "op": "add", @@ -157,17 +157,16 @@ "category": "fork", "name": "silver metal fork", "description": "four-tined silver stainless-steel fork with a plain handle", - "orientation_state": "lying" + "pose_description": "Lie flat on the support surface." } ] } The two orange additions intentionally share category and name. The bottle example represents an explicit user request to stand it upright, and the fork example represents an explicit user request to lay it flat. Only add operations -may introduce a new non-null orientation_state. A move may use null or repeat -its existing orientation_state from the supplied scene metadata, but it must not -change that state. Delete operations must use null. Do not add fields beyond the -required schema.""" +contains an explicit self-pose request. A move uses null to preserve its existing +pose_description, or replaces it only for an explicit self-pose request. Delete +operations must use null. Do not add fields beyond the required schema.""" def understand_scene_edit( @@ -224,7 +223,7 @@ def _build_updated_scene_graph( parent_id=node.parent_id, parent_relation=node.parent_relation, table_region=node.table_region, - orientation_state=node.orientation_state, + pose_description=node.pose_description, ) for node in scene_graph.nodes ], @@ -253,7 +252,8 @@ def _apply_scene_edit_plan_to_scene_graph( """Apply the target graph updates implied by add and move operations.""" deleted_object_ids: set[str] = set() added_object_ids: list[str] = [] - added_orientation_states_by_id: dict[str, OrientationState | None] = {} + added_pose_descriptions_by_id: dict[str, str | None] = {} + pose_description_updates_by_id: dict[str, str] = {} on_parent_updates: list[tuple[str, str, TableRegion | None]] = [] planar_relation_updates: list[tuple[str, PlanarRelationType, str]] = [] for operation in scene_edit_plan.operations: @@ -265,8 +265,12 @@ def _apply_scene_edit_plan_to_scene_graph( raise ValueError("Add and move operations must have an object_id.") if operation.op == "add": added_object_ids.append(operation.object_id) - added_orientation_states_by_id[operation.object_id] = ( - operation.orientation_state + added_pose_descriptions_by_id[operation.object_id] = ( + operation.pose_description + ) + elif operation.pose_description is not None: + pose_description_updates_by_id[operation.object_id] = ( + operation.pose_description ) if operation.target_id is None or operation.relation is None: continue @@ -287,7 +291,8 @@ def _apply_scene_edit_plan_to_scene_graph( scene_graph.apply_updates( deleted_object_ids=deleted_object_ids, added_object_ids=added_object_ids, - added_orientation_states_by_id=added_orientation_states_by_id, + added_pose_descriptions_by_id=added_pose_descriptions_by_id, + pose_description_updates_by_id=pose_description_updates_by_id, on_parent_updates=on_parent_updates, planar_relation_updates=planar_relation_updates, ) @@ -302,8 +307,8 @@ def _simplify_scene_info( table_regions_by_id = { node.object_id: node.table_region for node in scene_graph.nodes } - orientation_states_by_id = { - node.object_id: node.orientation_state for node in scene_graph.nodes + pose_descriptions_by_id = { + node.object_id: node.pose_description for node in scene_graph.nodes } return { "existing_object_ids": [scene_object.id for scene_object in scene.objects], @@ -315,7 +320,7 @@ def _simplify_scene_info( "description": scene_object.description, "center_xy": scene_object.center_xy, "table_region": table_regions_by_id.get(scene_object.id), - "orientation_state": orientation_states_by_id.get(scene_object.id), + "pose_description": pose_descriptions_by_id.get(scene_object.id), } for scene_object in scene.objects ], @@ -390,7 +395,7 @@ def _parse_scene_edit_operations( "category", "name", "description", - "orientation_state", + "pose_description", } # Get ids and counts of existing objects to assign new add IDs. assigned_object_ids = {scene_object.id for scene_object in scene.objects} @@ -411,7 +416,9 @@ def _parse_scene_edit_operations( raise ValueError("Scene edit operations must use the required schema.") object_id = _optional_string(value.get("object_id"), field_name="object_id") category = _optional_string(value.get("category"), field_name="category") - orientation_state = _optional_orientation_state(value.get("orientation_state")) + pose_description = _optional_string( + value.get("pose_description"), field_name="pose_description" + ) if op == "add": if object_id is not None: raise ValueError("VLM add operations must set object_id to null.") @@ -437,7 +444,7 @@ def _parse_scene_edit_operations( description=_optional_string( value.get("description"), field_name="description" ), - orientation_state=orientation_state, + pose_description=pose_description, ) ) return operations @@ -492,12 +499,3 @@ def _optional_table_region(value: object) -> TableRegion | None: if value not in TABLE_REGIONS: raise ValueError("Scene edit operation table_region is invalid.") return value - - -def _optional_orientation_state(value: object) -> OrientationState | None: - """Validate an optional explicit upright or lying edit intent.""" - if value is None: - return None - if value not in {"standing", "lying"}: - raise ValueError("Scene edit operation orientation_state is invalid.") - return value diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 2a7e526dd..986d7e18f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -132,12 +132,11 @@ def generate_scene_and_refine( coarse_layout_by_id = { layout_object["id"]: layout_object for layout_object in coarse_layout } - # Every graph asset, including null, receives a VLM check before SimReady - # canonicalization; null requests its natural stable tabletop pose. - orientation_states_by_id = { - node.object_id: node.orientation_state + # Only an explicit graph pose description requires a VLM pose adjustment. + pose_descriptions_by_id = { + node.object_id: node.pose_description for node in scene_graph.nodes - if node.object_id != TABLE_OBJECT_ID + if node.object_id != TABLE_OBJECT_ID and node.pose_description is not None } simready_processor = SimReadyProcessor( scene=scene, @@ -145,11 +144,11 @@ def generate_scene_and_refine( coarse_geometry_root=coarse_geometry_output_root, simready_geometry_root=simready_geometry_output_root, debug_output_root=debug_output_root, - # Keep geometry-server scale while applying explicit orientation semantics. + # Keep geometry-server scale while applying graph self-pose semantics. config=SimReadyProcessorConfig( use_vlm_scale=False, use_vlm_rotation=False, - orientation_states_by_id=orientation_states_by_id, + pose_descriptions_by_id=pose_descriptions_by_id, ), vlm_client=vlm_client, ) @@ -483,13 +482,6 @@ def _layout_refinement( ) ) - # 3. Correct image-observed standing containers before every geometry-based - # layout stage measures their footprint. - refined_assets_layout = _scene_graph_based_calibration( - scene_graph=scene_graph, - assets_layout=refined_assets_layout, - ) - # 4. Only direct on-table children participate in the table-level layout # stages. Their descendants follow each solved root transform until their # own parent-surface optimization is introduced in a later BFS pass. @@ -1198,104 +1190,6 @@ def _apply_root_layout_updates_to_descendant_subtrees( ] -def _scene_graph_based_calibration( - *, - scene_graph: SceneGraph, - assets_layout: list[dict[str, object]], -) -> list[dict[str, object]]: - """Minimally align graph-marked standing assets with the z-up table frame.""" - # This is the extension point for future image-conditioned scene generation - # calibration. The scene graph may later provide richer image-grounded - # constraints, but the current implementation deliberately consumes only - # ``orientation_state`` to correct standing container axes before layout. - y_up_to_z_up_matrix = np.eye(4) - y_up_to_z_up_matrix[:3, :3] = np.array( - [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] - ) - z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) - nodes_by_id = scene_graph.node_by_id() - calibrated_assets_layout: list[dict[str, object]] = [] - - for asset_layout in assets_layout: - asset_id = asset_layout.get("id") - if not isinstance(asset_id, str) or not asset_id: - raise ValueError("Each asset layout must contain a non-empty string id.") - node = nodes_by_id.get(asset_id) - if node is None: - raise ValueError(f"Scene graph does not contain asset {asset_id!r}.") - # Only correct the standing assets. - if node.orientation_state != "standing": - calibrated_assets_layout.append(asset_layout) - continue - - # Conjugate the y-up pose so the SimReady container axis is local z. - z_up_asset_to_table_matrix = ( - y_up_to_z_up_matrix - @ layout_object_to_transform_matrix(asset_layout) - @ z_up_to_y_up_matrix - ) - linear_matrix = z_up_asset_to_table_matrix[:3, :3] - # Layout transforms store rotation and per-axis scale in the same matrix. - scale = np.linalg.norm(linear_matrix, axis=0) - if np.any(scale <= 1e-8): - raise ValueError(f"Asset {asset_id!r} has a zero scale axis.") - rotation_matrix = linear_matrix / scale - if not np.allclose(rotation_matrix.T @ rotation_matrix, np.eye(3), atol=1e-6): - raise ValueError(f"Asset {asset_id!r} layout contains shear.") - - local_z_axis_in_table = rotation_matrix[:, 2] - # Treat the long axis as unsigned to avoid an unnecessary 180-degree flip. - target_z_axis = np.array( - [0.0, 0.0, 1.0 if local_z_axis_in_table[2] >= 0.0 else -1.0] - ) - # Left multiplication applies the correction in the table/world frame. - z_up_asset_to_table_matrix[:3, :3] = ( - _minimum_axis_alignment_rotation( - source_axis=local_z_axis_in_table, - target_axis=target_z_axis, - ) - @ rotation_matrix - @ np.diag(scale) - ) - calibrated_assets_layout.append( - transform_matrix_to_layout_object( - asset_id, - z_up_to_y_up_matrix @ z_up_asset_to_table_matrix @ y_up_to_z_up_matrix, - ) - ) - return calibrated_assets_layout - - -def _minimum_axis_alignment_rotation( - *, - source_axis: np.ndarray, - target_axis: np.ndarray, -) -> np.ndarray: - """Return the smallest proper rotation mapping one nonzero axis to another.""" - source = np.asarray(source_axis, dtype=float) - target = np.asarray(target_axis, dtype=float) - source_norm = np.linalg.norm(source) - target_norm = np.linalg.norm(target) - if source_norm <= 1e-8 or target_norm <= 1e-8: - raise ValueError("Axis alignment requires nonzero axes.") - source /= source_norm - target /= target_norm - - cross_product = np.cross(source, target) - sine = np.linalg.norm(cross_product) - cosine = float(np.clip(np.dot(source, target), -1.0, 1.0)) - if sine <= 1e-8: - if cosine > 0.0: - return np.eye(3) - basis_axis = np.eye(3)[np.argmin(np.abs(source))] - rotation_axis = np.cross(source, basis_axis) - rotation_axis /= np.linalg.norm(rotation_axis) - return Rotation.from_rotvec(np.pi * rotation_axis).as_matrix() - - rotation_axis = cross_product / sine - return Rotation.from_rotvec(np.arctan2(sine, cosine) * rotation_axis).as_matrix() - - def _measure_table_and_assets_in_z_up_world( *, table_layout: dict[str, object], diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py index f8fb48d50..706f1ca4a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_understanding.py @@ -162,7 +162,7 @@ text.""" _INITIAL_SCENE_GRAPH_SYSTEM_PROMPT = """You inspect an outlined tabletop-scene image. Each visible asset has an outline and an ID label. Build a support graph for the -listed assets and determine each asset's image-observed orientation state. +listed assets and determine each asset's image-observed self-pose description. Every asset must have exactly one direct support parent with relation "on". Use "table" when the asset directly rests on the table. Use another supplied @@ -171,20 +171,21 @@ When the direct support parent is uncertain, use "table". The table is a fixed support ID, not an output node: never include it in nodes. -Use a non-null orientation_state only for an elongated object with a clear -primary long axis. Use "standing" when that axis is approximately vertical to -the tabletop, and "lying" when it is approximately parallel to the tabletop. -Use null for every object without a clear primary long axis or when uncertain. -The orientation state describes the asset itself and is independent of its -support parent. +pose_description must be one short English sentence describing only the asset's +pose relative to its direct support. Do not describe left/right/front/back or +relations to other objects. Use null when the image does not provide a reliable +pose requirement. For a bottle, pen, or long tool, describe whether it stands +on its base or lies flat. For a pan, bowl, plate, or cup, describe functional +up/down semantics such as an opening or cooking surface facing upward. Do not +force every object into standing or lying. Examples: - A bottle directly on the table is upright: - {"nodes": [{"object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", "orientation_state": "standing"}]} + {"nodes": [{"object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", "pose_description": "Stand upright on its base."}]} - A pen lies flat on a book, and the book is on the table: - {"nodes": [{"object_id": "book_001", "parent_id": "table", "parent_relation": "on", "orientation_state": null}, {"object_id": "pen_001", "parent_id": "book_001", "parent_relation": "on", "orientation_state": "lying"}]} -- A round cup directly on the table has no reliable long axis: - {"nodes": [{"object_id": "cup_001", "parent_id": "table", "parent_relation": "on", "orientation_state": null}]} + {"nodes": [{"object_id": "book_001", "parent_id": "table", "parent_relation": "on", "pose_description": null}, {"object_id": "pen_001", "parent_id": "book_001", "parent_relation": "on", "pose_description": "Lie flat on the support surface."}]} +- A frying pan rests normally on the table: + {"nodes": [{"object_id": "frying_pan_001", "parent_id": "table", "parent_relation": "on", "pose_description": "Rest stably with the cooking surface facing upward."}]} Return JSON only, with exactly one key: nodes. Include every supplied asset ID exactly once and no unknown IDs. Do not include Markdown or any other text.""" @@ -349,32 +350,36 @@ def _parse_initial_scene_graph_response( "object_id", "parent_id", "parent_relation", - "orientation_state", + "pose_description", }: raise ValueError( "VLM JSON nodes[" f"{index}] must contain exactly object_id, parent_id, " - "parent_relation, and orientation_state." + "parent_relation, and pose_description." ) object_id = node_value["object_id"] parent_id = node_value["parent_id"] parent_relation = node_value["parent_relation"] - orientation_state = node_value["orientation_state"] + pose_description = node_value["pose_description"] if not isinstance(object_id, str) or not object_id: raise ValueError(f"VLM JSON nodes[{index}].object_id is invalid.") if not isinstance(parent_id, str) or not parent_id: raise ValueError(f"VLM JSON nodes[{index}].parent_id is invalid.") if parent_relation != "on": raise ValueError(f"VLM JSON nodes[{index}].parent_relation is invalid.") - if orientation_state not in {None, "standing", "lying"}: - raise ValueError(f"VLM JSON nodes[{index}].orientation_state is invalid.") + if pose_description is not None and ( + not isinstance(pose_description, str) + or not pose_description.strip() + or len(pose_description) > 240 + ): + raise ValueError(f"VLM JSON nodes[{index}].pose_description is invalid.") if object_id in nodes_by_id: raise ValueError(f"VLM JSON repeats scene graph node for {object_id!r}.") nodes_by_id[object_id] = SceneGraphNode( object_id=object_id, parent_id=parent_id, parent_relation=parent_relation, - orientation_state=orientation_state, + pose_description=pose_description, ) if set(nodes_by_id) != set(asset_ids): diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index bf5d9caea..030560858 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -192,14 +192,14 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: "parent_id", "parent_relation", "table_region", - "orientation_state", + "pose_description", }: raise ValueError("Scene graph nodes must use the serialized node schema.") object_id = value["object_id"] parent_id = value["parent_id"] parent_relation = value["parent_relation"] table_region = value["table_region"] - orientation_state = value["orientation_state"] + pose_description = value["pose_description"] if not isinstance(object_id, str) or not isinstance( parent_id, (str, type(None)) ): @@ -208,14 +208,18 @@ def _scene_graph_node_from_data(value: object) -> SceneGraphNode: raise ValueError("Scene graph parent_relation must be 'on' or null.") if table_region is not None and table_region not in TABLE_REGIONS: raise ValueError("Scene graph table_region is invalid.") - if orientation_state not in {None, "standing", "lying"}: - raise ValueError("Scene graph orientation_state is invalid.") + if pose_description is not None and ( + not isinstance(pose_description, str) + or not pose_description.strip() + or len(pose_description) > 240 + ): + raise ValueError("Scene graph pose_description is invalid.") return SceneGraphNode( object_id=object_id, parent_id=parent_id, parent_relation=parent_relation, table_region=table_region, - orientation_state=orientation_state, + pose_description=pose_description, ) @staticmethod diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index fb8d4c2d9..b18b97008 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -25,7 +25,6 @@ import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import OrientationState from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, @@ -35,8 +34,6 @@ ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( DEFAULT_NEEDED_LAYOUT, - LYING_NEEDED_LAYOUT, - STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, query_vlm_pose_switch_candidate, query_vlm_object_pose_and_target_size, @@ -73,10 +70,8 @@ class SimReadyProcessorConfig: use_vlm_scale: bool = False # Use the VLM-selected asset scale. use_vlm_rotation: bool = False # Use the VLM-selected asset rotation. - # Explicit graph orientation overrides the default stable tabletop pose. - orientation_states_by_id: dict[str, OrientationState | None] = field( - default_factory=dict - ) + # Explicit graph pose descriptions override the default stable tabletop pose. + pose_descriptions_by_id: dict[str, str | None] = field(default_factory=dict) class SimReadyProcessor: @@ -112,7 +107,7 @@ def __init__( if ( self.config.use_vlm_scale or self.config.use_vlm_rotation - or self.config.orientation_states_by_id + or self.config.pose_descriptions_by_id ) and vlm_client is None: raise ValueError("vlm_client is required when VLM transforms are enabled.") @@ -209,14 +204,14 @@ def _prepare_vlm_rotated_glb( ) -> tuple[Path, list[float] | None]: """Render, query, and optionally bake the VLM-selected x-axis rotation.""" coarse_path = self.coarse_geometry_root / f"{scene_object.id}.glb" - # A graph entry, including null, requests a VLM check of the desired pose. - orientation_pose_required = ( - scene_object.id in self.config.orientation_states_by_id + # Only an explicit graph pose description requests a VLM pose check. + pose_description_required = ( + scene_object.id in self.config.pose_descriptions_by_id ) if not ( self.config.use_vlm_scale or self.config.use_vlm_rotation - or orientation_pose_required + or pose_description_required ): return coarse_path, None decision = self._vlm_transform_for_object( @@ -273,20 +268,10 @@ def _prepare_vlm_rotated_glb( [vlm_scale, vlm_scale, vlm_scale] if self.config.use_vlm_scale else None, ) - def _orientation_state_for_object(self, object_id: str) -> OrientationState | None: - """Return the explicit graph orientation requested for one object.""" - return self.config.orientation_states_by_id.get(object_id) - def _needed_layout_for_object(self, object_id: str) -> str: """Return the VLM layout instruction for one object's graph semantics.""" return ( - STANDING_NEEDED_LAYOUT - if self._orientation_state_for_object(object_id) == "standing" - else ( - LYING_NEEDED_LAYOUT - if self._orientation_state_for_object(object_id) == "lying" - else DEFAULT_NEEDED_LAYOUT - ) + self.config.pose_descriptions_by_id.get(object_id) or DEFAULT_NEEDED_LAYOUT ) def _vlm_transform_for_object( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py index 3b57463ca..770a5a00a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor_utils.py @@ -117,19 +117,10 @@ _VISUAL_MIN_FLOOR_SIZE = 0.45 # Keep a contact cue for compact upright objects. DEFAULT_NEEDED_LAYOUT = ( - "Place this asset on the table in its natural, physically stable resting " - "orientation. For example, a fork should lie flat on the table rather " - "than stand on an edge." -) -STANDING_NEEDED_LAYOUT = ( - "The scene graph requires this asset to stand vertically on the table, " - "even when its natural stable pose would be lying down. For example, a " - "bottle should stand on its base and a fork should stand upright." -) -LYING_NEEDED_LAYOUT = ( - "The scene graph requires this asset to lie flat on the table, even when " - "its natural stable pose would be standing. For example, a bottle should " - "lie on its side and a fork should lie flat." + "Rest this asset naturally, physically stably, and functionally correctly " + "on its direct support. Preserve obvious up/down semantics: a pan, bowl, " + "plate, or cup opening faces upward; a bottle rests on its base; and a " + "long tool lies flat unless its requested pose says otherwise." ) diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 19df695ee..1445d3379 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -177,14 +177,14 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No "parent_id": None, "parent_relation": None, "table_region": None, - "orientation_state": None, + "pose_description": None, }, { "object_id": "cup", "parent_id": "table", "parent_relation": "on", "table_region": None, - "orientation_state": None, + "pose_description": None, }, ], "relations": [], @@ -200,7 +200,7 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert imported_graph.to_dict() == _scene_graph(scene).to_dict() -def test_scene_graph_importer_restores_node_orientation_state() -> None: +def test_scene_graph_importer_restores_node_pose_description() -> None: imported_graph = SceneExportImporter._scene_graph_from_data( { "nodes": [ @@ -209,21 +209,42 @@ def test_scene_graph_importer_restores_node_orientation_state() -> None: "parent_id": None, "parent_relation": None, "table_region": None, - "orientation_state": None, + "pose_description": None, }, { "object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", "table_region": None, - "orientation_state": "standing", + "pose_description": "Stand upright on its base.", }, ], "relations": [], } ) - assert imported_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert ( + imported_graph.node_by_id()["bottle_001"].pose_description + == "Stand upright on its base." + ) + + +def test_scene_graph_importer_rejects_the_removed_orientation_state_schema() -> None: + with pytest.raises(ValueError, match="serialized node schema"): + SceneExportImporter._scene_graph_from_data( + { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + } + ], + "relations": [], + } + ) def test_scene_export_overwrites_an_existing_scene_export(tmp_path: Path) -> None: diff --git a/tests/gen_sim/scene_engine/test_scene_edit_plan.py b/tests/gen_sim/scene_engine/test_scene_edit_plan.py index a68293668..9bdd2f606 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit_plan.py +++ b/tests/gen_sim/scene_engine/test_scene_edit_plan.py @@ -115,7 +115,7 @@ def test_scene_edit_plan_accepts_add_without_a_position() -> None: "category": "cup", "name": "green cup", "description": "A small green ceramic cup.", - "orientation_state": None, + "pose_description": None, } ] @@ -160,7 +160,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", - "orientation_state": None, + "pose_description": None, }, { "op": "add", @@ -171,7 +171,7 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "category": "orange", "name": "small_orange", "description": "A small round orange with a textured peel.", - "orientation_state": "lying", + "pose_description": "Lie flat on the support surface.", }, ] } @@ -184,9 +184,9 @@ def test_scene_edit_parser_assigns_ids_to_same_category_adds_in_order() -> None: "orange_002", "orange_003", ] - assert [operation.orientation_state for operation in operations] == [ + assert [operation.pose_description for operation in operations] == [ None, - "lying", + "Lie flat on the support surface.", ] @@ -210,7 +210,7 @@ def test_scene_edit_parser_rejects_path_traversal_add_categories( "category": unsafe_category, "name": "unsafe cup", "description": "A small cup.", - "orientation_state": None, + "pose_description": None, } ] } @@ -219,24 +219,25 @@ def test_scene_edit_parser_rejects_path_traversal_add_categories( _parse_scene_edit_operations(draft, scene=scene) -def test_scene_edit_plan_rejects_a_changed_move_orientation_state() -> None: +def test_scene_edit_plan_accepts_a_move_pose_description_override() -> None: scene, scene_graph = _scene_and_graph() - scene_graph.node_by_id()["book_001"].orientation_state = "lying" + scene_graph.node_by_id()["book_001"].pose_description = "Lie flat on the table." - with pytest.raises(ValueError, match="may only preserve"): - SceneEditPlan( - scene=scene, - scene_graph=scene_graph, - operations=[ - SceneEditOperation( - op="move", - object_id="book_001", - target_id="table", - relation="on", - orientation_state="standing", - ) - ], - ) + plan = SceneEditPlan( + scene=scene, + scene_graph=scene_graph, + operations=[ + SceneEditOperation( + op="move", + object_id="book_001", + target_id="table", + relation="on", + pose_description="Stand upright on its bottom edge.", + ) + ], + ) + + assert plan.operations[0].pose_description == "Stand upright on its bottom edge." def test_scene_edit_plan_rejects_targets_outside_the_input_scene() -> None: @@ -482,7 +483,7 @@ def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> No category="cup", name="green cup", description="A small green ceramic cup.", - orientation_state="standing", + pose_description="Stand upright on its base.", ) ], ) @@ -495,12 +496,12 @@ def test_scene_edit_graph_builder_adds_unpositioned_objects_on_the_table() -> No added_node = updated_scene_graph.node_by_id()["cup_001"] assert added_node.parent_id == "table" assert added_node.parent_relation == "on" - assert added_node.orientation_state == "standing" + assert added_node.pose_description == "Stand upright on its base." def test_scene_edit_graph_builder_updates_move_on_parent() -> None: scene, scene_graph = _scene_and_graph() - scene_graph.node_by_id()["orange_001"].orientation_state = "lying" + scene_graph.node_by_id()["orange_001"].pose_description = "Lie flat on the table." plan = SceneEditPlan( scene=scene, scene_graph=scene_graph, @@ -510,7 +511,7 @@ def test_scene_edit_graph_builder_updates_move_on_parent() -> None: object_id="orange_001", target_id="table", relation="on", - orientation_state="lying", + pose_description="Stand upright on its base.", ) ], ) @@ -521,7 +522,10 @@ def test_scene_edit_graph_builder_updates_move_on_parent() -> None: ) assert updated_scene_graph.node_by_id()["orange_001"].parent_id == "table" - assert updated_scene_graph.node_by_id()["orange_001"].orientation_state == "lying" + assert ( + updated_scene_graph.node_by_id()["orange_001"].pose_description + == "Stand upright on its base." + ) def test_scene_edit_graph_builder_adds_planar_relation_with_target_parent() -> None: diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py index a1a76f113..cd1c878a1 100644 --- a/tests/gen_sim/scene_engine/test_scene_generation.py +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -35,7 +35,6 @@ _optimize_simready_asset_visual_yaws, _project_child_aabb_centers_into_parent_aabb, _refine_on_children_bfs, - _scene_graph_based_calibration, _table_on_asset_ids, ) from embodichain.gen_sim.scene_engine.pipeline.utils.visual_yaw_optimizer import ( @@ -47,23 +46,6 @@ ) -def _y_up_layout_from_z_up_rotation( - object_id: str, - rotation_matrix: np.ndarray, -) -> dict[str, object]: - y_up_to_z_up_matrix = np.eye(4) - y_up_to_z_up_matrix[:3, :3] = np.array( - [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] - ) - z_up_to_y_up_matrix = np.linalg.inv(y_up_to_z_up_matrix) - z_up_transform = np.eye(4) - z_up_transform[:3, :3] = rotation_matrix - return transform_matrix_to_layout_object( - object_id, - z_up_to_y_up_matrix @ z_up_transform @ y_up_to_z_up_matrix, - ) - - def _z_up_rotation_from_y_up_layout(layout: dict[str, object]) -> np.ndarray: y_up_to_z_up_matrix = np.eye(4) y_up_to_z_up_matrix[:3, :3] = np.array( @@ -76,40 +58,6 @@ def _z_up_rotation_from_y_up_layout(layout: dict[str, object]) -> np.ndarray: )[:3, :3] -def test_scene_graph_calibration_makes_standing_asset_vertical() -> None: - scene_graph = SceneGraph( - nodes=[ - SceneGraphNode(object_id="table", parent_id=None), - SceneGraphNode( - object_id="bottle_001", - parent_id="table", - parent_relation="on", - orientation_state="standing", - ), - SceneGraphNode( - object_id="book_001", - parent_id="table", - parent_relation="on", - ), - ] - ) - lying_rotation = Rotation.from_euler("x", 90.0, degrees=True).as_matrix() - bottle_layout = _y_up_layout_from_z_up_rotation("bottle_001", lying_rotation) - book_layout = _y_up_layout_from_z_up_rotation("book_001", lying_rotation) - - calibrated_layouts = _scene_graph_based_calibration( - scene_graph=scene_graph, - assets_layout=[bottle_layout, book_layout], - ) - - bottle_axis = _z_up_rotation_from_y_up_layout(calibrated_layouts[0])[:, 2] - assert np.isclose(abs(bottle_axis[2]), 1.0) - assert np.allclose( - _z_up_rotation_from_y_up_layout(calibrated_layouts[1]), - lying_rotation, - ) - - def test_visual_yaw_optimizer_returns_zero_for_unobserved_asset( tmp_path, ) -> None: diff --git a/tests/gen_sim/scene_engine/test_scene_graph.py b/tests/gen_sim/scene_engine/test_scene_graph.py index c6cbd3243..ef403f2bd 100644 --- a/tests/gen_sim/scene_engine/test_scene_graph.py +++ b/tests/gen_sim/scene_engine/test_scene_graph.py @@ -34,7 +34,7 @@ def test_scene_graph_accepts_layered_on_relations() -> None: parent_id="table", parent_relation="on", table_region="center", - orientation_state="standing", + pose_description="Stand upright on its base.", ), SceneGraphNode( object_id="cup", @@ -302,7 +302,8 @@ def test_scene_graph_batch_planar_updates_preserve_chained_constraints() -> None graph.apply_updates( deleted_object_ids=set(), added_object_ids=[], - added_orientation_states_by_id={}, + added_pose_descriptions_by_id={}, + pose_description_updates_by_id={}, on_parent_updates=[], planar_relation_updates=[ ("plate", "left_of", "cup"), @@ -330,7 +331,7 @@ def test_scene_graph_to_dict_serializes_graph_state() -> None: parent_id="table", parent_relation="on", table_region="center", - orientation_state="standing", + pose_description="Stand upright on its base.", ), ], ) @@ -344,14 +345,14 @@ def test_scene_graph_to_dict_serializes_graph_state() -> None: "parent_id": None, "parent_relation": None, "table_region": None, - "orientation_state": None, + "pose_description": None, }, { "object_id": "plate", "parent_id": "table", "parent_relation": "on", "table_region": "center", - "orientation_state": "standing", + "pose_description": "Stand upright on its base.", }, ], "relations": [], diff --git a/tests/gen_sim/scene_engine/test_scene_understanding.py b/tests/gen_sim/scene_engine/test_scene_understanding.py index 24c8b6261..89e79c5fd 100644 --- a/tests/gen_sim/scene_engine/test_scene_understanding.py +++ b/tests/gen_sim/scene_engine/test_scene_understanding.py @@ -248,7 +248,7 @@ def complete(self, **_: object) -> str: "object_id": "cup_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": None, + "pose_description": None, }, ] } @@ -288,14 +288,14 @@ def complete(self, **_: object) -> str: "parent_id": None, "parent_relation": None, "table_region": None, - "orientation_state": None, + "pose_description": None, }, { "object_id": "cup_001", "parent_id": "table", "parent_relation": "on", "table_region": None, - "orientation_state": None, + "pose_description": None, }, ], "relations": [], @@ -318,13 +318,13 @@ def complete(self, **_: object) -> str: "object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": "standing", + "pose_description": "Stand upright on its base.", }, { "object_id": "book_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": "lying", + "pose_description": "Lie flat on the support surface.", }, ] } @@ -381,8 +381,14 @@ def complete(self, **_: object) -> str: }, ], } - assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" - assert scene_graph.node_by_id()["book_001"].orientation_state == "lying" + assert ( + scene_graph.node_by_id()["bottle_001"].pose_description + == "Stand upright on its base." + ) + assert ( + scene_graph.node_by_id()["book_001"].pose_description + == "Lie flat on the support surface." + ) def test_scene_graph_initialization_retries_a_response_containing_table( @@ -398,13 +404,13 @@ def __init__(self) -> None: "object_id": "table", "parent_id": "table", "parent_relation": "on", - "orientation_state": None, + "pose_description": None, }, { "object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": "standing", + "pose_description": "Stand upright on its base.", }, ] } @@ -416,7 +422,7 @@ def __init__(self) -> None: "object_id": "bottle_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": "standing", + "pose_description": "Stand upright on its base.", }, ] } @@ -454,7 +460,10 @@ def complete(self, **_: object) -> str: json_max_attempts=2, ) - assert scene_graph.node_by_id()["bottle_001"].orientation_state == "standing" + assert ( + scene_graph.node_by_id()["bottle_001"].pose_description + == "Stand upright on its base." + ) def test_scene_graph_initialization_retries_a_response_with_a_parent_cycle( @@ -470,13 +479,13 @@ def __init__(self) -> None: "object_id": "book_001", "parent_id": "pen_001", "parent_relation": "on", - "orientation_state": None, + "pose_description": None, }, { "object_id": "pen_001", "parent_id": "book_001", "parent_relation": "on", - "orientation_state": "lying", + "pose_description": "Lie flat on the support surface.", }, ] } @@ -488,13 +497,13 @@ def __init__(self) -> None: "object_id": "book_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": None, + "pose_description": None, }, { "object_id": "pen_001", "parent_id": "book_001", "parent_relation": "on", - "orientation_state": "lying", + "pose_description": "Lie flat on the support surface.", }, ] } @@ -626,13 +635,13 @@ def complete(self, **_: object) -> str: "object_id": "book_001", "parent_id": "table", "parent_relation": "on", - "orientation_state": None, + "pose_description": None, }, { "object_id": "pen_001", "parent_id": "book_001", "parent_relation": "on", - "orientation_state": "lying", + "pose_description": "Lie flat on the support surface.", }, ] } @@ -675,4 +684,4 @@ def complete(self, **_: object) -> str: pen = scene_graph.node_by_id()["pen_001"] assert pen.parent_id == "book_001" assert pen.parent_relation == "on" - assert pen.orientation_state == "lying" + assert pen.pose_description == "Lie flat on the support surface." diff --git a/tests/gen_sim/scene_engine/test_simready_processor_utils.py b/tests/gen_sim/scene_engine/test_simready_processor_utils.py index 959412af1..a4527ff35 100644 --- a/tests/gen_sim/scene_engine/test_simready_processor_utils.py +++ b/tests/gen_sim/scene_engine/test_simready_processor_utils.py @@ -28,15 +28,13 @@ ) from embodichain.gen_sim.scene_engine.pipeline.utils.simready_processor_utils import ( DEFAULT_NEEDED_LAYOUT, - LYING_NEEDED_LAYOUT, - STANDING_NEEDED_LAYOUT, compute_uniform_xy_scale_for_target, query_vlm_pose_switch_candidate, query_vlm_object_pose_and_target_size, ) -def test_simready_pose_layout_uses_graph_orientation_states( +def test_simready_pose_layout_uses_graph_pose_descriptions( tmp_path: Path, ) -> None: class VLM: @@ -49,20 +47,23 @@ def complete(self, **_: object) -> str: coarse_geometry_root=tmp_path / "coarse", simready_geometry_root=tmp_path / "simready", config=SimReadyProcessorConfig( - orientation_states_by_id={ - "bottle_001": "standing", - "fork_001": "lying", + pose_descriptions_by_id={ + "bottle_001": "Stand upright on its base.", + "fork_001": "Lie flat on the support surface.", "knife_001": None, }, ), vlm_client=VLM(), # type: ignore[arg-type] ) - assert processor._orientation_state_for_object("bottle_001") == "standing" - assert processor._orientation_state_for_object("fork_001") == "lying" - assert processor._orientation_state_for_object("knife_001") is None - assert processor._needed_layout_for_object("bottle_001") == STANDING_NEEDED_LAYOUT - assert processor._needed_layout_for_object("fork_001") == LYING_NEEDED_LAYOUT + assert ( + processor._needed_layout_for_object("bottle_001") + == "Stand upright on its base." + ) + assert ( + processor._needed_layout_for_object("fork_001") + == "Lie flat on the support surface." + ) assert processor._needed_layout_for_object("knife_001") == DEFAULT_NEEDED_LAYOUT @@ -81,7 +82,7 @@ def complete(self, **_: object) -> str: vlm_client = VLM() decision = query_vlm_object_pose_and_target_size( scene_object_description="small blue bottle", - needed_layout=STANDING_NEEDED_LAYOUT, + needed_layout="Stand upright on its base.", rendered_views_path=tmp_path / "views.png", vlm_client=vlm_client, # type: ignore[arg-type] ) @@ -101,7 +102,7 @@ def complete(self, **_: object) -> str: selected_rotation_degrees, reason = query_vlm_pose_switch_candidate( scene_object_description="small saucepan", - needed_layout=LYING_NEEDED_LAYOUT, + needed_layout="Rest stably with the cooking surface facing upward.", rendered_candidates_path=tmp_path / "candidates.png", vlm_client=VLM(), # type: ignore[arg-type] ) @@ -126,7 +127,11 @@ def test_simready_pose_switch_uses_the_vlm_selected_candidate( coarse_layout_by_id={"pan_001": {}}, coarse_geometry_root=tmp_path / "coarse", simready_geometry_root=tmp_path / "simready", - config=SimReadyProcessorConfig(orientation_states_by_id={"pan_001": "lying"}), + config=SimReadyProcessorConfig( + pose_descriptions_by_id={ + "pan_001": "Rest stably with the cooking surface facing upward." + } + ), vlm_client=object(), # type: ignore[arg-type] ) calls: dict[str, object] = {} @@ -175,7 +180,7 @@ def fake_rotate(**kwargs: object) -> Path: } assert calls["candidate_query"] == { "scene_object_description": "small saucepan", - "needed_layout": LYING_NEEDED_LAYOUT, + "needed_layout": "Rest stably with the cooking surface facing upward.", "rendered_candidates_path": tmp_path / "candidates.png", "vlm_client": processor.vlm_client, "debug_output_path": tmp_path @@ -190,7 +195,7 @@ def fake_rotate(**kwargs: object) -> Path: } -def test_simready_null_orientation_checks_the_default_stable_pose( +def test_simready_null_pose_description_checks_the_default_stable_pose( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: @@ -206,7 +211,7 @@ def test_simready_null_orientation_checks_the_default_stable_pose( coarse_layout_by_id={"book_001": {}}, coarse_geometry_root=tmp_path / "coarse", simready_geometry_root=tmp_path / "simready", - config=SimReadyProcessorConfig(orientation_states_by_id={"book_001": None}), + config=SimReadyProcessorConfig(pose_descriptions_by_id={"book_001": None}), vlm_client=object(), # type: ignore[arg-type] ) requested_layouts: list[str] = [] From 342f151ba5ace2a8b3ff34dc164daeb0f7baaf3c Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 28 Aug 2026 10:53:32 +0800 Subject: [PATCH 42/85] Transform each asset into the SimReady table frame to preserve its coarse relative position, then discard its unreliable coarse rotation and replace it with the canonical SimReady pose plus VLM-estimated yaw. --- .../pipeline/generation/scene_generation.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 986d7e18f..c820f7c50 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -161,10 +161,6 @@ def generate_scene_and_refine( vlm_client=vlm_client, debug_output_root=debug_output_root / "visual_yaw", ) - simready_assets_layout = _apply_visual_yaws_to_simready_asset_layouts( - simready_assets_layout=simready_assets_layout, - z_up_yaws_degrees_by_id=visual_yaws_by_id, - ) simready_table_layout = simready_processor.process_table() # Concat then save the table info and the assets info in one JSON file. simready_layout = [simready_table_layout, *simready_assets_layout] @@ -179,6 +175,7 @@ def generate_scene_and_refine( scene_graph=scene_graph, simready_geometry_output_root=simready_geometry_output_root, # Contains simready assets and their current coarse layout JSON. debug_output_root=debug_output_root, # Keep the table support surface info + optimized layout info (render with matplotlib) for debugging. + z_up_yaws_degrees_by_id=visual_yaws_by_id, ) # Write the Updated scene JSON for debugging. @@ -233,7 +230,7 @@ def _apply_visual_yaws_to_simready_asset_layouts( simready_assets_layout: list[dict[str, object]], z_up_yaws_degrees_by_id: dict[str, float], ) -> list[dict[str, object]]: - """Keep SimReady positions but replace each coarse rotation with canonical yaw.""" + """Keep table-frame positions but replace each coarse rotation with canonical yaw.""" layout_ids = [layout.get("id") for layout in simready_assets_layout] if not all(isinstance(layout_id, str) and layout_id for layout_id in layout_ids): raise ValueError("Every SimReady asset layout must contain a non-empty id.") @@ -418,6 +415,7 @@ def _layout_refinement( scene_graph: SceneGraph, simready_geometry_output_root: str | Path, debug_output_root: str | Path, + z_up_yaws_degrees_by_id: dict[str, float], ) -> tuple[dict[str, object], list[dict[str, object]]]: # 1. All layouts and geometries below are SimReady outputs. Do not mix a @@ -482,6 +480,14 @@ def _layout_refinement( ) ) + # Table-frame conversion retains the coarse relative positions but can also + # transfer an inverted coarse-table orientation; replace only that rotation + # with the canonical SimReady pose and its observed z-up yaw. + refined_assets_layout = _apply_visual_yaws_to_simready_asset_layouts( + simready_assets_layout=refined_assets_layout, + z_up_yaws_degrees_by_id=z_up_yaws_degrees_by_id, + ) + # 4. Only direct on-table children participate in the table-level layout # stages. Their descendants follow each solved root transform until their # own parent-surface optimization is introduced in a later BFS pass. From d4dd2113f99b4472d56fa9ff0651817cee5f7de5 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 28 Aug 2026 13:40:13 +0800 Subject: [PATCH 43/85] finished articulated replace: need to be tested --- .../gen_sim/scene_engine/cli/preview.py | 59 ++++++++- .../gen_sim/scene_engine/core/scene_object.py | 8 ++ .../gen_sim/scene_engine/pipeline/generate.py | 11 ++ .../pipeline/generation/scene_generation.py | 71 ++++++++++- .../pipeline/utils/scene_exporter.py | 112 +++++++++++++++++- .../pipeline/utils/scene_importer.py | 77 ++++++++++++ .../test_scene_core_and_export.py | 95 ++++++++++++++- .../scene_engine/test_scene_generation.py | 76 ++++++++++++ 8 files changed, 500 insertions(+), 9 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index f2d19c263..7cef620ab 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -26,7 +26,7 @@ from typing import Any from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import LightCfg, MeshCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import ArticulationCfg, LightCfg, MeshCfg, RigidObjectCfg from embodichain.lab.visualization import ( VisualizationCfg, add_viser_args_to_parser, @@ -94,6 +94,11 @@ def preview_scene_export( config_dir=config_path.parent, label="asset", ) + _add_articulations( + sim=sim, + entries=_config_entries(scene_config, "articulation"), + config_dir=config_path.parent, + ) is_viser = sim.sim_config.visualization.backend == "viser" if headless and not is_viser: @@ -203,6 +208,58 @@ def _add_objects( print(f"[{label}] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") +def _add_articulations( + *, + sim: SimulationManager, + entries: list[dict[str, Any]], + config_dir: Path, +) -> None: + """Add exported USDC articulations without also loading their GLB proxies.""" + resolved_config_dir = config_dir.resolve() + for entry in entries: + uid = entry.get("uid") + raw_fpath = entry.get("fpath") + if not isinstance(uid, str) or not uid: + raise ValueError("Articulation entry has no valid uid.") + if not isinstance(raw_fpath, str): + raise ValueError(f"Articulation entry {uid!r} has no fpath.") + fpath = Path(raw_fpath) + if fpath.is_absolute() or fpath.suffix.lower() != ".usdc": + raise ValueError( + f"Articulation entry {uid!r} fpath must be a relative USDC path." + ) + usdc_path = (resolved_config_dir / fpath).resolve() + if resolved_config_dir not in usdc_path.parents: + raise ValueError( + f"Articulation entry {uid!r} fpath must stay within " + f"{resolved_config_dir}." + ) + if not usdc_path.is_file(): + raise FileNotFoundError( + f"Articulation USDC for {uid!r} not found: {usdc_path}" + ) + init_pos = _vector3(entry.get("init_pos"), field_name=f"{uid}.init_pos") + init_rot = _vector3(entry.get("init_rot"), field_name=f"{uid}.init_rot") + body_scale = _vector3( + entry.get("body_scale", [1.0, 1.0, 1.0]), + field_name=f"{uid}.body_scale", + ) + if entry.get("fix_base", True) is not True: + raise ValueError(f"Articulation entry {uid!r} must set fix_base=true.") + # SimulationManager converts this y-up USDC to z-up with its bottom on XY. + sim.add_articulation( + ArticulationCfg( + uid=uid, + fpath=str(usdc_path), + init_pos=tuple(init_pos), + init_rot=tuple(init_rot), + body_scale=tuple(body_scale), + fix_base=True, + ) + ) + print(f"[articulation] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") + + def _vector3(value: object, *, field_name: str) -> list[float]: if not isinstance(value, list) or len(value) != 3: raise ValueError(f"Scene config field {field_name!r} must be a length-3 list.") diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index 251dd0762..c8279118e 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -65,6 +65,12 @@ class SceneObject: mask_path: str | None = None # Absolute path to the validated binary image mask. visible_rgba_path: str | None = None # None for future unsegmented objects. simready_glb_path: str | None = None # Absolute path to the canonical SimReady GLB. + articulated_usdc_path: str | None = ( + None # Generated articulation asset, unused by the GLB pipeline for now. + ) + articulated_usdc_scale: list[float] | None = ( + None # Y-up runtime scale retained because the GLB pipeline bakes coarse scale. + ) rot: list[float] | None = None # Final y-up Euler XYZ rotation in degrees. pos: list[float] | None = None # Final y-up world position in metres. scale: list[float] | None = None # Final y-up object scale. @@ -86,6 +92,8 @@ def to_dict(self) -> dict[str, object]: "mask_path": self.mask_path, "visible_rgba_path": self.visible_rgba_path, "simready_glb_path": self.simready_glb_path, + "articulated_usdc_path": self.articulated_usdc_path, + "articulated_usdc_scale": self.articulated_usdc_scale, "rot": self.rot, "pos": self.pos, "scale": self.scale, diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 144551c29..71155a40a 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -25,6 +25,9 @@ from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, ) +from embodichain.gen_sim.scene_engine.clients.articulated_generation import ( + ArticulatedGenerationClient, +) from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( ImageSegmentationClient, ) @@ -73,18 +76,26 @@ def generate_scene_from_image( log_info("Starting Objects + Coarse Layout Generation") # Load .env settings and fail if the Geometry Generation Server is unavailable. geometry_generation_client = GeometryGenerationClient.from_dotenv() + articulated_generation_client: ArticulatedGenerationClient | None = None + if any(scene_object.is_articulated for scene_object in scene.objects): + articulated_generation_client = ArticulatedGenerationClient.from_dotenv() try: geometry_generation_client.check_health() # Error raising will happen internally. + if articulated_generation_client is not None: + articulated_generation_client.check_health() scene = generate_scene_and_refine( image_path=image_path, output_root=resolved_output_root, scene=scene, scene_graph=scene_graph, geometry_generation_client=geometry_generation_client, + articulated_generation_client=articulated_generation_client, vlm_client=vlm_client, ) finally: geometry_generation_client.close() # Kill the session to avoid resource leaks. + if articulated_generation_client is not None: + articulated_generation_client.close() log_info("Completed Objects + Coarse Layout Generation") # 3. Scene Export diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index c820f7c50..75e9d2139 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -36,6 +36,9 @@ from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, ) +from embodichain.gen_sim.scene_engine.clients.articulated_generation import ( + ArticulatedGenerationClient, +) from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_graph import ( TABLE_OBJECT_ID, @@ -93,6 +96,7 @@ def generate_scene_and_refine( *, geometry_generation_client: GeometryGenerationClient, vlm_client: OpenAICompatibleVLM, + articulated_generation_client: ArticulatedGenerationClient | None = None, ) -> Scene: resolved_image_path = _validate_image_path(image_path) @@ -118,13 +122,20 @@ def generate_scene_and_refine( simready_geometry_output_root.mkdir() # Coarse geometry generation and coarse layout generation. - _generate_coarse_results_from_masks( + coarse_scales_y_up_by_id = _generate_coarse_results_from_masks( image_path=resolved_image_path, debug_output_root=debug_output_root, coarse_geometry_output_root=coarse_geometry_output_root, scene=scene, # Use the masks which are kept in the scene data structure. geometry_generation_client=geometry_generation_client, ) + # Generate articulated USDCs for every articulated object, if any. + _generate_articulated_usdcs( + scene=scene, + output_root=stage_output_root / "articulated_geometry", + coarse_scales_y_up_by_id=coarse_scales_y_up_by_id, + articulated_generation_client=articulated_generation_client, + ) # Simready all the assets(includes table). # Treat table and assets seperately. @@ -273,7 +284,7 @@ def _generate_coarse_results_from_masks( scene: Scene, *, geometry_generation_client: GeometryGenerationClient, -) -> None: +) -> dict[str, list[float]]: # Parse whether the scene has each assets' binary masks. # The original image has already been validated. @@ -326,8 +337,60 @@ def _generate_coarse_results_from_masks( json.dumps(coarse_layout, indent=2, ensure_ascii=False) + "\n", encoding="utf-8", ) - # Nothing to be returned. - return None + return { + str(layout_object["id"]): [float(value) for value in layout_object["scale"]] + for layout_object in coarse_layout + } + + +def _generate_articulated_usdcs( + *, + scene: Scene, + output_root: str | Path, + coarse_scales_y_up_by_id: dict[str, list[float]], + articulated_generation_client: ArticulatedGenerationClient | None, +) -> None: + """Generate and persist one articulation USDC for every articulated object.""" + articulated_objects = [ + scene_object for scene_object in scene.objects if scene_object.is_articulated + ] + if not articulated_objects: + return + if articulated_generation_client is None: + raise ValueError( + "Articulated scene objects require an articulated-generation client." + ) + + resolved_output_root = Path(output_root).expanduser().resolve() + resolved_output_root.mkdir(parents=True, exist_ok=True) + for scene_object in articulated_objects: + if scene_object.visible_rgba_path is None: + raise ValueError( + "Articulated scene object " + f"{scene_object.id!r} has no visible RGBA observation." + ) + coarse_scale_y_up = coarse_scales_y_up_by_id.get(scene_object.id) + if ( + not isinstance(coarse_scale_y_up, list) + or len(coarse_scale_y_up) != 3 + or not np.all(np.isfinite(coarse_scale_y_up)) + or any(scale <= 0.0 for scale in coarse_scale_y_up) + ): + raise ValueError( + "Articulated scene object " + f"{scene_object.id!r} has no valid coarse-layout scale." + ) + # Run serially so the stage ends only after every required USDC is saved. + scene_object.articulated_usdc_path = str( + articulated_generation_client.generate_articulated_usdc( + prompt=scene_object.description, + image_path=scene_object.visible_rgba_path, + output_path=resolved_output_root / f"{scene_object.id}.usdc", + ) + ) + # USDC is y-up like GLB; SimulationManager performs the shared z-up conversion. + scene_object.articulated_usdc_scale = list(coarse_scale_y_up) + log_info(f"Created articulated USDC: {scene_object.id!r}.") def _update_scene_final_y_up_layout_and_z_up_centers( diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py index c8f9a4591..2eba948f0 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_exporter.py @@ -59,7 +59,7 @@ def __init__( self.scene_json_path: Path | None = None def export(self) -> Path: - """Write a scene-only config and copy SimReady GLBs into ``mesh_assets``. + """Write a scene-only config and copy runtime and proxy scene assets. Scene layouts are y-up. The simulator automatically converts each y-up GLB to z-up, so this exporter copies each GLB unchanged and converts @@ -70,9 +70,13 @@ def export(self) -> Path: """ if self.scene.table is None: raise ValueError("Cannot export a scene without a table.") + if self.scene.table.is_articulated: + raise ValueError("The exported table must remain a rigid support object.") mesh_assets_root = self.export_root / "mesh_assets" + articulated_assets_root = self.export_root / "articulated_assets" mesh_assets_root.mkdir(parents=True, exist_ok=True) + articulated_assets_root.mkdir(parents=True, exist_ok=True) scene_objects = self.scene.objects object_ids = [scene_object.id for scene_object in scene_objects] if len(set(object_ids)) != len(object_ids): @@ -88,6 +92,14 @@ def export(self) -> Path: ) for scene_object in scene_objects } + exported_articulated_entries = { + asset.id: self._copy_articulated_usdc_to_assets( + scene_object=asset, + articulated_assets_root=articulated_assets_root, + ) + for asset in self.scene.assets + if asset.is_articulated + } scene_config = { "format": "embodichain.scene-export/v1", # This identifies the exported scene data only. It is deliberately not @@ -106,6 +118,16 @@ def export(self) -> Path: asset_relative_path=exported_entries[asset.id], ) for asset in self.scene.assets + if not asset.is_articulated + ], + "articulation": [ + self._articulation_config( + scene_object=asset, + articulated_relative_path=exported_articulated_entries[asset.id], + proxy_glb_relative_path=exported_entries[asset.id], + ) + for asset in self.scene.assets + if asset.is_articulated ], } self.scene_config_path = self.export_root / "scene_config.json" @@ -131,6 +153,10 @@ def export(self) -> Path: mesh_assets_root=mesh_assets_root, object_ids=set(object_ids), ) + self._remove_stale_mesh_assets( + mesh_assets_root=articulated_assets_root, + object_ids=set(exported_articulated_entries), + ) return self.scene_config_path @staticmethod @@ -167,6 +193,38 @@ def _copy_scene_object_to_assets( shutil.copy2(source_glb_path, destination_glb_path) return destination_glb_path.relative_to(mesh_assets_root.parent).as_posix() + @staticmethod + def _copy_articulated_usdc_to_assets( + *, + scene_object: SceneObject, + articulated_assets_root: Path, + ) -> str: + """Copy one runtime USDC and return its config-relative path.""" + object_id = scene_object.id + if scene_object.articulated_usdc_path is None: + raise ValueError( + f"Articulated scene object {object_id!r} has no USDC path." + ) + source_usdc_path = ( + Path(scene_object.articulated_usdc_path).expanduser().resolve() + ) + if not source_usdc_path.is_file(): + raise FileNotFoundError( + f"Articulated USDC for scene object {object_id!r} not found: " + f"{source_usdc_path}" + ) + destination_usdc_path = ( + articulated_assets_root / object_id / f"{object_id}.usdc" + ) + destination_usdc_path.parent.mkdir(parents=True, exist_ok=True) + if not destination_usdc_path.is_file() or not source_usdc_path.samefile( + destination_usdc_path + ): + shutil.copy2(source_usdc_path, destination_usdc_path) + return destination_usdc_path.relative_to( + articulated_assets_root.parent + ).as_posix() + @staticmethod def _remove_stale_mesh_assets( *, @@ -233,6 +291,58 @@ def _scene_object_config( "max_convex_hull_num": scene_object.physics.max_convex_hull_num, } + @staticmethod + def _articulation_config( + *, + scene_object: SceneObject, + articulated_relative_path: str, + proxy_glb_relative_path: str, + ) -> dict[str, object]: + """Build one y-up USDC articulation config from an optimized scene asset.""" + pos_y_up = SceneExporter._scene_vector(scene_object, "pos") + rot_y_up = SceneExporter._scene_vector(scene_object, "rot") + proxy_scale_y_up = SceneExporter._scene_vector(scene_object, "scale") + articulated_scale_y_up = SceneExporter._articulated_usdc_scale(scene_object) + pos_z_up = _Y_UP_TO_Z_UP_ROTATION @ np.asarray(pos_y_up, dtype=float) + rotation_y_up = Rotation.from_euler("xyz", rot_y_up, degrees=True).as_matrix() + rotation_z_up = ( + _Y_UP_TO_Z_UP_ROTATION @ rotation_y_up @ _Y_UP_TO_Z_UP_ROTATION.T + ) + rot_z_up = Rotation.from_matrix(rotation_z_up).as_euler("XYZ", degrees=True) + return { + "uid": scene_object.id, + "category": scene_object.category, + "name": scene_object.name, + "description": scene_object.description, + "is_articulated": True, + "fpath": articulated_relative_path, + # The proxy stays in the export so scene edit can keep using GLB AABBs. + "proxy_glb_fpath": proxy_glb_relative_path, + # Both y-up assets load with their bottom on z-up's XY plane. + "init_pos": pos_z_up.tolist(), + "init_rot": rot_z_up.tolist(), + "body_scale": articulated_scale_y_up, + "proxy_body_scale": proxy_scale_y_up, + "fix_base": True, + } + + @staticmethod + def _articulated_usdc_scale(scene_object: SceneObject) -> list[float]: + """Return the finite positive y-up scale for one USDC asset.""" + scale = scene_object.articulated_usdc_scale + if not isinstance(scale, list) or len(scale) != 3: + raise ValueError( + f"Articulated scene object {scene_object.id!r} has no USDC scale." + ) + typed_scale = [float(value) for value in scale] + if not np.all(np.isfinite(typed_scale)) or any( + value <= 0.0 for value in typed_scale + ): + raise ValueError( + f"Articulated scene object {scene_object.id!r} has invalid USDC scale." + ) + return typed_scale + @staticmethod def _scene_vector(scene_object: SceneObject, field_name: str) -> list[float]: """Read one finite final y-up layout vector from a scene object.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index 030560858..978143dda 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -154,6 +154,9 @@ def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: rigid_object_entries = scene_config.get("rigid_object", []) if not isinstance(rigid_object_entries, list): raise ValueError("Scene config rigid_object must be a list.") + articulation_entries = scene_config.get("articulation", []) + if not isinstance(articulation_entries, list): + raise ValueError("Scene config articulation must be a list.") return Scene( objects=[ @@ -162,6 +165,10 @@ def _scene_from_config(self, scene_config: dict[str, Any]) -> Scene: self._scene_object_from_export_entry(entry, kind="asset") for entry in rigid_object_entries ], + *[ + self._scene_object_from_articulation_entry(entry) + for entry in articulation_entries + ], ] ) @@ -325,6 +332,43 @@ def _scene_object_from_export_entry( ), ) + def _scene_object_from_articulation_entry(self, entry: object) -> SceneObject: + """Restore an editable GLB proxy and its runtime USDC articulation.""" + if not isinstance(entry, dict): + raise ValueError("Articulation entries must be objects.") + uid = entry.get("uid") + if not isinstance(uid, str) or not uid: + raise ValueError("Articulation entries must contain a valid uid.") + proxy_glb_fpath = entry.get("proxy_glb_fpath") + if not isinstance(proxy_glb_fpath, str): + raise ValueError( + f"Articulation entry {uid!r} must contain proxy_glb_fpath." + ) + proxy_body_scale = self._vector3( + entry.get("proxy_body_scale"), field_name=f"{uid}.proxy_body_scale" + ) + # Scene edit still measures and optimizes the canonical GLB proxy. + proxy_entry = { + **entry, + "shape": {"shape_type": "Mesh", "fpath": proxy_glb_fpath}, + "body_scale": proxy_body_scale, + } + scene_object = self._scene_object_from_export_entry(proxy_entry, kind="asset") + if not scene_object.is_articulated: + raise ValueError( + f"Articulation entry {uid!r} must set is_articulated=true." + ) + scene_object.articulated_usdc_path = str( + self._resolve_export_articulated_usdc_path(entry, uid=uid) + ) + articulated_usdc_scale = self._vector3( + entry.get("body_scale"), field_name=f"{uid}.body_scale" + ) + if any(value <= 0.0 for value in articulated_usdc_scale): + raise ValueError(f"Articulation entry {uid!r} body_scale must be positive.") + scene_object.articulated_usdc_scale = articulated_usdc_scale + return scene_object + def _resolve_export_glb_path( self, entry: dict[str, Any], @@ -355,6 +399,39 @@ def _resolve_export_glb_path( raise FileNotFoundError(f"Scene object {uid!r} GLB not found: {glb_path}") return glb_path + def _resolve_export_articulated_usdc_path( + self, + entry: dict[str, Any], + *, + uid: str, + ) -> Path: + """Validate one exported articulation reference and return its USDC path.""" + raw_fpath = entry.get("fpath") + if not isinstance(raw_fpath, str): + raise ValueError(f"Articulation entry {uid!r} must contain fpath.") + fpath = Path(raw_fpath) + if fpath.is_absolute() or fpath.suffix.lower() != ".usdc": + raise ValueError( + f"Articulation entry {uid!r} fpath must be a relative USDC path." + ) + expected_fpath = Path("articulated_assets") / uid / f"{uid}.usdc" + if fpath != expected_fpath: + raise ValueError( + f"Articulation entry {uid!r} fpath must be " + f"{expected_fpath.as_posix()!r}." + ) + usdc_path = (self.scene_export_root / fpath).resolve() + if self.scene_export_root.resolve() not in usdc_path.parents: + raise ValueError( + f"Articulation entry {uid!r} fpath must stay within " + f"{self.scene_export_root.resolve()}." + ) + if not usdc_path.is_file(): + raise FileNotFoundError( + f"Articulation USDC for {uid!r} not found: {usdc_path}" + ) + return usdc_path + @staticmethod def _vector3(value: object, *, field_name: str) -> list[float]: """Validate one length-3 numeric vector.""" diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 1445d3379..7ecf588d8 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -32,6 +32,7 @@ ObjectPhysics, SceneObject, ) +from embodichain.gen_sim.scene_engine.cli.preview import _add_articulations from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( SceneExportImporter, @@ -146,7 +147,6 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No physics=_physics("dynamic"), ) asset.center_xy = [0.25, -0.5] - asset.is_articulated = True scene = Scene(objects=[table, asset]) export_path = SceneExporter( @@ -164,7 +164,7 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert entry["uid"] == "cup" assert entry["category"] == "asset" assert entry["name"] == "cup" - assert entry["is_articulated"] is True + assert entry["is_articulated"] is False assert entry["body_type"] == "dynamic" assert entry["init_pos"] == [1.0, -3.0, 2.0] assert entry["body_scale"] == [1.0, 2.0, 3.0] @@ -196,10 +196,99 @@ def test_scene_export_copies_meshes_and_converts_y_up_pose(tmp_path: Path) -> No assert [asset.id for asset in imported_scene.assets] == ["cup"] assert imported_scene.assets[0].category == "asset" assert imported_scene.assets[0].name == "cup" - assert imported_scene.assets[0].is_articulated is True + assert imported_scene.assets[0].is_articulated is False assert imported_graph.to_dict() == _scene_graph(scene).to_dict() +def test_scene_export_uses_usdc_for_articulated_runtime_and_glb_for_editing( + tmp_path: Path, +) -> None: + table_glb = tmp_path / "table.glb" + drawer_glb = tmp_path / "drawer.glb" + drawer_usdc = tmp_path / "drawer.usdc" + table_glb.write_bytes(b"glTF-table") + drawer_glb.write_bytes(b"glTF-drawer") + drawer_usdc.write_bytes(b"USDC-drawer") + table = _scene_object( + object_id="table", + kind="table", + glb_path=table_glb, + physics=_physics("kinematic"), + ) + drawer = _scene_object( + object_id="drawer", + kind="asset", + glb_path=drawer_glb, + physics=_physics("dynamic"), + ) + drawer.is_articulated = True + drawer.articulated_usdc_path = str(drawer_usdc) + drawer.articulated_usdc_scale = [1.25, 2.5, 3.75] + scene = Scene(objects=[table, drawer]) + + export_path = SceneExporter( + scene=scene, + scene_graph=_scene_graph(scene), + output_root=tmp_path / "output", + ).export() + exported = json.loads(export_path.read_text(encoding="utf-8")) + + assert exported["rigid_object"] == [] + articulation = exported["articulation"][0] + assert articulation["fpath"] == "articulated_assets/drawer/drawer.usdc" + assert articulation["proxy_glb_fpath"] == "mesh_assets/drawer/drawer.glb" + assert articulation["body_scale"] == [1.25, 2.5, 3.75] + assert articulation["proxy_body_scale"] == [1.0, 2.0, 3.0] + assert (export_path.parent / articulation["fpath"]).read_bytes() == b"USDC-drawer" + assert ( + export_path.parent / articulation["proxy_glb_fpath"] + ).read_bytes() == b"glTF-drawer" + + imported_scene = SceneExportImporter(output_root=tmp_path / "output").import_scene() + imported_drawer = imported_scene.assets[0] + assert imported_drawer.simready_glb_path == str( + export_path.parent / "mesh_assets" / "drawer" / "drawer.glb" + ) + assert imported_drawer.articulated_usdc_path == str( + export_path.parent / "articulated_assets" / "drawer" / "drawer.usdc" + ) + assert imported_drawer.articulated_usdc_scale == [1.25, 2.5, 3.75] + + +def test_preview_loads_exported_usdc_as_an_articulation(tmp_path: Path) -> None: + class FakeSimulationManager: + def __init__(self) -> None: + self.articulation_cfgs: list[object] = [] + + def add_articulation(self, cfg: object) -> None: + self.articulation_cfgs.append(cfg) + + usdc_path = tmp_path / "articulated_assets" / "drawer" / "drawer.usdc" + usdc_path.parent.mkdir(parents=True) + usdc_path.write_bytes(b"USDC-drawer") + sim = FakeSimulationManager() + + _add_articulations( + sim=sim, # type: ignore[arg-type] + entries=[ + { + "uid": "drawer", + "fpath": "articulated_assets/drawer/drawer.usdc", + "init_pos": [1.0, 2.0, 3.0], + "init_rot": [10.0, 20.0, 30.0], + "body_scale": [1.25, 2.5, 3.75], + "fix_base": True, + } + ], + config_dir=tmp_path, + ) + + articulation_cfg = sim.articulation_cfgs[0] + assert articulation_cfg.uid == "drawer" # type: ignore[attr-defined] + assert articulation_cfg.fpath == str(usdc_path) # type: ignore[attr-defined] + assert articulation_cfg.body_scale == (1.25, 2.5, 3.75) # type: ignore[attr-defined] + + def test_scene_graph_importer_restores_node_pose_description() -> None: imported_graph = SceneExportImporter._scene_graph_from_data( { diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py index cd1c878a1..23e1ae498 100644 --- a/tests/gen_sim/scene_engine/test_scene_generation.py +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -30,6 +30,7 @@ from embodichain.gen_sim.scene_engine.core.scene import Scene, SceneObject from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( _align_table_roots_individually, + _generate_articulated_usdcs, _apply_visual_yaws_to_simready_asset_layouts, _apply_root_layout_updates_to_descendant_subtrees, _optimize_simready_asset_visual_yaws, @@ -218,6 +219,81 @@ def test_visual_yaws_replace_coarse_rotations_but_preserve_positions() -> None: assert np.allclose(yawed_layout["pos"], [0.1, 0.2, 0.3]) +def test_articulated_usdcs_use_visible_rgba_in_scene_order(tmp_path: Path) -> None: + class FakeArticulatedGenerationClient: + def __init__(self) -> None: + self.calls: list[tuple[str, str | Path, str | Path]] = [] + + def generate_articulated_usdc( + self, + *, + prompt: str, + image_path: str | Path, + output_path: str | Path, + ) -> Path: + self.calls.append((prompt, image_path, output_path)) + resolved_output_path = Path(output_path) + resolved_output_path.write_bytes(b"USDC") + return resolved_output_path + + drawer_rgba_path = tmp_path / "drawer_rgba.png" + microwave_rgba_path = tmp_path / "microwave_rgba.png" + drawer_rgba_path.write_bytes(b"PNG") + microwave_rgba_path.write_bytes(b"PNG") + drawer = SceneObject( + id="drawer_001", + kind="asset", + category="drawer", + name="drawer", + description="white drawer with a pull handle", + is_articulated=True, + visible_rgba_path=str(drawer_rgba_path), + ) + microwave = SceneObject( + id="microwave_001", + kind="asset", + category="microwave", + name="microwave", + description="black microwave with a hinged door", + is_articulated=True, + visible_rgba_path=str(microwave_rgba_path), + ) + static_mug = SceneObject( + id="mug_001", + kind="asset", + category="mug", + name="mug", + description="blue ceramic mug", + visible_rgba_path=str(tmp_path / "mug_rgba.png"), + ) + client = FakeArticulatedGenerationClient() + + _generate_articulated_usdcs( + scene=Scene(objects=[drawer, static_mug, microwave]), + output_root=tmp_path / "articulated_geometry", + coarse_scales_y_up_by_id={ + "drawer_001": [1.0, 2.0, 3.0], + "microwave_001": [4.0, 5.0, 6.0], + }, + articulated_generation_client=client, # type: ignore[arg-type] + ) + + assert [call[0] for call in client.calls] == [ + "white drawer with a pull handle", + "black microwave with a hinged door", + ] + assert drawer.articulated_usdc_path == str( + tmp_path / "articulated_geometry" / "drawer_001.usdc" + ) + assert microwave.articulated_usdc_path == str( + tmp_path / "articulated_geometry" / "microwave_001.usdc" + ) + assert drawer.to_dict()["articulated_usdc_path"] == drawer.articulated_usdc_path + assert drawer.articulated_usdc_scale == [1.0, 2.0, 3.0] + assert microwave.articulated_usdc_scale == [4.0, 5.0, 6.0] + assert static_mug.articulated_usdc_path is None + + def test_table_root_update_propagates_its_pose_delta_to_descendants() -> None: scene_graph = SceneGraph( nodes=[ From 6217e79fc4d4e7ef1fa3d4dae984904247112517 Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 28 Aug 2026 15:24:11 +0800 Subject: [PATCH 44/85] fix the preview problem after added usdc format articulated object --- embodichain/gen_sim/scene_engine/cli/preview.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index 7cef620ab..dcb744c2f 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -255,6 +255,8 @@ def _add_articulations( init_rot=tuple(init_rot), body_scale=tuple(body_scale), fix_base=True, + # Generated USDC is not URDF, so it cannot build a PK chain. + build_pk_chain=False, ) ) print(f"[articulation] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") From 6c156277f9cff0ca4068731a3b746412556a57ef Mon Sep 17 00:00:00 2001 From: Muzi Wong <178915912+MuziWong@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:13:30 +0800 Subject: [PATCH 45/85] fix bug in scene engine preview, delete the wrong-import --- .../gen_sim/scene_engine/cli/preview.py | 85 ++++++++++++++--- .../pipeline/generation/scene_generation.py | 15 ++- .../pipeline/utils/articulated_usdc_utils.py | 93 +++++++++++++++++++ .../pipeline/utils/simready_processor.py | 1 - .../test_scene_core_and_export.py | 48 +++++++++- .../scene_engine/test_scene_generation.py | 49 +++++++++- 6 files changed, 268 insertions(+), 23 deletions(-) create mode 100644 embodichain/gen_sim/scene_engine/pipeline/utils/articulated_usdc_utils.py diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index dcb744c2f..08c0b84fa 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -23,7 +23,7 @@ from pathlib import Path import time from collections.abc import Sequence -from typing import Any +from typing import TYPE_CHECKING, Any from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ArticulationCfg, LightCfg, MeshCfg, RigidObjectCfg @@ -33,6 +33,12 @@ visualization_cfg_from_args, ) +if TYPE_CHECKING: + from embodichain.lab.scripts.preview_joint_control import ( + ArticulationPreviewController, + ) + from embodichain.lab.sim.objects import Articulation + def preview_scene_export( *, @@ -40,6 +46,7 @@ def preview_scene_export( device: str = "cpu", headless: bool = False, visualization: VisualizationCfg | None = None, + joint_control: bool = True, ) -> None: """Load ``scene_export/scene_config.json`` and preview its table and assets. @@ -48,6 +55,7 @@ def preview_scene_export( device: Simulation device, for example ``"cpu"`` or ``"cuda"``. headless: Load and validate the scene without an interactive preview. visualization: Optional live-visualization configuration. + joint_control: Expose supported articulation joints in Viser. """ resolved_output_root = Path(output_root).expanduser().resolve() config_path = resolved_output_root / "scene_export" / "scene_config.json" @@ -94,27 +102,36 @@ def preview_scene_export( config_dir=config_path.parent, label="asset", ) - _add_articulations( + articulations = _add_articulations( sim=sim, entries=_config_entries(scene_config, "articulation"), config_dir=config_path.parent, ) is_viser = sim.sim_config.visualization.backend == "viser" + joint_controller = _setup_viser_joint_control( + sim=sim, + articulations=articulations, + enabled=is_viser and joint_control, + ) if headless and not is_viser: sim.update(step=1) print(f"Loaded scene export headlessly: {config_path}") return if is_viser: - sim.update(step=1) print(f"Previewing in Viser: {config_path}") else: print(f"Previewing: {config_path}") sim.open_window() print("Close with Ctrl-C.") while True: - time.sleep(0.1) + if is_viser: + # Browser commands are applied on the simulation thread before capture. + if joint_controller is not None: + joint_controller.update() + sim.update(step=1) + time.sleep(0.01) except KeyboardInterrupt: print("Stopping preview.") finally: @@ -213,9 +230,10 @@ def _add_articulations( sim: SimulationManager, entries: list[dict[str, Any]], config_dir: Path, -) -> None: +) -> list[Articulation]: """Add exported USDC articulations without also loading their GLB proxies.""" resolved_config_dir = config_dir.resolve() + articulations: list[Articulation] = [] for entry in entries: uid = entry.get("uid") raw_fpath = entry.get("fpath") @@ -247,19 +265,49 @@ def _add_articulations( if entry.get("fix_base", True) is not True: raise ValueError(f"Articulation entry {uid!r} must set fix_base=true.") # SimulationManager converts this y-up USDC to z-up with its bottom on XY. - sim.add_articulation( - ArticulationCfg( - uid=uid, - fpath=str(usdc_path), - init_pos=tuple(init_pos), - init_rot=tuple(init_rot), - body_scale=tuple(body_scale), - fix_base=True, - # Generated USDC is not URDF, so it cannot build a PK chain. - build_pk_chain=False, + articulations.append( + sim.add_articulation( + ArticulationCfg( + uid=uid, + fpath=str(usdc_path), + init_pos=tuple(init_pos), + init_rot=tuple(init_rot), + body_scale=tuple(body_scale), + fix_base=True, + # Generated USDC is not URDF, so it cannot build a PK chain. + build_pk_chain=False, + ) ) ) print(f"[articulation] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") + return articulations + + +def _setup_viser_joint_control( + *, + sim: SimulationManager, + articulations: list[Articulation], + enabled: bool, +) -> ArticulationPreviewController | None: + """Expose supported exported-articulation joints through the Viser runtime.""" + if not enabled or not articulations: + return None + runtime = sim.visualization_runtime + if runtime is None: + raise RuntimeError( + "Viser joint control requires an active visualization runtime." + ) + from embodichain.lab.scripts.preview_joint_control import ( + ArticulationPreviewController, + ) + + controller = ArticulationPreviewController(articulations, runtime) + if not controller.has_controls: + return None + controller.update() + runtime.set_joint_control_provider(controller) + print("Viser articulation joint controls enabled.") + return controller def _vector3(value: object, *, field_name: str) -> list[float]: @@ -292,6 +340,12 @@ def main(argv: Sequence[str] | None = None) -> None: action="store_true", help="Load and validate the exported scene without opening a window.", ) + parser.add_argument( + "--joint-control", + action=argparse.BooleanOptionalAction, + default=True, + help="Expose supported articulation joints in Viser (default: enabled).", + ) add_viser_args_to_parser(parser) args = parser.parse_args(argv) preview_scene_export( @@ -299,6 +353,7 @@ def main(argv: Sequence[str] | None = None) -> None: device=args.device, headless=args.headless, visualization=visualization_cfg_from_args(args), + joint_control=args.joint_control, ) diff --git a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py index 75e9d2139..a791ae47f 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generation/scene_generation.py @@ -57,6 +57,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.assets_group_layout_optimizer import ( AssetsSupportLayoutOptimizer, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.articulated_usdc_utils import ( + _canonicalize_articulated_usdc_bottom_center, +) from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( GravitySettleBody, GravitySettler, @@ -381,12 +384,14 @@ def _generate_articulated_usdcs( f"{scene_object.id!r} has no valid coarse-layout scale." ) # Run serially so the stage ends only after every required USDC is saved. + generated_usdc_path = articulated_generation_client.generate_articulated_usdc( + prompt=scene_object.description, + image_path=scene_object.visible_rgba_path, + output_path=resolved_output_root / f"{scene_object.id}.usdc", + ) + # Canonicalize the runtime USDC so it shares the SimReady GLB origin. scene_object.articulated_usdc_path = str( - articulated_generation_client.generate_articulated_usdc( - prompt=scene_object.description, - image_path=scene_object.visible_rgba_path, - output_path=resolved_output_root / f"{scene_object.id}.usdc", - ) + _canonicalize_articulated_usdc_bottom_center(generated_usdc_path) ) # USDC is y-up like GLB; SimulationManager performs the shared z-up conversion. scene_object.articulated_usdc_scale = list(coarse_scale_y_up) diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/articulated_usdc_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/articulated_usdc_utils.py new file mode 100644 index 000000000..af6780df8 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/articulated_usdc_utils.py @@ -0,0 +1,93 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + + +from __future__ import annotations + +from pathlib import Path + + +def _canonicalize_articulated_usdc_bottom_center(usdc_path: str | Path) -> Path: + """Place one y-up articulation's local origin at its bottom AABB center. + + The translation is authored on the default articulation root, so every link, + collider, and joint frame moves together without changing their relative + articulation structure. + """ + resolved_usdc_path = Path(usdc_path).expanduser().resolve() + if not resolved_usdc_path.is_file() or resolved_usdc_path.suffix.lower() != ".usdc": + raise FileNotFoundError( + "Articulated USDC canonicalization requires an existing .usdc file: " + f"{resolved_usdc_path}" + ) + try: + from pxr import Gf, Usd, UsdGeom, UsdPhysics + except ImportError as exc: + raise RuntimeError( + "Articulated USDC canonicalization requires the USD Python bindings." + ) from exc + + stage = Usd.Stage.Open(str(resolved_usdc_path)) + if stage is None: + raise ValueError(f"Cannot open articulated USDC: {resolved_usdc_path}") + root_prim = stage.GetDefaultPrim() + if not root_prim or not root_prim.IsValid(): + raise ValueError( + f"Articulated USDC has no valid default prim: {resolved_usdc_path}" + ) + if not root_prim.HasAPI(UsdPhysics.ArticulationRootAPI): + raise ValueError( + "Articulated USDC default prim must have ArticulationRootAPI: " + f"{root_prim.GetPath()}" + ) + + root_xformable = UsdGeom.Xformable(root_prim) + if not root_xformable: + raise ValueError( + "Articulated USDC default prim must be transformable: " + f"{root_prim.GetPath()}" + ) + op_name = "xformOp:translate:scene_engine_bottom_center" + if any(op.GetOpName() == op_name for op in root_xformable.GetOrderedXformOps()): + raise ValueError( + "Articulated USDC is already canonicalized at its bottom center: " + f"{resolved_usdc_path}" + ) + + bounds = ( + UsdGeom.BBoxCache( + Usd.TimeCode.Default(), + [UsdGeom.Tokens.default_, UsdGeom.Tokens.render, UsdGeom.Tokens.proxy], + ) + .ComputeLocalBound(root_prim) + .ComputeAlignedBox() + ) + if bounds.IsEmpty(): + raise ValueError( + f"Articulated USDC default prim has an empty bound: {root_prim.GetPath()}" + ) + minimum, maximum = bounds.GetMin(), bounds.GetMax() + bottom_center = Gf.Vec3d( + (minimum[0] + maximum[0]) / 2.0, + minimum[1], + (minimum[2] + maximum[2]) / 2.0, + ) + # Shift the complete articulation hierarchy into the shared SimReady origin. + root_xformable.AddTranslateOp(opSuffix="scene_engine_bottom_center").Set( + -bottom_center + ) + stage.GetRootLayer().Save() + return resolved_usdc_path diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index b66c5fae2..b18b97008 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -25,7 +25,6 @@ import trimesh from embodichain.gen_sim.scene_engine.core.scene import Scene -from embodichain.gen_sim.scene_engine.core.scene_graph import OrientationState from embodichain.gen_sim.scene_engine.core.scene_object import ( ObjectPhysics, SceneObject, diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 7ecf588d8..87d0043b3 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -32,7 +32,10 @@ ObjectPhysics, SceneObject, ) -from embodichain.gen_sim.scene_engine.cli.preview import _add_articulations +from embodichain.gen_sim.scene_engine.cli.preview import ( + _add_articulations, + _setup_viser_joint_control, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( SceneExportImporter, @@ -289,6 +292,49 @@ def add_articulation(self, cfg: object) -> None: assert articulation_cfg.body_scale == (1.25, 2.5, 3.75) # type: ignore[attr-defined] +def test_preview_registers_exported_articulation_joint_controls( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class FakeRuntime: + def __init__(self) -> None: + self.provider: object | None = None + + def set_joint_control_provider(self, provider: object) -> None: + self.provider = provider + + class FakeSimulationManager: + def __init__(self) -> None: + self.visualization_runtime = FakeRuntime() + + class FakeController: + def __init__(self, articulations: list[object], runtime: FakeRuntime) -> None: + self.articulations = articulations + self.runtime = runtime + self.has_controls = True + self.update_count = 0 + + def update(self) -> None: + self.update_count += 1 + + monkeypatch.setattr( + "embodichain.lab.scripts.preview_joint_control.ArticulationPreviewController", + FakeController, + ) + articulation = object() + sim = FakeSimulationManager() + + controller = _setup_viser_joint_control( + sim=sim, # type: ignore[arg-type] + articulations=[articulation], # type: ignore[list-item] + enabled=True, + ) + + assert controller is sim.visualization_runtime.provider + assert controller is not None + assert controller.articulations == [articulation] # type: ignore[attr-defined] + assert controller.update_count == 1 # type: ignore[attr-defined] + + def test_scene_graph_importer_restores_node_pose_description() -> None: imported_graph = SceneExportImporter._scene_graph_from_data( { diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py index 23e1ae498..b8c633884 100644 --- a/tests/gen_sim/scene_engine/test_scene_generation.py +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -21,6 +21,7 @@ import numpy as np import pytest from PIL import Image +from pxr import Gf, Usd, UsdGeom, UsdPhysics from scipy.spatial.transform import Rotation from embodichain.gen_sim.scene_engine.core.scene_graph import ( @@ -41,6 +42,9 @@ from embodichain.gen_sim.scene_engine.pipeline.utils.visual_yaw_optimizer import ( VisualYawOptimizer, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.articulated_usdc_utils import ( + _canonicalize_articulated_usdc_bottom_center, +) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, transform_matrix_to_layout_object, @@ -219,7 +223,10 @@ def test_visual_yaws_replace_coarse_rotations_but_preserve_positions() -> None: assert np.allclose(yawed_layout["pos"], [0.1, 0.2, 0.3]) -def test_articulated_usdcs_use_visible_rgba_in_scene_order(tmp_path: Path) -> None: +def test_articulated_usdcs_use_visible_rgba_in_scene_order( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: class FakeArticulatedGenerationClient: def __init__(self) -> None: self.calls: list[tuple[str, str | Path, str | Path]] = [] @@ -268,6 +275,14 @@ def generate_articulated_usdc( ) client = FakeArticulatedGenerationClient() + def fake_canonicalize(usdc_path: str | Path) -> Path: + return Path(usdc_path) + + monkeypatch.setattr( + "embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation._canonicalize_articulated_usdc_bottom_center", + fake_canonicalize, + ) + _generate_articulated_usdcs( scene=Scene(objects=[drawer, static_mug, microwave]), output_root=tmp_path / "articulated_geometry", @@ -294,6 +309,38 @@ def generate_articulated_usdc( assert static_mug.articulated_usdc_path is None +def test_articulated_usdc_canonicalization_moves_bottom_center_to_origin( + tmp_path: Path, +) -> None: + """A root-level translation preserves the hierarchy while normalizing its origin.""" + usdc_path = tmp_path / "offset_drawer.usdc" + stage = Usd.Stage.CreateNew(str(usdc_path)) + UsdGeom.Xform.Define(stage, "/World") + root = UsdGeom.Xform.Define(stage, "/World/drawer") + UsdPhysics.ArticulationRootAPI.Apply(root.GetPrim()) + cube = UsdGeom.Cube.Define(stage, "/World/drawer/housing") + cube.CreateSizeAttr(2.0) + UsdGeom.Xformable(cube).AddTranslateOp().Set(Gf.Vec3d(2.0, 3.0, 4.0)) + stage.SetDefaultPrim(root.GetPrim()) + stage.GetRootLayer().Save() + + _canonicalize_articulated_usdc_bottom_center(usdc_path) + + reopened_stage = Usd.Stage.Open(str(usdc_path)) + bounds = ( + UsdGeom.BBoxCache( + Usd.TimeCode.Default(), + [UsdGeom.Tokens.default_, UsdGeom.Tokens.render, UsdGeom.Tokens.proxy], + ) + .ComputeLocalBound(reopened_stage.GetDefaultPrim()) + .ComputeAlignedBox() + ) + minimum, maximum = bounds.GetMin(), bounds.GetMax() + assert minimum[1] == pytest.approx(0.0) + assert (minimum[0] + maximum[0]) / 2.0 == pytest.approx(0.0) + assert (minimum[2] + maximum[2]) / 2.0 == pytest.approx(0.0) + + def test_table_root_update_propagates_its_pose_delta_to_descendants() -> None: scene_graph = SceneGraph( nodes=[ From 245d7160f9106ad750556d4a853de02552fd0825 Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 12 Aug 2026 19:35:33 +0800 Subject: [PATCH 46/85] add articulation affordance | add turn knob --- .../atomic_actions/primitives/turn_knob.py | 326 ++++++++++++++++++ scripts/tutorials/atomic_action/turn_knob.py | 206 +++++++++++ 2 files changed, 532 insertions(+) create mode 100644 embodichain/lab/sim/atomic_actions/primitives/turn_knob.py create mode 100644 scripts/tutorials/atomic_action/turn_knob.py diff --git a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py new file mode 100644 index 000000000..c96fb5b53 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py @@ -0,0 +1,326 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""TurnKnob atomic action implementation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from embodichain.utils.math import axis_angle_to_rotation_matrix, pose_inv + +from ._helpers import arm_qpos_from_state +from ..affordance import TurnAffordance +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics +from ..effects import StateDelta +from ..goals import ObjectActionGoal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..state import PlanningContext +from ..trajectory_ops import ( + build_pose_plan_states, + interpolate_hand_qpos, + translate_pose_world, +) + + +@dataclass(frozen=True, slots=True, eq=False) +class TurnKnobGoal(ObjectActionGoal): + """Articulation-link knob described by a turn affordance.""" + + goal_kind: ClassVar[str] = "turn_knob" + + +@dataclass(frozen=True, slots=True, eq=False) +class TurnKnobOptions(ActionOptions): + """Per-invocation knob-turning behavior.""" + + hand_interp_steps: int = 5 + """Number of waypoints used for each close/open hand segment.""" + + turn_waypoint_count: int = 8 + """Number of Cartesian keyframes along the knob's circular turn arc.""" + + pre_grasp_distance: float = 0.1 + """Distance from the grasp pose along its negative z-axis.""" + + turn_angle: float = math.pi / 4 + """Requested knob rotation in radians.""" + + def __post_init__(self) -> None: + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if self.turn_waypoint_count < 1: + raise ValueError("turn_waypoint_count must be at least 1.") + if self.pre_grasp_distance < 0.0: + raise ValueError("pre_grasp_distance must be non-negative.") + if not math.isfinite(self.turn_angle): + raise ValueError("turn_angle must be finite.") + + +class TurnKnob(AtomicAction[TurnKnobGoal, TurnKnobOptions]): + """Approach, grasp, rotate, release, and retract from a knob.""" + + skill_id: ClassVar[str] = "turn_knob" + GoalType: ClassVar[type] = TurnKnobGoal + OptionsType: ClassVar[type] = TurnKnobOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__(self, default_options: TurnKnobOptions | None = None) -> None: + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve dimensions owned by the engine's robot.""" + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _plan( + self, + request: ResolvedActionRequest[TurnKnobGoal, TurnKnobOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan all six knob-turning segments without stepping simulation.""" + target = self.require_goal(request) + affordance = self._require_turn_affordance(target.semantics) + options = request.skill_options + manipulator = request.binding.manipulator() + end_effector = request.binding.end_effector() + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + hand_open_qpos = end_effector.joint_positions( + OPEN_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + + link_pose = affordance.get_link_pose().to( + device=self.device, dtype=torch.float32 + ) + if link_pose.shape != (self.n_envs, 4, 4): + raise ValueError( + "Articulation link pose must have shape " + f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." + ) + grasp_xpos = affordance.get_grasp_pose(link_pose).to( + device=self.device, dtype=torch.float32 + ) + + pre_grasp_xpos = translate_pose_world( + grasp_xpos, + -grasp_xpos[:, :3, 2] * options.pre_grasp_distance, + ) + turn_xpos = self._turned_grasp_poses( + link_pose, + grasp_xpos, + affordance.turn_axis, + options.turn_angle, + options.turn_waypoint_count, + ) + + n_approach, n_reach, n_turn, n_retract = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + ) + + approach_success, approach_arm = self._plan_pose_segment( + pre_grasp_xpos, + start_arm_qpos, + manipulator.name, + request, + n_approach, + ) + reach_success, reach_arm = self._plan_pose_segment( + grasp_xpos, + approach_arm[:, -1], + manipulator.name, + request, + n_reach, + ) + turn_success, turn_arm = self._plan_pose_segment( + turn_xpos, + reach_arm[:, -1], + manipulator.name, + request, + n_turn, + ) + retract_success, retract_arm = self._plan_pose_segment( + pre_grasp_xpos, + turn_arm[:, -1], + manipulator.name, + request, + n_retract, + ) + success = approach_success & reach_success & turn_success & retract_success + + hand_close = interpolate_hand_qpos( + hand_open_qpos, + hand_grasp_qpos, + n_waypoints=options.hand_interp_steps, + ) + hand_open = interpolate_hand_qpos( + hand_grasp_qpos, + hand_open_qpos, + n_waypoints=options.hand_interp_steps, + ) + parts = ( + approach_arm, + reach_arm, + hand_close, + turn_arm, + hand_open, + retract_arm, + ) + lengths = tuple(part.shape[1] for part in parts) + full = torch.empty( + (self.n_envs, sum(lengths), self.robot_dof), + dtype=context.robot.qpos.dtype, + device=self.device, + ) + full[:] = context.last_qpos.unsqueeze(1) + offset = 0 + arm_parts = (approach_arm, reach_arm, turn_arm, retract_arm) + arm_hands = ( + hand_open_qpos, + hand_open_qpos, + hand_grasp_qpos, + hand_open_qpos, + ) + for arm, hand in zip(arm_parts[:2], arm_hands[:2]): + stop = offset + arm.shape[1] + full[:, offset:stop, arm_joint_ids] = arm + full[:, offset:stop, hand_joint_ids] = hand.unsqueeze(1) + offset = stop + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = reach_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + stop = offset + turn_arm.shape[1] + full[:, offset:stop, arm_joint_ids] = turn_arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop + stop = offset + hand_open.shape[1] + full[:, offset:stop, arm_joint_ids] = turn_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_open + offset = stop + full[:, offset:, arm_joint_ids] = retract_arm + full[:, offset:, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + + return self.build_plan( + request, + context, + success=success, + trajectory=full, + expected_effects=StateDelta(), + segment_lengths={ + "approach": lengths[0], + "reach": lengths[1], + "close": lengths[2], + "turn": lengths[3], + "open": lengths[4], + "retract": lengths[5], + }, + ) + + @staticmethod + def _require_turn_affordance( + semantics: ObjectSemantics, + ) -> TurnAffordance: + affordance = semantics.affordance + if not isinstance(affordance, TurnAffordance): + raise ValueError("TurnKnob requires a TurnAffordance.") + return affordance + + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + ) -> tuple[int, int, int, int]: + motion_count = sample_count - 2 * hand_interp_steps + if motion_count < 8: + raise ValueError( + "Not enough waypoints for TurnKnob. Increase sample_count or " + "decrease hand_interp_steps." + ) + base, remainder = divmod(motion_count, 4) + values = [base + (index < remainder) for index in range(4)] + return values[0], values[1], values[2], values[3] + + def _plan_pose_segment( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, + request: ResolvedActionRequest[TurnKnobGoal, TurnKnobOptions], + sample_count: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.motion_generator.generate( + build_pose_plan_states(target_pose), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_qpos, + control_part=control_part, + sample_count=sample_count, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions + + def _turned_grasp_poses( + self, + link_pose: torch.Tensor, + grasp_xpos: torch.Tensor, + turn_axis: torch.Tensor, + turn_angle: float, + waypoint_count: int, + ) -> torch.Tensor: + """Build Cartesian EEF keyframes that follow the knob's circular arc.""" + axis = turn_axis.to(device=self.device, dtype=torch.float32) + axis = axis / torch.linalg.vector_norm(axis) + angles = torch.linspace( + turn_angle / waypoint_count, + turn_angle, + waypoint_count, + dtype=torch.float32, + device=self.device, + ) + rotations = ( + torch.eye(4, dtype=torch.float32, device=self.device) + .reshape(1, 4, 4) + .repeat(waypoint_count, 1, 1) + ) + rotations[:, :3, :3] = axis_angle_to_rotation_matrix(angles[:, None] * axis) + link_to_eef = torch.bmm(pose_inv(link_pose), grasp_xpos) + return torch.matmul( + torch.matmul(link_pose[:, None], rotations[None]), + link_to_eef[:, None], + ) + + +__all__ = ["TurnKnob", "TurnKnobGoal", "TurnKnobOptions"] diff --git a/scripts/tutorials/atomic_action/turn_knob.py b/scripts/tutorials/atomic_action/turn_knob.py new file mode 100644 index 000000000..db8f8c3bb --- /dev/null +++ b/scripts/tutorials/atomic_action/turn_knob.py @@ -0,0 +1,206 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Demonstrate TurnKnob on a microwave articulation.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + ControlPartCommandProfile, + MotionPolicy, + ObjectSemantics, + TurnAffordance, + TurnKnobGoal, + TurnKnobOptions, +) +from embodichain.lab.sim.cfg import ArticulationCfg +from embodichain.lab.sim.objects import Articulation, Robot +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + make_eef_pose_at, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +MICROWAVE_ASSET = "MicrowaveOven/microwave_oven.urdf" +KNOB_LINK_NAME = "cap_1" +MICROWAVE_POSITION = (-1.0, -0.30, 0.4) +MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees +TURN_SAMPLE_INTERVAL = 140 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the TurnKnob tutorial.""" + parser = create_tutorial_argument_parser( + "Demonstrate TurnKnob on a microwave power knob.", + features=("visualize_axes",), + ) + parser.add_argument("--turn_angle", type=float, default=-0.7853981634) + return parser.parse_args() + + +def create_microwave(sim) -> Articulation: + """Create the fixed-base microwave articulation used by the demo.""" + microwave = sim.add_articulation( + cfg=ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + init_pos=MICROWAVE_POSITION, + init_rot=MICROWAVE_ORIENTATION, + fix_base=True, + ) + ) + sim.update(step=10) + return microwave + + +def create_knob_semantics(microwave: Articulation) -> ObjectSemantics: + """Create turn semantics for the microwave power knob.""" + return ObjectSemantics( + label="microwave_power_knob", + geometry={}, + entity=microwave, + affordance=TurnAffordance( + articulation=microwave, + link_name=KNOB_LINK_NAME, + turn_axis=torch.tensor([0.0, 0.0, -1.0], device=microwave.device), + ), + ) + + +def initialize_robot_near_knob( + robot: Robot, + microwave: Articulation, + hand_open: torch.Tensor, +) -> None: + """Place the open gripper near the knob; values are intentionally coarse.""" + knob_position = microwave.get_link_pose(KNOB_LINK_NAME, to_matrix=True)[:, :3, 3] + start_position = knob_position.clone() + start_position[:, 1] += 0.25 + start_position[:, 2] += 0.05 + success, arm_qpos = robot.compute_ik( + pose=make_eef_pose_at(robot, start_position), + joint_seed=robot.get_qpos(name="arm"), + name="arm", + ) + if not torch.all(success): + logger.log_warning( + "The coarse microwave pre-turn pose is not reachable; keeping the " + "robot's configured initial arm pose." + ) + arm_qpos = robot.get_qpos(name="arm") + hand_qpos = hand_open.unsqueeze(0).expand(robot.get_qpos().shape[0], -1) + for target in (False, True): + robot.set_qpos(arm_qpos, name="arm", target=target) + robot.set_qpos(hand_qpos, name="hand", target=target) + robot.clear_dynamics() + + +def main() -> None: + """Plan and replay the microwave power-knob TurnKnob trajectory.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot( + sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] + ) + microwave = create_microwave(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot) + # initialize_robot_near_knob(robot, microwave, hand_open) + motion_gen = create_toppra_motion_generator(robot) + semantics = create_knob_semantics(microwave) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + if not args.no_vis_eef_axis: + draw_axis_marker( + sim, + "microwave_power_knob_axis", + microwave.get_link_pose(KNOB_LINK_NAME, to_matrix=True), + ) + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the microwave, then press Enter to plan TurnKnob...", + ) + + compiled = engine.compile( + ( + ActionInvocation( + skill_id="turn_knob", + goal=TurnKnobGoal(semantics), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=TURN_SAMPLE_INTERVAL), + skill_options=TurnKnobOptions( + hand_interp_steps=HAND_INTERP_STEPS, + pre_grasp_distance=0.12, + turn_angle=args.turn_angle, + ), + ), + ) + ) + if not compiled.plan_success.all(): + logger.log_warning("Failed to plan the TurnKnob demo trajectory.") + return + + if wait_for_user: + input("Press Enter to replay the TurnKnob demo...") + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix="turn_microwave_knob_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) From 3ec392c88e48d65906f2d062b22489e961ee5ed5 Mon Sep 17 00:00:00 2001 From: matafela Date: Fri, 14 Aug 2026 11:36:01 +0800 Subject: [PATCH 47/85] update --- scripts/tutorials/atomic_action/turn_knob.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/scripts/tutorials/atomic_action/turn_knob.py b/scripts/tutorials/atomic_action/turn_knob.py index db8f8c3bb..e24a766ef 100644 --- a/scripts/tutorials/atomic_action/turn_knob.py +++ b/scripts/tutorials/atomic_action/turn_knob.py @@ -40,7 +40,7 @@ TurnKnobGoal, TurnKnobOptions, ) -from embodichain.lab.sim.cfg import ArticulationCfg +from embodichain.lab.sim.cfg import ArticulationCfg, JointDrivePropertiesCfg from embodichain.lab.sim.objects import Articulation, Robot from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( @@ -56,7 +56,8 @@ run_tutorial, ) -MICROWAVE_ASSET = "MicrowaveOven/microwave_oven.urdf" +# MICROWAVE_ASSET = "MicrowaveOven/microwave_oven.urdf" +MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" KNOB_LINK_NAME = "cap_1" MICROWAVE_POSITION = (-1.0, -0.30, 0.4) MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees @@ -83,6 +84,9 @@ def create_microwave(sim) -> Articulation: fpath=get_data_path(MICROWAVE_ASSET), init_pos=MICROWAVE_POSITION, init_rot=MICROWAVE_ORIENTATION, + drive_pros=JointDrivePropertiesCfg( + stiffness=1e-3, damping=1e2, max_effort=1e-2 + ), fix_base=True, ) ) From 647b40681cf32574e44102642ff9f8076844f548 Mon Sep 17 00:00:00 2001 From: matafela Date: Mon, 17 Aug 2026 11:27:22 +0800 Subject: [PATCH 48/85] update --- .../atomic_actions/primitives/turn_knob.py | 29 +++++++++++++++++-- scripts/tutorials/atomic_action/turn_knob.py | 7 ----- 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py index c96fb5b53..5a07cbc4d 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py +++ b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py @@ -24,7 +24,11 @@ import torch -from embodichain.utils.math import axis_angle_to_rotation_matrix, pose_inv +from embodichain.utils.math import ( + axis_angle_to_rotation_matrix, + pose_inv, + get_relative_rotation, +) from ._helpers import arm_qpos_from_state from ..affordance import TurnAffordance @@ -93,6 +97,22 @@ def _on_bind(self) -> None: self.n_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof + def _find_symmetric_nearest_xpos( + self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor + ): + """Find the nearest symmetric pose to the reference pose.""" + symmetric_xpos = target_xpos.clone() + symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] + symmetric_xpos[:, :3, 1] = -symmetric_xpos[:, :3, 1] + angle_a = get_relative_rotation( + reference_xpos[:, :3, :3], target_xpos[:, :3, :3] + ) + angle_b = get_relative_rotation( + reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] + ) + target_xpos = torch.where(angle_a < angle_b, target_xpos, symmetric_xpos) + return target_xpos + def _plan( self, request: ResolvedActionRequest[TurnKnobGoal, TurnKnobOptions], @@ -131,7 +151,12 @@ def _plan( grasp_xpos = affordance.get_grasp_pose(link_pose).to( device=self.device, dtype=torch.float32 ) - + grasp_xpos = self._find_symmetric_nearest_xpos( + grasp_xpos, + reference_xpos=self.robot.compute_fk( + qpos=start_arm_qpos, name=manipulator.name, to_matrix=True + ), + ) pre_grasp_xpos = translate_pose_world( grasp_xpos, -grasp_xpos[:, :3, 2] * options.pre_grasp_distance, diff --git a/scripts/tutorials/atomic_action/turn_knob.py b/scripts/tutorials/atomic_action/turn_knob.py index e24a766ef..098cf17b6 100644 --- a/scripts/tutorials/atomic_action/turn_knob.py +++ b/scripts/tutorials/atomic_action/turn_knob.py @@ -56,7 +56,6 @@ run_tutorial, ) -# MICROWAVE_ASSET = "MicrowaveOven/microwave_oven.urdf" MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" KNOB_LINK_NAME = "cap_1" MICROWAVE_POSITION = (-1.0, -0.30, 0.4) @@ -158,12 +157,6 @@ def main() -> None: ) }, ) - if not args.no_vis_eef_axis: - draw_axis_marker( - sim, - "microwave_power_knob_axis", - microwave.get_link_pose(KNOB_LINK_NAME, to_matrix=True), - ) wait_for_user = prepare_tutorial_scene( sim, args, From 9d3c9cc9ec82a85e7db11d5d3b3c54cd71807a6a Mon Sep 17 00:00:00 2001 From: matafela Date: Mon, 17 Aug 2026 16:00:34 +0800 Subject: [PATCH 49/85] update press button --- .../atomic_actions/primitives/press_button.py | 269 ++++++++++++++++++ .../tutorials/atomic_action/press_button.py | 178 ++++++++++++ 2 files changed, 447 insertions(+) create mode 100644 embodichain/lab/sim/atomic_actions/primitives/press_button.py create mode 100644 scripts/tutorials/atomic_action/press_button.py diff --git a/embodichain/lab/sim/atomic_actions/primitives/press_button.py b/embodichain/lab/sim/atomic_actions/primitives/press_button.py new file mode 100644 index 000000000..db72ff063 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/press_button.py @@ -0,0 +1,269 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""PressButton atomic action implementation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from ._helpers import arm_qpos_from_state +from embodichain.utils.math import get_relative_rotation +from ..affordance import PressButtonAffordance +from ..control import GRASP_COMMAND +from ..core import AtomicAction, ObjectSemantics +from ..effects import StateDelta +from ..goals import ObjectActionGoal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan +from ..state import PlanningContext +from ..trajectory_ops import ( + build_pose_plan_states, + interpolate_hand_qpos, + translate_pose_world, +) + + +@dataclass(frozen=True, slots=True, eq=False) +class PressButtonGoal(ObjectActionGoal): + """Articulation-link button described by a press affordance.""" + + goal_kind: ClassVar[str] = "press_button" + + +@dataclass(frozen=True, slots=True, eq=False) +class PressButtonOptions(ActionOptions): + """Per-invocation button-pressing behavior.""" + + hand_interp_steps: int = 5 + """Number of waypoints used to close the hand.""" + + approach_distance: float = 0.1 + """Distance from the button surface opposite the press direction.""" + + press_distance: float = 0.05 + """Distance traveled into the button along its press axis.""" + + def __post_init__(self) -> None: + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if not math.isfinite(self.approach_distance): + raise ValueError("approach_distance must be finite.") + if self.approach_distance < 0.0: + raise ValueError("approach_distance must be non-negative.") + if not math.isfinite(self.press_distance): + raise ValueError("press_distance must be finite.") + if self.press_distance <= 0.0: + raise ValueError("press_distance must be positive.") + + +class PressButton(AtomicAction[PressButtonGoal, PressButtonOptions]): + """Close the gripper, approach and press a button, then retract.""" + + skill_id: ClassVar[str] = "press_button" + GoalType: ClassVar[type] = PressButtonGoal + OptionsType: ClassVar[type] = PressButtonOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__(self, default_options: PressButtonOptions | None = None) -> None: + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve dimensions owned by the engine's robot.""" + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _find_symmetric_nearest_xpos( + self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor + ): + """Find the nearest symmetric pose to the reference pose.""" + symmetric_xpos = target_xpos.clone() + symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] + symmetric_xpos[:, :3, 1] = -symmetric_xpos[:, :3, 1] + angle_a = get_relative_rotation( + reference_xpos[:, :3, :3], target_xpos[:, :3, :3] + ) + angle_b = get_relative_rotation( + reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] + ) + target_xpos = torch.where(angle_a < angle_b, target_xpos, symmetric_xpos) + return target_xpos + + def _plan( + self, + request: ResolvedActionRequest[PressButtonGoal, PressButtonOptions], + context: PlanningContext, + ) -> ActionPlan: + """Plan close, approach, press, and retract without stepping simulation.""" + target = self.require_goal(request) + affordance = self._require_press_button_affordance(target.semantics) + options = request.skill_options + manipulator = request.binding.manipulator() + end_effector = request.binding.end_effector() + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + start_hand_qpos = context.last_qpos[:, hand_joint_ids] + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + + link_pose = affordance.get_link_pose().to( + device=self.device, dtype=torch.float32 + ) + if link_pose.shape != (self.n_envs, 4, 4): + raise ValueError( + "Articulation link pose must have shape " + f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." + ) + contact_xpos = affordance.get_press_pose(link_pose).to( + device=self.device, dtype=torch.float32 + ) + contact_xpos = self._find_symmetric_nearest_xpos( + contact_xpos, + reference_xpos=self.robot.compute_fk( + qpos=start_arm_qpos, name=manipulator.name, to_matrix=True + ), + ) + approach_xpos = translate_pose_world( + contact_xpos, + -contact_xpos[:, :3, 2] * options.approach_distance, + ) + pressed_xpos = translate_pose_world( + contact_xpos, + contact_xpos[:, :3, 2] * options.press_distance, + ) + n_approach, n_press, n_retract = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + ) + hand_close = interpolate_hand_qpos( + start_hand_qpos, + hand_grasp_qpos, + n_waypoints=options.hand_interp_steps, + ) + approach_success, approach_arm = self._plan_pose_segment( + approach_xpos, + start_arm_qpos, + manipulator.name, + request, + n_approach, + ) + press_success, press_arm = self._plan_pose_segment( + pressed_xpos, + approach_arm[:, -1], + manipulator.name, + request, + n_press, + ) + retract_success, retract_arm = self._plan_pose_segment( + approach_xpos, + press_arm[:, -1], + manipulator.name, + request, + n_retract, + ) + success = approach_success & press_success & retract_success + + parts = (hand_close, approach_arm, press_arm, retract_arm) + lengths = tuple(part.shape[1] for part in parts) + full = torch.empty( + (self.n_envs, sum(lengths), self.robot_dof), + dtype=context.robot.qpos.dtype, + device=self.device, + ) + full[:] = context.last_qpos.unsqueeze(1) + offset = 0 + + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = start_arm_qpos.unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + + for arm in (approach_arm, press_arm, retract_arm): + stop = offset + arm.shape[1] + full[:, offset:stop, arm_joint_ids] = arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop + + return self.build_plan( + request, + context, + success=success, + trajectory=full, + expected_effects=StateDelta(), + segment_lengths={ + "close": lengths[0], + "approach": lengths[1], + "press": lengths[2], + "retract": lengths[3], + }, + ) + + @staticmethod + def _require_press_button_affordance( + semantics: ObjectSemantics, + ) -> PressButtonAffordance: + affordance = semantics.affordance + if not isinstance(affordance, PressButtonAffordance): + raise ValueError("PressButton requires a PressButtonAffordance.") + return affordance + + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + ) -> tuple[int, int, int]: + motion_count = sample_count - hand_interp_steps + if motion_count < 6: + raise ValueError( + "Not enough waypoints for PressButton. Increase sample_count or " + "decrease hand_interp_steps." + ) + base, remainder = divmod(motion_count, 3) + values = [base + (index < remainder) for index in range(3)] + return values[0], values[1], values[2] + + def _plan_pose_segment( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, + request: ResolvedActionRequest[PressButtonGoal, PressButtonOptions], + sample_count: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.motion_generator.generate( + build_pose_plan_states(target_pose), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_qpos, + control_part=control_part, + sample_count=sample_count, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions + + +__all__ = ["PressButton", "PressButtonGoal", "PressButtonOptions"] diff --git a/scripts/tutorials/atomic_action/press_button.py b/scripts/tutorials/atomic_action/press_button.py new file mode 100644 index 000000000..dd698062c --- /dev/null +++ b/scripts/tutorials/atomic_action/press_button.py @@ -0,0 +1,178 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Demonstrate PressButton on a microwave start button.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + ControlPartCommandProfile, + MotionPolicy, + ObjectSemantics, + PressButtonAffordance, + PressButtonGoal, + PressButtonOptions, +) +from embodichain.lab.sim.cfg import ArticulationCfg, JointDrivePropertiesCfg +from embodichain.lab.sim.objects import Articulation +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" +BUTTON_LINK_NAME = "button_cap" +MICROWAVE_POSITION = (-1.0, -0.30, 0.4) +MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees +PRESS_SAMPLE_INTERVAL = 140 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the PressButton tutorial.""" + parser = create_tutorial_argument_parser( + "Demonstrate PressButton on a microwave start button.", + features=("visualize_axes",), + ) + parser.add_argument("--press_distance", type=float, default=0.03) + return parser.parse_args() + + +def create_microwave(sim) -> Articulation: + """Create the fixed-base microwave articulation used by the demo.""" + microwave = sim.add_articulation( + cfg=ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + init_pos=MICROWAVE_POSITION, + init_qpos=(0, 0, 0, 0), + init_rot=MICROWAVE_ORIENTATION, + drive_pros=JointDrivePropertiesCfg( + stiffness=1e-3, damping=1e2, max_effort=1e-2 + ), + fix_base=True, + ) + ) + sim.update(step=10) + return microwave + + +def create_button_semantics(microwave: Articulation) -> ObjectSemantics: + """Create press semantics for the microwave start button.""" + return ObjectSemantics( + label="microwave_start_button", + geometry={}, + entity=microwave, + affordance=PressButtonAffordance( + articulation=microwave, + link_name=BUTTON_LINK_NAME, + # button_cap's local -z direction matches the prismatic joint's + # inward press direction in this asset. + press_axis=torch.tensor([0.0, 0.0, -1.0], device=microwave.device), + ), + ) + + +def main() -> None: + """Plan and replay the microwave start-button PressButton trajectory.""" + args = parse_arguments() + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot( + sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] + ) + microwave = create_microwave(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot, close_qpos=0.040) + motion_gen = create_toppra_motion_generator(robot) + semantics = create_button_semantics(microwave) + affordance = semantics.affordance + assert isinstance(affordance, PressButtonAffordance) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the microwave, then press Enter to plan PressButton...", + ) + + compiled = engine.compile( + ( + ActionInvocation( + skill_id="press_button", + goal=PressButtonGoal(semantics), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), + skill_options=PressButtonOptions( + hand_interp_steps=HAND_INTERP_STEPS, + approach_distance=0.12, + press_distance=args.press_distance, + ), + ), + ) + ) + if not compiled.plan_success.all(): + logger.log_warning("Failed to plan the PressButton demo trajectory.") + return + + if wait_for_user: + input("Press Enter to replay the PressButton demo...") + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix="press_microwave_button_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + ) + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) From 8c5d6e7d68ec6008780708f968b69cc6313fdb2e Mon Sep 17 00:00:00 2001 From: matafela Date: Mon, 17 Aug 2026 16:33:30 +0800 Subject: [PATCH 50/85] fix test --- .../lab/sim/atomic_actions/primitives/press_button.py | 5 +++-- embodichain/lab/sim/atomic_actions/primitives/turn_knob.py | 5 +++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/primitives/press_button.py b/embodichain/lab/sim/atomic_actions/primitives/press_button.py index db72ff063..bdd925fe6 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press_button.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press_button.py @@ -93,7 +93,7 @@ def _on_bind(self) -> None: def _find_symmetric_nearest_xpos( self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor - ): + ) -> torch.Tensor: """Find the nearest symmetric pose to the reference pose.""" symmetric_xpos = target_xpos.clone() symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] @@ -104,7 +104,8 @@ def _find_symmetric_nearest_xpos( angle_b = get_relative_rotation( reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] ) - target_xpos = torch.where(angle_a < angle_b, target_xpos, symmetric_xpos) + choose_target = (angle_a < angle_b)[..., None, None] + target_xpos = torch.where(choose_target, target_xpos, symmetric_xpos) return target_xpos def _plan( diff --git a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py index 5a07cbc4d..ef12b6523 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py +++ b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py @@ -99,7 +99,7 @@ def _on_bind(self) -> None: def _find_symmetric_nearest_xpos( self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor - ): + ) -> torch.Tensor: """Find the nearest symmetric pose to the reference pose.""" symmetric_xpos = target_xpos.clone() symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] @@ -110,7 +110,8 @@ def _find_symmetric_nearest_xpos( angle_b = get_relative_rotation( reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] ) - target_xpos = torch.where(angle_a < angle_b, target_xpos, symmetric_xpos) + choose_target = (angle_a < angle_b)[..., None, None] + target_xpos = torch.where(choose_target, target_xpos, symmetric_xpos) return target_xpos def _plan( From deb94c45bee9abc7d227a72cd5cf88c487917949 Mon Sep 17 00:00:00 2001 From: matafela Date: Tue, 18 Aug 2026 17:48:18 +0800 Subject: [PATCH 51/85] add pull push articulation part --- .../primitives/pull_push_articulated_part.py | 331 ++++++++++++++++++ .../pull_push_articulated_part.py | 291 +++++++++++++++ 2 files changed, 622 insertions(+) create mode 100644 embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py create mode 100644 scripts/tutorials/atomic_action/pull_push_articulated_part.py diff --git a/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py b/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py new file mode 100644 index 000000000..fa3e57984 --- /dev/null +++ b/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""PullPushArticulatedPart atomic action implementation.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from ._helpers import arm_qpos_from_state +from ..affordance import PullPushAffordance +from ..control import GRASP_COMMAND, OPEN_COMMAND +from ..core import AtomicAction, ObjectSemantics +from ..effects import StateDelta +from ..goals import ObjectActionGoal +from ..invocation import ActionOptions, ResolvedActionRequest +from ..plans import ActionPlan, normalize_success_mask +from ..state import PlanningContext +from ..trajectory_ops import ( + build_pose_plan_states, + interpolate_hand_qpos, + translate_pose_world, +) + + +@dataclass(frozen=True, slots=True, eq=False) +class PullPushArticulatedPartGoal(ObjectActionGoal): + """Translating articulation link described by a pull/push affordance.""" + + goal_kind: ClassVar[str] = "pull_push_articulated_part" + + +@dataclass(frozen=True, slots=True, eq=False) +class PullPushArticulatedPartOptions(ActionOptions): + """Per-invocation articulated-part pull/push behavior.""" + + is_pull: bool = True + """Whether to pull open; ``False`` pushes the part closed.""" + + hand_interp_steps: int = 5 + """Number of waypoints used for each close/open hand segment.""" + + approach_distance: float = 0.1 + """Pre-grasp distance opposite the approach/push axis.""" + + translation_distance: float = 0.15 + """Distance traveled along the pull or push direction.""" + + def __post_init__(self) -> None: + if not isinstance(self.is_pull, bool): + raise TypeError("is_pull must be a bool.") + if self.hand_interp_steps < 1: + raise ValueError("hand_interp_steps must be at least 1.") + if not math.isfinite(self.approach_distance): + raise ValueError("approach_distance must be finite.") + if self.approach_distance < 0.0: + raise ValueError("approach_distance must be non-negative.") + if not math.isfinite(self.translation_distance): + raise ValueError("translation_distance must be finite.") + if self.translation_distance <= 0.0: + raise ValueError("translation_distance must be positive.") + + +class PullPushArticulatedPart( + AtomicAction[PullPushArticulatedPartGoal, PullPushArticulatedPartOptions] +): + """Approach, grasp, and pull or push one translating articulation link.""" + + skill_id: ClassVar[str] = "pull_push_articulated_part" + GoalType: ClassVar[type] = PullPushArticulatedPartGoal + OptionsType: ClassVar[type] = PullPushArticulatedPartOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) + end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) + + def __init__( + self, + default_options: PullPushArticulatedPartOptions | None = None, + ) -> None: + super().__init__(default_options) + + def _on_bind(self) -> None: + """Resolve dimensions owned by the engine's robot.""" + self.n_envs = self.robot.get_qpos().shape[0] + self.robot_dof = self.robot.dof + + def _plan( + self, + request: ResolvedActionRequest[ + PullPushArticulatedPartGoal, + PullPushArticulatedPartOptions, + ], + context: PlanningContext, + ) -> ActionPlan: + """Plan the complete pull/push sequence without stepping simulation.""" + target = self.require_goal(request) + affordance = self._require_pull_push_affordance(target.semantics) + options = request.skill_options + manipulator = request.binding.manipulator() + end_effector = request.binding.end_effector() + arm_joint_ids = list(manipulator.joint_ids) + hand_joint_ids = list(end_effector.joint_ids) + start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) + hand_open_qpos = end_effector.joint_positions( + OPEN_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + hand_grasp_qpos = end_effector.joint_positions( + GRASP_COMMAND, + n_envs=context.batch_size, + device=self.device, + dtype=context.robot.qpos.dtype, + ) + + link_pose = affordance.get_articulation_link_pose().to( + device=self.device, dtype=torch.float32 + ) + if link_pose.shape != (self.n_envs, 4, 4): + raise ValueError( + "Articulation link pose must have shape " + f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." + ) + translation_axis = affordance.translation_axis.to( + device=self.device, dtype=torch.float32 + ) + translation_axis = translation_axis / torch.linalg.vector_norm(translation_axis) + translation_axis_world = torch.matmul(link_pose[:, :3, :3], translation_axis) + grasp_success, grasp_xpos, _ = affordance.get_best_grasp_poses( + obj_poses=link_pose, + approach_direction=translation_axis_world, + ) + grasp_xpos = grasp_xpos.to(device=self.device, dtype=torch.float32) + grasp_success = normalize_success_mask( + grasp_success, + n_envs=self.n_envs, + device=self.device, + name="Pull/push grasp-pose success", + ) + if not grasp_success.any(): + return self.failed_plan( + request, + context, + message="Failed to resolve an articulated-part grasp pose.", + ) + approach_xpos = translate_pose_world( + grasp_xpos, + -translation_axis_world * options.approach_distance, + ) + direction = -1.0 if options.is_pull else 1.0 + translated_xpos = translate_pose_world( + grasp_xpos, + translation_axis_world * (direction * options.translation_distance), + ) + + motion_lengths = self._motion_segment_lengths( + request.motion_policy.sample_count, + options.hand_interp_steps, + is_pull=options.is_pull, + ) + approach_success, approach_arm = self._plan_pose_segment( + approach_xpos, + start_arm_qpos, + manipulator.name, + request, + motion_lengths[0], + ) + reach_success, reach_arm = self._plan_pose_segment( + grasp_xpos, + approach_arm[:, -1], + manipulator.name, + request, + motion_lengths[1], + ) + translate_success, translate_arm = self._plan_pose_segment( + translated_xpos, + reach_arm[:, -1], + manipulator.name, + request, + motion_lengths[2], + ) + success = grasp_success & approach_success & reach_success & translate_success + + return_arm: torch.Tensor | None = None + if not options.is_pull: + return_success, return_arm = self._plan_pose_segment( + approach_xpos, + translate_arm[:, -1], + manipulator.name, + request, + motion_lengths[3], + ) + success = success & return_success + + hand_close = interpolate_hand_qpos( + hand_open_qpos, + hand_grasp_qpos, + n_waypoints=options.hand_interp_steps, + ) + hand_open = interpolate_hand_qpos( + hand_grasp_qpos, + hand_open_qpos, + n_waypoints=options.hand_interp_steps, + ) + named_parts: list[tuple[str, torch.Tensor]] = [ + ("approach", approach_arm), + ("reach", reach_arm), + ("close", hand_close), + ("pull" if options.is_pull else "push", translate_arm), + ("open", hand_open), + ] + if return_arm is not None: + named_parts.append(("return", return_arm)) + + segment_lengths = {name: part.shape[1] for name, part in named_parts} + full = torch.empty( + (self.n_envs, sum(segment_lengths.values()), self.robot_dof), + dtype=context.robot.qpos.dtype, + device=self.device, + ) + full[:] = context.last_qpos.unsqueeze(1) + offset = 0 + + for arm in (approach_arm, reach_arm): + stop = offset + arm.shape[1] + full[:, offset:stop, arm_joint_ids] = arm + full[:, offset:stop, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + offset = stop + + stop = offset + hand_close.shape[1] + full[:, offset:stop, arm_joint_ids] = reach_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_close + offset = stop + + stop = offset + translate_arm.shape[1] + full[:, offset:stop, arm_joint_ids] = translate_arm + full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) + offset = stop + + stop = offset + hand_open.shape[1] + full[:, offset:stop, arm_joint_ids] = translate_arm[:, -1].unsqueeze(1) + full[:, offset:stop, hand_joint_ids] = hand_open + offset = stop + + if return_arm is not None: + full[:, offset:, arm_joint_ids] = return_arm + full[:, offset:, hand_joint_ids] = hand_open_qpos.unsqueeze(1) + + return self.build_plan( + request, + context, + success=success, + trajectory=full, + expected_effects=StateDelta(), + segment_lengths=segment_lengths, + ) + + @staticmethod + def _require_pull_push_affordance( + semantics: ObjectSemantics, + ) -> PullPushAffordance: + affordance = semantics.affordance + if not isinstance(affordance, PullPushAffordance): + raise ValueError("PullPushArticulatedPart requires a PullPushAffordance.") + return affordance + + @staticmethod + def _motion_segment_lengths( + sample_count: int, + hand_interp_steps: int, + *, + is_pull: bool, + ) -> tuple[int, ...]: + motion_segment_count = 3 if is_pull else 4 + motion_count = sample_count - 2 * hand_interp_steps + if motion_count < 2 * motion_segment_count: + raise ValueError( + "Not enough waypoints for PullPushArticulatedPart. Increase " + "sample_count or decrease hand_interp_steps." + ) + base, remainder = divmod(motion_count, motion_segment_count) + return tuple( + base + (index < remainder) for index in range(motion_segment_count) + ) + + def _plan_pose_segment( + self, + target_pose: torch.Tensor, + start_qpos: torch.Tensor, + control_part: str, + request: ResolvedActionRequest[ + PullPushArticulatedPartGoal, + PullPushArticulatedPartOptions, + ], + sample_count: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.motion_generator.generate( + build_pose_plan_states(target_pose), + options=request.motion_policy.to_motion_gen_options( + start_qpos=start_qpos, + control_part=control_part, + sample_count=sample_count, + ), + ) + assert isinstance(result.success, torch.Tensor) + assert result.positions is not None + return result.success, result.positions + + +__all__ = [ + "PullPushArticulatedPart", + "PullPushArticulatedPartGoal", + "PullPushArticulatedPartOptions", +] diff --git a/scripts/tutorials/atomic_action/pull_push_articulated_part.py b/scripts/tutorials/atomic_action/pull_push_articulated_part.py new file mode 100644 index 000000000..11227cb25 --- /dev/null +++ b/scripts/tutorials/atomic_action/pull_push_articulated_part.py @@ -0,0 +1,291 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Demonstrate PullPushArticulatedPart on a translating drawer.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim import SimulationManager +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + AtomicActionEngine, + ControlPartCommandProfile, + MotionPolicy, + ObjectSemantics, + PullPushAffordance, + PullPushArticulatedPartGoal, + PullPushArticulatedPartOptions, +) +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + JointDrivePropertiesCfg, + RigidBodyAttributesCfg, +) +from embodichain.lab.sim.objects import Articulation +from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, +) +from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( + GripperCollisionCfg, +) +from embodichain.utils import logger +from scripts.tutorials.atomic_action.tutorial_utils import ( + add_ur5_gripper_robot, + create_toppra_motion_generator, + create_tutorial_argument_parser, + create_tutorial_simulation, + draw_axis_marker, + get_hand_open_close_qpos, + prepare_tutorial_scene, + replay_trajectory, + run_tutorial, +) + +DRAWER_ASSET = "Drawer/model_split_links_with_inertials.urdf" +HANDLE_LINK_NAME = "large_handle_bar" +DRAWER_POSITION = (-1.1, 0.0, 0.0) +DRAWER_ORIENTATION = (0.0, 0.0, 90.0) # degrees +TRANSLATION_AXIS = (0.0, 1.0, 0.0) # handle-link frame, approach/push direction +TRAJECTORY_SAMPLE_COUNT = 140 +HAND_INTERP_STEPS = 12 +POST_TRAJECTORY_STEPS = 240 + + +def parse_arguments() -> argparse.Namespace: + """Parse command-line arguments for the drawer pull/push tutorial.""" + parser = create_tutorial_argument_parser( + "Pull a drawer open, then push it closed with PullPushArticulatedPart.", + features=("grasp_sampling", "visualize_axes"), + ) + parser.add_argument("--translation_distance", type=float, default=0.18) + parser.add_argument("--approach_distance", type=float, default=0.10) + return parser.parse_args() + + +def create_drawer( + sim: SimulationManager, +) -> Articulation: + """Create the fixed-base drawer in its closed initial state.""" + drawer = sim.add_articulation( + cfg=ArticulationCfg( + uid="drawer", + fpath=get_data_path(DRAWER_ASSET), + init_pos=DRAWER_POSITION, + init_rot=DRAWER_ORIENTATION, + init_qpos=(0.0,), + drive_pros=JointDrivePropertiesCfg(drive_type="none"), + attrs=RigidBodyAttributesCfg( + static_friction=1.0, + dynamic_friction=1.0, + ), + fix_base=True, + ) + ) + sim.update(step=10) + return drawer + + +def create_drawer_semantics( + drawer: Articulation, + *, + n_sample: int, + force_reannotate: bool, +) -> ObjectSemantics: + """Create sampled-grasp translation semantics for the drawer handle. + + Args: + drawer: Drawer articulation that owns the target handle link. + n_sample: Number of antipodal surface samples. + force_reannotate: Whether to ignore a cached grasp annotation. + + Returns: + Object semantics backed by an articulation-link pull/push affordance. + """ + return ObjectSemantics( + label="drawer_large_handle", + geometry={}, + entity=drawer, + affordance=PullPushAffordance( + articulation=drawer, + link_name=HANDLE_LINK_NAME, + translation_axis=torch.tensor( + TRANSLATION_AXIS, + dtype=torch.float32, + device=drawer.device, + ), + generator_cfg=GraspGeneratorCfg( + antipodal_sampler_cfg=AntipodalSamplerCfg( + n_sample=n_sample, + max_length=0.1, + min_length=0.003, + ), + is_partial_annotate=False, + is_filter_ground_collision=False, + ), + gripper_collision_cfg=GripperCollisionCfg( + max_open_length=0.1, + finger_length=0.1, + y_thickness=0.04, + root_z_width=0.096, + open_check_margin=0.03, + point_sample_dense=0.012, + ), + force_reannotate=force_reannotate, + ), + ) + + +def create_invocation( + semantics: ObjectSemantics, + *, + is_pull: bool, + approach_distance: float, + translation_distance: float, +) -> ActionInvocation: + """Create one pull or push invocation for the shared drawer target. + + Args: + semantics: Drawer-handle semantics shared by both operations. + is_pull: Whether this invocation pulls open instead of pushing closed. + approach_distance: Pre-grasp offset opposite the approach axis. + translation_distance: Drawer travel distance for this operation. + + Returns: + A grounded pull/push invocation for the tutorial UR5. + """ + return ActionInvocation( + skill_id="pull_push_articulated_part", + goal=PullPushArticulatedPartGoal(semantics), + binding=ActionBinding( + manipulators={"primary": "arm"}, + end_effectors={"primary": "hand"}, + ), + motion_policy=MotionPolicy(sample_count=TRAJECTORY_SAMPLE_COUNT), + skill_options=PullPushArticulatedPartOptions( + is_pull=is_pull, + hand_interp_steps=HAND_INTERP_STEPS, + approach_distance=approach_distance, + translation_distance=translation_distance, + ), + ) + + +def main() -> None: + """Plan and replay a drawer pull followed by a push.""" + args = parse_arguments() + if args.translation_distance <= 0.0: + raise ValueError("--translation_distance must be positive.") + if args.translation_distance > 0.285: + raise ValueError( + "--translation_distance must not exceed the drawer limit 0.285." + ) + + sim = create_tutorial_simulation(args) + robot = add_ur5_gripper_robot( + sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0], tcp_z=0.15 + ) + drawer = create_drawer(sim) + hand_open, hand_close = get_hand_open_close_qpos(robot) + motion_gen = create_toppra_motion_generator(robot) + semantics = create_drawer_semantics( + drawer, + n_sample=args.n_sample, + force_reannotate=args.force_reannotate, + ) + affordance = semantics.affordance + assert isinstance(affordance, PullPushAffordance) + if not args.no_vis_eef_axis: + draw_axis_marker( + sim, + "drawer_handle_link_pose", + affordance.get_articulation_link_pose(), + ) + + engine = AtomicActionEngine( + motion_generator=motion_gen, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=hand_open, + grasp=hand_close, + ) + }, + ) + wait_for_user = prepare_tutorial_scene( + sim, + args, + "Inspect the closed drawer, then press Enter to plan the pull...", + ) + + for is_pull in (True, False): + operation_name = "pull" if is_pull else "push" + if not is_pull and wait_for_user: + input( + "Pull replay finished. Press Enter to read the moved handle " + "pose and plan the push..." + ) + + compiled = engine.compile( + ( + create_invocation( + semantics, + is_pull=is_pull, + approach_distance=args.approach_distance, + translation_distance=args.translation_distance, + ), + ) + ) + if not compiled.plan_success.all(): + logger.log_warning( + "Failed to plan the PullPushArticulatedPart " + f"{operation_name} trajectory." + ) + return + + if wait_for_user: + input(f"Press Enter to replay the drawer {operation_name}...") + replay_trajectory( + sim, + robot, + compiled.trajectory, + args, + video_prefix=f"{operation_name}_drawer_auto_play", + hold_steps=POST_TRAJECTORY_STEPS, + look_at=( + (-1.35, -1.15, 1.0), + (-0.55, -0.2, 0.45), + (0.0, 0.0, 1.0), + ), + ) + + if wait_for_user: + input("Press Enter to exit the simulation...") + + +if __name__ == "__main__": + run_tutorial(main) From 91d7137f8a75ce20b9a85cc304c80ae33570862e Mon Sep 17 00:00:00 2001 From: matafela Date: Tue, 18 Aug 2026 18:34:22 +0800 Subject: [PATCH 52/85] update --- .../atomic_actions/primitives/press_button.py | 25 ++++- .../atomic_actions/primitives/turn_knob.py | 4 +- .../tutorials/atomic_action/press_button.py | 93 ++++++++++++---- scripts/tutorials/atomic_action/turn_knob.py | 104 ++++++++++-------- 4 files changed, 150 insertions(+), 76 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/primitives/press_button.py b/embodichain/lab/sim/atomic_actions/primitives/press_button.py index bdd925fe6..59a782bec 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press_button.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press_button.py @@ -43,7 +43,7 @@ @dataclass(frozen=True, slots=True, eq=False) class PressButtonGoal(ObjectActionGoal): - """Articulation-link button described by a press affordance.""" + """Articulation-link or rigid button described by a press affordance.""" goal_kind: ClassVar[str] = "press_button" @@ -56,11 +56,14 @@ class PressButtonOptions(ActionOptions): """Number of waypoints used to close the hand.""" approach_distance: float = 0.1 - """Distance from the button surface opposite the press direction.""" + """Distance from the press position opposite the press direction.""" press_distance: float = 0.05 """Distance traveled into the button along its press axis.""" + press_position: tuple[float, float, float] | None = None + """Optional local-frame position overriding the affordance press position.""" + def __post_init__(self) -> None: if self.hand_interp_steps < 1: raise ValueError("hand_interp_steps must be at least 1.") @@ -72,6 +75,15 @@ def __post_init__(self) -> None: raise ValueError("press_distance must be finite.") if self.press_distance <= 0.0: raise ValueError("press_distance must be positive.") + if self.press_position is not None: + position = torch.as_tensor(self.press_position, dtype=torch.float32) + if position.shape != (3,) or not torch.isfinite(position).all(): + raise ValueError("press_position must be a finite (x, y, z) tuple.") + object.__setattr__( + self, + "press_position", + tuple(float(component) for component in position), + ) class PressButton(AtomicAction[PressButtonGoal, PressButtonOptions]): @@ -135,12 +147,13 @@ def _plan( ) if link_pose.shape != (self.n_envs, 4, 4): raise ValueError( - "Articulation link pose must have shape " + "Button target pose must have shape " f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." ) - contact_xpos = affordance.get_press_pose(link_pose).to( - device=self.device, dtype=torch.float32 - ) + contact_xpos = affordance.get_press_pose( + link_pose, + press_position=options.press_position, + ).to(device=self.device, dtype=torch.float32) contact_xpos = self._find_symmetric_nearest_xpos( contact_xpos, reference_xpos=self.robot.compute_fk( diff --git a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py index ef12b6523..a246a72a3 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py +++ b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py @@ -48,7 +48,7 @@ @dataclass(frozen=True, slots=True, eq=False) class TurnKnobGoal(ObjectActionGoal): - """Articulation-link knob described by a turn affordance.""" + """Articulation-link or rigid knob described by a turn affordance.""" goal_kind: ClassVar[str] = "turn_knob" @@ -146,7 +146,7 @@ def _plan( ) if link_pose.shape != (self.n_envs, 4, 4): raise ValueError( - "Articulation link pose must have shape " + "Knob target pose must have shape " f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." ) grasp_xpos = affordance.get_grasp_pose(link_pose).to( diff --git a/scripts/tutorials/atomic_action/press_button.py b/scripts/tutorials/atomic_action/press_button.py index dd698062c..5be762d28 100644 --- a/scripts/tutorials/atomic_action/press_button.py +++ b/scripts/tutorials/atomic_action/press_button.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate PressButton on a microwave start button.""" +"""Demonstrate PressButton on an articulation link or rigid button.""" from __future__ import annotations @@ -40,15 +40,19 @@ PressButtonGoal, PressButtonOptions, ) -from embodichain.lab.sim.cfg import ArticulationCfg, JointDrivePropertiesCfg -from embodichain.lab.sim.objects import Articulation +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + JointDrivePropertiesCfg, + RigidObjectCfg, +) +from embodichain.lab.sim.objects import Articulation, RigidObject +from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, create_toppra_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, - draw_axis_marker, get_hand_open_close_qpos, prepare_tutorial_scene, replay_trajectory, @@ -62,15 +66,30 @@ PRESS_SAMPLE_INTERVAL = 140 HAND_INTERP_STEPS = 12 POST_TRAJECTORY_STEPS = 240 +RIGID_BUTTON_POSITION = (-0.7, -0.00, 0.70) +RIGID_BUTTON_SIZE = (0.04, 0.02, 0.04) def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the PressButton tutorial.""" parser = create_tutorial_argument_parser( - "Demonstrate PressButton on a microwave start button.", + "Demonstrate PressButton on an articulation-link or rigid button.", features=("visualize_axes",), ) parser.add_argument("--press_distance", type=float, default=0.03) + parser.add_argument( + "--press_position", + type=float, + nargs=3, + default=None, + metavar=("X", "Y", "Z"), + help="Optional target-local press position overriding the affordance.", + ) + parser.add_argument( + "--rigid_object", + action="store_true", + help="Use a standalone rigid button instead of the microwave link.", + ) return parser.parse_args() @@ -93,33 +112,58 @@ def create_microwave(sim) -> Articulation: return microwave -def create_button_semantics(microwave: Articulation) -> ObjectSemantics: - """Create press semantics for the microwave start button.""" - return ObjectSemantics( - label="microwave_start_button", - geometry={}, - entity=microwave, - affordance=PressButtonAffordance( - articulation=microwave, +def create_rigid_button(sim) -> RigidObject: + """Create the standalone static rigid button used by the optional demo.""" + button = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="rigid_button", + shape=CubeCfg(size=list(RIGID_BUTTON_SIZE)), + body_type="static", + init_pos=RIGID_BUTTON_POSITION, + ) + ) + sim.update(step=10) + return button + + +def create_button_semantics( + target: Articulation | RigidObject, +) -> ObjectSemantics: + """Create press semantics for an articulation-link or rigid button.""" + if isinstance(target, Articulation): + affordance = PressButtonAffordance( + articulation=target, link_name=BUTTON_LINK_NAME, # button_cap's local -z direction matches the prismatic joint's # inward press direction in this asset. - press_axis=torch.tensor([0.0, 0.0, -1.0], device=microwave.device), - ), + press_axis=torch.tensor([0.0, 0.0, -1.0], device=target.device), + ) + label = "microwave_start_button" + else: + affordance = PressButtonAffordance( + rigid_object=target, + press_axis=torch.tensor([-1.0, 0.0, 0.0], device=target.device), + ) + label = "rigid_button" + return ObjectSemantics( + label=label, + geometry={}, + entity=target, + affordance=affordance, ) def main() -> None: - """Plan and replay the microwave start-button PressButton trajectory.""" + """Plan and replay PressButton for the selected target object type.""" args = parse_arguments() sim = create_tutorial_simulation(args) robot = add_ur5_gripper_robot( sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] ) - microwave = create_microwave(sim) + target = create_rigid_button(sim) if args.rigid_object else create_microwave(sim) hand_open, hand_close = get_hand_open_close_qpos(robot, close_qpos=0.040) motion_gen = create_toppra_motion_generator(robot) - semantics = create_button_semantics(microwave) + semantics = create_button_semantics(target) affordance = semantics.affordance assert isinstance(affordance, PressButtonAffordance) @@ -135,7 +179,7 @@ def main() -> None: wait_for_user = prepare_tutorial_scene( sim, args, - "Inspect the microwave, then press Enter to plan PressButton...", + "Inspect the button target, then press Enter to plan PressButton...", ) compiled = engine.compile( @@ -152,6 +196,11 @@ def main() -> None: hand_interp_steps=HAND_INTERP_STEPS, approach_distance=0.12, press_distance=args.press_distance, + press_position=( + None + if args.press_position is None + else tuple(args.press_position) + ), ), ), ) @@ -167,7 +216,11 @@ def main() -> None: robot, compiled.trajectory, args, - video_prefix="press_microwave_button_auto_play", + video_prefix=( + "press_rigid_button_auto_play" + if args.rigid_object + else "press_microwave_button_auto_play" + ), hold_steps=POST_TRAJECTORY_STEPS, ) if wait_for_user: diff --git a/scripts/tutorials/atomic_action/turn_knob.py b/scripts/tutorials/atomic_action/turn_knob.py index 098cf17b6..d18659b3f 100644 --- a/scripts/tutorials/atomic_action/turn_knob.py +++ b/scripts/tutorials/atomic_action/turn_knob.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Demonstrate TurnKnob on a microwave articulation.""" +"""Demonstrate TurnKnob on an articulation link or rigid knob.""" from __future__ import annotations @@ -40,17 +40,20 @@ TurnKnobGoal, TurnKnobOptions, ) -from embodichain.lab.sim.cfg import ArticulationCfg, JointDrivePropertiesCfg -from embodichain.lab.sim.objects import Articulation, Robot +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + JointDrivePropertiesCfg, + RigidObjectCfg, +) +from embodichain.lab.sim.objects import Articulation, RigidObject +from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, create_toppra_motion_generator, create_tutorial_argument_parser, create_tutorial_simulation, - draw_axis_marker, get_hand_open_close_qpos, - make_eef_pose_at, prepare_tutorial_scene, replay_trajectory, run_tutorial, @@ -63,15 +66,22 @@ TURN_SAMPLE_INTERVAL = 140 HAND_INTERP_STEPS = 12 POST_TRAJECTORY_STEPS = 240 +RIGID_KNOB_POSITION = (-0.7, -0.00, 0.70) +RIGID_KNOB_SIZE = (0.05, 0.05, 0.05) def parse_arguments() -> argparse.Namespace: """Parse command-line arguments for the TurnKnob tutorial.""" parser = create_tutorial_argument_parser( - "Demonstrate TurnKnob on a microwave power knob.", + "Demonstrate TurnKnob on an articulation-link or rigid knob.", features=("visualize_axes",), ) parser.add_argument("--turn_angle", type=float, default=-0.7853981634) + parser.add_argument( + "--rigid_object", + action="store_true", + help="Use a standalone rigid knob instead of the microwave link.", + ) return parser.parse_args() @@ -93,60 +103,54 @@ def create_microwave(sim) -> Articulation: return microwave -def create_knob_semantics(microwave: Articulation) -> ObjectSemantics: - """Create turn semantics for the microwave power knob.""" - return ObjectSemantics( - label="microwave_power_knob", - geometry={}, - entity=microwave, - affordance=TurnAffordance( - articulation=microwave, - link_name=KNOB_LINK_NAME, - turn_axis=torch.tensor([0.0, 0.0, -1.0], device=microwave.device), - ), +def create_rigid_knob(sim) -> RigidObject: + """Create the standalone static rigid knob used by the optional demo.""" + knob = sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="rigid_knob", + shape=CubeCfg(size=list(RIGID_KNOB_SIZE)), + body_type="static", + init_pos=RIGID_KNOB_POSITION, + ) ) + sim.update(step=10) + return knob -def initialize_robot_near_knob( - robot: Robot, - microwave: Articulation, - hand_open: torch.Tensor, -) -> None: - """Place the open gripper near the knob; values are intentionally coarse.""" - knob_position = microwave.get_link_pose(KNOB_LINK_NAME, to_matrix=True)[:, :3, 3] - start_position = knob_position.clone() - start_position[:, 1] += 0.25 - start_position[:, 2] += 0.05 - success, arm_qpos = robot.compute_ik( - pose=make_eef_pose_at(robot, start_position), - joint_seed=robot.get_qpos(name="arm"), - name="arm", - ) - if not torch.all(success): - logger.log_warning( - "The coarse microwave pre-turn pose is not reachable; keeping the " - "robot's configured initial arm pose." +def create_knob_semantics(target: Articulation | RigidObject) -> ObjectSemantics: + """Create turn semantics for an articulation-link or rigid knob.""" + if isinstance(target, Articulation): + affordance = TurnAffordance( + articulation=target, + link_name=KNOB_LINK_NAME, + turn_axis=torch.tensor([0.0, 0.0, -1.0], device=target.device), ) - arm_qpos = robot.get_qpos(name="arm") - hand_qpos = hand_open.unsqueeze(0).expand(robot.get_qpos().shape[0], -1) - for target in (False, True): - robot.set_qpos(arm_qpos, name="arm", target=target) - robot.set_qpos(hand_qpos, name="hand", target=target) - robot.clear_dynamics() + label = "microwave_power_knob" + else: + affordance = TurnAffordance( + rigid_object=target, + turn_axis=torch.tensor([-1.0, 0.0, 0.0], device=target.device), + ) + label = "rigid_knob" + return ObjectSemantics( + label=label, + geometry={}, + entity=target, + affordance=affordance, + ) def main() -> None: - """Plan and replay the microwave power-knob TurnKnob trajectory.""" + """Plan and replay TurnKnob for the selected target object type.""" args = parse_arguments() sim = create_tutorial_simulation(args) robot = add_ur5_gripper_robot( sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] ) - microwave = create_microwave(sim) + target = create_rigid_knob(sim) if args.rigid_object else create_microwave(sim) hand_open, hand_close = get_hand_open_close_qpos(robot) - # initialize_robot_near_knob(robot, microwave, hand_open) motion_gen = create_toppra_motion_generator(robot) - semantics = create_knob_semantics(microwave) + semantics = create_knob_semantics(target) engine = AtomicActionEngine( motion_generator=motion_gen, @@ -160,7 +164,7 @@ def main() -> None: wait_for_user = prepare_tutorial_scene( sim, args, - "Inspect the microwave, then press Enter to plan TurnKnob...", + "Inspect the knob target, then press Enter to plan TurnKnob...", ) compiled = engine.compile( @@ -192,7 +196,11 @@ def main() -> None: robot, compiled.trajectory, args, - video_prefix="turn_microwave_knob_auto_play", + video_prefix=( + "turn_rigid_knob_auto_play" + if args.rigid_object + else "turn_microwave_knob_auto_play" + ), hold_steps=POST_TRAJECTORY_STEPS, ) if wait_for_user: From c0962b618f5a0ac002bd38d3caa49869a7d399c8 Mon Sep 17 00:00:00 2001 From: matafela Date: Tue, 18 Aug 2026 18:57:10 +0800 Subject: [PATCH 53/85] fix module import --- .../atomic_actions/primitives/press_button.py | 33 +++++++++-------- .../primitives/pull_push_articulated_part.py | 37 ++++++++++--------- .../atomic_actions/primitives/turn_knob.py | 35 ++++++++++-------- 3 files changed, 57 insertions(+), 48 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/primitives/press_button.py b/embodichain/lab/sim/atomic_actions/primitives/press_button.py index 59a782bec..896e09160 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/press_button.py +++ b/embodichain/lab/sim/atomic_actions/primitives/press_button.py @@ -24,17 +24,20 @@ import torch -from ._helpers import arm_qpos_from_state from embodichain.utils.math import get_relative_rotation -from ..affordance import PressButtonAffordance -from ..control import GRASP_COMMAND -from ..core import AtomicAction, ObjectSemantics -from ..effects import StateDelta -from ..goals import ObjectActionGoal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..state import PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.affordance import PressButtonAffordance +from embodichain.lab.sim.atomic_actions.control import GRASP_COMMAND +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ObjectActionGoal +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_pose_plan_states, interpolate_hand_qpos, translate_pose_world, @@ -100,7 +103,7 @@ def __init__(self, default_options: PressButtonOptions | None = None) -> None: def _on_bind(self) -> None: """Resolve dimensions owned by the engine's robot.""" - self.n_envs = self.robot.get_qpos().shape[0] + self.num_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof def _find_symmetric_nearest_xpos( @@ -137,7 +140,7 @@ def _plan( start_hand_qpos = context.last_qpos[:, hand_joint_ids] hand_grasp_qpos = end_effector.joint_positions( GRASP_COMMAND, - n_envs=context.batch_size, + num_envs=context.batch_size, device=self.device, dtype=context.robot.qpos.dtype, ) @@ -145,10 +148,10 @@ def _plan( link_pose = affordance.get_link_pose().to( device=self.device, dtype=torch.float32 ) - if link_pose.shape != (self.n_envs, 4, 4): + if link_pose.shape != (self.num_envs, 4, 4): raise ValueError( "Button target pose must have shape " - f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." + f"({self.num_envs}, 4, 4), got {tuple(link_pose.shape)}." ) contact_xpos = affordance.get_press_pose( link_pose, @@ -203,7 +206,7 @@ def _plan( parts = (hand_close, approach_arm, press_arm, retract_arm) lengths = tuple(part.shape[1] for part in parts) full = torch.empty( - (self.n_envs, sum(lengths), self.robot_dof), + (self.num_envs, sum(lengths), self.robot_dof), dtype=context.robot.qpos.dtype, device=self.device, ) diff --git a/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py b/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py index fa3e57984..370f92c6d 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py +++ b/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py @@ -24,16 +24,19 @@ import torch -from ._helpers import arm_qpos_from_state -from ..affordance import PullPushAffordance -from ..control import GRASP_COMMAND, OPEN_COMMAND -from ..core import AtomicAction, ObjectSemantics -from ..effects import StateDelta -from ..goals import ObjectActionGoal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan, normalize_success_mask -from ..state import PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.affordance import PullPushAffordance +from embodichain.lab.sim.atomic_actions.control import GRASP_COMMAND, OPEN_COMMAND +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ObjectActionGoal +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_pose_plan_states, interpolate_hand_qpos, translate_pose_world, @@ -97,7 +100,7 @@ def __init__( def _on_bind(self) -> None: """Resolve dimensions owned by the engine's robot.""" - self.n_envs = self.robot.get_qpos().shape[0] + self.num_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof def _plan( @@ -119,13 +122,13 @@ def _plan( start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) hand_open_qpos = end_effector.joint_positions( OPEN_COMMAND, - n_envs=context.batch_size, + num_envs=context.batch_size, device=self.device, dtype=context.robot.qpos.dtype, ) hand_grasp_qpos = end_effector.joint_positions( GRASP_COMMAND, - n_envs=context.batch_size, + num_envs=context.batch_size, device=self.device, dtype=context.robot.qpos.dtype, ) @@ -133,10 +136,10 @@ def _plan( link_pose = affordance.get_articulation_link_pose().to( device=self.device, dtype=torch.float32 ) - if link_pose.shape != (self.n_envs, 4, 4): + if link_pose.shape != (self.num_envs, 4, 4): raise ValueError( "Articulation link pose must have shape " - f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." + f"({self.num_envs}, 4, 4), got {tuple(link_pose.shape)}." ) translation_axis = affordance.translation_axis.to( device=self.device, dtype=torch.float32 @@ -150,7 +153,7 @@ def _plan( grasp_xpos = grasp_xpos.to(device=self.device, dtype=torch.float32) grasp_success = normalize_success_mask( grasp_success, - n_envs=self.n_envs, + num_envs=self.num_envs, device=self.device, name="Pull/push grasp-pose success", ) @@ -231,7 +234,7 @@ def _plan( segment_lengths = {name: part.shape[1] for name, part in named_parts} full = torch.empty( - (self.n_envs, sum(segment_lengths.values()), self.robot_dof), + (self.num_envs, sum(segment_lengths.values()), self.robot_dof), dtype=context.robot.qpos.dtype, device=self.device, ) diff --git a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py index a246a72a3..12525b0b3 100644 --- a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py +++ b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py @@ -30,16 +30,19 @@ get_relative_rotation, ) -from ._helpers import arm_qpos_from_state -from ..affordance import TurnAffordance -from ..control import GRASP_COMMAND, OPEN_COMMAND -from ..core import AtomicAction, ObjectSemantics -from ..effects import StateDelta -from ..goals import ObjectActionGoal -from ..invocation import ActionOptions, ResolvedActionRequest -from ..plans import ActionPlan -from ..state import PlanningContext -from ..trajectory_ops import ( +from embodichain.lab.sim.atomic_actions.affordance import TurnAffordance +from embodichain.lab.sim.atomic_actions.control import GRASP_COMMAND, OPEN_COMMAND +from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics +from embodichain.lab.sim.atomic_actions.effects import StateDelta +from embodichain.lab.sim.atomic_actions.goals import ObjectActionGoal +from embodichain.lab.sim.atomic_actions.invocation import ( + ActionOptions, + ResolvedActionRequest, +) +from embodichain.lab.sim.atomic_actions.plans import ActionPlan +from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state +from embodichain.lab.sim.atomic_actions.state import PlanningContext +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( build_pose_plan_states, interpolate_hand_qpos, translate_pose_world, @@ -94,7 +97,7 @@ def __init__(self, default_options: TurnKnobOptions | None = None) -> None: def _on_bind(self) -> None: """Resolve dimensions owned by the engine's robot.""" - self.n_envs = self.robot.get_qpos().shape[0] + self.num_envs = self.robot.get_qpos().shape[0] self.robot_dof = self.robot.dof def _find_symmetric_nearest_xpos( @@ -130,13 +133,13 @@ def _plan( start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) hand_open_qpos = end_effector.joint_positions( OPEN_COMMAND, - n_envs=context.batch_size, + num_envs=context.batch_size, device=self.device, dtype=context.robot.qpos.dtype, ) hand_grasp_qpos = end_effector.joint_positions( GRASP_COMMAND, - n_envs=context.batch_size, + num_envs=context.batch_size, device=self.device, dtype=context.robot.qpos.dtype, ) @@ -144,10 +147,10 @@ def _plan( link_pose = affordance.get_link_pose().to( device=self.device, dtype=torch.float32 ) - if link_pose.shape != (self.n_envs, 4, 4): + if link_pose.shape != (self.num_envs, 4, 4): raise ValueError( "Knob target pose must have shape " - f"({self.n_envs}, 4, 4), got {tuple(link_pose.shape)}." + f"({self.num_envs}, 4, 4), got {tuple(link_pose.shape)}." ) grasp_xpos = affordance.get_grasp_pose(link_pose).to( device=self.device, dtype=torch.float32 @@ -225,7 +228,7 @@ def _plan( ) lengths = tuple(part.shape[1] for part in parts) full = torch.empty( - (self.n_envs, sum(lengths), self.robot_dof), + (self.num_envs, sum(lengths), self.robot_dof), dtype=context.robot.qpos.dtype, device=self.device, ) From 6926307ffbee2e7a0d9a0f359f33f22f83e751d0 Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 19 Aug 2026 11:39:11 +0800 Subject: [PATCH 54/85] rename --- .../atomic_actions/primitives/press_button.py | 286 -------------- .../primitives/pull_push_articulated_part.py | 334 ---------------- .../atomic_actions/primitives/turn_knob.py | 355 ------------------ .../tutorials/atomic_action/press_button.py | 231 ------------ .../pull_push_articulated_part.py | 291 -------------- scripts/tutorials/atomic_action/turn_knob.py | 211 ----------- 6 files changed, 1708 deletions(-) delete mode 100644 embodichain/lab/sim/atomic_actions/primitives/press_button.py delete mode 100644 embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py delete mode 100644 embodichain/lab/sim/atomic_actions/primitives/turn_knob.py delete mode 100644 scripts/tutorials/atomic_action/press_button.py delete mode 100644 scripts/tutorials/atomic_action/pull_push_articulated_part.py delete mode 100644 scripts/tutorials/atomic_action/turn_knob.py diff --git a/embodichain/lab/sim/atomic_actions/primitives/press_button.py b/embodichain/lab/sim/atomic_actions/primitives/press_button.py deleted file mode 100644 index 896e09160..000000000 --- a/embodichain/lab/sim/atomic_actions/primitives/press_button.py +++ /dev/null @@ -1,286 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""PressButton atomic action implementation.""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import ClassVar - -import torch - -from embodichain.utils.math import get_relative_rotation -from embodichain.lab.sim.atomic_actions.affordance import PressButtonAffordance -from embodichain.lab.sim.atomic_actions.control import GRASP_COMMAND -from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics -from embodichain.lab.sim.atomic_actions.effects import StateDelta -from embodichain.lab.sim.atomic_actions.goals import ObjectActionGoal -from embodichain.lab.sim.atomic_actions.invocation import ( - ActionOptions, - ResolvedActionRequest, -) -from embodichain.lab.sim.atomic_actions.plans import ActionPlan -from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state -from embodichain.lab.sim.atomic_actions.state import PlanningContext -from embodichain.lab.sim.atomic_actions.trajectory_ops import ( - build_pose_plan_states, - interpolate_hand_qpos, - translate_pose_world, -) - - -@dataclass(frozen=True, slots=True, eq=False) -class PressButtonGoal(ObjectActionGoal): - """Articulation-link or rigid button described by a press affordance.""" - - goal_kind: ClassVar[str] = "press_button" - - -@dataclass(frozen=True, slots=True, eq=False) -class PressButtonOptions(ActionOptions): - """Per-invocation button-pressing behavior.""" - - hand_interp_steps: int = 5 - """Number of waypoints used to close the hand.""" - - approach_distance: float = 0.1 - """Distance from the press position opposite the press direction.""" - - press_distance: float = 0.05 - """Distance traveled into the button along its press axis.""" - - press_position: tuple[float, float, float] | None = None - """Optional local-frame position overriding the affordance press position.""" - - def __post_init__(self) -> None: - if self.hand_interp_steps < 1: - raise ValueError("hand_interp_steps must be at least 1.") - if not math.isfinite(self.approach_distance): - raise ValueError("approach_distance must be finite.") - if self.approach_distance < 0.0: - raise ValueError("approach_distance must be non-negative.") - if not math.isfinite(self.press_distance): - raise ValueError("press_distance must be finite.") - if self.press_distance <= 0.0: - raise ValueError("press_distance must be positive.") - if self.press_position is not None: - position = torch.as_tensor(self.press_position, dtype=torch.float32) - if position.shape != (3,) or not torch.isfinite(position).all(): - raise ValueError("press_position must be a finite (x, y, z) tuple.") - object.__setattr__( - self, - "press_position", - tuple(float(component) for component in position), - ) - - -class PressButton(AtomicAction[PressButtonGoal, PressButtonOptions]): - """Close the gripper, approach and press a button, then retract.""" - - skill_id: ClassVar[str] = "press_button" - GoalType: ClassVar[type] = PressButtonGoal - OptionsType: ClassVar[type] = PressButtonOptions - manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) - end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) - - def __init__(self, default_options: PressButtonOptions | None = None) -> None: - super().__init__(default_options) - - def _on_bind(self) -> None: - """Resolve dimensions owned by the engine's robot.""" - self.num_envs = self.robot.get_qpos().shape[0] - self.robot_dof = self.robot.dof - - def _find_symmetric_nearest_xpos( - self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor - ) -> torch.Tensor: - """Find the nearest symmetric pose to the reference pose.""" - symmetric_xpos = target_xpos.clone() - symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] - symmetric_xpos[:, :3, 1] = -symmetric_xpos[:, :3, 1] - angle_a = get_relative_rotation( - reference_xpos[:, :3, :3], target_xpos[:, :3, :3] - ) - angle_b = get_relative_rotation( - reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] - ) - choose_target = (angle_a < angle_b)[..., None, None] - target_xpos = torch.where(choose_target, target_xpos, symmetric_xpos) - return target_xpos - - def _plan( - self, - request: ResolvedActionRequest[PressButtonGoal, PressButtonOptions], - context: PlanningContext, - ) -> ActionPlan: - """Plan close, approach, press, and retract without stepping simulation.""" - target = self.require_goal(request) - affordance = self._require_press_button_affordance(target.semantics) - options = request.skill_options - manipulator = request.binding.manipulator() - end_effector = request.binding.end_effector() - arm_joint_ids = list(manipulator.joint_ids) - hand_joint_ids = list(end_effector.joint_ids) - start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) - start_hand_qpos = context.last_qpos[:, hand_joint_ids] - hand_grasp_qpos = end_effector.joint_positions( - GRASP_COMMAND, - num_envs=context.batch_size, - device=self.device, - dtype=context.robot.qpos.dtype, - ) - - link_pose = affordance.get_link_pose().to( - device=self.device, dtype=torch.float32 - ) - if link_pose.shape != (self.num_envs, 4, 4): - raise ValueError( - "Button target pose must have shape " - f"({self.num_envs}, 4, 4), got {tuple(link_pose.shape)}." - ) - contact_xpos = affordance.get_press_pose( - link_pose, - press_position=options.press_position, - ).to(device=self.device, dtype=torch.float32) - contact_xpos = self._find_symmetric_nearest_xpos( - contact_xpos, - reference_xpos=self.robot.compute_fk( - qpos=start_arm_qpos, name=manipulator.name, to_matrix=True - ), - ) - approach_xpos = translate_pose_world( - contact_xpos, - -contact_xpos[:, :3, 2] * options.approach_distance, - ) - pressed_xpos = translate_pose_world( - contact_xpos, - contact_xpos[:, :3, 2] * options.press_distance, - ) - n_approach, n_press, n_retract = self._motion_segment_lengths( - request.motion_policy.sample_count, - options.hand_interp_steps, - ) - hand_close = interpolate_hand_qpos( - start_hand_qpos, - hand_grasp_qpos, - n_waypoints=options.hand_interp_steps, - ) - approach_success, approach_arm = self._plan_pose_segment( - approach_xpos, - start_arm_qpos, - manipulator.name, - request, - n_approach, - ) - press_success, press_arm = self._plan_pose_segment( - pressed_xpos, - approach_arm[:, -1], - manipulator.name, - request, - n_press, - ) - retract_success, retract_arm = self._plan_pose_segment( - approach_xpos, - press_arm[:, -1], - manipulator.name, - request, - n_retract, - ) - success = approach_success & press_success & retract_success - - parts = (hand_close, approach_arm, press_arm, retract_arm) - lengths = tuple(part.shape[1] for part in parts) - full = torch.empty( - (self.num_envs, sum(lengths), self.robot_dof), - dtype=context.robot.qpos.dtype, - device=self.device, - ) - full[:] = context.last_qpos.unsqueeze(1) - offset = 0 - - stop = offset + hand_close.shape[1] - full[:, offset:stop, arm_joint_ids] = start_arm_qpos.unsqueeze(1) - full[:, offset:stop, hand_joint_ids] = hand_close - offset = stop - - for arm in (approach_arm, press_arm, retract_arm): - stop = offset + arm.shape[1] - full[:, offset:stop, arm_joint_ids] = arm - full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) - offset = stop - - return self.build_plan( - request, - context, - success=success, - trajectory=full, - expected_effects=StateDelta(), - segment_lengths={ - "close": lengths[0], - "approach": lengths[1], - "press": lengths[2], - "retract": lengths[3], - }, - ) - - @staticmethod - def _require_press_button_affordance( - semantics: ObjectSemantics, - ) -> PressButtonAffordance: - affordance = semantics.affordance - if not isinstance(affordance, PressButtonAffordance): - raise ValueError("PressButton requires a PressButtonAffordance.") - return affordance - - @staticmethod - def _motion_segment_lengths( - sample_count: int, - hand_interp_steps: int, - ) -> tuple[int, int, int]: - motion_count = sample_count - hand_interp_steps - if motion_count < 6: - raise ValueError( - "Not enough waypoints for PressButton. Increase sample_count or " - "decrease hand_interp_steps." - ) - base, remainder = divmod(motion_count, 3) - values = [base + (index < remainder) for index in range(3)] - return values[0], values[1], values[2] - - def _plan_pose_segment( - self, - target_pose: torch.Tensor, - start_qpos: torch.Tensor, - control_part: str, - request: ResolvedActionRequest[PressButtonGoal, PressButtonOptions], - sample_count: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - result = self.motion_generator.generate( - build_pose_plan_states(target_pose), - options=request.motion_policy.to_motion_gen_options( - start_qpos=start_qpos, - control_part=control_part, - sample_count=sample_count, - ), - ) - assert isinstance(result.success, torch.Tensor) - assert result.positions is not None - return result.success, result.positions - - -__all__ = ["PressButton", "PressButtonGoal", "PressButtonOptions"] diff --git a/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py b/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py deleted file mode 100644 index 370f92c6d..000000000 --- a/embodichain/lab/sim/atomic_actions/primitives/pull_push_articulated_part.py +++ /dev/null @@ -1,334 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""PullPushArticulatedPart atomic action implementation.""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import ClassVar - -import torch - -from embodichain.lab.sim.atomic_actions.affordance import PullPushAffordance -from embodichain.lab.sim.atomic_actions.control import GRASP_COMMAND, OPEN_COMMAND -from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics -from embodichain.lab.sim.atomic_actions.effects import StateDelta -from embodichain.lab.sim.atomic_actions.goals import ObjectActionGoal -from embodichain.lab.sim.atomic_actions.invocation import ( - ActionOptions, - ResolvedActionRequest, -) -from embodichain.lab.sim.atomic_actions.plans import ActionPlan, normalize_success_mask -from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state -from embodichain.lab.sim.atomic_actions.state import PlanningContext -from embodichain.lab.sim.atomic_actions.trajectory_ops import ( - build_pose_plan_states, - interpolate_hand_qpos, - translate_pose_world, -) - - -@dataclass(frozen=True, slots=True, eq=False) -class PullPushArticulatedPartGoal(ObjectActionGoal): - """Translating articulation link described by a pull/push affordance.""" - - goal_kind: ClassVar[str] = "pull_push_articulated_part" - - -@dataclass(frozen=True, slots=True, eq=False) -class PullPushArticulatedPartOptions(ActionOptions): - """Per-invocation articulated-part pull/push behavior.""" - - is_pull: bool = True - """Whether to pull open; ``False`` pushes the part closed.""" - - hand_interp_steps: int = 5 - """Number of waypoints used for each close/open hand segment.""" - - approach_distance: float = 0.1 - """Pre-grasp distance opposite the approach/push axis.""" - - translation_distance: float = 0.15 - """Distance traveled along the pull or push direction.""" - - def __post_init__(self) -> None: - if not isinstance(self.is_pull, bool): - raise TypeError("is_pull must be a bool.") - if self.hand_interp_steps < 1: - raise ValueError("hand_interp_steps must be at least 1.") - if not math.isfinite(self.approach_distance): - raise ValueError("approach_distance must be finite.") - if self.approach_distance < 0.0: - raise ValueError("approach_distance must be non-negative.") - if not math.isfinite(self.translation_distance): - raise ValueError("translation_distance must be finite.") - if self.translation_distance <= 0.0: - raise ValueError("translation_distance must be positive.") - - -class PullPushArticulatedPart( - AtomicAction[PullPushArticulatedPartGoal, PullPushArticulatedPartOptions] -): - """Approach, grasp, and pull or push one translating articulation link.""" - - skill_id: ClassVar[str] = "pull_push_articulated_part" - GoalType: ClassVar[type] = PullPushArticulatedPartGoal - OptionsType: ClassVar[type] = PullPushArticulatedPartOptions - manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) - end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) - - def __init__( - self, - default_options: PullPushArticulatedPartOptions | None = None, - ) -> None: - super().__init__(default_options) - - def _on_bind(self) -> None: - """Resolve dimensions owned by the engine's robot.""" - self.num_envs = self.robot.get_qpos().shape[0] - self.robot_dof = self.robot.dof - - def _plan( - self, - request: ResolvedActionRequest[ - PullPushArticulatedPartGoal, - PullPushArticulatedPartOptions, - ], - context: PlanningContext, - ) -> ActionPlan: - """Plan the complete pull/push sequence without stepping simulation.""" - target = self.require_goal(request) - affordance = self._require_pull_push_affordance(target.semantics) - options = request.skill_options - manipulator = request.binding.manipulator() - end_effector = request.binding.end_effector() - arm_joint_ids = list(manipulator.joint_ids) - hand_joint_ids = list(end_effector.joint_ids) - start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) - hand_open_qpos = end_effector.joint_positions( - OPEN_COMMAND, - num_envs=context.batch_size, - device=self.device, - dtype=context.robot.qpos.dtype, - ) - hand_grasp_qpos = end_effector.joint_positions( - GRASP_COMMAND, - num_envs=context.batch_size, - device=self.device, - dtype=context.robot.qpos.dtype, - ) - - link_pose = affordance.get_articulation_link_pose().to( - device=self.device, dtype=torch.float32 - ) - if link_pose.shape != (self.num_envs, 4, 4): - raise ValueError( - "Articulation link pose must have shape " - f"({self.num_envs}, 4, 4), got {tuple(link_pose.shape)}." - ) - translation_axis = affordance.translation_axis.to( - device=self.device, dtype=torch.float32 - ) - translation_axis = translation_axis / torch.linalg.vector_norm(translation_axis) - translation_axis_world = torch.matmul(link_pose[:, :3, :3], translation_axis) - grasp_success, grasp_xpos, _ = affordance.get_best_grasp_poses( - obj_poses=link_pose, - approach_direction=translation_axis_world, - ) - grasp_xpos = grasp_xpos.to(device=self.device, dtype=torch.float32) - grasp_success = normalize_success_mask( - grasp_success, - num_envs=self.num_envs, - device=self.device, - name="Pull/push grasp-pose success", - ) - if not grasp_success.any(): - return self.failed_plan( - request, - context, - message="Failed to resolve an articulated-part grasp pose.", - ) - approach_xpos = translate_pose_world( - grasp_xpos, - -translation_axis_world * options.approach_distance, - ) - direction = -1.0 if options.is_pull else 1.0 - translated_xpos = translate_pose_world( - grasp_xpos, - translation_axis_world * (direction * options.translation_distance), - ) - - motion_lengths = self._motion_segment_lengths( - request.motion_policy.sample_count, - options.hand_interp_steps, - is_pull=options.is_pull, - ) - approach_success, approach_arm = self._plan_pose_segment( - approach_xpos, - start_arm_qpos, - manipulator.name, - request, - motion_lengths[0], - ) - reach_success, reach_arm = self._plan_pose_segment( - grasp_xpos, - approach_arm[:, -1], - manipulator.name, - request, - motion_lengths[1], - ) - translate_success, translate_arm = self._plan_pose_segment( - translated_xpos, - reach_arm[:, -1], - manipulator.name, - request, - motion_lengths[2], - ) - success = grasp_success & approach_success & reach_success & translate_success - - return_arm: torch.Tensor | None = None - if not options.is_pull: - return_success, return_arm = self._plan_pose_segment( - approach_xpos, - translate_arm[:, -1], - manipulator.name, - request, - motion_lengths[3], - ) - success = success & return_success - - hand_close = interpolate_hand_qpos( - hand_open_qpos, - hand_grasp_qpos, - n_waypoints=options.hand_interp_steps, - ) - hand_open = interpolate_hand_qpos( - hand_grasp_qpos, - hand_open_qpos, - n_waypoints=options.hand_interp_steps, - ) - named_parts: list[tuple[str, torch.Tensor]] = [ - ("approach", approach_arm), - ("reach", reach_arm), - ("close", hand_close), - ("pull" if options.is_pull else "push", translate_arm), - ("open", hand_open), - ] - if return_arm is not None: - named_parts.append(("return", return_arm)) - - segment_lengths = {name: part.shape[1] for name, part in named_parts} - full = torch.empty( - (self.num_envs, sum(segment_lengths.values()), self.robot_dof), - dtype=context.robot.qpos.dtype, - device=self.device, - ) - full[:] = context.last_qpos.unsqueeze(1) - offset = 0 - - for arm in (approach_arm, reach_arm): - stop = offset + arm.shape[1] - full[:, offset:stop, arm_joint_ids] = arm - full[:, offset:stop, hand_joint_ids] = hand_open_qpos.unsqueeze(1) - offset = stop - - stop = offset + hand_close.shape[1] - full[:, offset:stop, arm_joint_ids] = reach_arm[:, -1].unsqueeze(1) - full[:, offset:stop, hand_joint_ids] = hand_close - offset = stop - - stop = offset + translate_arm.shape[1] - full[:, offset:stop, arm_joint_ids] = translate_arm - full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) - offset = stop - - stop = offset + hand_open.shape[1] - full[:, offset:stop, arm_joint_ids] = translate_arm[:, -1].unsqueeze(1) - full[:, offset:stop, hand_joint_ids] = hand_open - offset = stop - - if return_arm is not None: - full[:, offset:, arm_joint_ids] = return_arm - full[:, offset:, hand_joint_ids] = hand_open_qpos.unsqueeze(1) - - return self.build_plan( - request, - context, - success=success, - trajectory=full, - expected_effects=StateDelta(), - segment_lengths=segment_lengths, - ) - - @staticmethod - def _require_pull_push_affordance( - semantics: ObjectSemantics, - ) -> PullPushAffordance: - affordance = semantics.affordance - if not isinstance(affordance, PullPushAffordance): - raise ValueError("PullPushArticulatedPart requires a PullPushAffordance.") - return affordance - - @staticmethod - def _motion_segment_lengths( - sample_count: int, - hand_interp_steps: int, - *, - is_pull: bool, - ) -> tuple[int, ...]: - motion_segment_count = 3 if is_pull else 4 - motion_count = sample_count - 2 * hand_interp_steps - if motion_count < 2 * motion_segment_count: - raise ValueError( - "Not enough waypoints for PullPushArticulatedPart. Increase " - "sample_count or decrease hand_interp_steps." - ) - base, remainder = divmod(motion_count, motion_segment_count) - return tuple( - base + (index < remainder) for index in range(motion_segment_count) - ) - - def _plan_pose_segment( - self, - target_pose: torch.Tensor, - start_qpos: torch.Tensor, - control_part: str, - request: ResolvedActionRequest[ - PullPushArticulatedPartGoal, - PullPushArticulatedPartOptions, - ], - sample_count: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - result = self.motion_generator.generate( - build_pose_plan_states(target_pose), - options=request.motion_policy.to_motion_gen_options( - start_qpos=start_qpos, - control_part=control_part, - sample_count=sample_count, - ), - ) - assert isinstance(result.success, torch.Tensor) - assert result.positions is not None - return result.success, result.positions - - -__all__ = [ - "PullPushArticulatedPart", - "PullPushArticulatedPartGoal", - "PullPushArticulatedPartOptions", -] diff --git a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py b/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py deleted file mode 100644 index 12525b0b3..000000000 --- a/embodichain/lab/sim/atomic_actions/primitives/turn_knob.py +++ /dev/null @@ -1,355 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""TurnKnob atomic action implementation.""" - -from __future__ import annotations - -import math -from dataclasses import dataclass -from typing import ClassVar - -import torch - -from embodichain.utils.math import ( - axis_angle_to_rotation_matrix, - pose_inv, - get_relative_rotation, -) - -from embodichain.lab.sim.atomic_actions.affordance import TurnAffordance -from embodichain.lab.sim.atomic_actions.control import GRASP_COMMAND, OPEN_COMMAND -from embodichain.lab.sim.atomic_actions.core import AtomicAction, ObjectSemantics -from embodichain.lab.sim.atomic_actions.effects import StateDelta -from embodichain.lab.sim.atomic_actions.goals import ObjectActionGoal -from embodichain.lab.sim.atomic_actions.invocation import ( - ActionOptions, - ResolvedActionRequest, -) -from embodichain.lab.sim.atomic_actions.plans import ActionPlan -from embodichain.lab.sim.atomic_actions.primitives._helpers import arm_qpos_from_state -from embodichain.lab.sim.atomic_actions.state import PlanningContext -from embodichain.lab.sim.atomic_actions.trajectory_ops import ( - build_pose_plan_states, - interpolate_hand_qpos, - translate_pose_world, -) - - -@dataclass(frozen=True, slots=True, eq=False) -class TurnKnobGoal(ObjectActionGoal): - """Articulation-link or rigid knob described by a turn affordance.""" - - goal_kind: ClassVar[str] = "turn_knob" - - -@dataclass(frozen=True, slots=True, eq=False) -class TurnKnobOptions(ActionOptions): - """Per-invocation knob-turning behavior.""" - - hand_interp_steps: int = 5 - """Number of waypoints used for each close/open hand segment.""" - - turn_waypoint_count: int = 8 - """Number of Cartesian keyframes along the knob's circular turn arc.""" - - pre_grasp_distance: float = 0.1 - """Distance from the grasp pose along its negative z-axis.""" - - turn_angle: float = math.pi / 4 - """Requested knob rotation in radians.""" - - def __post_init__(self) -> None: - if self.hand_interp_steps < 1: - raise ValueError("hand_interp_steps must be at least 1.") - if self.turn_waypoint_count < 1: - raise ValueError("turn_waypoint_count must be at least 1.") - if self.pre_grasp_distance < 0.0: - raise ValueError("pre_grasp_distance must be non-negative.") - if not math.isfinite(self.turn_angle): - raise ValueError("turn_angle must be finite.") - - -class TurnKnob(AtomicAction[TurnKnobGoal, TurnKnobOptions]): - """Approach, grasp, rotate, release, and retract from a knob.""" - - skill_id: ClassVar[str] = "turn_knob" - GoalType: ClassVar[type] = TurnKnobGoal - OptionsType: ClassVar[type] = TurnKnobOptions - manipulator_roles: ClassVar[tuple[str, ...]] = ("primary",) - end_effector_roles: ClassVar[tuple[str, ...]] = ("primary",) - - def __init__(self, default_options: TurnKnobOptions | None = None) -> None: - super().__init__(default_options) - - def _on_bind(self) -> None: - """Resolve dimensions owned by the engine's robot.""" - self.num_envs = self.robot.get_qpos().shape[0] - self.robot_dof = self.robot.dof - - def _find_symmetric_nearest_xpos( - self, target_xpos: torch.Tensor, reference_xpos: torch.Tensor - ) -> torch.Tensor: - """Find the nearest symmetric pose to the reference pose.""" - symmetric_xpos = target_xpos.clone() - symmetric_xpos[:, :3, 0] = -symmetric_xpos[:, :3, 0] - symmetric_xpos[:, :3, 1] = -symmetric_xpos[:, :3, 1] - angle_a = get_relative_rotation( - reference_xpos[:, :3, :3], target_xpos[:, :3, :3] - ) - angle_b = get_relative_rotation( - reference_xpos[:, :3, :3], symmetric_xpos[:, :3, :3] - ) - choose_target = (angle_a < angle_b)[..., None, None] - target_xpos = torch.where(choose_target, target_xpos, symmetric_xpos) - return target_xpos - - def _plan( - self, - request: ResolvedActionRequest[TurnKnobGoal, TurnKnobOptions], - context: PlanningContext, - ) -> ActionPlan: - """Plan all six knob-turning segments without stepping simulation.""" - target = self.require_goal(request) - affordance = self._require_turn_affordance(target.semantics) - options = request.skill_options - manipulator = request.binding.manipulator() - end_effector = request.binding.end_effector() - arm_joint_ids = list(manipulator.joint_ids) - hand_joint_ids = list(end_effector.joint_ids) - start_arm_qpos = arm_qpos_from_state(context, arm_joint_ids) - hand_open_qpos = end_effector.joint_positions( - OPEN_COMMAND, - num_envs=context.batch_size, - device=self.device, - dtype=context.robot.qpos.dtype, - ) - hand_grasp_qpos = end_effector.joint_positions( - GRASP_COMMAND, - num_envs=context.batch_size, - device=self.device, - dtype=context.robot.qpos.dtype, - ) - - link_pose = affordance.get_link_pose().to( - device=self.device, dtype=torch.float32 - ) - if link_pose.shape != (self.num_envs, 4, 4): - raise ValueError( - "Knob target pose must have shape " - f"({self.num_envs}, 4, 4), got {tuple(link_pose.shape)}." - ) - grasp_xpos = affordance.get_grasp_pose(link_pose).to( - device=self.device, dtype=torch.float32 - ) - grasp_xpos = self._find_symmetric_nearest_xpos( - grasp_xpos, - reference_xpos=self.robot.compute_fk( - qpos=start_arm_qpos, name=manipulator.name, to_matrix=True - ), - ) - pre_grasp_xpos = translate_pose_world( - grasp_xpos, - -grasp_xpos[:, :3, 2] * options.pre_grasp_distance, - ) - turn_xpos = self._turned_grasp_poses( - link_pose, - grasp_xpos, - affordance.turn_axis, - options.turn_angle, - options.turn_waypoint_count, - ) - - n_approach, n_reach, n_turn, n_retract = self._motion_segment_lengths( - request.motion_policy.sample_count, - options.hand_interp_steps, - ) - - approach_success, approach_arm = self._plan_pose_segment( - pre_grasp_xpos, - start_arm_qpos, - manipulator.name, - request, - n_approach, - ) - reach_success, reach_arm = self._plan_pose_segment( - grasp_xpos, - approach_arm[:, -1], - manipulator.name, - request, - n_reach, - ) - turn_success, turn_arm = self._plan_pose_segment( - turn_xpos, - reach_arm[:, -1], - manipulator.name, - request, - n_turn, - ) - retract_success, retract_arm = self._plan_pose_segment( - pre_grasp_xpos, - turn_arm[:, -1], - manipulator.name, - request, - n_retract, - ) - success = approach_success & reach_success & turn_success & retract_success - - hand_close = interpolate_hand_qpos( - hand_open_qpos, - hand_grasp_qpos, - n_waypoints=options.hand_interp_steps, - ) - hand_open = interpolate_hand_qpos( - hand_grasp_qpos, - hand_open_qpos, - n_waypoints=options.hand_interp_steps, - ) - parts = ( - approach_arm, - reach_arm, - hand_close, - turn_arm, - hand_open, - retract_arm, - ) - lengths = tuple(part.shape[1] for part in parts) - full = torch.empty( - (self.num_envs, sum(lengths), self.robot_dof), - dtype=context.robot.qpos.dtype, - device=self.device, - ) - full[:] = context.last_qpos.unsqueeze(1) - offset = 0 - arm_parts = (approach_arm, reach_arm, turn_arm, retract_arm) - arm_hands = ( - hand_open_qpos, - hand_open_qpos, - hand_grasp_qpos, - hand_open_qpos, - ) - for arm, hand in zip(arm_parts[:2], arm_hands[:2]): - stop = offset + arm.shape[1] - full[:, offset:stop, arm_joint_ids] = arm - full[:, offset:stop, hand_joint_ids] = hand.unsqueeze(1) - offset = stop - stop = offset + hand_close.shape[1] - full[:, offset:stop, arm_joint_ids] = reach_arm[:, -1].unsqueeze(1) - full[:, offset:stop, hand_joint_ids] = hand_close - offset = stop - stop = offset + turn_arm.shape[1] - full[:, offset:stop, arm_joint_ids] = turn_arm - full[:, offset:stop, hand_joint_ids] = hand_grasp_qpos.unsqueeze(1) - offset = stop - stop = offset + hand_open.shape[1] - full[:, offset:stop, arm_joint_ids] = turn_arm[:, -1].unsqueeze(1) - full[:, offset:stop, hand_joint_ids] = hand_open - offset = stop - full[:, offset:, arm_joint_ids] = retract_arm - full[:, offset:, hand_joint_ids] = hand_open_qpos.unsqueeze(1) - - return self.build_plan( - request, - context, - success=success, - trajectory=full, - expected_effects=StateDelta(), - segment_lengths={ - "approach": lengths[0], - "reach": lengths[1], - "close": lengths[2], - "turn": lengths[3], - "open": lengths[4], - "retract": lengths[5], - }, - ) - - @staticmethod - def _require_turn_affordance( - semantics: ObjectSemantics, - ) -> TurnAffordance: - affordance = semantics.affordance - if not isinstance(affordance, TurnAffordance): - raise ValueError("TurnKnob requires a TurnAffordance.") - return affordance - - @staticmethod - def _motion_segment_lengths( - sample_count: int, - hand_interp_steps: int, - ) -> tuple[int, int, int, int]: - motion_count = sample_count - 2 * hand_interp_steps - if motion_count < 8: - raise ValueError( - "Not enough waypoints for TurnKnob. Increase sample_count or " - "decrease hand_interp_steps." - ) - base, remainder = divmod(motion_count, 4) - values = [base + (index < remainder) for index in range(4)] - return values[0], values[1], values[2], values[3] - - def _plan_pose_segment( - self, - target_pose: torch.Tensor, - start_qpos: torch.Tensor, - control_part: str, - request: ResolvedActionRequest[TurnKnobGoal, TurnKnobOptions], - sample_count: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - result = self.motion_generator.generate( - build_pose_plan_states(target_pose), - options=request.motion_policy.to_motion_gen_options( - start_qpos=start_qpos, - control_part=control_part, - sample_count=sample_count, - ), - ) - assert isinstance(result.success, torch.Tensor) - assert result.positions is not None - return result.success, result.positions - - def _turned_grasp_poses( - self, - link_pose: torch.Tensor, - grasp_xpos: torch.Tensor, - turn_axis: torch.Tensor, - turn_angle: float, - waypoint_count: int, - ) -> torch.Tensor: - """Build Cartesian EEF keyframes that follow the knob's circular arc.""" - axis = turn_axis.to(device=self.device, dtype=torch.float32) - axis = axis / torch.linalg.vector_norm(axis) - angles = torch.linspace( - turn_angle / waypoint_count, - turn_angle, - waypoint_count, - dtype=torch.float32, - device=self.device, - ) - rotations = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .reshape(1, 4, 4) - .repeat(waypoint_count, 1, 1) - ) - rotations[:, :3, :3] = axis_angle_to_rotation_matrix(angles[:, None] * axis) - link_to_eef = torch.bmm(pose_inv(link_pose), grasp_xpos) - return torch.matmul( - torch.matmul(link_pose[:, None], rotations[None]), - link_to_eef[:, None], - ) - - -__all__ = ["TurnKnob", "TurnKnobGoal", "TurnKnobOptions"] diff --git a/scripts/tutorials/atomic_action/press_button.py b/scripts/tutorials/atomic_action/press_button.py deleted file mode 100644 index 5be762d28..000000000 --- a/scripts/tutorials/atomic_action/press_button.py +++ /dev/null @@ -1,231 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Demonstrate PressButton on an articulation link or rigid button.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -import torch - -from embodichain.data import get_data_path -from embodichain.lab.sim.atomic_actions import ( - ActionBinding, - ActionInvocation, - AtomicActionEngine, - ControlPartCommandProfile, - MotionPolicy, - ObjectSemantics, - PressButtonAffordance, - PressButtonGoal, - PressButtonOptions, -) -from embodichain.lab.sim.cfg import ( - ArticulationCfg, - JointDrivePropertiesCfg, - RigidObjectCfg, -) -from embodichain.lab.sim.objects import Articulation, RigidObject -from embodichain.lab.sim.shapes import CubeCfg -from embodichain.utils import logger -from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - create_toppra_motion_generator, - create_tutorial_argument_parser, - create_tutorial_simulation, - get_hand_open_close_qpos, - prepare_tutorial_scene, - replay_trajectory, - run_tutorial, -) - -MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" -BUTTON_LINK_NAME = "button_cap" -MICROWAVE_POSITION = (-1.0, -0.30, 0.4) -MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees -PRESS_SAMPLE_INTERVAL = 140 -HAND_INTERP_STEPS = 12 -POST_TRAJECTORY_STEPS = 240 -RIGID_BUTTON_POSITION = (-0.7, -0.00, 0.70) -RIGID_BUTTON_SIZE = (0.04, 0.02, 0.04) - - -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the PressButton tutorial.""" - parser = create_tutorial_argument_parser( - "Demonstrate PressButton on an articulation-link or rigid button.", - features=("visualize_axes",), - ) - parser.add_argument("--press_distance", type=float, default=0.03) - parser.add_argument( - "--press_position", - type=float, - nargs=3, - default=None, - metavar=("X", "Y", "Z"), - help="Optional target-local press position overriding the affordance.", - ) - parser.add_argument( - "--rigid_object", - action="store_true", - help="Use a standalone rigid button instead of the microwave link.", - ) - return parser.parse_args() - - -def create_microwave(sim) -> Articulation: - """Create the fixed-base microwave articulation used by the demo.""" - microwave = sim.add_articulation( - cfg=ArticulationCfg( - uid="microwave", - fpath=get_data_path(MICROWAVE_ASSET), - init_pos=MICROWAVE_POSITION, - init_qpos=(0, 0, 0, 0), - init_rot=MICROWAVE_ORIENTATION, - drive_pros=JointDrivePropertiesCfg( - stiffness=1e-3, damping=1e2, max_effort=1e-2 - ), - fix_base=True, - ) - ) - sim.update(step=10) - return microwave - - -def create_rigid_button(sim) -> RigidObject: - """Create the standalone static rigid button used by the optional demo.""" - button = sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="rigid_button", - shape=CubeCfg(size=list(RIGID_BUTTON_SIZE)), - body_type="static", - init_pos=RIGID_BUTTON_POSITION, - ) - ) - sim.update(step=10) - return button - - -def create_button_semantics( - target: Articulation | RigidObject, -) -> ObjectSemantics: - """Create press semantics for an articulation-link or rigid button.""" - if isinstance(target, Articulation): - affordance = PressButtonAffordance( - articulation=target, - link_name=BUTTON_LINK_NAME, - # button_cap's local -z direction matches the prismatic joint's - # inward press direction in this asset. - press_axis=torch.tensor([0.0, 0.0, -1.0], device=target.device), - ) - label = "microwave_start_button" - else: - affordance = PressButtonAffordance( - rigid_object=target, - press_axis=torch.tensor([-1.0, 0.0, 0.0], device=target.device), - ) - label = "rigid_button" - return ObjectSemantics( - label=label, - geometry={}, - entity=target, - affordance=affordance, - ) - - -def main() -> None: - """Plan and replay PressButton for the selected target object type.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot( - sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] - ) - target = create_rigid_button(sim) if args.rigid_object else create_microwave(sim) - hand_open, hand_close = get_hand_open_close_qpos(robot, close_qpos=0.040) - motion_gen = create_toppra_motion_generator(robot) - semantics = create_button_semantics(target) - affordance = semantics.affordance - assert isinstance(affordance, PressButtonAffordance) - - engine = AtomicActionEngine( - motion_generator=motion_gen, - control_profiles={ - "hand": ControlPartCommandProfile.joint_positions( - open=hand_open, - grasp=hand_close, - ) - }, - ) - wait_for_user = prepare_tutorial_scene( - sim, - args, - "Inspect the button target, then press Enter to plan PressButton...", - ) - - compiled = engine.compile( - ( - ActionInvocation( - skill_id="press_button", - goal=PressButtonGoal(semantics), - binding=ActionBinding( - manipulators={"primary": "arm"}, - end_effectors={"primary": "hand"}, - ), - motion_policy=MotionPolicy(sample_count=PRESS_SAMPLE_INTERVAL), - skill_options=PressButtonOptions( - hand_interp_steps=HAND_INTERP_STEPS, - approach_distance=0.12, - press_distance=args.press_distance, - press_position=( - None - if args.press_position is None - else tuple(args.press_position) - ), - ), - ), - ) - ) - if not compiled.plan_success.all(): - logger.log_warning("Failed to plan the PressButton demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the PressButton demo...") - replay_trajectory( - sim, - robot, - compiled.trajectory, - args, - video_prefix=( - "press_rigid_button_auto_play" - if args.rigid_object - else "press_microwave_button_auto_play" - ), - hold_steps=POST_TRAJECTORY_STEPS, - ) - if wait_for_user: - input("Press Enter to exit the simulation...") - - -if __name__ == "__main__": - run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/pull_push_articulated_part.py b/scripts/tutorials/atomic_action/pull_push_articulated_part.py deleted file mode 100644 index 11227cb25..000000000 --- a/scripts/tutorials/atomic_action/pull_push_articulated_part.py +++ /dev/null @@ -1,291 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Demonstrate PullPushArticulatedPart on a translating drawer.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -import torch - -from embodichain.data import get_data_path -from embodichain.lab.sim import SimulationManager -from embodichain.lab.sim.atomic_actions import ( - ActionBinding, - ActionInvocation, - AtomicActionEngine, - ControlPartCommandProfile, - MotionPolicy, - ObjectSemantics, - PullPushAffordance, - PullPushArticulatedPartGoal, - PullPushArticulatedPartOptions, -) -from embodichain.lab.sim.cfg import ( - ArticulationCfg, - JointDrivePropertiesCfg, - RigidBodyAttributesCfg, -) -from embodichain.lab.sim.objects import Articulation -from embodichain.toolkits.graspkit.pg_grasp.antipodal_generator import ( - AntipodalSamplerCfg, - GraspGeneratorCfg, -) -from embodichain.toolkits.graspkit.pg_grasp.gripper_collision_checker import ( - GripperCollisionCfg, -) -from embodichain.utils import logger -from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - create_toppra_motion_generator, - create_tutorial_argument_parser, - create_tutorial_simulation, - draw_axis_marker, - get_hand_open_close_qpos, - prepare_tutorial_scene, - replay_trajectory, - run_tutorial, -) - -DRAWER_ASSET = "Drawer/model_split_links_with_inertials.urdf" -HANDLE_LINK_NAME = "large_handle_bar" -DRAWER_POSITION = (-1.1, 0.0, 0.0) -DRAWER_ORIENTATION = (0.0, 0.0, 90.0) # degrees -TRANSLATION_AXIS = (0.0, 1.0, 0.0) # handle-link frame, approach/push direction -TRAJECTORY_SAMPLE_COUNT = 140 -HAND_INTERP_STEPS = 12 -POST_TRAJECTORY_STEPS = 240 - - -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the drawer pull/push tutorial.""" - parser = create_tutorial_argument_parser( - "Pull a drawer open, then push it closed with PullPushArticulatedPart.", - features=("grasp_sampling", "visualize_axes"), - ) - parser.add_argument("--translation_distance", type=float, default=0.18) - parser.add_argument("--approach_distance", type=float, default=0.10) - return parser.parse_args() - - -def create_drawer( - sim: SimulationManager, -) -> Articulation: - """Create the fixed-base drawer in its closed initial state.""" - drawer = sim.add_articulation( - cfg=ArticulationCfg( - uid="drawer", - fpath=get_data_path(DRAWER_ASSET), - init_pos=DRAWER_POSITION, - init_rot=DRAWER_ORIENTATION, - init_qpos=(0.0,), - drive_pros=JointDrivePropertiesCfg(drive_type="none"), - attrs=RigidBodyAttributesCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - fix_base=True, - ) - ) - sim.update(step=10) - return drawer - - -def create_drawer_semantics( - drawer: Articulation, - *, - n_sample: int, - force_reannotate: bool, -) -> ObjectSemantics: - """Create sampled-grasp translation semantics for the drawer handle. - - Args: - drawer: Drawer articulation that owns the target handle link. - n_sample: Number of antipodal surface samples. - force_reannotate: Whether to ignore a cached grasp annotation. - - Returns: - Object semantics backed by an articulation-link pull/push affordance. - """ - return ObjectSemantics( - label="drawer_large_handle", - geometry={}, - entity=drawer, - affordance=PullPushAffordance( - articulation=drawer, - link_name=HANDLE_LINK_NAME, - translation_axis=torch.tensor( - TRANSLATION_AXIS, - dtype=torch.float32, - device=drawer.device, - ), - generator_cfg=GraspGeneratorCfg( - antipodal_sampler_cfg=AntipodalSamplerCfg( - n_sample=n_sample, - max_length=0.1, - min_length=0.003, - ), - is_partial_annotate=False, - is_filter_ground_collision=False, - ), - gripper_collision_cfg=GripperCollisionCfg( - max_open_length=0.1, - finger_length=0.1, - y_thickness=0.04, - root_z_width=0.096, - open_check_margin=0.03, - point_sample_dense=0.012, - ), - force_reannotate=force_reannotate, - ), - ) - - -def create_invocation( - semantics: ObjectSemantics, - *, - is_pull: bool, - approach_distance: float, - translation_distance: float, -) -> ActionInvocation: - """Create one pull or push invocation for the shared drawer target. - - Args: - semantics: Drawer-handle semantics shared by both operations. - is_pull: Whether this invocation pulls open instead of pushing closed. - approach_distance: Pre-grasp offset opposite the approach axis. - translation_distance: Drawer travel distance for this operation. - - Returns: - A grounded pull/push invocation for the tutorial UR5. - """ - return ActionInvocation( - skill_id="pull_push_articulated_part", - goal=PullPushArticulatedPartGoal(semantics), - binding=ActionBinding( - manipulators={"primary": "arm"}, - end_effectors={"primary": "hand"}, - ), - motion_policy=MotionPolicy(sample_count=TRAJECTORY_SAMPLE_COUNT), - skill_options=PullPushArticulatedPartOptions( - is_pull=is_pull, - hand_interp_steps=HAND_INTERP_STEPS, - approach_distance=approach_distance, - translation_distance=translation_distance, - ), - ) - - -def main() -> None: - """Plan and replay a drawer pull followed by a push.""" - args = parse_arguments() - if args.translation_distance <= 0.0: - raise ValueError("--translation_distance must be positive.") - if args.translation_distance > 0.285: - raise ValueError( - "--translation_distance must not exceed the drawer limit 0.285." - ) - - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot( - sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0], tcp_z=0.15 - ) - drawer = create_drawer(sim) - hand_open, hand_close = get_hand_open_close_qpos(robot) - motion_gen = create_toppra_motion_generator(robot) - semantics = create_drawer_semantics( - drawer, - n_sample=args.n_sample, - force_reannotate=args.force_reannotate, - ) - affordance = semantics.affordance - assert isinstance(affordance, PullPushAffordance) - if not args.no_vis_eef_axis: - draw_axis_marker( - sim, - "drawer_handle_link_pose", - affordance.get_articulation_link_pose(), - ) - - engine = AtomicActionEngine( - motion_generator=motion_gen, - control_profiles={ - "hand": ControlPartCommandProfile.joint_positions( - open=hand_open, - grasp=hand_close, - ) - }, - ) - wait_for_user = prepare_tutorial_scene( - sim, - args, - "Inspect the closed drawer, then press Enter to plan the pull...", - ) - - for is_pull in (True, False): - operation_name = "pull" if is_pull else "push" - if not is_pull and wait_for_user: - input( - "Pull replay finished. Press Enter to read the moved handle " - "pose and plan the push..." - ) - - compiled = engine.compile( - ( - create_invocation( - semantics, - is_pull=is_pull, - approach_distance=args.approach_distance, - translation_distance=args.translation_distance, - ), - ) - ) - if not compiled.plan_success.all(): - logger.log_warning( - "Failed to plan the PullPushArticulatedPart " - f"{operation_name} trajectory." - ) - return - - if wait_for_user: - input(f"Press Enter to replay the drawer {operation_name}...") - replay_trajectory( - sim, - robot, - compiled.trajectory, - args, - video_prefix=f"{operation_name}_drawer_auto_play", - hold_steps=POST_TRAJECTORY_STEPS, - look_at=( - (-1.35, -1.15, 1.0), - (-0.55, -0.2, 0.45), - (0.0, 0.0, 1.0), - ), - ) - - if wait_for_user: - input("Press Enter to exit the simulation...") - - -if __name__ == "__main__": - run_tutorial(main) diff --git a/scripts/tutorials/atomic_action/turn_knob.py b/scripts/tutorials/atomic_action/turn_knob.py deleted file mode 100644 index d18659b3f..000000000 --- a/scripts/tutorials/atomic_action/turn_knob.py +++ /dev/null @@ -1,211 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Demonstrate TurnKnob on an articulation link or rigid knob.""" - -from __future__ import annotations - -import argparse -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -import torch - -from embodichain.data import get_data_path -from embodichain.lab.sim.atomic_actions import ( - ActionBinding, - ActionInvocation, - AtomicActionEngine, - ControlPartCommandProfile, - MotionPolicy, - ObjectSemantics, - TurnAffordance, - TurnKnobGoal, - TurnKnobOptions, -) -from embodichain.lab.sim.cfg import ( - ArticulationCfg, - JointDrivePropertiesCfg, - RigidObjectCfg, -) -from embodichain.lab.sim.objects import Articulation, RigidObject -from embodichain.lab.sim.shapes import CubeCfg -from embodichain.utils import logger -from scripts.tutorials.atomic_action.tutorial_utils import ( - add_ur5_gripper_robot, - create_toppra_motion_generator, - create_tutorial_argument_parser, - create_tutorial_simulation, - get_hand_open_close_qpos, - prepare_tutorial_scene, - replay_trajectory, - run_tutorial, -) - -MICROWAVE_ASSET = "MicrowaveOven/microwave_oven_with_inertials.urdf" -KNOB_LINK_NAME = "cap_1" -MICROWAVE_POSITION = (-1.0, -0.30, 0.4) -MICROWAVE_ORIENTATION = (0.0, 0.0, 90) # degrees -TURN_SAMPLE_INTERVAL = 140 -HAND_INTERP_STEPS = 12 -POST_TRAJECTORY_STEPS = 240 -RIGID_KNOB_POSITION = (-0.7, -0.00, 0.70) -RIGID_KNOB_SIZE = (0.05, 0.05, 0.05) - - -def parse_arguments() -> argparse.Namespace: - """Parse command-line arguments for the TurnKnob tutorial.""" - parser = create_tutorial_argument_parser( - "Demonstrate TurnKnob on an articulation-link or rigid knob.", - features=("visualize_axes",), - ) - parser.add_argument("--turn_angle", type=float, default=-0.7853981634) - parser.add_argument( - "--rigid_object", - action="store_true", - help="Use a standalone rigid knob instead of the microwave link.", - ) - return parser.parse_args() - - -def create_microwave(sim) -> Articulation: - """Create the fixed-base microwave articulation used by the demo.""" - microwave = sim.add_articulation( - cfg=ArticulationCfg( - uid="microwave", - fpath=get_data_path(MICROWAVE_ASSET), - init_pos=MICROWAVE_POSITION, - init_rot=MICROWAVE_ORIENTATION, - drive_pros=JointDrivePropertiesCfg( - stiffness=1e-3, damping=1e2, max_effort=1e-2 - ), - fix_base=True, - ) - ) - sim.update(step=10) - return microwave - - -def create_rigid_knob(sim) -> RigidObject: - """Create the standalone static rigid knob used by the optional demo.""" - knob = sim.add_rigid_object( - cfg=RigidObjectCfg( - uid="rigid_knob", - shape=CubeCfg(size=list(RIGID_KNOB_SIZE)), - body_type="static", - init_pos=RIGID_KNOB_POSITION, - ) - ) - sim.update(step=10) - return knob - - -def create_knob_semantics(target: Articulation | RigidObject) -> ObjectSemantics: - """Create turn semantics for an articulation-link or rigid knob.""" - if isinstance(target, Articulation): - affordance = TurnAffordance( - articulation=target, - link_name=KNOB_LINK_NAME, - turn_axis=torch.tensor([0.0, 0.0, -1.0], device=target.device), - ) - label = "microwave_power_knob" - else: - affordance = TurnAffordance( - rigid_object=target, - turn_axis=torch.tensor([-1.0, 0.0, 0.0], device=target.device), - ) - label = "rigid_knob" - return ObjectSemantics( - label=label, - geometry={}, - entity=target, - affordance=affordance, - ) - - -def main() -> None: - """Plan and replay TurnKnob for the selected target object type.""" - args = parse_arguments() - sim = create_tutorial_simulation(args) - robot = add_ur5_gripper_robot( - sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] - ) - target = create_rigid_knob(sim) if args.rigid_object else create_microwave(sim) - hand_open, hand_close = get_hand_open_close_qpos(robot) - motion_gen = create_toppra_motion_generator(robot) - semantics = create_knob_semantics(target) - - engine = AtomicActionEngine( - motion_generator=motion_gen, - control_profiles={ - "hand": ControlPartCommandProfile.joint_positions( - open=hand_open, - grasp=hand_close, - ) - }, - ) - wait_for_user = prepare_tutorial_scene( - sim, - args, - "Inspect the knob target, then press Enter to plan TurnKnob...", - ) - - compiled = engine.compile( - ( - ActionInvocation( - skill_id="turn_knob", - goal=TurnKnobGoal(semantics), - binding=ActionBinding( - manipulators={"primary": "arm"}, - end_effectors={"primary": "hand"}, - ), - motion_policy=MotionPolicy(sample_count=TURN_SAMPLE_INTERVAL), - skill_options=TurnKnobOptions( - hand_interp_steps=HAND_INTERP_STEPS, - pre_grasp_distance=0.12, - turn_angle=args.turn_angle, - ), - ), - ) - ) - if not compiled.plan_success.all(): - logger.log_warning("Failed to plan the TurnKnob demo trajectory.") - return - - if wait_for_user: - input("Press Enter to replay the TurnKnob demo...") - replay_trajectory( - sim, - robot, - compiled.trajectory, - args, - video_prefix=( - "turn_rigid_knob_auto_play" - if args.rigid_object - else "turn_microwave_knob_auto_play" - ), - hold_steps=POST_TRAJECTORY_STEPS, - ) - if wait_for_user: - input("Press Enter to exit the simulation...") - - -if __name__ == "__main__": - run_tutorial(main) From 1b51dd055a318d41fcab07e1c38dd1e08faecd6e Mon Sep 17 00:00:00 2001 From: ACRL Date: Wed, 19 Aug 2026 11:21:05 +0800 Subject: [PATCH 55/85] Add unified policy evaluation for RL checkpoints (#510) Co-authored-by: acrlw <13927622+acrlw@users.noreply.github.com> --- ...odichain.learning.rl.policy_evaluation.rst | 7 + .../embodichain/embodichain.learning.rl.rst | 1 + .../features/toolkits/grasp_generator.rst | 2 +- docs/source/guides/cli.md | 52 +- docs/source/guides/index.rst | 1 + docs/source/guides/policy_evaluation.md | 196 +++++++ docs/source/guides/preview_asset.md | 2 +- docs/source/overview/sim/sim_manager.md | 2 +- docs/source/tutorial/rl.rst | 18 +- embodichain/lab/gym/utils/gym_utils.py | 4 +- embodichain/lab/scripts/analyze_workspace.py | 2 +- embodichain/lab/scripts/preview_asset.py | 2 +- embodichain/lab/sim/cfg.py | 14 +- embodichain/lab/sim/sim_manager.py | 4 +- embodichain/lab/sim/utility/render_utils.py | 3 +- embodichain/learning/rl/evaluation.py | 45 +- .../learning/rl/policy_evaluation/__init__.py | 33 ++ .../learning/rl/policy_evaluation/bridge.py | 189 +++++++ .../learning/rl/policy_evaluation/cli.py | 523 ++++++++++++++++++ .../learning/rl/policy_evaluation/manifest.py | 193 +++++++ .../learning/rl/policy_evaluation/profile.py | 128 +++++ .../learning/rl/policy_evaluation/report.py | 78 +++ .../learning/rl/policy_evaluation/viewer.py | 480 ++++++++++++++++ embodichain/learning/rl/runtime.py | 312 +++++++++++ embodichain/learning/rl/train.py | 240 ++++---- .../cart_pole/agents/grpo.json | 3 +- .../cart_pole/agents/grpo.yaml | 1 - .../classic_control/cart_pole/agents/ppo.json | 3 +- .../classic_control/cart_pole/agents/ppo.yaml | 1 - .../manipulation/push_cube/agents/grpo.json | 3 +- .../manipulation/push_cube/agents/ppo.json | 3 +- examples/learning/policy_evaluation/README.md | 137 +++++ .../policy_evaluation/anymal_c/__init__.py | 30 + .../policy_evaluation/anymal_c/profile.py | 278 ++++++++++ .../learning/policy_evaluation/eval_policy.py | 72 +++ .../policy_evaluation/prepare_resources.py | 209 +++++++ scripts/benchmark/atomic_action/common.py | 2 +- .../benchmark/atomic_action/run_benchmark.py | 2 +- tests/gym/utils/test_gym_utils.py | 10 +- .../test_anymal_c_example.py | 181 ++++++ .../rl/policy_evaluation/test_bridge.py | 84 +++ .../learning/rl/policy_evaluation/test_cli.py | 130 +++++ .../rl/policy_evaluation/test_viewer.py | 221 ++++++++ tests/learning/test_point_mass.py | 8 + tests/learning/test_runtime.py | 125 +++++ tests/learning/test_train_profile.py | 11 + tests/sim/test_cfg.py | 3 +- 47 files changed, 3859 insertions(+), 189 deletions(-) create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst create mode 100644 docs/source/guides/policy_evaluation.md create mode 100644 embodichain/learning/rl/policy_evaluation/__init__.py create mode 100644 embodichain/learning/rl/policy_evaluation/bridge.py create mode 100644 embodichain/learning/rl/policy_evaluation/cli.py create mode 100644 embodichain/learning/rl/policy_evaluation/manifest.py create mode 100644 embodichain/learning/rl/policy_evaluation/profile.py create mode 100644 embodichain/learning/rl/policy_evaluation/report.py create mode 100644 embodichain/learning/rl/policy_evaluation/viewer.py create mode 100644 embodichain/learning/rl/runtime.py create mode 100644 examples/learning/policy_evaluation/README.md create mode 100644 examples/learning/policy_evaluation/anymal_c/__init__.py create mode 100644 examples/learning/policy_evaluation/anymal_c/profile.py create mode 100644 examples/learning/policy_evaluation/eval_policy.py create mode 100644 examples/learning/policy_evaluation/prepare_resources.py create mode 100644 tests/learning/rl/policy_evaluation/test_anymal_c_example.py create mode 100644 tests/learning/rl/policy_evaluation/test_bridge.py create mode 100644 tests/learning/rl/policy_evaluation/test_cli.py create mode 100644 tests/learning/rl/policy_evaluation/test_viewer.py create mode 100644 tests/learning/test_runtime.py diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst new file mode 100644 index 000000000..b5aab6c0a --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst @@ -0,0 +1,7 @@ +embodichain.learning.rl.policy_evaluation +========================================= + +.. automodule:: embodichain.learning.rl.policy_evaluation + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst index bf1b5b01e..e2e51b69d 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst @@ -18,6 +18,7 @@ collection logic, policy/model builders, and training entry points. buffer collector models + policy_evaluation train utils diff --git a/docs/source/features/toolkits/grasp_generator.rst b/docs/source/features/toolkits/grasp_generator.rst index 18cb8ec24..fa8fb110b 100644 --- a/docs/source/features/toolkits/grasp_generator.rst +++ b/docs/source/features/toolkits/grasp_generator.rst @@ -151,7 +151,7 @@ You can customize the run with additional arguments: .. code-block:: bash - python scripts/tutorials/grasp/grasp_generator.py --num_envs --device --renderer --headless + python scripts/tutorials/grasp/grasp_generator.py --num_envs --device --renderer --headless The script computes a grasp pose, prints the elapsed time, and then waits for you to press **Enter** before executing the full grasp trajectory. Press diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 44e860e87..c1330f90c 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -159,7 +159,7 @@ embodichain run-env --gym_config config.yaml \ | ``--num_envs`` | ``1`` | Number of parallel environments | | ``--device`` | ``cpu`` | Device (``cpu`` or ``cuda``) | | ``--headless`` | ``False`` | Run in headless mode | -| ``--renderer`` | ``auto`` | Renderer backend: ``auto``, ``hybrid``, ``fast-rt`` or ``rt`` | +| ``--renderer`` | ``auto`` | Renderer backend: ``auto``, ``hybrid``, ``fast-rt`` or ``offline-rt`` | | ``--arena_space`` | ``5.0`` | Arena space size | | ``--gpu_id`` | ``0`` | GPU ID to use | | ``--preview`` | ``False`` | Enter interactive preview mode | @@ -386,6 +386,56 @@ See the Profiling section under Run Env for report format. Outputs are written t --- +## Policy Evaluation + +Evaluate the latest checkpoint from an EmbodiChain training run: + +```bash +embodichain eval-policy outputs/my_policy_ +``` + +Open a simulator task in the Viewer: + +```bash +embodichain eval-policy outputs/my_policy_ \ + --checkpoint best \ + --viewer \ + --renderer hybrid +``` + +Evaluate an explicit EmbodiChain checkpoint: + +```bash +embodichain eval-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --gym-config /path/to/gym.yaml +``` + +### Main arguments + +| Argument | Default | Description | +|---|---|---| +| ``RUN`` | *(optional)* | Training run containing ``run-manifest.json`` | +| ``--checkpoint`` | ``latest`` with RUN | ``latest``, ``best``, or a checkpoint path | +| ``--config`` | RUN manifest | Training configuration override | +| ``--gym-config`` | RUN manifest | Simulator task configuration override | +| ``--episodes`` | Training configuration | Number of completed task episodes | +| ``--num-envs`` | Training configuration | Number of parallel Headless environments | +| ``--viewer`` | Headless | Open the original simulator task in the DexSim Viewer | +| ``--control-steps`` | Viewer runs continuously | Exact number of Policy actions | +| ``--duration`` | *(optional)* | Duration converted to integer control steps | +| ``--renderer`` | Training configuration or ``hybrid`` | Viewer renderer | +| ``--device`` | Training configuration | PyTorch inference device | +| ``--sim-device`` | Inference device | Simulation device | +| ``--output`` | RUN or checkpoint evaluations | Evaluation output parent directory | + +External Motion Profiles use the same command with `--profile`. See +{doc}`policy_evaluation` for training-run layout, execution paths, Viewer +controls, output reports, and the complete ANYmal-C example. + +--- + ## Annotate Grasp Launch the browser-based grasp-region annotation tool. diff --git a/docs/source/guides/index.rst b/docs/source/guides/index.rst index 220cbb89d..a5816d4fd 100644 --- a/docs/source/guides/index.rst +++ b/docs/source/guides/index.rst @@ -12,4 +12,5 @@ Practical guides for common tasks in EmbodiChain. add_robot preview_asset run_env + policy_evaluation cli diff --git a/docs/source/guides/policy_evaluation.md b/docs/source/guides/policy_evaluation.md new file mode 100644 index 000000000..392da7cc2 --- /dev/null +++ b/docs/source/guides/policy_evaluation.md @@ -0,0 +1,196 @@ +# Policy Evaluation + +`embodichain eval-policy` evaluates a saved EmbodiChain policy after training. +It reconstructs the policy and environment from the training configuration, +loads the selected checkpoint, and writes a standalone evaluation report. + +The command runs Headless by default. Add `--viewer` to open the original +simulator task in the DexSim Viewer. + +## Training output + +`train-rl` records the files required by a later evaluation: + +```text +outputs/_/ +├── checkpoints/ +│ └── policy_*.pt +├── configs/ +│ ├── train.yaml +│ └── gym.yaml +├── logs/ +├── videos/ +│ ├── train/ +│ └── eval/ +└── run-manifest.json +``` + +`configs/gym.yaml` is present for simulator tasks. The first evaluation adds: + +```text +evaluations/ +└── -policy/ + └── evaluation.json +``` + +`run-manifest.json` connects the run directory to its configuration snapshots +and checkpoints: + +```json +{ + "schema_version": 1, + "configs": { + "train": "configs/train.yaml", + "gym": "configs/gym.yaml" + }, + "checkpoints": { + "best": "checkpoints/cart_pole_grpo_best.pt", + "latest": "checkpoints/cart_pole_grpo_step_4096.pt" + } +} +``` + +All paths in the manifest are relative to the run directory. `best` is `null` +when training did not select a best checkpoint. + +## Evaluate a training run + +The shortest command selects `latest` and runs the configured number of +Headless evaluation episodes: + +```bash +embodichain eval-policy outputs/_ +``` + +Select the best checkpoint and override the episode count: + +```bash +embodichain eval-policy outputs/_ \ + --checkpoint best \ + --episodes 10 +``` + +Open the original simulator task in the Viewer: + +```bash +embodichain eval-policy outputs/_ \ + --checkpoint best \ + --viewer \ + --renderer hybrid \ + --device cuda:0 \ + --sim-device gpu +``` + +The Viewer uses one environment and keeps running until the window closes. Use +`--episodes`, `--control-steps`, or `--duration` to select another stopping +condition. `--renderer` accepts `hybrid`, `fast-rt`, and `offline-rt`. + +For a checkpoint created before `run-manifest.json` was introduced, provide +its training configuration directly: + +```bash +embodichain eval-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --gym-config /path/to/gym.yaml \ + --viewer +``` + +`--gym-config` can be omitted when the training configuration already refers to +the task configuration. + +## Execution paths + +```mermaid +flowchart LR + Run[Training run] --> Manifest[run-manifest.json] + Manifest --> Config[Training config] + Manifest --> Checkpoint[Checkpoint] + Config --> Runtime[EmbodiChain RL runtime] + Checkpoint --> Runtime + Runtime --> Headless[Headless episode evaluation] + Runtime --> Viewer[DexSim MotionPolicyEvaluator] + Headless --> Report[evaluation.json] + Viewer --> Report + Profile[External Motion Profile] --> Viewer +``` + +Headless evaluation calls the existing `evaluate_episodes()` path. Viewer +evaluation keeps the task's original observation, action processing, reset, +reward, termination, objects, and sensors: + +```mermaid +sequenceDiagram + participant Evaluator as MotionPolicyEvaluator + participant Adapter as EmbodiChainTaskPolicyAdapter + participant Task as EmbodiChainTaskEnvironment + participant Policy as EmbodiChain Policy + participant Env as Original task Environment + + Evaluator->>Task: reset() + Task->>Env: reset() + Env-->>Task: observation and task state + Task-->>Evaluator: EvaluationFrame + loop Each control step + Evaluator->>Adapter: infer(frame) + Adapter->>Policy: deterministic inference + Policy-->>Adapter: action + Adapter-->>Evaluator: PolicyOutput + Evaluator->>Task: step(action) + Task->>Env: action processing and env.step() + Env-->>Task: observation, reward, termination and info + Task-->>Evaluator: EnvironmentStep + end +``` + +| Input | Headless | Viewer | +|---|---:|---:| +| EmbodiChain lightweight RL environment | Yes | — | +| EmbodiChain simulator RL environment | Yes | Yes | +| Registered external Motion Profile | Yes | Yes | + +Policy reconstruction follows the model definition stored in the training +configuration. The Viewer path has been validated with CartPole GRPO and +PushCube PPO checkpoints. + +## Viewer controls + +| Key | Action | +|---|---| +| `Backspace` | Reset the task and camera framing | +| `T` | Switch between tracking and free camera modes when the Environment provides a tracking target | +| `R` | Start or stop recording | +| `Esc` | Close the Viewer | + +While tracking is active, drag with the left mouse button to orbit and use the +mouse wheel to zoom. + +## External policy example + +The repository includes a concrete ANYmal-C velocity example under +`examples/learning/policy_evaluation/`. It prepares a public TorchScript +checkpoint and robot assets, registers an adjacent Motion Profile, and forwards +the remaining arguments to `eval-policy`. + +```bash +python examples/learning/policy_evaluation/prepare_resources.py +python examples/learning/policy_evaluation/eval_policy.py \ + --viewer \ + --renderer hybrid \ + --sim-device gpu +``` + +Use W/S for `vx`, A/D for `vy`, Q/E for `yaw`, and M to zero the command. See +the [example README](https://github.com/DexForce/EmbodiChain/tree/main/examples/learning/policy_evaluation) +for the resource layout, observation construction, action conversion, and +Profile implementation. + +This example tracks the robot root in the ground plane. Press `T` to switch +between tracking and free view. + +## Evaluation report + +`evaluation.json` records the selected checkpoint and configs, task and device +information, episode results, and aggregated metrics. Reports are written to +`/evaluations/` for a training run and next to an explicit checkpoint by +default. Use `--output` to select another parent directory. diff --git a/docs/source/guides/preview_asset.md b/docs/source/guides/preview_asset.md index a83fa49bc..7049382fe 100644 --- a/docs/source/guides/preview_asset.md +++ b/docs/source/guides/preview_asset.md @@ -149,7 +149,7 @@ asset.set_local_pose(pose) | `--use_usd_properties` | disabled | Use physical properties stored in the USD file. | | `--fix_base` / `--no-fix_base` | fixed | Fix or unfix articulation bases. | | `--sim_device` | `cpu` | Simulation device. | -| `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `rt`. | +| `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `offline-rt`. | | `--env_map` | none | Built-in IBL resource name or absolute `.hdr`, `.png`, or `.exr` path. | | `--headless` | disabled | Run without the native window. | | `--preview` | disabled | Enter the interactive terminal after loading. | diff --git a/docs/source/overview/sim/sim_manager.md b/docs/source/overview/sim/sim_manager.md index cd280c26e..fa57cf93c 100644 --- a/docs/source/overview/sim/sim_manager.md +++ b/docs/source/overview/sim/sim_manager.md @@ -68,7 +68,7 @@ The {class}`~cfg.RenderCfg` class controls the rendering backend and quality set | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `renderer` | `str` | `"auto"` | Renderer backend to use. Options are `'auto'` (pick a default based on the detected GPU), `'hybrid'` (ray tracing for shadows/reflections + rasterization), `'fast-rt'` (full ray tracing), and `'rt'` (offline ray-traced renderer for maximum visual fidelity). | +| `renderer` | `str` | `"auto"` | Renderer backend to use. Options are `'auto'` (pick a default based on the detected GPU), `'hybrid'` (ray tracing for shadows/reflections + rasterization), `'fast-rt'` (full ray tracing), and `'offline-rt'` (offline ray-traced renderer for maximum visual fidelity). | | `spp` | `int` | `1` | Samples per pixel for ray-traced rendering. Must be at least 1. | | `tone_mapping_enabled` | `bool` | `False` | Whether to map HDR RGB output with the modified Reinhard curve. | | `tone_mapping_exposure` | `float` | `1.0` | Non-negative fixed linear exposure multiplier applied before tone mapping. | diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index 6470ef0e6..a5a57d8bd 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -264,6 +264,19 @@ All outputs are written to ``./outputs/_/``: - **logs/**: TensorBoard logs - **checkpoints/**: Model checkpoints +- **configs/**: Training config and referenced gym config snapshots +- **evaluations/**: Timestamped policy evaluation reports +- **run-manifest.json**: Training configs and best/latest checkpoint index used by ``eval-policy`` + +A training run can be evaluated Headless or opened in its simulator task: + +.. code-block:: bash + + embodichain eval-policy outputs/_ + embodichain eval-policy outputs/_ --viewer + +See :doc:`../guides/policy_evaluation` for EmbodiChain ``.pt`` training +runs and the external Motion Profile example. Training Process ~~~~~~~~~~~~~~~~ @@ -452,9 +465,9 @@ Best Practices - **Configuration**: Use JSON for all hyperparameters. This makes experiments reproducible and easy to track. -- **Logging**: Metrics are automatically logged to TensorBoard and Weights & Biases. Check ``outputs//logs/`` for TensorBoard logs. +- **Logging**: Metrics are automatically logged to TensorBoard and Weights & Biases. Check ``outputs/_/logs/`` for TensorBoard logs. -- **Checkpoints**: Regular checkpoints are saved to ``outputs//checkpoints/``. Use these to resume training or evaluate policies. +- **Checkpoints**: Regular checkpoints are saved to ``outputs/_/checkpoints/``. Use these to resume training or evaluate policies. See Also -------- @@ -464,3 +477,4 @@ See Also - :doc:`basic_env` — Creating basic Gymnasium environments - :doc:`modular_env` — Advanced modular environments with managers - :doc:`/resources/task/index` — List of available RL task environments +- :doc:`/guides/policy_evaluation` — Headless and Viewer evaluation of EmbodiChain ``.pt`` checkpoints and external Motion Profiles diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 8b2010b00..84c4314c5 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -974,7 +974,7 @@ def add_env_launcher_args_to_parser( --num_envs: Number of environments to run in parallel (default: 1) --device: Device to run the environment on (default: 'cpu') --headless: Whether to perform the simulation in headless mode (default: False) - --renderer: Renderer backend to use for the simulation. Options are 'hybrid', 'fast-rt', and 'rt'. (default: 'hybrid') + --renderer: Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'offline-rt'. --gpu_id: The GPU ID to use for the simulation (default: 0) --gym_config: Path to gym config file (default: '') --action_config: Path to action config file (default: None) @@ -1015,7 +1015,7 @@ def add_env_launcher_args_to_parser( parser.add_argument( "--renderer", type=str, - choices=["auto", "hybrid", "fast-rt", "rt"], + choices=["auto", "hybrid", "fast-rt", "offline-rt"], default=None if require_gym_config else "auto", help="Renderer backend to use for the simulation. When loading a gym " "config, the configured render_cfg.renderer is used unless this option " diff --git a/embodichain/lab/scripts/analyze_workspace.py b/embodichain/lab/scripts/analyze_workspace.py index 7f1649c14..e26a17746 100644 --- a/embodichain/lab/scripts/analyze_workspace.py +++ b/embodichain/lab/scripts/analyze_workspace.py @@ -1047,7 +1047,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: sim.add_argument( "--renderer", type=str, - choices=["hybrid", "fast-rt", "rt"], + choices=["hybrid", "fast-rt", "offline-rt"], default="hybrid", help="Renderer backend (default: hybrid).", ) diff --git a/embodichain/lab/scripts/preview_asset.py b/embodichain/lab/scripts/preview_asset.py index 0bb4f1416..8660dfb5d 100644 --- a/embodichain/lab/scripts/preview_asset.py +++ b/embodichain/lab/scripts/preview_asset.py @@ -434,7 +434,7 @@ def _create_parser() -> argparse.ArgumentParser: parser.add_argument( "--renderer", type=str, - choices=["hybrid", "fast-rt", "rt"], + choices=["hybrid", "fast-rt", "offline-rt"], default="hybrid", help="Renderer backend (default: hybrid).", ) diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index c36fabbcd..4fef9a2b0 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -56,13 +56,13 @@ # :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`). Assigning a # concrete renderer here (e.g. in test fixtures) forces that renderer and takes # precedence over auto-selection. -DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" +DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "offline-rt"] = "auto" @configclass class RenderCfg: - renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" - """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'rt'. + renderer: Literal["auto", "hybrid", "fast-rt", "offline-rt"] = "auto" + """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'offline-rt'. Note: - 'auto' selects a default renderer based on the detected GPU: RTX-series cards use @@ -71,11 +71,11 @@ class RenderCfg: - 'hybrid' uses ray tracing for shadows and reflections while keeping rasterization for primary rendering, providing a balance between performance and visual quality. - 'fast-rt' is a fully ray-traced renderer for maximum visual fidelity, but may have higher computational cost. - - 'rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. + - 'offline-rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. """ spp: int = 1 - """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'rt'.""" + """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'offline-rt'.""" tone_mapping_enabled: bool = False """Whether to map HDR RGB output with the modified Reinhard curve.""" @@ -98,7 +98,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID elif self.renderer == "fast-rt": return Renderer.FASTRT - elif self.renderer == "rt": + elif self.renderer == "offline-rt": return Renderer.OFFLINERT elif self.renderer == "auto": # 'auto' is normally resolved by the SimulationManager before this is @@ -110,7 +110,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID else: logger.log_error( - f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'rt'." + f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'offline-rt'." ) def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index b044575f6..707da39be 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -440,7 +440,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: Args: renderer: The renderer to set. One of ``"auto"``, ``"hybrid"``, - ``"fast-rt"``, or ``"rt"``. When ``"auto"``, the renderer is + ``"fast-rt"``, or ``"offline-rt"``. When ``"auto"``, the renderer is resolved immediately from the detected GPU via :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`. gpu_id: The CUDA device index to query when ``renderer="auto"``. @@ -451,7 +451,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: from embodichain.lab.sim import cfg from embodichain.lab.sim.utility.render_utils import select_default_renderer - valid = {"auto", "hybrid", "fast-rt", "rt"} + valid = {"auto", "hybrid", "fast-rt", "offline-rt"} if renderer not in valid: logger.log_error( f"Invalid renderer '{renderer}'. Must be one of {sorted(valid)}." diff --git a/embodichain/lab/sim/utility/render_utils.py b/embodichain/lab/sim/utility/render_utils.py index d82bb2644..8469ad767 100644 --- a/embodichain/lab/sim/utility/render_utils.py +++ b/embodichain/lab/sim/utility/render_utils.py @@ -47,7 +47,8 @@ def select_default_renderer(gpu_id: int = 0) -> str: gpu_id: The CUDA device index to query for selecting the renderer. Returns: - The resolved renderer name, one of ``"hybrid"``, ``"fast-rt"``, or ``"rt"``. + The resolved renderer name, one of ``"hybrid"``, ``"fast-rt"``, or + ``"offline-rt"``. """ from embodichain.lab.sim import cfg diff --git a/embodichain/learning/rl/evaluation.py b/embodichain/learning/rl/evaluation.py index 974a2ab7c..3a54715b2 100644 --- a/embodichain/learning/rl/evaluation.py +++ b/embodichain/learning/rl/evaluation.py @@ -30,15 +30,43 @@ flatten_dict_observation, ) -__all__ = ["evaluate_episodes"] +__all__ = [ + "convert_policy_action_for_env", + "evaluate_episodes", + "infer_policy_action", + "prepare_policy_observation", +] -def _flat_observation(observation: Any, device: torch.device) -> torch.Tensor: +def prepare_policy_observation( + observation: Any, + device: torch.device | str, +) -> torch.Tensor: + """Flatten one Environment observation in the training input order.""" + device = torch.device(device) tensor_dict = dict_to_tensordict(observation, device) return flatten_dict_observation(tensor_dict) -def _action_for_env(env: Any, action: torch.Tensor) -> Any: +def infer_policy_action( + policy: torch.nn.Module, + observation: Any, + *, + device: torch.device | str, + num_envs: int, +) -> torch.Tensor: + """Run the same deterministic Policy call used by RL evaluation.""" + device = torch.device(device) + policy_input = TensorDict( + {"obs": prepare_policy_observation(observation, device)}, + batch_size=[num_envs], + device=device, + ) + return policy.get_action(policy_input, deterministic=True)["action"] + + +def convert_policy_action_for_env(env: Any, action: torch.Tensor) -> Any: + """Convert a flat Policy action to the task Environment input layout.""" action_manager = getattr(env, "action_manager", None) if action_manager is None and hasattr(env, "get_wrapper_attr"): try: @@ -106,15 +134,14 @@ def evaluate_episodes( try: observation, _ = env.reset(seed=seed) while len(returns) < num_episodes: - flat_observation = _flat_observation(observation, device) - policy_input = TensorDict( - {"obs": flat_observation}, - batch_size=[num_envs], + action = infer_policy_action( + policy, + observation, device=device, + num_envs=num_envs, ) - policy_output = policy.get_action(policy_input, deterministic=True) observation, reward, terminated, truncated, info = env.step( - _action_for_env(env, policy_output["action"]) + convert_policy_action_for_env(env, action) ) reward = torch.as_tensor(reward, device=device).reshape(num_envs) done = ( diff --git a/embodichain/learning/rl/policy_evaluation/__init__.py b/embodichain/learning/rl/policy_evaluation/__init__.py new file mode 100644 index 000000000..4dbc590b4 --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/__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. +# ---------------------------------------------------------------------------- + +"""External Policy Profiles for ``embodichain eval-policy``.""" + +from __future__ import annotations + +from .profile import ( + MotionProfile, + MotionProfileRequest, + build_motion_profile, + register_motion_profile, +) + +__all__ = [ + "MotionProfile", + "MotionProfileRequest", + "build_motion_profile", + "register_motion_profile", +] diff --git a/embodichain/learning/rl/policy_evaluation/bridge.py b/embodichain/learning/rl/policy_evaluation/bridge.py new file mode 100644 index 000000000..a336b9ccc --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/bridge.py @@ -0,0 +1,189 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Run an external Policy Profile through DexSim Motion Policy Kit.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from dexsim.kit.motion_policy import ( + PolicySpec, + ResolvedPolicy, + ResourceResolver, + RunOptions, + load_scene_config, + parse_policy_spec, + policy_spec_to_dict, + resolve_policy_spec, + run_motion_policy, + scene_config_to_dict, +) + +from .profile import MotionProfile + +__all__ = [ + "MotionEvaluationResult", + "evaluate_motion_profile", +] + + +@dataclass(frozen=True) +class MotionEvaluationResult: + """Normalized inputs and per-episode motion evaluation results.""" + + profile: MotionProfile + policy_spec: Mapping[str, Any] + scene_config: Mapping[str, Any] + episodes: tuple[Mapping[str, Any], ...] + summary: Mapping[str, Any] + viewer: bool + + +def evaluate_motion_profile( + profile: MotionProfile, + *, + episodes: int = 1, + viewer: bool = False, + control_steps: int | None = None, + duration: float | None = None, + command: tuple[float, ...] | None = None, + scene_config: str | Path = "standard", + physics_backend: str | None = None, + simulation_device: str = "cpu", + renderer: str = "hybrid", + gpu_id: int = 0, + termination_behavior: str | None = None, + cache_dir: str | Path | None = None, + offline: bool = False, +) -> MotionEvaluationResult: + """Resolve one Motion Profile and run its visual evaluation. + + Args: + profile: Provider-built profile containing the DexSim Policy Spec. + episodes: Number of independent runs. + viewer: Open the DexSim Viewer. + control_steps: Exact number of applied policy commands per run. + duration: Convenience duration converted by DexSim to policy steps. + command: Optional task command override. + scene_config: Built-in scene style or custom YAML path. + physics_backend: Optional DexSim physics backend override. + simulation_device: ``cpu`` or ``gpu``. + renderer: DexSim renderer. + gpu_id: Selected GPU index. + termination_behavior: Policy termination handling override. + cache_dir: Motion Policy Kit resource cache. + offline: Use resources already available in the cache. + + Returns: + Normalized inputs, episode results, and aggregate metrics. + """ + if episodes <= 0: + raise ValueError("episodes must be positive") + if viewer and episodes != 1: + raise ValueError("Viewer evaluation supports one episode") + parsed, resolved = _resolve_profile(profile, cache_dir, offline) + resolved_scene = load_scene_config(scene_config) + options = RunOptions( + physics_backend=physics_backend, + simulation_device=simulation_device, + renderer=renderer, + gpu_id=gpu_id, + headless=not viewer, + control_steps=control_steps, + duration=duration, + command=command, + termination_behavior=termination_behavior, + scene_config=resolved_scene, + ) + results = tuple( + _episode( + index, + run_motion_policy( + resolved, + options, + ), + ) + for index in range(episodes) + ) + return MotionEvaluationResult( + profile=profile, + policy_spec=policy_spec_to_dict(parsed), + scene_config=scene_config_to_dict(resolved_scene), + episodes=results, + summary=_summary(results), + viewer=viewer, + ) + + +def _resolve_profile( + profile: MotionProfile, + cache_dir: str | Path | None, + offline: bool, +) -> tuple[PolicySpec, ResolvedPolicy]: + parsed = parse_policy_spec(profile.policy_spec) + resolved = resolve_policy_spec( + parsed, + ResourceResolver( + None if cache_dir is None else Path(cache_dir), + offline=offline, + ), + ) + return parsed, resolved + + +def _episode(index: int, result: Any) -> dict[str, Any]: + return { + "index": index, + "reason": str(result.reason), + "simulation_time": float(result.simulation_time), + "simulation_steps": int(result.simulation_steps), + "control_steps": int(result.control_steps), + "physics_backend": str(result.physics_backend), + "requested_duration": ( + None + if result.requested_duration is None + else float(result.requested_duration) + ), + "effective_duration": float(result.effective_duration), + "metrics": {name: float(value) for name, value in result.metrics.items()}, + } + + +def _summary(episodes: tuple[Mapping[str, Any], ...]) -> dict[str, Any]: + count = len(episodes) + metric_names = set.intersection(*(set(episode["metrics"]) for episode in episodes)) + metrics = { + name: sum(episode["metrics"][name] for episode in episodes) / count + for name in sorted(metric_names) + } + result: dict[str, Any] = { + "episodes": count, + "avg_simulation_time": sum(episode["simulation_time"] for episode in episodes) + / count, + "avg_control_steps": sum(episode["control_steps"] for episode in episodes) + / count, + "avg_effective_duration": sum( + episode["effective_duration"] for episode in episodes + ) + / count, + } + if metrics: + result["metrics"] = metrics + return result diff --git a/embodichain/learning/rl/policy_evaluation/cli.py b/embodichain/learning/rl/policy_evaluation/cli.py new file mode 100644 index 000000000..8c9b9c2d2 --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/cli.py @@ -0,0 +1,523 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Unified policy evaluation for EmbodiChain training runs.""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from embodichain import __version__ +from embodichain.lab.gym.utils.registration import ( + discover_task_packages, + execute_init_hooks, +) +from embodichain.learning.rl.evaluation import evaluate_episodes +from embodichain.learning.rl.runtime import ( + PolicyRuntime, + build_gym_policy_runtime, + build_learning_policy_runtime, +) +from embodichain.utils.utility import load_config + +from .manifest import RunManifest +from .report import write_evaluation_report + +__all__ = ["cli", "parse_args", "run"] + + +@dataclass(frozen=True) +class EvaluationInput: + """Checkpoint and configuration selected for one evaluation.""" + + checkpoint: Path + profile: str | None + configs: Mapping[str, Path] + run: Path | None + requested_checkpoint: str + selected_checkpoint: str + + +@dataclass(frozen=True) +class NativeRuntime: + """Reconstructed EmbodiChain task and its runtime choices.""" + + runtime: PolicyRuntime + device: torch.device + simulation_device: torch.device + seed: int + renderer: str + uses_simulator: bool + trainer: Mapping[str, Any] + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + """Parse ``embodichain eval-policy`` arguments.""" + parser = argparse.ArgumentParser( + prog="embodichain eval-policy", + description="Evaluate an EmbodiChain or external policy checkpoint.", + ) + parser.add_argument("run", nargs="?", help="EmbodiChain training run directory.") + parser.add_argument("--profile", help="Registered external Policy Profile.") + parser.add_argument( + "--checkpoint", + help="latest, best, or a checkpoint path; defaults to latest with RUN.", + ) + parser.add_argument("--config", help="Training config for an explicit checkpoint.") + parser.add_argument("--gym-config", help="Task config override.") + parser.add_argument("--resource-root", help="External Profile resource root.") + parser.add_argument("--episodes", type=int) + parser.add_argument("--num-envs", type=int) + count = parser.add_mutually_exclusive_group() + count.add_argument("--control-steps", type=int) + count.add_argument("--duration", type=float) + parser.add_argument("--command", nargs="+", type=float) + parser.add_argument("--device", help="PyTorch inference device.") + parser.add_argument("--sim-device", choices=("cpu", "gpu")) + parser.add_argument("--seed", type=int) + parser.add_argument("--physics-backend") + parser.add_argument( + "--renderer", + choices=("hybrid", "fast-rt", "offline-rt"), + ) + parser.add_argument("--gpu-id", type=int, default=0) + parser.add_argument("--scene-config") + parser.add_argument( + "--termination-behavior", + choices=("pause", "continue", "auto_reset"), + ) + parser.add_argument("--viewer", action="store_true") + parser.add_argument("--cache-dir") + parser.add_argument("--offline", action="store_true") + parser.add_argument("--output", help="Evaluation output parent directory.") + return parser.parse_args(argv) + + +def run(args: argparse.Namespace) -> Path: + """Run Headless or Viewer evaluation and write ``evaluation.json``.""" + resolved = _resolve_input(args) + if resolved.profile is not None: + return _run_profile(args, resolved) + discover_task_packages() + execute_init_hooks() + _validate_native_options(args) + if args.viewer: + return _run_native_viewer(args, resolved) + return _run_native_headless(args, resolved) + + +def cli(argv: Sequence[str] | None = None) -> None: + """Run policy evaluation from the unified EmbodiChain CLI.""" + try: + report = run(parse_args(argv)) + except ( + FileNotFoundError, + ImportError, + KeyError, + RuntimeError, + TypeError, + ValueError, + ) as error: + raise SystemExit(f"eval-policy: {error}") from error + print(f"Evaluation report: {report}") + + +def _resolve_input(args: argparse.Namespace) -> EvaluationInput: + if args.run is not None: + manifest = RunManifest.load(args.run) + requested = args.checkpoint or "latest" + if requested in {"best", "latest"}: + selected, checkpoint = manifest.select_checkpoint(requested) + else: + selected = "explicit" + candidate = Path(requested).expanduser() + checkpoint = ( + candidate.resolve() + if candidate.is_absolute() + else (manifest.root / candidate).resolve() + ) + configs = dict(manifest.configs) + if args.config is not None: + configs["train"] = Path(args.config).expanduser().resolve() + if args.gym_config is not None: + configs["gym"] = Path(args.gym_config).expanduser().resolve() + return EvaluationInput( + checkpoint=checkpoint, + profile=args.profile, + configs=configs, + run=manifest.root, + requested_checkpoint=requested, + selected_checkpoint=selected, + ) + if args.checkpoint is None: + raise ValueError("--checkpoint is required without RUN") + configs = {} + if args.config is not None: + configs["train"] = Path(args.config).expanduser().resolve() + if args.gym_config is not None: + configs["gym"] = Path(args.gym_config).expanduser().resolve() + if args.profile is None and "train" not in configs: + raise ValueError("--config is required for an EmbodiChain checkpoint") + return EvaluationInput( + checkpoint=Path(args.checkpoint).expanduser().resolve(), + profile=args.profile, + configs=configs, + run=None, + requested_checkpoint=args.checkpoint, + selected_checkpoint="explicit", + ) + + +def _run_native_headless( + args: argparse.Namespace, + resolved: EvaluationInput, +) -> Path: + native = _build_native_runtime(args, resolved, viewer=False) + episodes = ( + args.episodes + if args.episodes is not None + else int(native.trainer.get("num_eval_episodes", 5)) + ) + try: + metrics = evaluate_episodes( + policy=native.runtime.policy, + env=native.runtime.env, + num_episodes=episodes, + device=native.device, + seed=native.seed, + ) + finally: + native.runtime.close() + _flush_simulator(native.uses_simulator) + return write_evaluation_report( + _output_parent(args.output, resolved), + _headless_report(native, resolved, episodes, metrics), + ) + + +def _run_native_viewer( + args: argparse.Namespace, + resolved: EvaluationInput, +) -> Path: + from .viewer import evaluate_native_viewer + + native = _build_native_runtime(args, resolved, viewer=True) + try: + result = evaluate_native_viewer( + native.runtime, + seed=native.seed, + episodes=args.episodes, + control_steps=args.control_steps, + duration=args.duration, + termination_behavior=args.termination_behavior or "auto_reset", + ) + finally: + _flush_simulator(True) + return write_evaluation_report( + _output_parent(args.output, resolved), + _viewer_report(result, native, resolved), + ) + + +def _run_profile(args: argparse.Namespace, resolved: EvaluationInput) -> Path: + from .bridge import evaluate_motion_profile + from .profile import MotionProfileRequest, build_motion_profile + + device = _torch_device(args.device or "cpu") + renderer = args.renderer or "hybrid" + profile = build_motion_profile( + resolved.profile, + MotionProfileRequest( + checkpoint=resolved.checkpoint, + device=device, + configs=resolved.configs, + resource_root=( + None if args.resource_root is None else Path(args.resource_root) + ), + renderer=renderer, + ), + ) + for warning in profile.warnings: + print(f"Warning: {warning}", file=sys.stderr) + result = evaluate_motion_profile( + profile, + episodes=args.episodes if args.episodes is not None else 1, + viewer=args.viewer, + control_steps=args.control_steps, + duration=args.duration, + command=None if args.command is None else tuple(args.command), + scene_config=args.scene_config or "standard", + physics_backend=args.physics_backend, + simulation_device=args.sim_device or "cpu", + renderer=renderer, + gpu_id=args.gpu_id, + termination_behavior=args.termination_behavior, + cache_dir=args.cache_dir, + offline=args.offline, + ) + return write_evaluation_report( + _output_parent(args.output, resolved), + _profile_report(result, resolved, device), + ) + + +def _build_native_runtime( + args: argparse.Namespace, + resolved: EvaluationInput, + *, + viewer: bool, +) -> NativeRuntime: + train_config = resolved.configs.get("train") + if train_config is None: + raise ValueError("Training config is required for an EmbodiChain checkpoint") + config = load_config(train_config) + config["trainer"] = dict(config["trainer"]) + gym_config = resolved.configs.get("gym") + if gym_config is not None: + config["trainer"]["gym_config"] = str(gym_config) + trainer = config["trainer"] + device = _torch_device(args.device or trainer.get("device", "cpu")) + simulation_device = _simulation_device(args, device) + seed = int( + args.seed + if args.seed is not None + else trainer.get("eval_seed", int(trainer.get("seed", 1)) + 10_000) + ) + np.random.seed(seed) + torch.manual_seed(seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(seed) + uses_simulator = "gym_config" in trainer + if viewer and not uses_simulator: + raise ValueError("--viewer requires a simulator training task") + renderer = args.renderer or str(trainer.get("renderer", "hybrid")) + num_envs = ( + 1 + if viewer + else int( + args.num_envs + if args.num_envs is not None + else trainer.get("num_eval_envs", 4) + ) + ) + if uses_simulator: + runtime = build_gym_policy_runtime( + config, + device=device, + simulation_device=simulation_device, + num_envs=num_envs, + headless=not viewer, + renderer=renderer, + gpu_id=args.gpu_id, + config_dir=train_config.parent, + ) + else: + runtime = build_learning_policy_runtime( + config, + device=device, + num_envs=num_envs, + ) + try: + runtime.policy.load_state_dict(_load_policy_state_dict(resolved.checkpoint)) + except Exception: + runtime.close() + _flush_simulator(uses_simulator) + raise + return NativeRuntime( + runtime=runtime, + device=device, + simulation_device=simulation_device, + seed=seed, + renderer=renderer, + uses_simulator=uses_simulator, + trainer=trainer, + ) + + +def _validate_native_options(args: argparse.Namespace) -> None: + profile_options = { + "--resource-root": args.resource_root, + "--command": args.command, + "--physics-backend": args.physics_backend, + "--scene-config": args.scene_config, + "--cache-dir": args.cache_dir, + "--offline": args.offline, + } + selected = [ + name for name, value in profile_options.items() if value not in (None, False) + ] + if selected: + raise ValueError(f"{', '.join(selected)} requires --profile") + if not args.viewer and ( + args.control_steps is not None + or args.duration is not None + or args.termination_behavior is not None + ): + raise ValueError( + "--control-steps, --duration, and --termination-behavior require --viewer" + ) + + +def _load_policy_state_dict(checkpoint: Path) -> Mapping[str, Any]: + payload = torch.load(checkpoint, map_location="cpu", weights_only=True) + if not isinstance(payload, Mapping) or not isinstance( + payload.get("policy"), Mapping + ): + raise TypeError("Checkpoint must contain a 'policy' state mapping") + return payload["policy"] + + +def _torch_device(value: str) -> torch.device: + device = torch.device(value) + if device.type == "cuda": + index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + torch.cuda.set_device(index) + return torch.device(f"cuda:{index}") + if device.type != "cpu": + raise ValueError(f"Unsupported device type: {device.type}") + return device + + +def _simulation_device( + args: argparse.Namespace, + inference_device: torch.device, +) -> torch.device: + if args.sim_device == "gpu": + return _torch_device(f"cuda:{args.gpu_id}") + if args.sim_device == "cpu": + return torch.device("cpu") + return inference_device + + +def _flush_simulator(enabled: bool) -> None: + if enabled: + from embodichain.lab.sim.sim_manager import SimulationManager + + SimulationManager.flush_cleanup_queue() + + +def _output_parent(configured: str | None, resolved: EvaluationInput) -> Path: + if configured is not None: + return Path(configured) + if resolved.run is not None: + return resolved.run / "evaluations" + return resolved.checkpoint.parent / "evaluations" + + +def _checkpoint_inputs(resolved: EvaluationInput) -> dict[str, Any]: + return { + "run": resolved.run, + "checkpoint": { + "path": resolved.checkpoint, + "requested": resolved.requested_checkpoint, + "selected": resolved.selected_checkpoint, + }, + "configs": resolved.configs, + } + + +def _headless_report( + native: NativeRuntime, + resolved: EvaluationInput, + episodes: int, + metrics: Mapping[str, float], +) -> dict[str, Any]: + return { + "mode": "headless", + "inputs": { + **_checkpoint_inputs(resolved), + "task_id": native.runtime.env_id, + "seed": native.seed, + "num_envs": int(native.runtime.env.num_envs), + "device": str(native.device), + "embodichain_version": __version__, + }, + "result": {"episodes": episodes, "metrics": metrics}, + } + + +def _viewer_report( + result: Any, + native: NativeRuntime, + resolved: EvaluationInput, +) -> dict[str, Any]: + import dexsim + + return { + "mode": "viewer", + "inputs": { + **_checkpoint_inputs(resolved), + "task_id": result.task_id, + "seed": native.seed, + "inference_device": str(native.device), + "simulation_device": str(native.simulation_device), + "renderer": native.renderer, + "embodichain_version": __version__, + "dexsim_version": getattr(dexsim, "__version__", None), + "dexsim_commit": getattr(dexsim, "__commit_id__", None), + }, + "result": { + "reason": result.reason, + "simulation_time": result.simulation_time, + "simulation_steps": result.simulation_steps, + "control_steps": result.control_steps, + "requested_duration": result.requested_duration, + "effective_duration": result.effective_duration, + "episodes": result.episodes, + "metrics": result.metrics, + }, + } + + +def _profile_report( + result: Any, + resolved: EvaluationInput, + device: torch.device, +) -> dict[str, Any]: + import dexsim + + return { + "mode": "viewer" if result.viewer else "headless", + "inputs": { + **_checkpoint_inputs(resolved), + "profile": { + "id": result.profile.profile_id, + "provider_version": result.profile.provider_version, + "provenance": result.profile.provenance, + "warnings": result.profile.warnings, + }, + "policy_spec": result.policy_spec, + "scene_config": result.scene_config, + "inference_device": str(device), + "embodichain_version": __version__, + "dexsim_version": getattr(dexsim, "__version__", None), + "dexsim_commit": getattr(dexsim, "__commit_id__", None), + }, + "result": { + "episodes": result.episodes, + "summary": result.summary, + }, + } diff --git a/embodichain/learning/rl/policy_evaluation/manifest.py b/embodichain/learning/rl/policy_evaluation/manifest.py new file mode 100644 index 000000000..cc8baa75d --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/manifest.py @@ -0,0 +1,193 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Index a training run for standalone policy evaluation.""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +__all__ = ["RUN_MANIFEST_NAME", "RunManifest", "write_run_manifest"] + +RUN_MANIFEST_NAME = "run-manifest.json" + + +@dataclass(frozen=True) +class RunManifest: + """Resolved paths from one EmbodiChain training run.""" + + root: Path + configs: Mapping[str, Path] + checkpoints: Mapping[str, Path | None] + + def __post_init__(self) -> None: + object.__setattr__(self, "root", Path(self.root).resolve()) + object.__setattr__(self, "configs", dict(self.configs)) + object.__setattr__(self, "checkpoints", dict(self.checkpoints)) + + @classmethod + def load(cls, run: str | Path) -> RunManifest: + """Load ``run-manifest.json`` and resolve its referenced files. + + Args: + run: EmbodiChain training run directory. + + Returns: + Resolved manifest. + """ + root = Path(run).expanduser().resolve() + path = root / RUN_MANIFEST_NAME + if not path.is_file(): + raise FileNotFoundError(f"Run manifest does not exist: {path}") + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, Mapping) or value.get("schema_version") != 1: + raise ValueError(f"Unsupported run manifest: {path}") + configs = _resolve_group(root, value.get("configs"), "configs") + checkpoints = _resolve_group( + root, + value.get("checkpoints"), + "checkpoints", + allow_none=True, + ) + return cls(root, configs, checkpoints) + + def select_checkpoint(self, requested: str = "latest") -> tuple[str, Path]: + """Select ``best`` or ``latest`` and return its resolved path. + + Args: + requested: Checkpoint role. + + Returns: + Selected role and checkpoint path. ``best`` uses ``latest`` when + the training run has no best checkpoint. + """ + if requested not in {"best", "latest"}: + raise ValueError("checkpoint role must be best or latest") + selected = requested + checkpoint = self.checkpoints.get(selected) + if checkpoint is None and requested == "best": + selected = "latest" + checkpoint = self.checkpoints.get(selected) + if checkpoint is None: + raise FileNotFoundError( + f"Run manifest has no {requested} checkpoint: {self.root}" + ) + return selected, checkpoint + + +def write_run_manifest( + run: str | Path, + *, + train_config: str | Path, + latest_checkpoint: str | Path, + best_checkpoint: str | Path | None = None, + gym_config: str | Path | None = None, +) -> Path: + """Snapshot training configs and write the minimal run manifest. + + Args: + run: Training run directory containing the checkpoints. + train_config: Training config used for the run. + latest_checkpoint: Final saved checkpoint. + best_checkpoint: Best checkpoint when evaluation selected one. + gym_config: Referenced task config when the trainer uses one. + + Returns: + Written manifest path. + """ + root = Path(run).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + config_dir = root / "configs" + config_dir.mkdir(exist_ok=True) + configs = { + "train": _snapshot_config(train_config, config_dir, "train"), + } + if gym_config is not None: + configs["gym"] = _snapshot_config(gym_config, config_dir, "gym") + checkpoints = { + "best": _relative_file(root, best_checkpoint), + "latest": _relative_file(root, latest_checkpoint), + } + value: dict[str, Any] = { + "schema_version": 1, + "configs": configs, + "checkpoints": checkpoints, + } + path = root / RUN_MANIFEST_NAME + path.write_text( + json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return path + + +def _snapshot_config(source: str | Path, target: Path, name: str) -> str: + path = Path(source).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Training config does not exist: {path}") + suffix = path.suffix.lower() if path.suffix else ".yaml" + destination = target / f"{name}{suffix}" + shutil.copyfile(path, destination) + return destination.relative_to(target.parent).as_posix() + + +def _relative_file(root: Path, value: str | Path | None) -> str | None: + if value is None: + return None + path = Path(value).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Training checkpoint does not exist: {path}") + try: + return path.relative_to(root).as_posix() + except ValueError as error: + raise ValueError(f"Training checkpoint is outside its run: {path}") from error + + +def _resolve_group( + root: Path, + value: object, + field: str, + *, + allow_none: bool = False, +) -> dict[str, Path | None]: + if not isinstance(value, Mapping): + raise TypeError(f"Run manifest {field} must be a mapping") + result: dict[str, Path | None] = {} + for name, reference in value.items(): + if reference is None and allow_none: + result[str(name)] = None + continue + if not isinstance(reference, str) or not reference: + raise TypeError(f"Run manifest {field}.{name} must be a path") + relative = Path(reference) + if relative.is_absolute(): + raise ValueError(f"Run manifest {field}.{name} must be relative") + path = (root / relative).resolve() + try: + path.relative_to(root) + except ValueError as error: + raise ValueError( + f"Run manifest {field}.{name} escapes the run directory" + ) from error + if not path.is_file(): + raise FileNotFoundError(f"Run manifest file does not exist: {path}") + result[str(name)] = path + return result diff --git a/embodichain/learning/rl/policy_evaluation/profile.py b/embodichain/learning/rl/policy_evaluation/profile.py new file mode 100644 index 000000000..a215a084b --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/profile.py @@ -0,0 +1,128 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""External Policy Profile registration and construction.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch + +__all__ = [ + "MotionProfile", + "MotionProfileRequest", + "build_motion_profile", + "register_motion_profile", +] + + +@dataclass(frozen=True) +class MotionProfileRequest: + """Checkpoint, configs, and runtime choices supplied to a provider.""" + + checkpoint: Path + device: torch.device + configs: Mapping[str, Path] = field(default_factory=dict) + resource_root: Path | None = None + renderer: str = "hybrid" + + def __post_init__(self) -> None: + checkpoint = Path(self.checkpoint).expanduser().resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"Motion checkpoint does not exist: {checkpoint}") + configs = { + name: Path(path).expanduser().resolve() + for name, path in self.configs.items() + } + for name, path in configs.items(): + if not path.is_file(): + raise FileNotFoundError( + f"Motion config {name!r} does not exist: {path}" + ) + root = ( + None + if self.resource_root is None + else Path(self.resource_root).expanduser().resolve() + ) + object.__setattr__(self, "checkpoint", checkpoint) + object.__setattr__(self, "configs", configs) + object.__setattr__(self, "resource_root", root) + + +@dataclass(frozen=True) +class MotionProfile: + """DexSim Policy Spec and report metadata built by one provider.""" + + profile_id: str + policy_spec: Mapping[str, Any] + provider_version: int = 1 + provenance: Mapping[str, Any] = field(default_factory=dict) + warnings: tuple[str, ...] = () + + def __post_init__(self) -> None: + object.__setattr__(self, "policy_spec", dict(self.policy_spec)) + object.__setattr__(self, "provenance", dict(self.provenance)) + object.__setattr__(self, "warnings", tuple(self.warnings)) + + +MotionProfileProvider = Callable[[MotionProfileRequest], MotionProfile] +_PROVIDERS: dict[str, MotionProfileProvider] = {} + + +def register_motion_profile(name: str, provider: MotionProfileProvider) -> None: + """Register a Motion Profile provider under its CLI name. + + Args: + name: Stable profile name. + provider: Callable that builds one :class:`MotionProfile`. + """ + if not name: + raise ValueError("Motion profile name must not be empty") + if name in _PROVIDERS: + raise ValueError(f"Motion profile is already registered: {name}") + _PROVIDERS[name] = provider + + +def build_motion_profile( + name: str, + request: MotionProfileRequest, +) -> MotionProfile: + """Build one profile with its registered provider. + + Args: + name: Registered profile name. + request: Checkpoint, configs, and runtime choices. + + Returns: + Provider-built Motion Profile. + """ + try: + provider = _PROVIDERS[name] + except KeyError: + available = ", ".join(sorted(_PROVIDERS)) or "none" + raise ValueError( + f"Unknown motion profile {name!r}; available: {available}" + ) from None + profile = provider(request) + if profile.profile_id != name: + raise ValueError( + f"Motion provider {name!r} returned profile {profile.profile_id!r}" + ) + return profile diff --git a/embodichain/learning/rl/policy_evaluation/report.py b/embodichain/learning/rl/policy_evaluation/report.py new file mode 100644 index 000000000..545b41f2e --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/report.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Write timestamped policy evaluation reports.""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +__all__ = ["write_evaluation_report"] + + +def write_evaluation_report( + parent: str | Path, + payload: Mapping[str, Any], +) -> Path: + """Write ``evaluation.json`` under a new timestamped directory. + + Args: + parent: Output parent directory. + payload: Evaluation inputs and results. + + Returns: + Written report path. + """ + output = Path(parent).expanduser().resolve() + output.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + directory = output / f"{stamp}-policy" + directory.mkdir() + report = { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + **dict(payload), + } + path = directory / "evaluation.json" + path.write_text( + json.dumps( + _json_value(report), + indent=2, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + ) + + "\n", + encoding="utf-8", + ) + return path + + +def _json_value(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, Mapping): + return {str(name): _json_value(item) for name, item in value.items()} + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [_json_value(item) for item in value] + return value diff --git a/embodichain/learning/rl/policy_evaluation/viewer.py b/embodichain/learning/rl/policy_evaluation/viewer.py new file mode 100644 index 000000000..509176684 --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/viewer.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. +# ---------------------------------------------------------------------------- + +"""Connect an EmbodiChain task Viewer to Motion Policy Evaluator.""" + +from __future__ import annotations + +import math +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import torch + +from dexsim.kit.motion_policy import ( + EvaluationFrame, + PolicyContext, + PolicyOutput, + RunOptions, + create_motion_policy_evaluator, +) +from dexsim.kit.motion_policy.types import EnvironmentStep + +from embodichain.learning.rl.evaluation import ( + convert_policy_action_for_env, + infer_policy_action, +) +from embodichain.learning.rl.runtime import PolicyRuntime + +__all__ = [ + "EmbodiChainTaskEnvironment", + "EmbodiChainTaskPolicyAdapter", + "NativeViewerResult", + "evaluate_native_viewer", +] + +_MISSING = object() + + +@dataclass(frozen=True) +class NativeViewerResult: + """Result of visualizing one Policy in its EmbodiChain task.""" + + task_id: str + reason: str + simulation_time: float + simulation_steps: int + control_steps: int + effective_duration: float + requested_duration: float | None + episodes: tuple[Mapping[str, float | int | bool | str], ...] + metrics: Mapping[str, float] + + +class EmbodiChainTaskPolicyAdapter: + """Run an EmbodiChain Policy from the task observation in each frame.""" + + def __init__(self, policy: torch.nn.Module, device: torch.device): + self.policy = policy + self.device = device + self._previous_training = policy.training + + def setup(self, context: PolicyContext) -> None: + """Select deterministic inference for this evaluation.""" + del context + self.policy.eval() + + def reset(self, frame: EvaluationFrame) -> None: + """Validate that the Environment supplied the next observation.""" + if frame.observation is None: + raise RuntimeError("EmbodiChain task frame has no observation") + + @torch.no_grad() + def infer(self, frame: EvaluationFrame) -> PolicyOutput: + """Run the same observation and deterministic Policy path as RL evaluation.""" + if frame.observation is None: + raise RuntimeError("EmbodiChain task frame has no observation") + action = infer_policy_action( + self.policy, + frame.observation, + device=self.device, + num_envs=1, + ) + return PolicyOutput(action=action) + + def metrics(self) -> dict[str, float]: + """Return Policy-side metrics.""" + return {} + + def close(self) -> None: + """Restore the Policy mode used before evaluation.""" + self.policy.train(self._previous_training) + + +class EmbodiChainTaskEnvironment: + """Expose one original EmbodiChain RL Environment to the Evaluator.""" + + def __init__( + self, + env: Any, + *, + seed: int, + ) -> None: + if int(env.num_envs) != 1: + raise ValueError("Visual task evaluation requires num_envs=1") + self.env = env + self._base_env = getattr(env, "unwrapped", env) + world = self._world() + if world is None or not world.is_window_initialized(): + raise ValueError( + "Viewer evaluation requires an initialized simulator window" + ) + self._seed = seed + self._first_reset = True + self._reset_key_down = False + self._control_step = 0 + self._frame: EvaluationFrame | None = None + self._episode_return = 0.0 + self._episode_length = 0 + self._episodes: list[dict[str, float | int | bool | str]] = [] + self._reported_metrics: dict[str, float] = {} + self._closed = False + self._policy_context = _policy_context_from_env(self._base_env) + self._previous_no_auto_reset = getattr( + self._base_env, + "_demo_no_auto_reset", + _MISSING, + ) + self._base_env._demo_no_auto_reset = True + + @property + def policy_context(self) -> PolicyContext: + """Return the timing used by the original task Environment.""" + return self._policy_context + + @property + def physics_backend(self) -> str: + """Return the backend selected by the original task Environment.""" + return "default" + + @property + def viewer_is_open(self) -> bool: + """Return whether the original task Viewer remains open.""" + world = self._world() + return bool(world is not None and world.is_window_initialized()) + + @property + def current_frame(self) -> EvaluationFrame: + """Return the latest observation and task state.""" + if self._frame is None: + raise RuntimeError("Environment has not been reset") + return self._frame + + @property + def episodes(self) -> tuple[Mapping[str, float | int | bool | str], ...]: + """Return completed episode summaries.""" + return tuple(self._episodes) + + def open_viewer(self, title: str) -> None: + """Apply the evaluation title to the task Viewer.""" + self._world().get_windows().set_window_title(title) + + def reset(self) -> EvaluationFrame: + """Run the task's original reset and return its observation.""" + kwargs = {"seed": self._seed} if self._first_reset else {} + observation, info = self.env.reset(**kwargs) + self._first_reset = False + self._control_step = 0 + self._episode_return = 0.0 + self._episode_length = 0 + self._frame = self._make_frame(observation, {"info": info}) + return self._frame + + def poll(self) -> str | None: + """Report when the native Viewer is closed or Escape is pressed.""" + world = self._world() + if world is None or not world.is_window_initialized(): + return "viewer closed" + from dexsim.types import InputKey + + native = world.get_windows().native() + if native.key_state(InputKey.SCANCODE_ESCAPE): + return "viewer closed" + reset_down = bool(native.key_state(InputKey.SCANCODE_BACKSPACE)) + reset_pressed = reset_down and not self._reset_key_down + self._reset_key_down = reset_down + if reset_pressed: + return "manual reset" + return None + + def step(self, action: object) -> EnvironmentStep: + """Apply one raw Policy action through the task's original action path.""" + if not isinstance(action, torch.Tensor): + raise TypeError("EmbodiChain Policy action must be a torch.Tensor") + started = time.perf_counter() + env_action = convert_policy_action_for_env(self.env, action) + observation, reward, terminated, truncated, info = self.env.step(env_action) + reward_value = _single_float(reward, "reward") + terminated_value = _single_bool(terminated, "terminated") + truncated_value = _single_bool(truncated, "truncated") + self._control_step += 1 + self._episode_return += reward_value + self._episode_length += 1 + task_state = { + "reward": reward, + "terminated": terminated, + "truncated": truncated, + "info": info, + } + self._frame = self._make_frame(observation, task_state) + reason = _termination_reason(info, terminated_value, truncated_value) + metrics = _step_metrics(info, reward_value) + self._reported_metrics.update(metrics) + if reason is not None: + success = _info_bool(info, "success") + self._episodes.append( + { + "index": len(self._episodes), + "reason": reason, + "reward": self._episode_return, + "length": self._episode_length, + "success": success, + } + ) + remaining = self._policy_context.policy_dt - (time.perf_counter() - started) + if remaining > 0.0: + time.sleep(remaining) + return EnvironmentStep( + frame=self._frame, + termination_reason=reason, + metrics=metrics, + ) + + def metrics(self) -> dict[str, float]: + """Return task metrics and completed episode aggregates.""" + result = dict(self._reported_metrics) + if self._episodes: + count = len(self._episodes) + result.update( + { + "eval/avg_reward": sum( + float(episode["reward"]) for episode in self._episodes + ) + / count, + "eval/avg_length": sum( + float(episode["length"]) for episode in self._episodes + ) + / count, + "eval/success_rate": sum( + bool(episode["success"]) for episode in self._episodes + ) + / count, + } + ) + return result + + def wait_for_reset_or_close(self) -> str: + """Keep a paused Viewer responsive until it is closed. + + ``MotionPolicyEvaluator`` calls this method after a task termination + when the selected behavior is ``pause``. + """ + while self.viewer_is_open: + event = self.poll() + if event is not None: + return event + world = self._world() + if world is not None: + world.update(0.0) + time.sleep(0.01) + return "viewer closed" + + def close(self) -> None: + """Close the original task Environment.""" + if self._closed: + return + if self._previous_no_auto_reset is _MISSING: + delattr(self._base_env, "_demo_no_auto_reset") + else: + self._base_env._demo_no_auto_reset = self._previous_no_auto_reset + if getattr(self._base_env, "sim", None) is not None: + self._base_env.close(exit_process=False) + else: + self.env.close() + self._closed = True + + def _make_frame( + self, + observation: object, + task_state: Mapping[str, object], + ) -> EvaluationFrame: + simulation_step = ( + self._control_step * self._policy_context.sim_steps_per_control + ) + return EvaluationFrame( + control_step=self._control_step, + policy_time=self._control_step * self._policy_context.policy_dt, + simulation_step=simulation_step, + simulation_time=simulation_step * self._policy_context.physics_dt, + observation=observation, + task_state=task_state, + ) + + def _world(self) -> Any | None: + sim = getattr(self._base_env, "sim", None) + return None if sim is None else sim.get_world() + + +def evaluate_native_viewer( + runtime: PolicyRuntime, + *, + seed: int, + episodes: int | None, + control_steps: int | None, + duration: float | None, + termination_behavior: str = "auto_reset", +) -> NativeViewerResult: + """Visualize an EmbodiChain Policy in the task used for training.""" + if episodes is not None and episodes <= 0: + raise ValueError("episodes must be positive") + if control_steps is not None and control_steps <= 0: + raise ValueError("control_steps must be positive") + if duration is not None and (duration <= 0.0 or not math.isfinite(duration)): + raise ValueError("duration must be finite and positive") + if control_steps is not None and duration is not None: + raise ValueError("control_steps and duration are mutually exclusive") + if termination_behavior == "continue": + raise ValueError("Native task evaluation supports pause or auto_reset") + + environment = None + adapter = None + evaluator = None + try: + environment = EmbodiChainTaskEnvironment( + runtime.env, + seed=seed, + ) + adapter = EmbodiChainTaskPolicyAdapter( + runtime.policy, + runtime.device, + ) + if duration is not None: + control_steps = math.ceil( + duration / environment.policy_context.policy_dt - 1e-12 + ) + total_steps = 0 + reason = "viewer closed" + options = RunOptions( + headless=False, + termination_behavior=( + "continue" if termination_behavior == "auto_reset" else "pause" + ), + ) + evaluator = create_motion_policy_evaluator( + options=options, + adapter=adapter, + environment=environment, + title=f"{runtime.env_id} - EmbodiChain", + ) + evaluator.reset() + while True: + if control_steps is not None and total_steps >= control_steps: + reason = "control steps reached" + break + if episodes is not None and len(environment.episodes) >= episodes: + reason = "episode target reached" + break + completed_before = len(environment.episodes) + result = evaluator.step() + if result.advanced: + total_steps += 1 + if len(environment.episodes) > completed_before: + if episodes is not None and len(environment.episodes) >= episodes: + reason = "episode target reached" + break + if termination_behavior == "auto_reset": + evaluator.reset() + continue + if result.reason is not None and not result.reset_performed: + reason = result.reason + break + episode_results = environment.episodes + metrics = environment.metrics() + context = environment.policy_context + finally: + if evaluator is not None: + evaluator.close() + elif environment is not None: + if adapter is not None: + adapter.close() + environment.close() + else: + runtime.close() + + simulation_steps = total_steps * context.sim_steps_per_control + return NativeViewerResult( + task_id=runtime.env_id, + reason=reason, + simulation_time=simulation_steps * context.physics_dt, + simulation_steps=simulation_steps, + control_steps=total_steps, + effective_duration=total_steps * context.policy_dt, + requested_duration=duration, + episodes=episode_results, + metrics=metrics, + ) + + +def _single_float(value: object, name: str) -> float: + tensor = torch.as_tensor(value).reshape(-1) + if tensor.numel() != 1: + raise ValueError(f"Native task {name} must contain one value") + return float(tensor.item()) + + +def _single_bool(value: object, name: str) -> bool: + tensor = torch.as_tensor(value, dtype=torch.bool).reshape(-1) + if tensor.numel() != 1: + raise ValueError(f"Native task {name} must contain one value") + return bool(tensor.item()) + + +def _info_bool(info: object, name: str) -> bool: + if not isinstance(info, Mapping) or name not in info: + return False + return _single_bool(info[name], f"info.{name}") + + +def _termination_reason( + info: object, + terminated: bool, + truncated: bool, +) -> str | None: + if _info_bool(info, "success"): + return "success" + if _info_bool(info, "fail"): + return "failure" + if truncated: + return "time limit" + if terminated: + return "terminated" + return None + + +def _step_metrics(info: object, reward: float) -> dict[str, float]: + result = {"reward": reward} + if not isinstance(info, Mapping): + return result + metrics = info.get("metrics") + if not isinstance(metrics, Mapping): + return result + for name, value in metrics.items(): + tensor = torch.as_tensor(value).reshape(-1) + if tensor.numel() == 1: + result[str(name)] = float(tensor.item()) + return result + + +def _policy_context_from_env(env: Any) -> PolicyContext: + """Read timing from the simulator task.""" + return PolicyContext( + robot=None, + physics_dt=float(env.physics_dt), + sim_steps_per_control=int(env.cfg.sim_steps_per_control), + policy_dt=float(env.step_dt), + ) diff --git a/embodichain/learning/rl/runtime.py b/embodichain/learning/rl/runtime.py new file mode 100644 index 000000000..ed8f0c7f7 --- /dev/null +++ b/embodichain/learning/rl/runtime.py @@ -0,0 +1,312 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Shared environment and Policy construction for RL training and evaluation.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules +from embodichain.lab.gym.utils.profiler import EnvProfilerCfg +from embodichain.lab.gym.utils.registration import build_env +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import RenderCfg +from embodichain.learning.rl.env import build_learning_env +from embodichain.learning.rl.models import build_mlp_from_cfg, build_policy +from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation +from embodichain.utils.utility import load_config + +__all__ = [ + "PolicyRuntime", + "build_gym_policy_runtime", + "build_learning_policy_runtime", +] + + +@dataclass(frozen=True) +class _GymEnvironmentRuntime: + """A simulator task reconstructed from one training configuration.""" + + env: Any + env_id: str + env_cfg: Any + gym_config: dict[str, Any] + gym_config_path: Path + + +@dataclass(frozen=True) +class PolicyRuntime: + """An Environment and Policy reconstructed from one training configuration.""" + + env: Any + policy: torch.nn.Module + device: torch.device + env_id: str + env_cfg: Any | None = None + gym_config: dict[str, Any] | None = None + gym_config_path: Path | None = None + + def close(self) -> None: + """Close the Environment without terminating the current process.""" + _close_environment(self.env) + + +def _resolve_config_reference( + value: str | Path, + *, + base_dir: str | Path | None = None, +) -> Path: + """Resolve a referenced config relative to its containing config file.""" + path = Path(value).expanduser() + if path.is_absolute(): + return path + if base_dir is not None: + candidate = Path(base_dir).expanduser().resolve() / path + if candidate.exists(): + return candidate + return path + + +def _build_learning_environment( + config: dict[str, Any], + *, + device: torch.device, + num_envs: int, +) -> tuple[str, Any]: + """Build the lightweight Environment declared by a training config.""" + env_block = config["trainer"]["learning_env"] + if isinstance(env_block, str): + env_name = env_block + env_config: dict[str, Any] = {} + else: + env_name = env_block["name"] + env_config = dict(env_block.get("cfg", {})) + return str(env_name), build_learning_env( + str(env_name), + num_envs=num_envs, + device=device, + **env_config, + ) + + +def build_learning_policy_runtime( + config: dict[str, Any], + *, + device: torch.device, + num_envs: int, +) -> PolicyRuntime: + """Build a lightweight Environment and its configured Policy.""" + env_name, env = _build_learning_environment( + config, + device=device, + num_envs=num_envs, + ) + try: + policy = _build_learning_policy(config["policy"], env, device) + except Exception: + env.close() + raise + return PolicyRuntime(env, policy, device, env_name) + + +def _build_gym_environment( + config: dict[str, Any], + *, + simulation_device: torch.device, + num_envs: int | None, + headless: bool, + renderer: str, + gpu_id: int, + config_dir: str | Path | None = None, + profiler: EnvProfilerCfg | None = None, +) -> _GymEnvironmentRuntime: + """Build the simulator Environment declared by a training config.""" + trainer_cfg = config["trainer"] + gym_config_path = _resolve_config_reference( + trainer_cfg["gym_config"], + base_dir=config_dir, + ) + gym_config = load_config(gym_config_path) + env_cfg = config_to_cfg(gym_config, manager_modules=get_manager_modules()) + if num_envs is not None: + env_cfg.num_envs = int(num_envs) + if env_cfg.sim_cfg is None: + env_cfg.sim_cfg = SimulationManagerCfg() + env_cfg.sim_cfg.sim_device = simulation_device + env_cfg.sim_cfg.headless = headless + env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) + env_cfg.sim_cfg.gpu_id = ( + simulation_device.index + if simulation_device.type == "cuda" and simulation_device.index is not None + else gpu_id + ) + env_cfg.profiler = profiler + env = build_env(gym_config["id"], base_env_cfg=env_cfg) + return _GymEnvironmentRuntime( + env=env, + env_id=str(gym_config["id"]), + env_cfg=env_cfg, + gym_config=gym_config, + gym_config_path=gym_config_path.resolve(), + ) + + +def build_gym_policy_runtime( + config: dict[str, Any], + *, + device: torch.device, + num_envs: int | None, + headless: bool, + renderer: str, + gpu_id: int, + config_dir: str | Path | None = None, + profiler: EnvProfilerCfg | None = None, + simulation_device: torch.device | None = None, +) -> PolicyRuntime: + """Build a simulator task and the Policy declared by its training config.""" + task = _build_gym_environment( + config, + simulation_device=simulation_device or device, + num_envs=num_envs, + headless=headless, + renderer=renderer, + gpu_id=gpu_id, + config_dir=config_dir, + profiler=profiler, + ) + env = task.env + try: + sample_observation, _ = env.reset() + sample_observation_td = dict_to_tensordict(sample_observation, device) + observation_dim = int(flatten_dict_observation(sample_observation_td).shape[-1]) + action_manager = env.get_wrapper_attr("action_manager") + environment_action_dim = ( + action_manager.total_action_dim + if action_manager is not None + else len(env.get_wrapper_attr("active_joint_ids")) + ) + policy = _build_gym_policy( + config["policy"], + env=env, + device=device, + observation_dim=observation_dim, + action_dim=environment_action_dim, + ) + except Exception: + _close_environment(env) + raise + return PolicyRuntime( + env=env, + policy=policy, + device=device, + env_id=task.env_id, + env_cfg=task.env_cfg, + gym_config=task.gym_config, + gym_config_path=task.gym_config_path, + ) + + +def _build_gym_policy( + policy_block: dict[str, Any], + *, + env: Any, + device: torch.device, + observation_dim: int, + action_dim: int, +) -> torch.nn.Module: + configured_action_dim = int(policy_block.get("action_dim", action_dim)) + if configured_action_dim != action_dim: + raise ValueError( + f"Configured policy.action_dim={configured_action_dim} does not match " + f"env action dim {action_dim}." + ) + policy_name = str(policy_block["name"]).lower() + if policy_name == "actor_critic": + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + if actor_cfg is None or critic_cfg is None: + raise ValueError( + "ActorCritic requires policy.actor and policy.critic definitions." + ) + return build_policy( + policy_block, + env.flattened_observation_space, + env.action_space, + device, + actor=build_mlp_from_cfg(actor_cfg, observation_dim, action_dim), + critic=build_mlp_from_cfg(critic_cfg, observation_dim, 1), + ) + if policy_name == "actor_only": + actor_cfg = policy_block.get("actor") + if actor_cfg is None: + raise ValueError("ActorOnly requires a policy.actor definition.") + return build_policy( + policy_block, + env.flattened_observation_space, + env.action_space, + device, + actor=build_mlp_from_cfg(actor_cfg, observation_dim, action_dim), + ) + return build_policy( + policy_block, + env.observation_space, + env.action_space, + device, + ) + + +def _build_learning_policy( + policy_block: dict[str, Any], + env: Any, + device: torch.device, +) -> torch.nn.Module: + observation_dim = int(env.single_observation_space.shape[-1]) + action_dim = int(env.single_action_space.shape[-1]) + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + policy = build_policy( + policy_block, + env.single_observation_space, + env.single_action_space, + device, + actor=( + build_mlp_from_cfg(actor_cfg, observation_dim, action_dim) + if actor_cfg is not None + else None + ), + critic=( + build_mlp_from_cfg(critic_cfg, observation_dim, 1) + if critic_cfg is not None + else None + ), + ) + if "initial_log_std" in policy_block and hasattr(policy, "log_std"): + with torch.no_grad(): + policy.log_std.fill_(float(policy_block["initial_log_std"])) + return policy + + +def _close_environment(env: Any) -> None: + target = getattr(env, "unwrapped", env) + if getattr(target, "sim", None) is not None: + target.close(exit_process=False) + else: + env.close() diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 41c44a3f1..9556efd31 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -20,16 +20,18 @@ import os import time from collections.abc import Sequence +from copy import deepcopy from pathlib import Path import numpy as np import torch import wandb from torch.utils.tensorboard import SummaryWriter -from copy import deepcopy -from embodichain.learning.rl.models import build_policy, get_registered_policy_names -from embodichain.learning.rl.models import build_mlp_from_cfg +from embodichain.learning.rl.models import get_registered_policy_names +from embodichain.learning.rl.policy_evaluation.manifest import ( + write_run_manifest, +) from embodichain.learning.rl.algo import ( RolloutKind, build_algo, @@ -39,9 +41,12 @@ DifferentiableTrainer, DifferentiableTrainerCfg, ) -from embodichain.learning.rl.env import build_learning_env +from embodichain.learning.rl.runtime import ( + _build_learning_environment, + build_gym_policy_runtime, + build_learning_policy_runtime, +) from embodichain.learning.rl.routing import get_trainer_class -from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation from embodichain.learning.rl.utils.trainer import Trainer from embodichain.utils import logger from embodichain.lab.gym.utils.registration import ( @@ -49,14 +54,13 @@ discover_task_packages, execute_init_hooks, ) -from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules from embodichain.lab.gym.utils.profiler import EnvProfilerCfg from embodichain.utils.utility import load_config from embodichain.utils.module_utils import find_function_from_modules -from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.sim.cfg import RenderCfg from embodichain.lab.gym.envs.managers.cfg import EventCfg +_CAMERA_RECORDERS = {"record_camera_data", "record_camera_data_async"} + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: """Parse command-line arguments. @@ -114,51 +118,33 @@ def _resolve_profile_output( return str(output.with_name(f"{output.stem}_rank{rank}{output.suffix}")) -def _build_learning_policy( - policy_block: dict, - env, - device: torch.device, -): - obs_dim = int(env.single_observation_space.shape[-1]) - action_dim = int(env.single_action_space.shape[-1]) - policy_name = policy_block["name"].lower() - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - actor = ( - build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) - if actor_cfg is not None - else None - ) - critic = ( - build_mlp_from_cfg(critic_cfg, obs_dim, 1) if critic_cfg is not None else None - ) - policy = build_policy( - policy_block, - env.single_observation_space, - env.single_action_space, - device, - actor=actor, - critic=critic, - ) - if "initial_log_std" in policy_block and hasattr(policy, "log_std"): - with torch.no_grad(): - policy.log_std.fill_(float(policy_block["initial_log_std"])) - return policy +def _event_params( + event_info: dict, + *, + run_base: str | Path, + phase: str, +) -> dict: + """Place default camera recordings under the current training run.""" + params = dict(event_info.get("params", {})) + function_name = str(event_info.get("func", "")).rsplit(".", 1)[-1] + if function_name in _CAMERA_RECORDERS: + params.setdefault("save_path", str(Path(run_base) / "videos" / phase)) + return params def _train_learning_env( cfg_data: dict, *, + config_path: str | Path, distributed: bool | None, profile: bool = False, -): +) -> dict[str, object]: """Train a lightweight registered environment through the unified CLI.""" if profile: raise ValueError( "--profile requires trainer.gym_config; learning_env is unsupported." ) trainer_cfg = cfg_data["trainer"] - policy_block = cfg_data["policy"] algorithm_block = cfg_data["algorithm"] distributed = ( bool(trainer_cfg.get("distributed", False)) @@ -182,32 +168,25 @@ def _train_learning_env( np.random.seed(seed) torch.manual_seed(seed) - env_block = trainer_cfg["learning_env"] - if isinstance(env_block, str): - env_name = env_block - env_cfg = {} - else: - env_name = env_block["name"] - env_cfg = dict(env_block.get("cfg", {})) num_envs = int(trainer_cfg.get("num_envs", 64)) - env = build_learning_env( - env_name, + runtime = build_learning_policy_runtime( + cfg_data, num_envs=num_envs, device=device, - **env_cfg, ) + env = runtime.env + policy = runtime.policy + env_name = runtime.env_id enable_eval = bool(trainer_cfg.get("enable_eval", False)) eval_env = None if enable_eval: - eval_env = build_learning_env( - env_name, + _eval_name, eval_env = _build_learning_environment( + cfg_data, num_envs=int(trainer_cfg.get("num_eval_envs", 16)), device=device, - **env_cfg, ) - policy = _build_learning_policy(policy_block, env, device) algorithm = build_algo( algorithm_block["name"], dict(algorithm_block.get("cfg", {})), @@ -291,7 +270,7 @@ def _train_learning_env( total_timesteps = int(trainer_cfg.get("total_timesteps", default_steps)) trainer.train(total_timesteps) trainer.save_checkpoint() - return trainer.get_summary() + summary = trainer.get_summary() finally: writer.close() if use_wandb: @@ -299,6 +278,8 @@ def _train_learning_env( env.close() if eval_env is not None: eval_env.close() + _write_policy_run_manifest(run_base, config_path, summary) + return summary def train_from_config( @@ -307,7 +288,7 @@ def train_from_config( *, profile: bool = False, profile_output: str | None = None, -): +) -> dict[str, object] | None: """Run training from a config file path. Args: @@ -316,6 +297,9 @@ def train_from_config( If None, use trainer.distributed from config. profile: Enable gym ``EnvProfiler`` on the training environment. profile_output: Optional JSON dump path for the profiling report. + + Returns: + The lightweight trainer summary, or ``None`` for simulator training. """ if profile_output is not None and not profile: raise ValueError("--profile_output requires --profile.") @@ -326,6 +310,7 @@ def train_from_config( if "learning_env" in trainer_cfg: return _train_learning_env( cfg_data, + config_path=config_path, distributed=distributed, profile=profile, ) @@ -441,32 +426,11 @@ def train_from_config( if use_wandb and rank == 0: wandb.init(project=wandb_project_name, name=exp_name, config=cfg_data) - gym_config_path = Path(trainer_cfg["gym_config"]) if rank == 0: logger.log_info(f"Current working directory: {Path.cwd()}") - gym_config_data = load_config(str(gym_config_path)) - gym_env_cfg = config_to_cfg(gym_config_data, manager_modules=get_manager_modules()) - if num_envs is not None: - gym_env_cfg.num_envs = int(num_envs) - - # Ensure sim configuration mirrors runtime overrides - if gym_env_cfg.sim_cfg is None: - gym_env_cfg.sim_cfg = SimulationManagerCfg() - if device.type == "cuda": - gpu_index = device.index - if gpu_index is None: - gpu_index = torch.cuda.current_device() - gym_env_cfg.sim_cfg.sim_device = torch.device(f"cuda:{gpu_index}") - if hasattr(gym_env_cfg.sim_cfg, "gpu_id"): - gym_env_cfg.sim_cfg.gpu_id = gpu_index - else: - gym_env_cfg.sim_cfg.sim_device = torch.device("cpu") - gym_env_cfg.sim_cfg.headless = headless - gym_env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) - gym_env_cfg.sim_cfg.gpu_id = gpu_id - if profile: - gym_env_cfg.profiler = EnvProfilerCfg( + profiler = ( + EnvProfilerCfg( enable_time=True, output_path=_resolve_profile_output( profile_output, @@ -474,23 +438,37 @@ def train_from_config( world_size=world_size, ), ) + if profile + else None + ) + runtime = build_gym_policy_runtime( + cfg_data, + device=device, + num_envs=num_envs, + headless=headless, + renderer=renderer, + gpu_id=gpu_id, + config_dir=Path(config_path).expanduser().resolve().parent, + profiler=profiler, + ) + env = runtime.env + policy = runtime.policy + gym_config_path = runtime.gym_config_path + gym_config_data = runtime.gym_config + gym_env_cfg = runtime.env_cfg + if gym_config_path is None or gym_config_data is None or gym_env_cfg is None: + raise RuntimeError("Simulator Policy runtime is missing task configuration") if rank == 0: logger.log_info( f"Loaded gym_config from {gym_config_path} (env_id={gym_config_data['id']}, num_envs={gym_env_cfg.num_envs}, headless={gym_env_cfg.sim_cfg.headless}, renderer={gym_env_cfg.sim_cfg.render_cfg.renderer}, sim_device={gym_env_cfg.sim_cfg.sim_device})" ) - env = build_env(gym_config_data["id"], base_env_cfg=gym_env_cfg) - sample_obs, _ = env.reset() - sample_obs_td = dict_to_tensordict(sample_obs, device) - obs_dim = flatten_dict_observation(sample_obs_td).shape[-1] - flat_obs_space = env.flattened_observation_space - # Create evaluation environment only if enabled eval_env = None num_eval_envs = trainer_cfg.get("num_eval_envs", 4) if enable_eval and rank == 0: eval_gym_env_cfg = deepcopy(gym_env_cfg) - eval_gym_env_cfg.num_envs = num_eval_envs + eval_gym_env_cfg.num_envs = int(num_eval_envs) eval_gym_env_cfg.sim_cfg.headless = True eval_gym_env_cfg.profiler = None eval_env = build_env(gym_config_data["id"], base_env_cfg=eval_gym_env_cfg) @@ -498,59 +476,7 @@ def train_from_config( f"Evaluation environment created (num_envs={num_eval_envs}, headless=True)" ) - # Build Policy via registry policy_name = policy_block["name"] - env_action_dim = ( - env.get_wrapper_attr("action_manager").total_action_dim - if env.get_wrapper_attr("action_manager") is not None - else len(env.get_wrapper_attr("active_joint_ids")) - ) - action_dim = policy_block.get("action_dim", env_action_dim) - action_dim = int(action_dim) - if action_dim != env_action_dim: - raise ValueError( - f"Configured policy.action_dim={action_dim} does not match env action dim {env_action_dim}." - ) - # Build Policy via registry (actor/critic must be explicitly defined in JSON when using actor_critic/actor_only) - if policy_name.lower() == "actor_critic": - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - if actor_cfg is None or critic_cfg is None: - raise ValueError( - "ActorCritic requires 'actor' and 'critic' definitions in JSON (policy.actor / policy.critic)." - ) - - actor = build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) - critic = build_mlp_from_cfg(critic_cfg, obs_dim, 1) - - policy = build_policy( - policy_block, - flat_obs_space, - env.action_space, - device, - actor=actor, - critic=critic, - ) - elif policy_name.lower() == "actor_only": - actor_cfg = policy_block.get("actor") - if actor_cfg is None: - raise ValueError( - "ActorOnly requires 'actor' definition in JSON (policy.actor)." - ) - - actor = build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) - - policy = build_policy( - policy_block, - flat_obs_space, - env.action_space, - device, - actor=actor, - ) - else: - policy = build_policy( - policy_block, env.observation_space, env.action_space, device - ) # Build Algorithm via factory algo_name = algo_block["name"].lower() @@ -581,7 +507,7 @@ def train_from_config( for event_name, event_info in events_dict.get("train", {}).items(): event_func_str = event_info.get("func") mode = event_info.get("mode", "interval") - params = event_info.get("params", {}) + params = _event_params(event_info, run_base=run_base, phase="train") interval_step = event_info.get("interval_step", 1) event_func = find_function_from_modules( event_func_str, event_modules, raise_if_not_found=True @@ -597,7 +523,7 @@ def train_from_config( for event_name, event_info in events_dict.get("eval", {}).items(): event_func_str = event_info.get("func") mode = event_info.get("mode", "interval") - params = event_info.get("params", {}) + params = _event_params(event_info, run_base=run_base, phase="eval") interval_step = event_info.get("interval_step", 1) event_func = find_function_from_modules( event_func_str, event_modules, raise_if_not_found=True @@ -648,6 +574,7 @@ def train_from_config( f"Total steps: {total_steps} (iterations≈{iterations}, world_size={world_size})" ) + summary = None try: trainer.train(total_steps) except KeyboardInterrupt: @@ -655,6 +582,8 @@ def train_from_config( logger.log_info("Training interrupted by user") finally: trainer.save_checkpoint() + if rank == 0: + summary = trainer.get_summary() if writer is not None: writer.close() if use_wandb and rank == 0: @@ -681,8 +610,35 @@ def train_from_config( if distributed and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() - if rank == 0: - logger.log_info("Training finished") + if summary is not None: + _write_policy_run_manifest( + run_base, + config_path, + summary, + gym_config=gym_config_path, + ) + if rank == 0: + logger.log_info("Training finished") + + +def _write_policy_run_manifest( + run_base: str | Path, + config_path: str | Path, + summary: dict, + *, + gym_config: str | Path | None = None, +) -> Path: + """Write the checkpoint and configuration index for policy evaluation.""" + latest = summary.get("latest_checkpoint_path") + if latest is None: + raise RuntimeError("Training finished without a checkpoint") + return write_run_manifest( + run_base, + train_config=config_path, + gym_config=gym_config, + latest_checkpoint=latest, + best_checkpoint=summary.get("best_checkpoint_path"), + ) def cli(argv: Sequence[str] | None = None) -> None: diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json index d44e26f67..f7d951e61 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json @@ -46,8 +46,7 @@ 600, 320, 240 - ], - "save_path": "./outputs/videos/eval" + ] } } } diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml index d64961242..c3c77b95a 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml @@ -40,7 +40,6 @@ trainer: - 600 - 320 - 240 - save_path: ./outputs/videos/eval renderer: hybrid policy: name: actor_only diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json index e6598e818..62fd4c5a6 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json @@ -45,8 +45,7 @@ 600, 320, 240 - ], - "save_path": "./outputs/videos/eval" + ] } } } diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml index 5b5935b3d..05f10ba96 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml @@ -39,7 +39,6 @@ trainer: - 600 - 320 - 240 - save_path: ./outputs/videos/eval renderer: fast-rt policy: name: actor_critic diff --git a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json index 21f6cd767..e697ce0d4 100644 --- a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json +++ b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json @@ -28,8 +28,7 @@ "eye": [-1.4, 1.4, 2.0], "target": [0, 0, 0], "up": [0, 0, 1], - "intrinsics": [600, 600, 320, 240], - "save_path": "./outputs/videos/eval" + "intrinsics": [600, 600, 320, 240] } } } diff --git a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json index ae9adbbaa..25a1928df 100644 --- a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json +++ b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json @@ -28,8 +28,7 @@ "eye": [-1.4, 1.4, 2.0], "target": [0, 0, 0], "up": [0, 0, 1], - "intrinsics": [600, 600, 320, 240], - "save_path": "./outputs/videos_ppo1/eval" + "intrinsics": [600, 600, 320, 240] } } } diff --git a/examples/learning/policy_evaluation/README.md b/examples/learning/policy_evaluation/README.md new file mode 100644 index 000000000..8e3155b00 --- /dev/null +++ b/examples/learning/policy_evaluation/README.md @@ -0,0 +1,137 @@ +# ANYmal-C Velocity Policy Evaluation + +This example connects Newton's public ANYmal-C velocity TorchScript `.pt` to +EmbodiChain and opens it in the DexSim Viewer through Motion Policy Kit. The +model accepts `vx`, `vy`, and `yaw` commands. W/A/S/D and Q/E update these +commands while the Viewer is running. + +The model, configuration, and robot resources come from newton-assets commit +`261cd1f429619d8ef4f546bd788ab9dea906b5e1`. The Policy is distributed under +Apache-2.0, and the robot resources use BSD-3-Clause. The Adapter follows +Newton v1.2.1 +[`example_robot_policy.py`](https://github.com/newton-physics/newton/blob/v1.2.1/newton/examples/robot/example_robot_policy.py) +to reproduce the 48-dimensional observation, TorchScript inference, and joint +target processing. + +## Directory layout + +```text +policy_evaluation/ +├── README.md +├── prepare_resources.py +├── eval_policy.py # Register the local Profile and run the example +└── anymal_c/ + ├── __init__.py # Register newton-anymal-c-velocity + └── profile.py # Policy Spec and AnymalCVelocityAdapter +``` + +Resource preparation creates this local cache: + +```text +~/.cache/embodichain/examples/anymal_c_velocity/ +└── upstream/ + └── anybotics_anymal_c/ + ├── rl_policies/ + │ ├── mjw_anymal.pt + │ ├── anymal.yaml + │ └── LICENSE + ├── urdf/anymal.urdf + ├── meshes/... + └── LICENSE +``` + +## Run the example + +Run these commands from the EmbodiChain repository root. The preparation script +prints the model, asset, checkout, and digest verification progress. Re-running +the command continues an existing Git checkout after an interrupted download. + +```bash +python examples/learning/policy_evaluation/prepare_resources.py +python examples/learning/policy_evaluation/eval_policy.py \ + --viewer \ + --renderer hybrid +``` + +Viewer controls: + +| Key | Command | +|---|---| +| W / S | Increase / decrease `vx` | +| A / D | Increase / decrease `vy` | +| Q / E | Increase / decrease `yaw` | +| M | Set all three commands to zero | +| Backspace | Reset the robot, Policy history, and camera framing | +| T | Switch between tracking and free view | + +The camera follows the robot root in the ground plane. Hold the left mouse +button to change the orbit angle and use the mouse wheel to change the viewing +distance. Right-button panning is locked while tracking is active. Tracking +continues while the orbit angle is being adjusted. Switching back to tracking +centers the camera on the current robot position. + +The terminal prints the path to `evaluation.json` when the Viewer closes. Run a +Headless smoke test with: + +```bash +python examples/learning/policy_evaluation/eval_policy.py \ + --device cpu \ + --sim-device cpu \ + --control-steps 20 +``` + +`eval_policy.py` reads the checkpoint and robot assets from the default cache, +imports the adjacent `anymal_c/profile.py`, and registers the Profile in the +current process. Run the script directly from the repository root. To use +another cache directory: + +```bash +python examples/learning/policy_evaluation/prepare_resources.py \ + --output /tmp/anymal_c_velocity + +ANYMAL_C_EXAMPLE_CACHE=/tmp/anymal_c_velocity \ + python examples/learning/policy_evaluation/eval_policy.py --viewer +``` + +## Execution pipeline + +```mermaid +flowchart LR + CLI[eval-policy] --> Profile[build_profile] + Profile --> Spec[Policy Spec
assets, control parameters, frequency] + Spec --> Setup[Adapter.setup
load TorchScript and joint mapping] + Setup --> State[read RobotState] + Command[WASD + QE command] --> Obs[build 48-dimensional observation] + State --> Obs + Obs --> Actor[TorchScript actor] + Actor --> Action[map to 12 joint targets] + Action --> Sim[advance the Environment] +``` + +`AnymalCVelocityAdapter` restores the upstream data path: + +| Stage | Processing | +|---|---| +| `setup()` | Build a `JointMap` for the 12 ANYmal-C joints, then load and validate the TorchScript inputs and outputs | +| observation | 3 body linear velocity, 3 body angular velocity, 3 projected gravity, 3 command, 12 joint position, 12 joint velocity, and 12 previous action values | +| command | Read `vx`, `vy`, and `yaw` from `frame.controls["command"]`, with ranges ±1.0, ±0.5, and ±1.0 | +| actor | Pass a `[1, 48]` tensor through the model's normalizer and actor to produce `[1, 12]` | +| action | Apply `default_position + 0.5 * action` | +| control | Run simulation at 200 Hz and infer once every four simulation steps for a 50 Hz Policy rate | + +The Adapter clears the previous action during reset. After each inference call, +it stores the current action for the next observation. + +## Integrate another external Policy + +Copy this directory and replace: + +1. the fixed revisions, paths, and digests for the model and robot resources in `prepare_resources.py`; +2. the initial pose, joint control parameters, simulation step, and `sim_steps_per_control` in `build_profile()`; +3. the model format and training joint order in `Adapter.setup()`; +4. observation construction, normalization, network forward pass, action clipping, scale, and offset in `Adapter.infer()`; +5. `PROFILE_ID` and the Profile name used by `eval_policy.py`. + +An Adapter can call an existing project data reader from `__init__()` or +`setup()`. To let Policy Spec resolve a data file path, declare it under +`policy.resources` and read the resolved path from `AdapterRequest.resources`. diff --git a/examples/learning/policy_evaluation/anymal_c/__init__.py b/examples/learning/policy_evaluation/anymal_c/__init__.py new file mode 100644 index 000000000..0303a1211 --- /dev/null +++ b/examples/learning/policy_evaluation/anymal_c/__init__.py @@ -0,0 +1,30 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Register the ANYmal-C velocity Motion Profile.""" + +from __future__ import annotations + +from embodichain.learning.rl.policy_evaluation import register_motion_profile + +from .profile import PROFILE_ID, build_profile + +__all__ = ["register"] + + +def register() -> None: + """Register the ANYmal-C velocity Profile for the example process.""" + register_motion_profile(PROFILE_ID, build_profile) diff --git a/examples/learning/policy_evaluation/anymal_c/profile.py b/examples/learning/policy_evaluation/anymal_c/profile.py new file mode 100644 index 000000000..2f0e89e73 --- /dev/null +++ b/examples/learning/policy_evaluation/anymal_c/profile.py @@ -0,0 +1,278 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""ANYmal-C velocity Profile for a public TorchScript policy.""" + +from __future__ import annotations + +import math +from pathlib import Path + +import numpy as np +import torch +from dexsim.kit.motion_policy import ( + AdapterRequest, + EvaluationFrame, + JointMap, + PolicyContext, + PolicyOutput, + require_finite, +) + +from embodichain.learning.rl.policy_evaluation import ( + MotionProfile, + MotionProfileRequest, +) + +__all__ = ["AnymalCVelocityAdapter", "PROFILE_ID", "build_profile"] + +PROFILE_ID = "newton-anymal-c-velocity" + +_SOURCE_REVISION = "7249270ab41be1c2d4c809aa87536bab3a1a26f4" +_ASSET_REVISION = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" +_ROBOT_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") +_JOINT_NAMES = ( + "LF_HAA", + "LF_HFE", + "LF_KFE", + "LH_HAA", + "LH_HFE", + "LH_KFE", + "RF_HAA", + "RF_HFE", + "RF_KFE", + "RH_HAA", + "RH_HFE", + "RH_KFE", +) +_DEFAULT_POSITION = np.asarray( + (0.0, 0.4, -0.8, 0.0, -0.4, 0.8, 0.0, 0.4, -0.8, 0.0, -0.4, 0.8), + dtype=np.float32, +) + + +def build_profile(request: MotionProfileRequest) -> MotionProfile: + """Build the Policy Spec for the public ANYmal-C checkpoint. + + Args: + request: Checkpoint, resource checkout, device, and renderer. + + Returns: + A Motion Profile ready for DexSim Motion Policy Kit. + """ + if request.resource_root is None: + raise ValueError( + "The ANYmal-C example requires --resource-root from prepare_resources.py" + ) + robot_asset = request.resource_root / _ROBOT_PATH + if not robot_asset.is_file(): + raise FileNotFoundError( + f"ANYmal-C robot asset does not exist: {robot_asset}. " + "Run prepare_resources.py first." + ) + + return MotionProfile( + profile_id=PROFILE_ID, + policy_spec={ + "schema_version": 1, + "kind": "policy", + "id": PROFILE_ID, + "metadata": { + "title": "Newton ANYmal-C velocity policy", + "description": "Public 48-D command locomotion TorchScript policy.", + "status": "example", + "tags": ["external", "quadruped", "velocity", "torchscript"], + }, + "robot": { + "asset": {"path": str(robot_asset)}, + "use_urdf_material": True, + "initial": { + "root_height": 0.76, + "joint_positions": { + "default": 0.0, + "overrides": dict( + zip( + _JOINT_NAMES, + _DEFAULT_POSITION.tolist(), + strict=True, + ) + ), + }, + }, + "control": { + "defaults": { + "stiffness": 300.0, + "damping": 10.0, + "effort_limit": 80.0, + "armature": 0.06, + }, + }, + }, + "policy": { + "models": {"actor": {"path": str(request.checkpoint)}}, + "adapter": { + "type": "python", + "entrypoint": ("anymal_c.profile:AnymalCVelocityAdapter"), + "config": { + "inference_device": str(request.device), + "joint_names": list(_JOINT_NAMES), + }, + }, + }, + "evaluation": { + "initial_command": [0.0, 0.0, 0.0], + "termination": {"behavior": "pause"}, + }, + "runtime": { + "physics_dt": 0.005, + "sim_steps_per_control": 4, + "physics_backend": "default", + "simulation_device": "cpu", + "inference_provider": ( + "cuda" if request.device.type == "cuda" else "cpu" + ), + "renderer": request.renderer, + }, + }, + provenance={ + "source": "Newton ANYmal-C keyboard policy example", + "source_revision": _SOURCE_REVISION, + "source_example": "newton/examples/robot/example_robot_policy.py", + "asset_revision": _ASSET_REVISION, + "model_format": "torchscript", + "observation_size": 48, + "action_size": 12, + }, + ) + + +class AnymalCVelocityAdapter: + """Reproduce the upstream command locomotion observation and action path.""" + + command_enabled = True + command_step = (0.1, 0.05, 0.1) + command_limits = (1.0, 0.5, 1.0) + + def __init__(self, request: AdapterRequest) -> None: + config = dict(request.config) + self.device = torch.device(str(config["inference_device"])) + self.joint_names = tuple(config["joint_names"]) + self.checkpoint = request.models["actor"] + self.previous_action = np.zeros(12, dtype=np.float32) + self.joints: JointMap | None = None + self.model: torch.jit.ScriptModule | None = None + + def setup(self, context: PolicyContext) -> None: + """Load the model and bind the runtime joint order.""" + robot = context.robot + if robot is None: + raise RuntimeError("ANYmal-C robot description is required") + self.joints = JointMap.from_joint_names( + robot.joint_names, + self.joint_names, + ) + self.model = torch.jit.load( + str(self.checkpoint), + map_location=self.device, + ).eval() + with torch.inference_mode(): + output = self.model(torch.zeros((1, 48), device=self.device)) + if not isinstance(output, torch.Tensor) or tuple(output.shape) != (1, 12): + shape = ( + None if not isinstance(output, torch.Tensor) else tuple(output.shape) + ) + raise ValueError(f"ANYmal-C policy output must be (1, 12), got {shape}") + + def reset(self, frame: EvaluationFrame) -> None: + """Reset the previous action used by the policy observation.""" + self.previous_action.fill(0.0) + + def infer(self, frame: EvaluationFrame) -> PolicyOutput: + """Build one 48-D observation and return 12 joint targets.""" + observation = self._build_observation(frame) + tensor = torch.from_numpy(observation).to(self.device).unsqueeze(0) + with torch.inference_mode(): + output = self._model()(tensor) + action = require_finite( + "ANYmal-C action", + output[0].detach().cpu().numpy(), + ) + self.previous_action = action.copy() + return PolicyOutput( + action=self._joints().command( + position=_DEFAULT_POSITION + 0.5 * action, + ), + termination_reason=_fall_reason(frame.robot_state.root_pose), + ) + + def metrics(self) -> dict[str, float]: + """Return the metrics produced by this velocity example.""" + return {} + + def close(self) -> None: + """Release the loaded TorchScript model.""" + self.model = None + + def _build_observation(self, frame: EvaluationFrame) -> np.ndarray: + state = frame.robot_state + if state is None: + raise RuntimeError("ANYmal-C robot state is required") + pose = np.asarray(state.root_pose, dtype=np.float32) + velocity = np.asarray(state.root_velocity, dtype=np.float32) + rotation = pose[:3, :3] + qpos = self._joints().to_model(state.qpos) + qvel = self._joints().to_model(state.qvel) + command = require_finite( + "ANYmal-C command", + frame.controls["command"], + ) + if command.shape != (3,): + raise ValueError("ANYmal-C command must contain vx, vy, and yaw rate") + observation = np.concatenate( + ( + rotation.T @ velocity[:3], + rotation.T @ velocity[3:], + rotation.T @ np.asarray((0.0, 0.0, -1.0), dtype=np.float32), + command, + qpos - _DEFAULT_POSITION, + qvel, + self.previous_action, + ), + dtype=np.float32, + ) + return require_finite("ANYmal-C observation", observation) + + def _joints(self) -> JointMap: + if self.joints is None: + raise RuntimeError("ANYmal-C Adapter is not set up") + return self.joints + + def _model(self) -> torch.jit.ScriptModule: + if self.model is None: + raise RuntimeError("ANYmal-C Adapter is not set up") + return self.model + + +def _fall_reason(root_pose: np.ndarray) -> str | None: + pose = np.asarray(root_pose, dtype=np.float64) + height = float(pose[2, 3]) + tilt = math.acos(float(np.clip(pose[2, 2], -1.0, 1.0))) + reasons = [] + if height < 0.25: + reasons.append(f"base_height_below_minimum: {height:.3f} m") + if tilt > math.pi * 0.4: + reasons.append(f"bad_orientation: {tilt:.3f} rad") + return "; ".join(reasons) or None diff --git a/examples/learning/policy_evaluation/eval_policy.py b/examples/learning/policy_evaluation/eval_policy.py new file mode 100644 index 000000000..b01048b30 --- /dev/null +++ b/examples/learning/policy_evaluation/eval_policy.py @@ -0,0 +1,72 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Evaluate the public ANYmal-C checkpoint from the example directory.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +__all__ = ["example_arguments", "main"] + +EXAMPLE_ROOT = Path(__file__).resolve().parent +REPOSITORY_ROOT = EXAMPLE_ROOT.parents[2] +DEFAULT_CACHE = Path.home() / ".cache/embodichain/examples/anymal_c_velocity" + + +def example_arguments(argv: list[str]) -> list[str]: + """Add the example Profile, checkpoint, and resource paths. + + Args: + argv: Evaluation options accepted by ``eval-policy``. + + Returns: + Arguments ready for the EmbodiChain evaluation CLI. + """ + cache = Path(os.environ.get("ANYMAL_C_EXAMPLE_CACHE", DEFAULT_CACHE)) + resource_root = cache / "upstream" + checkpoint = resource_root / "anybotics_anymal_c/rl_policies/mjw_anymal.pt" + return [ + "--profile", + "newton-anymal-c-velocity", + "--checkpoint", + str(checkpoint), + "--resource-root", + str(resource_root), + *argv, + ] + + +def main(argv: list[str] | None = None) -> None: + """Register the local Profile and run visual policy evaluation. + + Args: + argv: Evaluation options. Uses command-line arguments when omitted. + """ + if str(REPOSITORY_ROOT) not in sys.path: + sys.path.insert(0, str(REPOSITORY_ROOT)) + + from anymal_c import register + from embodichain.learning.rl.policy_evaluation.cli import cli + + register() + cli(example_arguments(sys.argv[1:] if argv is None else argv)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/policy_evaluation/prepare_resources.py b/examples/learning/policy_evaluation/prepare_resources.py new file mode 100644 index 000000000..7e000e2a4 --- /dev/null +++ b/examples/learning/policy_evaluation/prepare_resources.py @@ -0,0 +1,209 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Prepare the pinned public policy and assets for the ANYmal-C example.""" + +from __future__ import annotations + +import argparse +import hashlib +import os +import subprocess +from pathlib import Path + +__all__ = ["main", "prepare_resources"] + +UPSTREAM_URL = "https://github.com/newton-physics/newton-assets.git" +UPSTREAM_REVISION = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" +MODEL_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/mjw_anymal.pt") +POLICY_CONFIG_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/anymal.yaml") +POLICY_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/LICENSE") +ROBOT_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/LICENSE") +ROBOT_RELATIVE_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") +MESH_RELATIVE_PATH = Path("anybotics_anymal_c/meshes/base.dae") +SHA256 = { + MODEL_RELATIVE_PATH: "00765c1c07e497be3825672b05f9cefff9238f2df72fb0bcb5ac9541155b945f", + POLICY_CONFIG_RELATIVE_PATH: "b5a463ac418c7f40ebe494c7bcf0d8031f021db70a0625dbcc28a718de8ee817", + POLICY_LICENSE_RELATIVE_PATH: "59899c6091b540582ed617e8eeaac4919dc985ccfc35459ee9752b699be5205b", + ROBOT_LICENSE_RELATIVE_PATH: "cef384faae108293b03b5e16a00bc3db8212d44575f69df6296438a3f901700b", + ROBOT_RELATIVE_PATH: "d6bd20292cdd4873ffdeeb6f8ca3f96c4a0096565d78d8b6204f6edf0d19fb83", + MESH_RELATIVE_PATH: "785bea9b33831f8c741fc0ca070162e73cbf560ea9b03c53abf8978be877fc48", +} + + +def prepare_resources(output: Path) -> tuple[Path, Path]: + """Fetch and verify the pinned upstream files. + + Args: + output: Cache directory that will contain the sparse Git checkout. + + Returns: + The local checkpoint and resource-root paths. + """ + output = output.expanduser().resolve() + checkout = output / "upstream" + _status("Preparing the ANYmal-C command policy and robot assets") + _prepare_checkout( + checkout, + UPSTREAM_URL, + UPSTREAM_REVISION, + ( + f"/{MODEL_RELATIVE_PATH}", + f"/{POLICY_CONFIG_RELATIVE_PATH}", + f"/{POLICY_LICENSE_RELATIVE_PATH}", + f"/{ROBOT_LICENSE_RELATIVE_PATH}", + "/anybotics_anymal_c/urdf/**", + "/anybotics_anymal_c/meshes/**", + ), + ) + + checkpoint = checkout / MODEL_RELATIVE_PATH + for relative, digest in SHA256.items(): + _verify_sha256(checkout / relative, digest) + _status("Resource verification completed") + return checkpoint, checkout + + +def main() -> None: + """Prepare resources and print the paths used by the evaluation command.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--output", + type=Path, + default=Path.home() / ".cache/embodichain/examples/anymal_c_velocity", + help="Directory used for the pinned upstream checkout", + ) + args = parser.parse_args() + checkpoint, resource_root = prepare_resources(args.output) + print(f"Checkpoint: {checkpoint}") + print(f"Resource root: {resource_root}") + + +def _git( + checkout: Path, + *args: str, + capture_output: bool = False, +) -> subprocess.CompletedProcess[str]: + command = ["git", "-C", str(checkout), *args] + environment = os.environ.copy() + environment["GIT_TERMINAL_PROMPT"] = "0" + try: + return subprocess.run( + command, + check=True, + text=True, + capture_output=capture_output, + env=environment, + timeout=600, + ) + except subprocess.TimeoutExpired as error: + raise RuntimeError( + f"Git command did not finish within 10 minutes: {' '.join(command)}" + ) from error + + +def _git_output(checkout: Path, *args: str) -> str | None: + try: + return _git(checkout, *args, capture_output=True).stdout.strip() + except subprocess.CalledProcessError: + return None + + +def _prepare_checkout( + checkout: Path, + url: str, + revision: str, + includes: tuple[str, ...], +) -> None: + if checkout.exists() and not (checkout / ".git").is_dir(): + raise RuntimeError( + f"Resource path exists but is not a Git checkout: {checkout}" + ) + + if checkout.exists(): + remote_url = _git_output(checkout, "remote", "get-url", "origin") + if remote_url != url: + raise RuntimeError( + f"Resource checkout uses an unexpected remote: {remote_url}" + ) + if _git_output(checkout, "rev-parse", "HEAD") == revision: + tracked_changes = _git_output( + checkout, "status", "--porcelain", "--untracked-files=no" + ) + if tracked_changes == "": + _status(f"Using cached revision {revision[:8]} from {checkout}") + return + else: + checkout.parent.mkdir(parents=True, exist_ok=True) + checkout.mkdir() + _git(checkout, "init", "--quiet") + _git(checkout, "remote", "add", "origin", url) + + _git(checkout, "sparse-checkout", "init", "--no-cone") + _git(checkout, "sparse-checkout", "set", *includes) + + _status(f"Fetching revision {revision[:8]} from {url}") + _git( + checkout, + "fetch", + "--progress", + "--filter=blob:none", + "--depth", + "1", + "origin", + revision, + ) + + _status(f"Checking out required files in {checkout}") + _git( + checkout, + "-c", + "advice.detachedHead=false", + "checkout", + "--progress", + "--force", + "--detach", + "FETCH_HEAD", + ) + actual_revision = _git_output(checkout, "rev-parse", "HEAD") + if actual_revision != revision: + raise RuntimeError( + f"Checkout revision mismatch: expected {revision}, got {actual_revision}" + ) + + +def _status(message: str) -> None: + print(f"[resources] {message}", flush=True) + + +def _verify_sha256(path: Path, expected: str) -> None: + actual = _sha256(path) + if actual != expected: + raise RuntimeError( + f"SHA256 mismatch for {path}: expected {expected}, got {actual}" + ) + + +def _sha256(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() + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py index a7b7e6db9..6fb13a58f 100644 --- a/scripts/benchmark/atomic_action/common.py +++ b/scripts/benchmark/atomic_action/common.py @@ -240,7 +240,7 @@ def add_common_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "rt"), + choices=("auto", "hybrid", "fast-rt", "offline-rt"), default="auto", help="Renderer backend used by SimulationManager.", ) diff --git a/scripts/benchmark/atomic_action/run_benchmark.py b/scripts/benchmark/atomic_action/run_benchmark.py index d50107756..d6cbc462e 100644 --- a/scripts/benchmark/atomic_action/run_benchmark.py +++ b/scripts/benchmark/atomic_action/run_benchmark.py @@ -102,7 +102,7 @@ def add_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "rt"), + choices=("auto", "hybrid", "fast-rt", "offline-rt"), default="auto", help="Renderer backend forwarded to each selected benchmark.", ) diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index ae0b855b3..763cf5e08 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -385,12 +385,12 @@ def test_launcher_preserves_gym_renderer_when_cli_omits_override(): add_env_launcher_args_to_parser(parser, require_gym_config=True) args = parser.parse_args(["--gym_config", "gym_config.yaml"]) - gym_config = {"id": "Dummy-v0", "render_cfg": {"renderer": "rt"}} + gym_config = {"id": "Dummy-v0", "render_cfg": {"renderer": "offline-rt"}} merged_config = merge_args_with_gym_config(args, gym_config) assert args.renderer is None assert "renderer" not in merged_config - assert merged_config["render_cfg"]["renderer"] == "rt" + assert merged_config["render_cfg"]["renderer"] == "offline-rt" def test_env_launcher_includes_viser_arguments(): @@ -1155,7 +1155,7 @@ def test_yaml_gym_config_parses_to_cfg(self, tmp_path): "speed_tolerance": 0.1, }, "render_cfg": { - "renderer": "rt", + "renderer": "offline-rt", "spp": 4, "tone_mapping_enabled": True, "tone_mapping_exposure": 1.25, @@ -1208,7 +1208,7 @@ def test_yaml_gym_config_parses_to_cfg(self, tmp_path): assert cfg.sim_cfg.physics_config.enable_ccd is True assert cfg.sim_cfg.physics_config.length_tolerance == 0.02 assert cfg.sim_cfg.physics_config.speed_tolerance == 0.1 - assert cfg.sim_cfg.render_cfg.renderer == "rt" + assert cfg.sim_cfg.render_cfg.renderer == "offline-rt" assert cfg.sim_cfg.render_cfg.spp == 4 assert cfg.sim_cfg.render_cfg.tone_mapping_enabled is True assert cfg.sim_cfg.render_cfg.tone_mapping_exposure == 1.25 @@ -1267,7 +1267,7 @@ def test_build_env_cfg_applies_modifier_before_parsing(self, tmp_path): "enable_ccd": True, }, "render_cfg": { - "renderer": "rt", + "renderer": "offline-rt", "spp": 8, "tone_mapping_enabled": True, }, diff --git a/tests/learning/rl/policy_evaluation/test_anymal_c_example.py b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py new file mode 100644 index 000000000..c07d6ef9f --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py @@ -0,0 +1,181 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch + +pytest.importorskip("dexsim.kit.motion_policy.evaluator") + +from dexsim.kit.motion_policy import ( + AdapterRequest, + EvaluationFrame, + PolicyContext, + RobotDescription, + RobotState, + parse_policy_spec, +) + +from embodichain.learning.rl.policy_evaluation import ( + MotionProfileRequest, + build_motion_profile, +) + +_JOINT_NAMES = ( + "LF_HAA", + "LF_HFE", + "LF_KFE", + "LH_HAA", + "LH_HFE", + "LH_KFE", + "RF_HAA", + "RF_HFE", + "RF_KFE", + "RH_HAA", + "RH_HFE", + "RH_KFE", +) +_DEFAULT_POSITION = np.asarray( + (0.0, 0.4, -0.8, 0.0, -0.4, 0.8, 0.0, 0.4, -0.8, 0.0, -0.4, 0.8), + dtype=np.float32, +) +_COMMAND = np.asarray((0.4, -0.2, 0.6), dtype=np.float32) + + +class _CommandPolicy(torch.nn.Module): + def forward(self, observation: torch.Tensor) -> torch.Tensor: + padding = torch.zeros( + (observation.shape[0], 9), + dtype=observation.dtype, + device=observation.device, + ) + return torch.cat((observation[:, 9:12], padding), dim=1) + + +def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): + example_root = ( + Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" + ) + monkeypatch.syspath_prepend(str(example_root)) + from anymal_c.profile import ( + AnymalCVelocityAdapter, + ) + from anymal_c import register + + checkpoint = tmp_path / "mjw_anymal.pt" + traced = torch.jit.trace(_CommandPolicy().eval(), torch.zeros((1, 48))) + torch.jit.save(traced, checkpoint) + + robot_asset = tmp_path / "anybotics_anymal_c/urdf/anymal.urdf" + robot_asset.parent.mkdir(parents=True) + robot_asset.write_text("\n", encoding="utf-8") + + request = MotionProfileRequest( + checkpoint=checkpoint, + device=torch.device("cpu"), + resource_root=tmp_path, + ) + register() + profile = build_motion_profile("newton-anymal-c-velocity", request) + spec = parse_policy_spec(profile.policy_spec) + assert spec.environment.entrypoint is None + config = profile.policy_spec["policy"]["adapter"]["config"] + adapter = AnymalCVelocityAdapter( + AdapterRequest( + asset_path=robot_asset, + models={"actor": checkpoint}, + resources={}, + config=config, + ) + ) + context = PolicyContext( + robot=RobotDescription( + _JOINT_NAMES, + ("base",), + "base", + ), + physics_dt=0.005, + sim_steps_per_control=4, + policy_dt=0.02, + ) + pose = np.eye(4, dtype=np.float32) + pose[2, 3] = 0.76 + frame = EvaluationFrame( + control_step=0, + policy_time=0.0, + simulation_time=0.0, + simulation_step=0, + robot_state=RobotState( + joint_names=_JOINT_NAMES, + qpos=_DEFAULT_POSITION.copy(), + qvel=np.zeros(12, dtype=np.float32), + target_qpos=_DEFAULT_POSITION.copy(), + target_qvel=np.zeros(12, dtype=np.float32), + joint_effort=np.zeros(12, dtype=np.float32), + root_name="base", + root_pose=pose, + root_velocity=np.zeros(6, dtype=np.float32), + link_names=("base",), + link_poses=pose[None, ...], + link_velocities=np.zeros((1, 6), dtype=np.float32), + ), + controls={"command": _COMMAND}, + ) + + adapter.setup(context) + adapter.reset(frame) + output = adapter.infer(frame) + + assert output.action.joint_names == _JOINT_NAMES + expected_position = _DEFAULT_POSITION.copy() + expected_position[:3] += 0.5 * _COMMAND + np.testing.assert_allclose(output.action.position, expected_position) + np.testing.assert_allclose( + adapter.previous_action, + np.concatenate((_COMMAND, np.zeros(9, dtype=np.float32))), + ) + assert adapter.command_enabled + assert adapter.command_limits == (1.0, 0.5, 1.0) + assert output.termination_reason is None + adapter.reset(frame) + np.testing.assert_array_equal(adapter.previous_action, np.zeros(12)) + adapter.close() + + +def test_example_script_supplies_default_resource_paths(tmp_path, monkeypatch): + example_root = ( + Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" + ) + monkeypatch.syspath_prepend(str(example_root)) + monkeypatch.setenv("ANYMAL_C_EXAMPLE_CACHE", str(tmp_path)) + from eval_policy import example_arguments + + arguments = example_arguments(["--control-steps", "5"]) + + assert arguments == [ + "--profile", + "newton-anymal-c-velocity", + "--checkpoint", + str(tmp_path / "upstream/anybotics_anymal_c/rl_policies/mjw_anymal.pt"), + "--resource-root", + str(tmp_path / "upstream"), + "--control-steps", + "5", + ] diff --git a/tests/learning/rl/policy_evaluation/test_bridge.py b/tests/learning/rl/policy_evaluation/test_bridge.py new file mode 100644 index 000000000..eb3f120f5 --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_bridge.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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("dexsim.kit.motion_policy.evaluator") + +from embodichain.learning.rl.policy_evaluation.bridge import ( + evaluate_motion_profile, +) +from embodichain.learning.rl.policy_evaluation.profile import MotionProfile + + +def test_bridge_forwards_physics_backend_and_exact_control_steps( + monkeypatch, +): + profile = MotionProfile( + profile_id="example", + policy_spec={"schema_version": 1}, + ) + options = [] + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.parse_policy_spec", + lambda value: "parsed", + ) + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.resolve_policy_spec", + lambda spec, resolver: "resolved", + ) + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.policy_spec_to_dict", + lambda value: {"policy_id": "example"}, + ) + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.scene_config_to_dict", + lambda value: {"style": "standard"}, + ) + + def run_policy(resolved, run_options): + options.append(run_options) + return SimpleNamespace( + reason="control steps reached", + simulation_time=0.2, + simulation_steps=40, + control_steps=10, + physics_backend="default", + requested_duration=None, + effective_duration=0.2, + metrics={"tracking/error": 0.25}, + ) + + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.run_motion_policy", + run_policy, + ) + + result = evaluate_motion_profile( + profile, + control_steps=10, + physics_backend="default", + ) + + assert options[0].control_steps == 10 + assert options[0].physics_backend == "default" + assert result.episodes[0]["control_steps"] == 10 + assert result.episodes[0]["effective_duration"] == 0.2 + assert result.summary["metrics"]["tracking/error"] == 0.25 diff --git a/tests/learning/rl/policy_evaluation/test_cli.py b/tests/learning/rl/policy_evaluation/test_cli.py new file mode 100644 index 000000000..d36059fee --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_cli.py @@ -0,0 +1,130 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 importlib + +import pytest + +from embodichain.learning.rl.policy_evaluation.cli import ( + _resolve_input, + _validate_native_options, + parse_args, +) +from embodichain.learning.rl.policy_evaluation.manifest import write_run_manifest + + +def _run(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + write_run_manifest( + run, + train_config=train, + latest_checkpoint=checkpoint, + ) + return run, checkpoint + + +def test_run_defaults_to_latest_checkpoint(tmp_path): + run, checkpoint = _run(tmp_path) + + resolved = _resolve_input(parse_args((str(run),))) + + assert resolved.checkpoint == checkpoint + assert resolved.requested_checkpoint == "latest" + assert resolved.selected_checkpoint == "latest" + + +@pytest.mark.parametrize( + "arguments, handler", + [ + ((), "_run_native_headless"), + (("--viewer",), "_run_native_viewer"), + (("--profile", "example"), "_run_profile"), + ], +) +def test_cli_routes_one_command_to_the_selected_evaluation( + tmp_path, + monkeypatch, + arguments, + handler, +): + module = importlib.import_module("embodichain.learning.rl.policy_evaluation.cli") + run, _checkpoint = _run(tmp_path) + expected = tmp_path / "evaluation.json" + calls = [] + + monkeypatch.setattr(module, "discover_task_packages", lambda: None) + monkeypatch.setattr(module, "execute_init_hooks", lambda: None) + for name in ("_run_native_headless", "_run_native_viewer", "_run_profile"): + monkeypatch.setattr( + module, + name, + lambda args, resolved, name=name: calls.append(name) or expected, + ) + + report = module.run(parse_args((str(run), *arguments))) + + assert report == expected + assert calls == [handler] + + +def test_explicit_checkpoint_requires_training_config(tmp_path): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + + with pytest.raises(ValueError, match="--config is required"): + _resolve_input(parse_args(("--checkpoint", str(checkpoint)))) + + +@pytest.mark.parametrize("renderer", ("hybrid", "fast-rt", "offline-rt")) +def test_cli_accepts_dexsim_renderer_names(renderer): + args = parse_args(("--renderer", renderer)) + + assert args.renderer == renderer + + +def test_native_options_keep_profile_and_viewer_inputs_explicit(): + profile_args = parse_args( + ( + "--checkpoint", + "policy.pt", + "--config", + "train.yaml", + "--command", + "0.5", + ) + ) + viewer_args = parse_args( + ( + "--checkpoint", + "policy.pt", + "--config", + "train.yaml", + "--control-steps", + "10", + ) + ) + + with pytest.raises(ValueError, match="--command requires --profile"): + _validate_native_options(profile_args) + with pytest.raises(ValueError, match="--control-steps.*require --viewer"): + _validate_native_options(viewer_args) diff --git a/tests/learning/rl/policy_evaluation/test_viewer.py b/tests/learning/rl/policy_evaluation/test_viewer.py new file mode 100644 index 000000000..bde501fe0 --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_viewer.py @@ -0,0 +1,221 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from tensordict import TensorDict + +pytest.importorskip("dexsim.kit.motion_policy.evaluator") + +from dexsim.kit.motion_policy import EvaluationFrame, PolicyContext + +from embodichain.learning.rl.evaluation import infer_policy_action +from embodichain.learning.rl.policy_evaluation.viewer import ( + EmbodiChainTaskEnvironment, + EmbodiChainTaskPolicyAdapter, + evaluate_native_viewer, +) +from embodichain.learning.rl.runtime import PolicyRuntime + + +class Policy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([[2.0], [-1.0]])) + + def get_action( + self, + tensordict: TensorDict, + deterministic: bool = False, + ) -> TensorDict: + assert deterministic + tensordict["action"] = tensordict["obs"] @ self.weight + return tensordict + + +class Window: + def __init__(self) -> None: + self.titles = [] + self.keys = set() + + def set_window_title(self, title): + self.titles.append(title) + + def native(self): + return self + + def key_state(self, key): + return key in self.keys + + +class World: + def __init__(self) -> None: + self.window = Window() + self.open = True + + def is_window_initialized(self): + return self.open + + def get_windows(self): + return self.window + + +class ActionManager: + def __init__(self) -> None: + self.calls = 0 + + def convert_policy_action_to_env_action(self, action: torch.Tensor): + self.calls += 1 + return action + 0.5 + + +class Environment: + num_envs = 1 + physics_dt = 0.005 + step_dt = 0.02 + + def __init__(self) -> None: + self.unwrapped = self + self.cfg = SimpleNamespace(sim_steps_per_control=4) + self.action_manager = ActionManager() + self.world = World() + self.sim = SimpleNamespace(get_world=lambda: self.world) + self.actions = [] + self.episode_step = 0 + self.reset_seeds = [] + self.exit_process_values = [] + + def reset(self, seed=None): + self.reset_seeds.append(seed) + self.episode_step = 0 + return self._observation(), {} + + def step(self, action): + self.actions.append(action.clone()) + self.episode_step += 1 + done = self.episode_step == 2 + return ( + self._observation(), + torch.tensor([1.25]), + torch.tensor([done]), + torch.tensor([False]), + { + "success": torch.tensor([done]), + "metrics": {"task_progress": torch.tensor([self.episode_step])}, + }, + ) + + def close(self, *, exit_process=None): + self.exit_process_values.append(exit_process) + + def _observation(self): + return { + "policy": torch.tensor( + [[float(self.episode_step), 1.0]], + dtype=torch.float32, + ) + } + + +def _runtime(env: Environment, policy: Policy | None = None) -> PolicyRuntime: + return PolicyRuntime( + env=env, + policy=policy or Policy(), + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + +def test_viewer_adapter_uses_the_shared_deterministic_inference_chain(): + policy = Policy() + observation = {"policy": torch.tensor([[0.25, 0.75]])} + expected = infer_policy_action(policy, observation, device="cpu", num_envs=1) + adapter = EmbodiChainTaskPolicyAdapter(policy, torch.device("cpu")) + adapter.setup(PolicyContext(None, 0.005, 4, 0.02)) + + output = adapter.infer(EvaluationFrame(0, 0.0, 0.0, 0, observation=observation)) + + assert torch.equal(output.action, expected) + adapter.close() + + +def test_viewer_reuses_task_actions_resets_and_metrics(): + env = Environment() + + result = evaluate_native_viewer( + _runtime(env), + seed=17, + episodes=2, + control_steps=None, + duration=None, + ) + + assert result.control_steps == 4 + assert result.simulation_steps == 16 + assert len(result.episodes) == 2 + assert result.metrics == pytest.approx( + { + "reward": 1.25, + "task_progress": 2.0, + "eval/avg_reward": 2.5, + "eval/avg_length": 2.0, + "eval/success_rate": 1.0, + } + ) + assert env.action_manager.calls == 4 + assert env.reset_seeds == [17, None] + assert env.exit_process_values == [False] + + +def test_backspace_requests_one_reset_per_key_press(): + from dexsim.types import InputKey + + env = Environment() + task = EmbodiChainTaskEnvironment(env, seed=1) + env.world.window.keys.add(InputKey.SCANCODE_BACKSPACE) + + assert task.poll() == "manual reset" + assert task.poll() is None + env.world.window.keys.clear() + assert task.poll() is None + env.world.window.keys.add(InputKey.SCANCODE_BACKSPACE) + assert task.poll() == "manual reset" + task.close() + + +def test_viewer_closes_resources_when_evaluator_creation_fails(monkeypatch): + env = Environment() + policy = Policy() + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.viewer.create_motion_policy_evaluator", + lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("setup failed")), + ) + + with pytest.raises(RuntimeError, match="setup failed"): + evaluate_native_viewer( + _runtime(env, policy), + seed=1, + episodes=1, + control_steps=None, + duration=None, + ) + + assert env.exit_process_values == [False] + assert policy.training is True diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py index 21012c4e6..43ceb9da9 100644 --- a/tests/learning/test_point_mass.py +++ b/tests/learning/test_point_mass.py @@ -31,6 +31,7 @@ build_learning_env, ) from embodichain.learning.rl.models import ActorCritic +from embodichain.learning.rl.policy_evaluation.manifest import RunManifest from embodichain.learning.rl.train import train_from_config from embodichain.learning.rl.utils import OptimizerCfg from embodichain.learning.rl.utils.trainer import Trainer @@ -183,6 +184,13 @@ def test_unified_train_entry_runs_apg_and_ppo( assert summary["global_step"] == 8 assert summary["latest_checkpoint_path"] is not None + checkpoint = Path(summary["latest_checkpoint_path"]).resolve() + run = checkpoint.parents[1] + manifest = RunManifest.load(run) + assert ( + manifest.configs["train"] == run / "configs" / f"train.{config_path.suffix[1:]}" + ) + assert manifest.select_checkpoint("latest")[1] == checkpoint def test_sync_collector_accepts_tensor_point_mass_observations() -> None: diff --git a/tests/learning/test_runtime.py b/tests/learning/test_runtime.py new file mode 100644 index 000000000..ca35bab74 --- /dev/null +++ b/tests/learning/test_runtime.py @@ -0,0 +1,125 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import gymnasium as gym +import torch + +from embodichain.learning.rl.runtime import ( + _GymEnvironmentRuntime, + _build_gym_environment, + build_gym_policy_runtime, +) + + +def _policy_config() -> dict: + network = { + "type": "mlp", + "network_cfg": {"hidden_sizes": [8], "activation": "relu"}, + } + return { + "name": "actor_critic", + "actor": network, + "critic": network, + } + + +class GymEnvironment: + flattened_observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,)) + observation_space = gym.spaces.Dict( + {"policy": gym.spaces.Box(-1.0, 1.0, shape=(1, 3))} + ) + action_space = gym.spaces.Box(-1.0, 1.0, shape=(1, 2)) + + def __init__(self) -> None: + self.closed = 0 + self.action_manager = SimpleNamespace(total_action_dim=2) + + def reset(self): + return {"policy": torch.zeros(1, 3)}, {} + + def get_wrapper_attr(self, name): + return getattr(self, name) + + def close(self) -> None: + self.closed += 1 + + +def test_gym_environment_applies_runtime_overrides(tmp_path, monkeypatch): + gym_config = tmp_path / "gym.yaml" + gym_config.write_text("id: Example\n", encoding="utf-8") + env_cfg = SimpleNamespace( + num_envs=8, + sim_cfg=None, + profiler=None, + ) + built = GymEnvironment() + monkeypatch.setattr( + "embodichain.learning.rl.runtime.config_to_cfg", + lambda value, manager_modules: env_cfg, + ) + monkeypatch.setattr( + "embodichain.learning.rl.runtime.build_env", + lambda env_id, base_env_cfg: built, + ) + + runtime = _build_gym_environment( + {"trainer": {"gym_config": "gym.yaml"}}, + simulation_device=torch.device("cpu"), + num_envs=1, + headless=True, + renderer="hybrid", + gpu_id=0, + config_dir=tmp_path, + ) + + assert runtime.env is built + assert runtime.env_id == "Example" + assert runtime.env_cfg.num_envs == 1 + assert runtime.env_cfg.sim_cfg.sim_device == torch.device("cpu") + assert runtime.env_cfg.sim_cfg.headless is True + assert runtime.env_cfg.sim_cfg.render_cfg.renderer == "hybrid" + + +def test_gym_runtime_uses_the_same_task_spaces_for_policy_build(monkeypatch): + env = GymEnvironment() + task = _GymEnvironmentRuntime( + env=env, + env_id="Example", + env_cfg=SimpleNamespace(), + gym_config={"id": "Example"}, + gym_config_path=SimpleNamespace(resolve=lambda: None), + ) + monkeypatch.setattr( + "embodichain.learning.rl.runtime._build_gym_environment", + lambda *args, **kwargs: task, + ) + + runtime = build_gym_policy_runtime( + {"trainer": {"gym_config": "gym.yaml"}, "policy": _policy_config()}, + device=torch.device("cpu"), + num_envs=1, + headless=True, + renderer="hybrid", + gpu_id=0, + ) + + assert runtime.env is env + assert runtime.policy.actor[0].in_features == 3 + assert runtime.policy.actor[-1].out_features == 2 diff --git a/tests/learning/test_train_profile.py b/tests/learning/test_train_profile.py index 8b30b74f2..7934fb7e2 100644 --- a/tests/learning/test_train_profile.py +++ b/tests/learning/test_train_profile.py @@ -21,6 +21,7 @@ import pytest from embodichain.learning.rl.train import ( + _event_params, _resolve_profile_output, parse_args, train_from_config, @@ -65,3 +66,13 @@ def test_learning_env_rejects_profile(tmp_path): with pytest.raises(ValueError, match="--profile_output requires --profile"): train_from_config(str(config_path), profile_output="prof.json") + + +def test_camera_recording_defaults_to_the_run_directory(tmp_path): + params = _event_params( + {"func": "record_camera_data_async", "params": {"name": "main"}}, + run_base=tmp_path / "run", + phase="eval", + ) + + assert params["save_path"] == str(tmp_path / "run" / "videos" / "eval") diff --git a/tests/sim/test_cfg.py b/tests/sim/test_cfg.py index c9cfc28fe..aef03c050 100644 --- a/tests/sim/test_cfg.py +++ b/tests/sim/test_cfg.py @@ -94,13 +94,14 @@ def test_render_cfg_applies_tone_mapping_and_fixed_exposure() -> None: expected_exposure = 1.25 world_config = dexsim.WorldConfig() render_cfg = RenderCfg( - renderer="rt", + renderer="offline-rt", tone_mapping_enabled=True, tone_mapping_exposure=expected_exposure, ) render_cfg.apply_to_dexsim_config(world_config) + assert world_config.renderer == Renderer.OFFLINERT assert world_config.postprocess_config.tone_mapping_enabled is True assert ( world_config.postprocess_config.tone_mapping_type From a4636fae218337fa40a5a45d78b38639b48e2c42 Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:27:00 +0800 Subject: [PATCH 56/85] Revert "Add unified policy evaluation for RL checkpoints (#510)" This reverts commit 3df0bed1b8b312c7d686cae7bea3c88e6beb49b9. --- ...odichain.learning.rl.policy_evaluation.rst | 7 - .../embodichain/embodichain.learning.rl.rst | 1 - .../features/toolkits/grasp_generator.rst | 2 +- docs/source/guides/cli.md | 52 +- docs/source/guides/index.rst | 1 - docs/source/guides/policy_evaluation.md | 196 ------- docs/source/guides/preview_asset.md | 2 +- docs/source/overview/sim/sim_manager.md | 2 +- docs/source/tutorial/rl.rst | 18 +- embodichain/lab/gym/utils/gym_utils.py | 4 +- embodichain/lab/scripts/analyze_workspace.py | 2 +- embodichain/lab/scripts/preview_asset.py | 2 +- embodichain/lab/sim/cfg.py | 14 +- embodichain/lab/sim/sim_manager.py | 4 +- embodichain/lab/sim/utility/render_utils.py | 3 +- embodichain/learning/rl/evaluation.py | 45 +- .../learning/rl/policy_evaluation/__init__.py | 33 -- .../learning/rl/policy_evaluation/bridge.py | 189 ------- .../learning/rl/policy_evaluation/cli.py | 523 ------------------ .../learning/rl/policy_evaluation/manifest.py | 193 ------- .../learning/rl/policy_evaluation/profile.py | 128 ----- .../learning/rl/policy_evaluation/report.py | 78 --- .../learning/rl/policy_evaluation/viewer.py | 480 ---------------- embodichain/learning/rl/runtime.py | 312 ----------- embodichain/learning/rl/train.py | 240 ++++---- .../cart_pole/agents/grpo.json | 5 +- .../cart_pole/agents/grpo.yaml | 1 + .../classic_control/cart_pole/agents/ppo.json | 5 +- .../classic_control/cart_pole/agents/ppo.yaml | 1 + .../manipulation/push_cube/agents/grpo.json | 3 +- .../manipulation/push_cube/agents/ppo.json | 5 +- examples/learning/policy_evaluation/README.md | 137 ----- .../policy_evaluation/anymal_c/__init__.py | 30 - .../policy_evaluation/anymal_c/profile.py | 278 ---------- .../learning/policy_evaluation/eval_policy.py | 72 --- .../policy_evaluation/prepare_resources.py | 209 ------- scripts/benchmark/atomic_action/common.py | 2 +- .../benchmark/atomic_action/run_benchmark.py | 2 +- tests/gym/utils/test_gym_utils.py | 10 +- .../test_anymal_c_example.py | 181 ------ .../rl/policy_evaluation/test_bridge.py | 84 --- .../learning/rl/policy_evaluation/test_cli.py | 130 ----- .../rl/policy_evaluation/test_viewer.py | 221 -------- tests/learning/test_point_mass.py | 8 - tests/learning/test_runtime.py | 125 ----- tests/learning/test_train_profile.py | 11 - tests/sim/test_cfg.py | 3 +- 47 files changed, 192 insertions(+), 3862 deletions(-) delete mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst delete mode 100644 docs/source/guides/policy_evaluation.md delete mode 100644 embodichain/learning/rl/policy_evaluation/__init__.py delete mode 100644 embodichain/learning/rl/policy_evaluation/bridge.py delete mode 100644 embodichain/learning/rl/policy_evaluation/cli.py delete mode 100644 embodichain/learning/rl/policy_evaluation/manifest.py delete mode 100644 embodichain/learning/rl/policy_evaluation/profile.py delete mode 100644 embodichain/learning/rl/policy_evaluation/report.py delete mode 100644 embodichain/learning/rl/policy_evaluation/viewer.py delete mode 100644 embodichain/learning/rl/runtime.py delete mode 100644 examples/learning/policy_evaluation/README.md delete mode 100644 examples/learning/policy_evaluation/anymal_c/__init__.py delete mode 100644 examples/learning/policy_evaluation/anymal_c/profile.py delete mode 100644 examples/learning/policy_evaluation/eval_policy.py delete mode 100644 examples/learning/policy_evaluation/prepare_resources.py delete mode 100644 tests/learning/rl/policy_evaluation/test_anymal_c_example.py delete mode 100644 tests/learning/rl/policy_evaluation/test_bridge.py delete mode 100644 tests/learning/rl/policy_evaluation/test_cli.py delete mode 100644 tests/learning/rl/policy_evaluation/test_viewer.py delete mode 100644 tests/learning/test_runtime.py diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst deleted file mode 100644 index b5aab6c0a..000000000 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst +++ /dev/null @@ -1,7 +0,0 @@ -embodichain.learning.rl.policy_evaluation -========================================= - -.. automodule:: embodichain.learning.rl.policy_evaluation - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst index e2e51b69d..bf1b5b01e 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst @@ -18,7 +18,6 @@ collection logic, policy/model builders, and training entry points. buffer collector models - policy_evaluation train utils diff --git a/docs/source/features/toolkits/grasp_generator.rst b/docs/source/features/toolkits/grasp_generator.rst index fa8fb110b..18cb8ec24 100644 --- a/docs/source/features/toolkits/grasp_generator.rst +++ b/docs/source/features/toolkits/grasp_generator.rst @@ -151,7 +151,7 @@ You can customize the run with additional arguments: .. code-block:: bash - python scripts/tutorials/grasp/grasp_generator.py --num_envs --device --renderer --headless + python scripts/tutorials/grasp/grasp_generator.py --num_envs --device --renderer --headless The script computes a grasp pose, prints the elapsed time, and then waits for you to press **Enter** before executing the full grasp trajectory. Press diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index c1330f90c..44e860e87 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -159,7 +159,7 @@ embodichain run-env --gym_config config.yaml \ | ``--num_envs`` | ``1`` | Number of parallel environments | | ``--device`` | ``cpu`` | Device (``cpu`` or ``cuda``) | | ``--headless`` | ``False`` | Run in headless mode | -| ``--renderer`` | ``auto`` | Renderer backend: ``auto``, ``hybrid``, ``fast-rt`` or ``offline-rt`` | +| ``--renderer`` | ``auto`` | Renderer backend: ``auto``, ``hybrid``, ``fast-rt`` or ``rt`` | | ``--arena_space`` | ``5.0`` | Arena space size | | ``--gpu_id`` | ``0`` | GPU ID to use | | ``--preview`` | ``False`` | Enter interactive preview mode | @@ -386,56 +386,6 @@ See the Profiling section under Run Env for report format. Outputs are written t --- -## Policy Evaluation - -Evaluate the latest checkpoint from an EmbodiChain training run: - -```bash -embodichain eval-policy outputs/my_policy_ -``` - -Open a simulator task in the Viewer: - -```bash -embodichain eval-policy outputs/my_policy_ \ - --checkpoint best \ - --viewer \ - --renderer hybrid -``` - -Evaluate an explicit EmbodiChain checkpoint: - -```bash -embodichain eval-policy \ - --checkpoint /path/to/policy.pt \ - --config /path/to/train.yaml \ - --gym-config /path/to/gym.yaml -``` - -### Main arguments - -| Argument | Default | Description | -|---|---|---| -| ``RUN`` | *(optional)* | Training run containing ``run-manifest.json`` | -| ``--checkpoint`` | ``latest`` with RUN | ``latest``, ``best``, or a checkpoint path | -| ``--config`` | RUN manifest | Training configuration override | -| ``--gym-config`` | RUN manifest | Simulator task configuration override | -| ``--episodes`` | Training configuration | Number of completed task episodes | -| ``--num-envs`` | Training configuration | Number of parallel Headless environments | -| ``--viewer`` | Headless | Open the original simulator task in the DexSim Viewer | -| ``--control-steps`` | Viewer runs continuously | Exact number of Policy actions | -| ``--duration`` | *(optional)* | Duration converted to integer control steps | -| ``--renderer`` | Training configuration or ``hybrid`` | Viewer renderer | -| ``--device`` | Training configuration | PyTorch inference device | -| ``--sim-device`` | Inference device | Simulation device | -| ``--output`` | RUN or checkpoint evaluations | Evaluation output parent directory | - -External Motion Profiles use the same command with `--profile`. See -{doc}`policy_evaluation` for training-run layout, execution paths, Viewer -controls, output reports, and the complete ANYmal-C example. - ---- - ## Annotate Grasp Launch the browser-based grasp-region annotation tool. diff --git a/docs/source/guides/index.rst b/docs/source/guides/index.rst index a5816d4fd..220cbb89d 100644 --- a/docs/source/guides/index.rst +++ b/docs/source/guides/index.rst @@ -12,5 +12,4 @@ Practical guides for common tasks in EmbodiChain. add_robot preview_asset run_env - policy_evaluation cli diff --git a/docs/source/guides/policy_evaluation.md b/docs/source/guides/policy_evaluation.md deleted file mode 100644 index 392da7cc2..000000000 --- a/docs/source/guides/policy_evaluation.md +++ /dev/null @@ -1,196 +0,0 @@ -# Policy Evaluation - -`embodichain eval-policy` evaluates a saved EmbodiChain policy after training. -It reconstructs the policy and environment from the training configuration, -loads the selected checkpoint, and writes a standalone evaluation report. - -The command runs Headless by default. Add `--viewer` to open the original -simulator task in the DexSim Viewer. - -## Training output - -`train-rl` records the files required by a later evaluation: - -```text -outputs/_/ -├── checkpoints/ -│ └── policy_*.pt -├── configs/ -│ ├── train.yaml -│ └── gym.yaml -├── logs/ -├── videos/ -│ ├── train/ -│ └── eval/ -└── run-manifest.json -``` - -`configs/gym.yaml` is present for simulator tasks. The first evaluation adds: - -```text -evaluations/ -└── -policy/ - └── evaluation.json -``` - -`run-manifest.json` connects the run directory to its configuration snapshots -and checkpoints: - -```json -{ - "schema_version": 1, - "configs": { - "train": "configs/train.yaml", - "gym": "configs/gym.yaml" - }, - "checkpoints": { - "best": "checkpoints/cart_pole_grpo_best.pt", - "latest": "checkpoints/cart_pole_grpo_step_4096.pt" - } -} -``` - -All paths in the manifest are relative to the run directory. `best` is `null` -when training did not select a best checkpoint. - -## Evaluate a training run - -The shortest command selects `latest` and runs the configured number of -Headless evaluation episodes: - -```bash -embodichain eval-policy outputs/_ -``` - -Select the best checkpoint and override the episode count: - -```bash -embodichain eval-policy outputs/_ \ - --checkpoint best \ - --episodes 10 -``` - -Open the original simulator task in the Viewer: - -```bash -embodichain eval-policy outputs/_ \ - --checkpoint best \ - --viewer \ - --renderer hybrid \ - --device cuda:0 \ - --sim-device gpu -``` - -The Viewer uses one environment and keeps running until the window closes. Use -`--episodes`, `--control-steps`, or `--duration` to select another stopping -condition. `--renderer` accepts `hybrid`, `fast-rt`, and `offline-rt`. - -For a checkpoint created before `run-manifest.json` was introduced, provide -its training configuration directly: - -```bash -embodichain eval-policy \ - --checkpoint /path/to/policy.pt \ - --config /path/to/train.yaml \ - --gym-config /path/to/gym.yaml \ - --viewer -``` - -`--gym-config` can be omitted when the training configuration already refers to -the task configuration. - -## Execution paths - -```mermaid -flowchart LR - Run[Training run] --> Manifest[run-manifest.json] - Manifest --> Config[Training config] - Manifest --> Checkpoint[Checkpoint] - Config --> Runtime[EmbodiChain RL runtime] - Checkpoint --> Runtime - Runtime --> Headless[Headless episode evaluation] - Runtime --> Viewer[DexSim MotionPolicyEvaluator] - Headless --> Report[evaluation.json] - Viewer --> Report - Profile[External Motion Profile] --> Viewer -``` - -Headless evaluation calls the existing `evaluate_episodes()` path. Viewer -evaluation keeps the task's original observation, action processing, reset, -reward, termination, objects, and sensors: - -```mermaid -sequenceDiagram - participant Evaluator as MotionPolicyEvaluator - participant Adapter as EmbodiChainTaskPolicyAdapter - participant Task as EmbodiChainTaskEnvironment - participant Policy as EmbodiChain Policy - participant Env as Original task Environment - - Evaluator->>Task: reset() - Task->>Env: reset() - Env-->>Task: observation and task state - Task-->>Evaluator: EvaluationFrame - loop Each control step - Evaluator->>Adapter: infer(frame) - Adapter->>Policy: deterministic inference - Policy-->>Adapter: action - Adapter-->>Evaluator: PolicyOutput - Evaluator->>Task: step(action) - Task->>Env: action processing and env.step() - Env-->>Task: observation, reward, termination and info - Task-->>Evaluator: EnvironmentStep - end -``` - -| Input | Headless | Viewer | -|---|---:|---:| -| EmbodiChain lightweight RL environment | Yes | — | -| EmbodiChain simulator RL environment | Yes | Yes | -| Registered external Motion Profile | Yes | Yes | - -Policy reconstruction follows the model definition stored in the training -configuration. The Viewer path has been validated with CartPole GRPO and -PushCube PPO checkpoints. - -## Viewer controls - -| Key | Action | -|---|---| -| `Backspace` | Reset the task and camera framing | -| `T` | Switch between tracking and free camera modes when the Environment provides a tracking target | -| `R` | Start or stop recording | -| `Esc` | Close the Viewer | - -While tracking is active, drag with the left mouse button to orbit and use the -mouse wheel to zoom. - -## External policy example - -The repository includes a concrete ANYmal-C velocity example under -`examples/learning/policy_evaluation/`. It prepares a public TorchScript -checkpoint and robot assets, registers an adjacent Motion Profile, and forwards -the remaining arguments to `eval-policy`. - -```bash -python examples/learning/policy_evaluation/prepare_resources.py -python examples/learning/policy_evaluation/eval_policy.py \ - --viewer \ - --renderer hybrid \ - --sim-device gpu -``` - -Use W/S for `vx`, A/D for `vy`, Q/E for `yaw`, and M to zero the command. See -the [example README](https://github.com/DexForce/EmbodiChain/tree/main/examples/learning/policy_evaluation) -for the resource layout, observation construction, action conversion, and -Profile implementation. - -This example tracks the robot root in the ground plane. Press `T` to switch -between tracking and free view. - -## Evaluation report - -`evaluation.json` records the selected checkpoint and configs, task and device -information, episode results, and aggregated metrics. Reports are written to -`/evaluations/` for a training run and next to an explicit checkpoint by -default. Use `--output` to select another parent directory. diff --git a/docs/source/guides/preview_asset.md b/docs/source/guides/preview_asset.md index 7049382fe..a83fa49bc 100644 --- a/docs/source/guides/preview_asset.md +++ b/docs/source/guides/preview_asset.md @@ -149,7 +149,7 @@ asset.set_local_pose(pose) | `--use_usd_properties` | disabled | Use physical properties stored in the USD file. | | `--fix_base` / `--no-fix_base` | fixed | Fix or unfix articulation bases. | | `--sim_device` | `cpu` | Simulation device. | -| `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `offline-rt`. | +| `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `rt`. | | `--env_map` | none | Built-in IBL resource name or absolute `.hdr`, `.png`, or `.exr` path. | | `--headless` | disabled | Run without the native window. | | `--preview` | disabled | Enter the interactive terminal after loading. | diff --git a/docs/source/overview/sim/sim_manager.md b/docs/source/overview/sim/sim_manager.md index fa57cf93c..cd280c26e 100644 --- a/docs/source/overview/sim/sim_manager.md +++ b/docs/source/overview/sim/sim_manager.md @@ -68,7 +68,7 @@ The {class}`~cfg.RenderCfg` class controls the rendering backend and quality set | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `renderer` | `str` | `"auto"` | Renderer backend to use. Options are `'auto'` (pick a default based on the detected GPU), `'hybrid'` (ray tracing for shadows/reflections + rasterization), `'fast-rt'` (full ray tracing), and `'offline-rt'` (offline ray-traced renderer for maximum visual fidelity). | +| `renderer` | `str` | `"auto"` | Renderer backend to use. Options are `'auto'` (pick a default based on the detected GPU), `'hybrid'` (ray tracing for shadows/reflections + rasterization), `'fast-rt'` (full ray tracing), and `'rt'` (offline ray-traced renderer for maximum visual fidelity). | | `spp` | `int` | `1` | Samples per pixel for ray-traced rendering. Must be at least 1. | | `tone_mapping_enabled` | `bool` | `False` | Whether to map HDR RGB output with the modified Reinhard curve. | | `tone_mapping_exposure` | `float` | `1.0` | Non-negative fixed linear exposure multiplier applied before tone mapping. | diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index a5a57d8bd..6470ef0e6 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -264,19 +264,6 @@ All outputs are written to ``./outputs/_/``: - **logs/**: TensorBoard logs - **checkpoints/**: Model checkpoints -- **configs/**: Training config and referenced gym config snapshots -- **evaluations/**: Timestamped policy evaluation reports -- **run-manifest.json**: Training configs and best/latest checkpoint index used by ``eval-policy`` - -A training run can be evaluated Headless or opened in its simulator task: - -.. code-block:: bash - - embodichain eval-policy outputs/_ - embodichain eval-policy outputs/_ --viewer - -See :doc:`../guides/policy_evaluation` for EmbodiChain ``.pt`` training -runs and the external Motion Profile example. Training Process ~~~~~~~~~~~~~~~~ @@ -465,9 +452,9 @@ Best Practices - **Configuration**: Use JSON for all hyperparameters. This makes experiments reproducible and easy to track. -- **Logging**: Metrics are automatically logged to TensorBoard and Weights & Biases. Check ``outputs/_/logs/`` for TensorBoard logs. +- **Logging**: Metrics are automatically logged to TensorBoard and Weights & Biases. Check ``outputs//logs/`` for TensorBoard logs. -- **Checkpoints**: Regular checkpoints are saved to ``outputs/_/checkpoints/``. Use these to resume training or evaluate policies. +- **Checkpoints**: Regular checkpoints are saved to ``outputs//checkpoints/``. Use these to resume training or evaluate policies. See Also -------- @@ -477,4 +464,3 @@ See Also - :doc:`basic_env` — Creating basic Gymnasium environments - :doc:`modular_env` — Advanced modular environments with managers - :doc:`/resources/task/index` — List of available RL task environments -- :doc:`/guides/policy_evaluation` — Headless and Viewer evaluation of EmbodiChain ``.pt`` checkpoints and external Motion Profiles diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 84c4314c5..8b2010b00 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -974,7 +974,7 @@ def add_env_launcher_args_to_parser( --num_envs: Number of environments to run in parallel (default: 1) --device: Device to run the environment on (default: 'cpu') --headless: Whether to perform the simulation in headless mode (default: False) - --renderer: Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'offline-rt'. + --renderer: Renderer backend to use for the simulation. Options are 'hybrid', 'fast-rt', and 'rt'. (default: 'hybrid') --gpu_id: The GPU ID to use for the simulation (default: 0) --gym_config: Path to gym config file (default: '') --action_config: Path to action config file (default: None) @@ -1015,7 +1015,7 @@ def add_env_launcher_args_to_parser( parser.add_argument( "--renderer", type=str, - choices=["auto", "hybrid", "fast-rt", "offline-rt"], + choices=["auto", "hybrid", "fast-rt", "rt"], default=None if require_gym_config else "auto", help="Renderer backend to use for the simulation. When loading a gym " "config, the configured render_cfg.renderer is used unless this option " diff --git a/embodichain/lab/scripts/analyze_workspace.py b/embodichain/lab/scripts/analyze_workspace.py index e26a17746..7f1649c14 100644 --- a/embodichain/lab/scripts/analyze_workspace.py +++ b/embodichain/lab/scripts/analyze_workspace.py @@ -1047,7 +1047,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: sim.add_argument( "--renderer", type=str, - choices=["hybrid", "fast-rt", "offline-rt"], + choices=["hybrid", "fast-rt", "rt"], default="hybrid", help="Renderer backend (default: hybrid).", ) diff --git a/embodichain/lab/scripts/preview_asset.py b/embodichain/lab/scripts/preview_asset.py index 8660dfb5d..0bb4f1416 100644 --- a/embodichain/lab/scripts/preview_asset.py +++ b/embodichain/lab/scripts/preview_asset.py @@ -434,7 +434,7 @@ def _create_parser() -> argparse.ArgumentParser: parser.add_argument( "--renderer", type=str, - choices=["hybrid", "fast-rt", "offline-rt"], + choices=["hybrid", "fast-rt", "rt"], default="hybrid", help="Renderer backend (default: hybrid).", ) diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 4fef9a2b0..c36fabbcd 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -56,13 +56,13 @@ # :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`). Assigning a # concrete renderer here (e.g. in test fixtures) forces that renderer and takes # precedence over auto-selection. -DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "offline-rt"] = "auto" +DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" @configclass class RenderCfg: - renderer: Literal["auto", "hybrid", "fast-rt", "offline-rt"] = "auto" - """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'offline-rt'. + renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" + """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'rt'. Note: - 'auto' selects a default renderer based on the detected GPU: RTX-series cards use @@ -71,11 +71,11 @@ class RenderCfg: - 'hybrid' uses ray tracing for shadows and reflections while keeping rasterization for primary rendering, providing a balance between performance and visual quality. - 'fast-rt' is a fully ray-traced renderer for maximum visual fidelity, but may have higher computational cost. - - 'offline-rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. + - 'rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. """ spp: int = 1 - """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'offline-rt'.""" + """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'rt'.""" tone_mapping_enabled: bool = False """Whether to map HDR RGB output with the modified Reinhard curve.""" @@ -98,7 +98,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID elif self.renderer == "fast-rt": return Renderer.FASTRT - elif self.renderer == "offline-rt": + elif self.renderer == "rt": return Renderer.OFFLINERT elif self.renderer == "auto": # 'auto' is normally resolved by the SimulationManager before this is @@ -110,7 +110,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID else: logger.log_error( - f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'offline-rt'." + f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'rt'." ) def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 707da39be..b044575f6 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -440,7 +440,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: Args: renderer: The renderer to set. One of ``"auto"``, ``"hybrid"``, - ``"fast-rt"``, or ``"offline-rt"``. When ``"auto"``, the renderer is + ``"fast-rt"``, or ``"rt"``. When ``"auto"``, the renderer is resolved immediately from the detected GPU via :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`. gpu_id: The CUDA device index to query when ``renderer="auto"``. @@ -451,7 +451,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: from embodichain.lab.sim import cfg from embodichain.lab.sim.utility.render_utils import select_default_renderer - valid = {"auto", "hybrid", "fast-rt", "offline-rt"} + valid = {"auto", "hybrid", "fast-rt", "rt"} if renderer not in valid: logger.log_error( f"Invalid renderer '{renderer}'. Must be one of {sorted(valid)}." diff --git a/embodichain/lab/sim/utility/render_utils.py b/embodichain/lab/sim/utility/render_utils.py index 8469ad767..d82bb2644 100644 --- a/embodichain/lab/sim/utility/render_utils.py +++ b/embodichain/lab/sim/utility/render_utils.py @@ -47,8 +47,7 @@ def select_default_renderer(gpu_id: int = 0) -> str: gpu_id: The CUDA device index to query for selecting the renderer. Returns: - The resolved renderer name, one of ``"hybrid"``, ``"fast-rt"``, or - ``"offline-rt"``. + The resolved renderer name, one of ``"hybrid"``, ``"fast-rt"``, or ``"rt"``. """ from embodichain.lab.sim import cfg diff --git a/embodichain/learning/rl/evaluation.py b/embodichain/learning/rl/evaluation.py index 3a54715b2..974a2ab7c 100644 --- a/embodichain/learning/rl/evaluation.py +++ b/embodichain/learning/rl/evaluation.py @@ -30,43 +30,15 @@ flatten_dict_observation, ) -__all__ = [ - "convert_policy_action_for_env", - "evaluate_episodes", - "infer_policy_action", - "prepare_policy_observation", -] +__all__ = ["evaluate_episodes"] -def prepare_policy_observation( - observation: Any, - device: torch.device | str, -) -> torch.Tensor: - """Flatten one Environment observation in the training input order.""" - device = torch.device(device) +def _flat_observation(observation: Any, device: torch.device) -> torch.Tensor: tensor_dict = dict_to_tensordict(observation, device) return flatten_dict_observation(tensor_dict) -def infer_policy_action( - policy: torch.nn.Module, - observation: Any, - *, - device: torch.device | str, - num_envs: int, -) -> torch.Tensor: - """Run the same deterministic Policy call used by RL evaluation.""" - device = torch.device(device) - policy_input = TensorDict( - {"obs": prepare_policy_observation(observation, device)}, - batch_size=[num_envs], - device=device, - ) - return policy.get_action(policy_input, deterministic=True)["action"] - - -def convert_policy_action_for_env(env: Any, action: torch.Tensor) -> Any: - """Convert a flat Policy action to the task Environment input layout.""" +def _action_for_env(env: Any, action: torch.Tensor) -> Any: action_manager = getattr(env, "action_manager", None) if action_manager is None and hasattr(env, "get_wrapper_attr"): try: @@ -134,14 +106,15 @@ def evaluate_episodes( try: observation, _ = env.reset(seed=seed) while len(returns) < num_episodes: - action = infer_policy_action( - policy, - observation, + flat_observation = _flat_observation(observation, device) + policy_input = TensorDict( + {"obs": flat_observation}, + batch_size=[num_envs], device=device, - num_envs=num_envs, ) + policy_output = policy.get_action(policy_input, deterministic=True) observation, reward, terminated, truncated, info = env.step( - convert_policy_action_for_env(env, action) + _action_for_env(env, policy_output["action"]) ) reward = torch.as_tensor(reward, device=device).reshape(num_envs) done = ( diff --git a/embodichain/learning/rl/policy_evaluation/__init__.py b/embodichain/learning/rl/policy_evaluation/__init__.py deleted file mode 100644 index 4dbc590b4..000000000 --- a/embodichain/learning/rl/policy_evaluation/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""External Policy Profiles for ``embodichain eval-policy``.""" - -from __future__ import annotations - -from .profile import ( - MotionProfile, - MotionProfileRequest, - build_motion_profile, - register_motion_profile, -) - -__all__ = [ - "MotionProfile", - "MotionProfileRequest", - "build_motion_profile", - "register_motion_profile", -] diff --git a/embodichain/learning/rl/policy_evaluation/bridge.py b/embodichain/learning/rl/policy_evaluation/bridge.py deleted file mode 100644 index a336b9ccc..000000000 --- a/embodichain/learning/rl/policy_evaluation/bridge.py +++ /dev/null @@ -1,189 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Run an external Policy Profile through DexSim Motion Policy Kit.""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -from dexsim.kit.motion_policy import ( - PolicySpec, - ResolvedPolicy, - ResourceResolver, - RunOptions, - load_scene_config, - parse_policy_spec, - policy_spec_to_dict, - resolve_policy_spec, - run_motion_policy, - scene_config_to_dict, -) - -from .profile import MotionProfile - -__all__ = [ - "MotionEvaluationResult", - "evaluate_motion_profile", -] - - -@dataclass(frozen=True) -class MotionEvaluationResult: - """Normalized inputs and per-episode motion evaluation results.""" - - profile: MotionProfile - policy_spec: Mapping[str, Any] - scene_config: Mapping[str, Any] - episodes: tuple[Mapping[str, Any], ...] - summary: Mapping[str, Any] - viewer: bool - - -def evaluate_motion_profile( - profile: MotionProfile, - *, - episodes: int = 1, - viewer: bool = False, - control_steps: int | None = None, - duration: float | None = None, - command: tuple[float, ...] | None = None, - scene_config: str | Path = "standard", - physics_backend: str | None = None, - simulation_device: str = "cpu", - renderer: str = "hybrid", - gpu_id: int = 0, - termination_behavior: str | None = None, - cache_dir: str | Path | None = None, - offline: bool = False, -) -> MotionEvaluationResult: - """Resolve one Motion Profile and run its visual evaluation. - - Args: - profile: Provider-built profile containing the DexSim Policy Spec. - episodes: Number of independent runs. - viewer: Open the DexSim Viewer. - control_steps: Exact number of applied policy commands per run. - duration: Convenience duration converted by DexSim to policy steps. - command: Optional task command override. - scene_config: Built-in scene style or custom YAML path. - physics_backend: Optional DexSim physics backend override. - simulation_device: ``cpu`` or ``gpu``. - renderer: DexSim renderer. - gpu_id: Selected GPU index. - termination_behavior: Policy termination handling override. - cache_dir: Motion Policy Kit resource cache. - offline: Use resources already available in the cache. - - Returns: - Normalized inputs, episode results, and aggregate metrics. - """ - if episodes <= 0: - raise ValueError("episodes must be positive") - if viewer and episodes != 1: - raise ValueError("Viewer evaluation supports one episode") - parsed, resolved = _resolve_profile(profile, cache_dir, offline) - resolved_scene = load_scene_config(scene_config) - options = RunOptions( - physics_backend=physics_backend, - simulation_device=simulation_device, - renderer=renderer, - gpu_id=gpu_id, - headless=not viewer, - control_steps=control_steps, - duration=duration, - command=command, - termination_behavior=termination_behavior, - scene_config=resolved_scene, - ) - results = tuple( - _episode( - index, - run_motion_policy( - resolved, - options, - ), - ) - for index in range(episodes) - ) - return MotionEvaluationResult( - profile=profile, - policy_spec=policy_spec_to_dict(parsed), - scene_config=scene_config_to_dict(resolved_scene), - episodes=results, - summary=_summary(results), - viewer=viewer, - ) - - -def _resolve_profile( - profile: MotionProfile, - cache_dir: str | Path | None, - offline: bool, -) -> tuple[PolicySpec, ResolvedPolicy]: - parsed = parse_policy_spec(profile.policy_spec) - resolved = resolve_policy_spec( - parsed, - ResourceResolver( - None if cache_dir is None else Path(cache_dir), - offline=offline, - ), - ) - return parsed, resolved - - -def _episode(index: int, result: Any) -> dict[str, Any]: - return { - "index": index, - "reason": str(result.reason), - "simulation_time": float(result.simulation_time), - "simulation_steps": int(result.simulation_steps), - "control_steps": int(result.control_steps), - "physics_backend": str(result.physics_backend), - "requested_duration": ( - None - if result.requested_duration is None - else float(result.requested_duration) - ), - "effective_duration": float(result.effective_duration), - "metrics": {name: float(value) for name, value in result.metrics.items()}, - } - - -def _summary(episodes: tuple[Mapping[str, Any], ...]) -> dict[str, Any]: - count = len(episodes) - metric_names = set.intersection(*(set(episode["metrics"]) for episode in episodes)) - metrics = { - name: sum(episode["metrics"][name] for episode in episodes) / count - for name in sorted(metric_names) - } - result: dict[str, Any] = { - "episodes": count, - "avg_simulation_time": sum(episode["simulation_time"] for episode in episodes) - / count, - "avg_control_steps": sum(episode["control_steps"] for episode in episodes) - / count, - "avg_effective_duration": sum( - episode["effective_duration"] for episode in episodes - ) - / count, - } - if metrics: - result["metrics"] = metrics - return result diff --git a/embodichain/learning/rl/policy_evaluation/cli.py b/embodichain/learning/rl/policy_evaluation/cli.py deleted file mode 100644 index 8c9b9c2d2..000000000 --- a/embodichain/learning/rl/policy_evaluation/cli.py +++ /dev/null @@ -1,523 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Unified policy evaluation for EmbodiChain training runs.""" - -from __future__ import annotations - -import argparse -import sys -from collections.abc import Mapping, Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from embodichain import __version__ -from embodichain.lab.gym.utils.registration import ( - discover_task_packages, - execute_init_hooks, -) -from embodichain.learning.rl.evaluation import evaluate_episodes -from embodichain.learning.rl.runtime import ( - PolicyRuntime, - build_gym_policy_runtime, - build_learning_policy_runtime, -) -from embodichain.utils.utility import load_config - -from .manifest import RunManifest -from .report import write_evaluation_report - -__all__ = ["cli", "parse_args", "run"] - - -@dataclass(frozen=True) -class EvaluationInput: - """Checkpoint and configuration selected for one evaluation.""" - - checkpoint: Path - profile: str | None - configs: Mapping[str, Path] - run: Path | None - requested_checkpoint: str - selected_checkpoint: str - - -@dataclass(frozen=True) -class NativeRuntime: - """Reconstructed EmbodiChain task and its runtime choices.""" - - runtime: PolicyRuntime - device: torch.device - simulation_device: torch.device - seed: int - renderer: str - uses_simulator: bool - trainer: Mapping[str, Any] - - -def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - """Parse ``embodichain eval-policy`` arguments.""" - parser = argparse.ArgumentParser( - prog="embodichain eval-policy", - description="Evaluate an EmbodiChain or external policy checkpoint.", - ) - parser.add_argument("run", nargs="?", help="EmbodiChain training run directory.") - parser.add_argument("--profile", help="Registered external Policy Profile.") - parser.add_argument( - "--checkpoint", - help="latest, best, or a checkpoint path; defaults to latest with RUN.", - ) - parser.add_argument("--config", help="Training config for an explicit checkpoint.") - parser.add_argument("--gym-config", help="Task config override.") - parser.add_argument("--resource-root", help="External Profile resource root.") - parser.add_argument("--episodes", type=int) - parser.add_argument("--num-envs", type=int) - count = parser.add_mutually_exclusive_group() - count.add_argument("--control-steps", type=int) - count.add_argument("--duration", type=float) - parser.add_argument("--command", nargs="+", type=float) - parser.add_argument("--device", help="PyTorch inference device.") - parser.add_argument("--sim-device", choices=("cpu", "gpu")) - parser.add_argument("--seed", type=int) - parser.add_argument("--physics-backend") - parser.add_argument( - "--renderer", - choices=("hybrid", "fast-rt", "offline-rt"), - ) - parser.add_argument("--gpu-id", type=int, default=0) - parser.add_argument("--scene-config") - parser.add_argument( - "--termination-behavior", - choices=("pause", "continue", "auto_reset"), - ) - parser.add_argument("--viewer", action="store_true") - parser.add_argument("--cache-dir") - parser.add_argument("--offline", action="store_true") - parser.add_argument("--output", help="Evaluation output parent directory.") - return parser.parse_args(argv) - - -def run(args: argparse.Namespace) -> Path: - """Run Headless or Viewer evaluation and write ``evaluation.json``.""" - resolved = _resolve_input(args) - if resolved.profile is not None: - return _run_profile(args, resolved) - discover_task_packages() - execute_init_hooks() - _validate_native_options(args) - if args.viewer: - return _run_native_viewer(args, resolved) - return _run_native_headless(args, resolved) - - -def cli(argv: Sequence[str] | None = None) -> None: - """Run policy evaluation from the unified EmbodiChain CLI.""" - try: - report = run(parse_args(argv)) - except ( - FileNotFoundError, - ImportError, - KeyError, - RuntimeError, - TypeError, - ValueError, - ) as error: - raise SystemExit(f"eval-policy: {error}") from error - print(f"Evaluation report: {report}") - - -def _resolve_input(args: argparse.Namespace) -> EvaluationInput: - if args.run is not None: - manifest = RunManifest.load(args.run) - requested = args.checkpoint or "latest" - if requested in {"best", "latest"}: - selected, checkpoint = manifest.select_checkpoint(requested) - else: - selected = "explicit" - candidate = Path(requested).expanduser() - checkpoint = ( - candidate.resolve() - if candidate.is_absolute() - else (manifest.root / candidate).resolve() - ) - configs = dict(manifest.configs) - if args.config is not None: - configs["train"] = Path(args.config).expanduser().resolve() - if args.gym_config is not None: - configs["gym"] = Path(args.gym_config).expanduser().resolve() - return EvaluationInput( - checkpoint=checkpoint, - profile=args.profile, - configs=configs, - run=manifest.root, - requested_checkpoint=requested, - selected_checkpoint=selected, - ) - if args.checkpoint is None: - raise ValueError("--checkpoint is required without RUN") - configs = {} - if args.config is not None: - configs["train"] = Path(args.config).expanduser().resolve() - if args.gym_config is not None: - configs["gym"] = Path(args.gym_config).expanduser().resolve() - if args.profile is None and "train" not in configs: - raise ValueError("--config is required for an EmbodiChain checkpoint") - return EvaluationInput( - checkpoint=Path(args.checkpoint).expanduser().resolve(), - profile=args.profile, - configs=configs, - run=None, - requested_checkpoint=args.checkpoint, - selected_checkpoint="explicit", - ) - - -def _run_native_headless( - args: argparse.Namespace, - resolved: EvaluationInput, -) -> Path: - native = _build_native_runtime(args, resolved, viewer=False) - episodes = ( - args.episodes - if args.episodes is not None - else int(native.trainer.get("num_eval_episodes", 5)) - ) - try: - metrics = evaluate_episodes( - policy=native.runtime.policy, - env=native.runtime.env, - num_episodes=episodes, - device=native.device, - seed=native.seed, - ) - finally: - native.runtime.close() - _flush_simulator(native.uses_simulator) - return write_evaluation_report( - _output_parent(args.output, resolved), - _headless_report(native, resolved, episodes, metrics), - ) - - -def _run_native_viewer( - args: argparse.Namespace, - resolved: EvaluationInput, -) -> Path: - from .viewer import evaluate_native_viewer - - native = _build_native_runtime(args, resolved, viewer=True) - try: - result = evaluate_native_viewer( - native.runtime, - seed=native.seed, - episodes=args.episodes, - control_steps=args.control_steps, - duration=args.duration, - termination_behavior=args.termination_behavior or "auto_reset", - ) - finally: - _flush_simulator(True) - return write_evaluation_report( - _output_parent(args.output, resolved), - _viewer_report(result, native, resolved), - ) - - -def _run_profile(args: argparse.Namespace, resolved: EvaluationInput) -> Path: - from .bridge import evaluate_motion_profile - from .profile import MotionProfileRequest, build_motion_profile - - device = _torch_device(args.device or "cpu") - renderer = args.renderer or "hybrid" - profile = build_motion_profile( - resolved.profile, - MotionProfileRequest( - checkpoint=resolved.checkpoint, - device=device, - configs=resolved.configs, - resource_root=( - None if args.resource_root is None else Path(args.resource_root) - ), - renderer=renderer, - ), - ) - for warning in profile.warnings: - print(f"Warning: {warning}", file=sys.stderr) - result = evaluate_motion_profile( - profile, - episodes=args.episodes if args.episodes is not None else 1, - viewer=args.viewer, - control_steps=args.control_steps, - duration=args.duration, - command=None if args.command is None else tuple(args.command), - scene_config=args.scene_config or "standard", - physics_backend=args.physics_backend, - simulation_device=args.sim_device or "cpu", - renderer=renderer, - gpu_id=args.gpu_id, - termination_behavior=args.termination_behavior, - cache_dir=args.cache_dir, - offline=args.offline, - ) - return write_evaluation_report( - _output_parent(args.output, resolved), - _profile_report(result, resolved, device), - ) - - -def _build_native_runtime( - args: argparse.Namespace, - resolved: EvaluationInput, - *, - viewer: bool, -) -> NativeRuntime: - train_config = resolved.configs.get("train") - if train_config is None: - raise ValueError("Training config is required for an EmbodiChain checkpoint") - config = load_config(train_config) - config["trainer"] = dict(config["trainer"]) - gym_config = resolved.configs.get("gym") - if gym_config is not None: - config["trainer"]["gym_config"] = str(gym_config) - trainer = config["trainer"] - device = _torch_device(args.device or trainer.get("device", "cpu")) - simulation_device = _simulation_device(args, device) - seed = int( - args.seed - if args.seed is not None - else trainer.get("eval_seed", int(trainer.get("seed", 1)) + 10_000) - ) - np.random.seed(seed) - torch.manual_seed(seed) - if device.type == "cuda": - torch.cuda.manual_seed_all(seed) - uses_simulator = "gym_config" in trainer - if viewer and not uses_simulator: - raise ValueError("--viewer requires a simulator training task") - renderer = args.renderer or str(trainer.get("renderer", "hybrid")) - num_envs = ( - 1 - if viewer - else int( - args.num_envs - if args.num_envs is not None - else trainer.get("num_eval_envs", 4) - ) - ) - if uses_simulator: - runtime = build_gym_policy_runtime( - config, - device=device, - simulation_device=simulation_device, - num_envs=num_envs, - headless=not viewer, - renderer=renderer, - gpu_id=args.gpu_id, - config_dir=train_config.parent, - ) - else: - runtime = build_learning_policy_runtime( - config, - device=device, - num_envs=num_envs, - ) - try: - runtime.policy.load_state_dict(_load_policy_state_dict(resolved.checkpoint)) - except Exception: - runtime.close() - _flush_simulator(uses_simulator) - raise - return NativeRuntime( - runtime=runtime, - device=device, - simulation_device=simulation_device, - seed=seed, - renderer=renderer, - uses_simulator=uses_simulator, - trainer=trainer, - ) - - -def _validate_native_options(args: argparse.Namespace) -> None: - profile_options = { - "--resource-root": args.resource_root, - "--command": args.command, - "--physics-backend": args.physics_backend, - "--scene-config": args.scene_config, - "--cache-dir": args.cache_dir, - "--offline": args.offline, - } - selected = [ - name for name, value in profile_options.items() if value not in (None, False) - ] - if selected: - raise ValueError(f"{', '.join(selected)} requires --profile") - if not args.viewer and ( - args.control_steps is not None - or args.duration is not None - or args.termination_behavior is not None - ): - raise ValueError( - "--control-steps, --duration, and --termination-behavior require --viewer" - ) - - -def _load_policy_state_dict(checkpoint: Path) -> Mapping[str, Any]: - payload = torch.load(checkpoint, map_location="cpu", weights_only=True) - if not isinstance(payload, Mapping) or not isinstance( - payload.get("policy"), Mapping - ): - raise TypeError("Checkpoint must contain a 'policy' state mapping") - return payload["policy"] - - -def _torch_device(value: str) -> torch.device: - device = torch.device(value) - if device.type == "cuda": - index = ( - device.index if device.index is not None else torch.cuda.current_device() - ) - torch.cuda.set_device(index) - return torch.device(f"cuda:{index}") - if device.type != "cpu": - raise ValueError(f"Unsupported device type: {device.type}") - return device - - -def _simulation_device( - args: argparse.Namespace, - inference_device: torch.device, -) -> torch.device: - if args.sim_device == "gpu": - return _torch_device(f"cuda:{args.gpu_id}") - if args.sim_device == "cpu": - return torch.device("cpu") - return inference_device - - -def _flush_simulator(enabled: bool) -> None: - if enabled: - from embodichain.lab.sim.sim_manager import SimulationManager - - SimulationManager.flush_cleanup_queue() - - -def _output_parent(configured: str | None, resolved: EvaluationInput) -> Path: - if configured is not None: - return Path(configured) - if resolved.run is not None: - return resolved.run / "evaluations" - return resolved.checkpoint.parent / "evaluations" - - -def _checkpoint_inputs(resolved: EvaluationInput) -> dict[str, Any]: - return { - "run": resolved.run, - "checkpoint": { - "path": resolved.checkpoint, - "requested": resolved.requested_checkpoint, - "selected": resolved.selected_checkpoint, - }, - "configs": resolved.configs, - } - - -def _headless_report( - native: NativeRuntime, - resolved: EvaluationInput, - episodes: int, - metrics: Mapping[str, float], -) -> dict[str, Any]: - return { - "mode": "headless", - "inputs": { - **_checkpoint_inputs(resolved), - "task_id": native.runtime.env_id, - "seed": native.seed, - "num_envs": int(native.runtime.env.num_envs), - "device": str(native.device), - "embodichain_version": __version__, - }, - "result": {"episodes": episodes, "metrics": metrics}, - } - - -def _viewer_report( - result: Any, - native: NativeRuntime, - resolved: EvaluationInput, -) -> dict[str, Any]: - import dexsim - - return { - "mode": "viewer", - "inputs": { - **_checkpoint_inputs(resolved), - "task_id": result.task_id, - "seed": native.seed, - "inference_device": str(native.device), - "simulation_device": str(native.simulation_device), - "renderer": native.renderer, - "embodichain_version": __version__, - "dexsim_version": getattr(dexsim, "__version__", None), - "dexsim_commit": getattr(dexsim, "__commit_id__", None), - }, - "result": { - "reason": result.reason, - "simulation_time": result.simulation_time, - "simulation_steps": result.simulation_steps, - "control_steps": result.control_steps, - "requested_duration": result.requested_duration, - "effective_duration": result.effective_duration, - "episodes": result.episodes, - "metrics": result.metrics, - }, - } - - -def _profile_report( - result: Any, - resolved: EvaluationInput, - device: torch.device, -) -> dict[str, Any]: - import dexsim - - return { - "mode": "viewer" if result.viewer else "headless", - "inputs": { - **_checkpoint_inputs(resolved), - "profile": { - "id": result.profile.profile_id, - "provider_version": result.profile.provider_version, - "provenance": result.profile.provenance, - "warnings": result.profile.warnings, - }, - "policy_spec": result.policy_spec, - "scene_config": result.scene_config, - "inference_device": str(device), - "embodichain_version": __version__, - "dexsim_version": getattr(dexsim, "__version__", None), - "dexsim_commit": getattr(dexsim, "__commit_id__", None), - }, - "result": { - "episodes": result.episodes, - "summary": result.summary, - }, - } diff --git a/embodichain/learning/rl/policy_evaluation/manifest.py b/embodichain/learning/rl/policy_evaluation/manifest.py deleted file mode 100644 index cc8baa75d..000000000 --- a/embodichain/learning/rl/policy_evaluation/manifest.py +++ /dev/null @@ -1,193 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Index a training run for standalone policy evaluation.""" - -from __future__ import annotations - -import json -import shutil -from collections.abc import Mapping -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -__all__ = ["RUN_MANIFEST_NAME", "RunManifest", "write_run_manifest"] - -RUN_MANIFEST_NAME = "run-manifest.json" - - -@dataclass(frozen=True) -class RunManifest: - """Resolved paths from one EmbodiChain training run.""" - - root: Path - configs: Mapping[str, Path] - checkpoints: Mapping[str, Path | None] - - def __post_init__(self) -> None: - object.__setattr__(self, "root", Path(self.root).resolve()) - object.__setattr__(self, "configs", dict(self.configs)) - object.__setattr__(self, "checkpoints", dict(self.checkpoints)) - - @classmethod - def load(cls, run: str | Path) -> RunManifest: - """Load ``run-manifest.json`` and resolve its referenced files. - - Args: - run: EmbodiChain training run directory. - - Returns: - Resolved manifest. - """ - root = Path(run).expanduser().resolve() - path = root / RUN_MANIFEST_NAME - if not path.is_file(): - raise FileNotFoundError(f"Run manifest does not exist: {path}") - value = json.loads(path.read_text(encoding="utf-8")) - if not isinstance(value, Mapping) or value.get("schema_version") != 1: - raise ValueError(f"Unsupported run manifest: {path}") - configs = _resolve_group(root, value.get("configs"), "configs") - checkpoints = _resolve_group( - root, - value.get("checkpoints"), - "checkpoints", - allow_none=True, - ) - return cls(root, configs, checkpoints) - - def select_checkpoint(self, requested: str = "latest") -> tuple[str, Path]: - """Select ``best`` or ``latest`` and return its resolved path. - - Args: - requested: Checkpoint role. - - Returns: - Selected role and checkpoint path. ``best`` uses ``latest`` when - the training run has no best checkpoint. - """ - if requested not in {"best", "latest"}: - raise ValueError("checkpoint role must be best or latest") - selected = requested - checkpoint = self.checkpoints.get(selected) - if checkpoint is None and requested == "best": - selected = "latest" - checkpoint = self.checkpoints.get(selected) - if checkpoint is None: - raise FileNotFoundError( - f"Run manifest has no {requested} checkpoint: {self.root}" - ) - return selected, checkpoint - - -def write_run_manifest( - run: str | Path, - *, - train_config: str | Path, - latest_checkpoint: str | Path, - best_checkpoint: str | Path | None = None, - gym_config: str | Path | None = None, -) -> Path: - """Snapshot training configs and write the minimal run manifest. - - Args: - run: Training run directory containing the checkpoints. - train_config: Training config used for the run. - latest_checkpoint: Final saved checkpoint. - best_checkpoint: Best checkpoint when evaluation selected one. - gym_config: Referenced task config when the trainer uses one. - - Returns: - Written manifest path. - """ - root = Path(run).expanduser().resolve() - root.mkdir(parents=True, exist_ok=True) - config_dir = root / "configs" - config_dir.mkdir(exist_ok=True) - configs = { - "train": _snapshot_config(train_config, config_dir, "train"), - } - if gym_config is not None: - configs["gym"] = _snapshot_config(gym_config, config_dir, "gym") - checkpoints = { - "best": _relative_file(root, best_checkpoint), - "latest": _relative_file(root, latest_checkpoint), - } - value: dict[str, Any] = { - "schema_version": 1, - "configs": configs, - "checkpoints": checkpoints, - } - path = root / RUN_MANIFEST_NAME - path.write_text( - json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n", - encoding="utf-8", - ) - return path - - -def _snapshot_config(source: str | Path, target: Path, name: str) -> str: - path = Path(source).expanduser().resolve() - if not path.is_file(): - raise FileNotFoundError(f"Training config does not exist: {path}") - suffix = path.suffix.lower() if path.suffix else ".yaml" - destination = target / f"{name}{suffix}" - shutil.copyfile(path, destination) - return destination.relative_to(target.parent).as_posix() - - -def _relative_file(root: Path, value: str | Path | None) -> str | None: - if value is None: - return None - path = Path(value).expanduser().resolve() - if not path.is_file(): - raise FileNotFoundError(f"Training checkpoint does not exist: {path}") - try: - return path.relative_to(root).as_posix() - except ValueError as error: - raise ValueError(f"Training checkpoint is outside its run: {path}") from error - - -def _resolve_group( - root: Path, - value: object, - field: str, - *, - allow_none: bool = False, -) -> dict[str, Path | None]: - if not isinstance(value, Mapping): - raise TypeError(f"Run manifest {field} must be a mapping") - result: dict[str, Path | None] = {} - for name, reference in value.items(): - if reference is None and allow_none: - result[str(name)] = None - continue - if not isinstance(reference, str) or not reference: - raise TypeError(f"Run manifest {field}.{name} must be a path") - relative = Path(reference) - if relative.is_absolute(): - raise ValueError(f"Run manifest {field}.{name} must be relative") - path = (root / relative).resolve() - try: - path.relative_to(root) - except ValueError as error: - raise ValueError( - f"Run manifest {field}.{name} escapes the run directory" - ) from error - if not path.is_file(): - raise FileNotFoundError(f"Run manifest file does not exist: {path}") - result[str(name)] = path - return result diff --git a/embodichain/learning/rl/policy_evaluation/profile.py b/embodichain/learning/rl/policy_evaluation/profile.py deleted file mode 100644 index a215a084b..000000000 --- a/embodichain/learning/rl/policy_evaluation/profile.py +++ /dev/null @@ -1,128 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""External Policy Profile registration and construction.""" - -from __future__ import annotations - -from collections.abc import Callable, Mapping -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import torch - -__all__ = [ - "MotionProfile", - "MotionProfileRequest", - "build_motion_profile", - "register_motion_profile", -] - - -@dataclass(frozen=True) -class MotionProfileRequest: - """Checkpoint, configs, and runtime choices supplied to a provider.""" - - checkpoint: Path - device: torch.device - configs: Mapping[str, Path] = field(default_factory=dict) - resource_root: Path | None = None - renderer: str = "hybrid" - - def __post_init__(self) -> None: - checkpoint = Path(self.checkpoint).expanduser().resolve() - if not checkpoint.is_file(): - raise FileNotFoundError(f"Motion checkpoint does not exist: {checkpoint}") - configs = { - name: Path(path).expanduser().resolve() - for name, path in self.configs.items() - } - for name, path in configs.items(): - if not path.is_file(): - raise FileNotFoundError( - f"Motion config {name!r} does not exist: {path}" - ) - root = ( - None - if self.resource_root is None - else Path(self.resource_root).expanduser().resolve() - ) - object.__setattr__(self, "checkpoint", checkpoint) - object.__setattr__(self, "configs", configs) - object.__setattr__(self, "resource_root", root) - - -@dataclass(frozen=True) -class MotionProfile: - """DexSim Policy Spec and report metadata built by one provider.""" - - profile_id: str - policy_spec: Mapping[str, Any] - provider_version: int = 1 - provenance: Mapping[str, Any] = field(default_factory=dict) - warnings: tuple[str, ...] = () - - def __post_init__(self) -> None: - object.__setattr__(self, "policy_spec", dict(self.policy_spec)) - object.__setattr__(self, "provenance", dict(self.provenance)) - object.__setattr__(self, "warnings", tuple(self.warnings)) - - -MotionProfileProvider = Callable[[MotionProfileRequest], MotionProfile] -_PROVIDERS: dict[str, MotionProfileProvider] = {} - - -def register_motion_profile(name: str, provider: MotionProfileProvider) -> None: - """Register a Motion Profile provider under its CLI name. - - Args: - name: Stable profile name. - provider: Callable that builds one :class:`MotionProfile`. - """ - if not name: - raise ValueError("Motion profile name must not be empty") - if name in _PROVIDERS: - raise ValueError(f"Motion profile is already registered: {name}") - _PROVIDERS[name] = provider - - -def build_motion_profile( - name: str, - request: MotionProfileRequest, -) -> MotionProfile: - """Build one profile with its registered provider. - - Args: - name: Registered profile name. - request: Checkpoint, configs, and runtime choices. - - Returns: - Provider-built Motion Profile. - """ - try: - provider = _PROVIDERS[name] - except KeyError: - available = ", ".join(sorted(_PROVIDERS)) or "none" - raise ValueError( - f"Unknown motion profile {name!r}; available: {available}" - ) from None - profile = provider(request) - if profile.profile_id != name: - raise ValueError( - f"Motion provider {name!r} returned profile {profile.profile_id!r}" - ) - return profile diff --git a/embodichain/learning/rl/policy_evaluation/report.py b/embodichain/learning/rl/policy_evaluation/report.py deleted file mode 100644 index 545b41f2e..000000000 --- a/embodichain/learning/rl/policy_evaluation/report.py +++ /dev/null @@ -1,78 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Write timestamped policy evaluation reports.""" - -from __future__ import annotations - -import json -import math -from collections.abc import Mapping, Sequence -from datetime import datetime, timezone -from pathlib import Path -from typing import Any - -__all__ = ["write_evaluation_report"] - - -def write_evaluation_report( - parent: str | Path, - payload: Mapping[str, Any], -) -> Path: - """Write ``evaluation.json`` under a new timestamped directory. - - Args: - parent: Output parent directory. - payload: Evaluation inputs and results. - - Returns: - Written report path. - """ - output = Path(parent).expanduser().resolve() - output.mkdir(parents=True, exist_ok=True) - stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") - directory = output / f"{stamp}-policy" - directory.mkdir() - report = { - "schema_version": 1, - "created_at": datetime.now(timezone.utc).isoformat(), - **dict(payload), - } - path = directory / "evaluation.json" - path.write_text( - json.dumps( - _json_value(report), - indent=2, - sort_keys=True, - ensure_ascii=False, - allow_nan=False, - ) - + "\n", - encoding="utf-8", - ) - return path - - -def _json_value(value: Any) -> Any: - if isinstance(value, Path): - return str(value) - if isinstance(value, float): - return value if math.isfinite(value) else None - if isinstance(value, Mapping): - return {str(name): _json_value(item) for name, item in value.items()} - if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): - return [_json_value(item) for item in value] - return value diff --git a/embodichain/learning/rl/policy_evaluation/viewer.py b/embodichain/learning/rl/policy_evaluation/viewer.py deleted file mode 100644 index 509176684..000000000 --- a/embodichain/learning/rl/policy_evaluation/viewer.py +++ /dev/null @@ -1,480 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Connect an EmbodiChain task Viewer to Motion Policy Evaluator.""" - -from __future__ import annotations - -import math -import time -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Any - -import torch - -from dexsim.kit.motion_policy import ( - EvaluationFrame, - PolicyContext, - PolicyOutput, - RunOptions, - create_motion_policy_evaluator, -) -from dexsim.kit.motion_policy.types import EnvironmentStep - -from embodichain.learning.rl.evaluation import ( - convert_policy_action_for_env, - infer_policy_action, -) -from embodichain.learning.rl.runtime import PolicyRuntime - -__all__ = [ - "EmbodiChainTaskEnvironment", - "EmbodiChainTaskPolicyAdapter", - "NativeViewerResult", - "evaluate_native_viewer", -] - -_MISSING = object() - - -@dataclass(frozen=True) -class NativeViewerResult: - """Result of visualizing one Policy in its EmbodiChain task.""" - - task_id: str - reason: str - simulation_time: float - simulation_steps: int - control_steps: int - effective_duration: float - requested_duration: float | None - episodes: tuple[Mapping[str, float | int | bool | str], ...] - metrics: Mapping[str, float] - - -class EmbodiChainTaskPolicyAdapter: - """Run an EmbodiChain Policy from the task observation in each frame.""" - - def __init__(self, policy: torch.nn.Module, device: torch.device): - self.policy = policy - self.device = device - self._previous_training = policy.training - - def setup(self, context: PolicyContext) -> None: - """Select deterministic inference for this evaluation.""" - del context - self.policy.eval() - - def reset(self, frame: EvaluationFrame) -> None: - """Validate that the Environment supplied the next observation.""" - if frame.observation is None: - raise RuntimeError("EmbodiChain task frame has no observation") - - @torch.no_grad() - def infer(self, frame: EvaluationFrame) -> PolicyOutput: - """Run the same observation and deterministic Policy path as RL evaluation.""" - if frame.observation is None: - raise RuntimeError("EmbodiChain task frame has no observation") - action = infer_policy_action( - self.policy, - frame.observation, - device=self.device, - num_envs=1, - ) - return PolicyOutput(action=action) - - def metrics(self) -> dict[str, float]: - """Return Policy-side metrics.""" - return {} - - def close(self) -> None: - """Restore the Policy mode used before evaluation.""" - self.policy.train(self._previous_training) - - -class EmbodiChainTaskEnvironment: - """Expose one original EmbodiChain RL Environment to the Evaluator.""" - - def __init__( - self, - env: Any, - *, - seed: int, - ) -> None: - if int(env.num_envs) != 1: - raise ValueError("Visual task evaluation requires num_envs=1") - self.env = env - self._base_env = getattr(env, "unwrapped", env) - world = self._world() - if world is None or not world.is_window_initialized(): - raise ValueError( - "Viewer evaluation requires an initialized simulator window" - ) - self._seed = seed - self._first_reset = True - self._reset_key_down = False - self._control_step = 0 - self._frame: EvaluationFrame | None = None - self._episode_return = 0.0 - self._episode_length = 0 - self._episodes: list[dict[str, float | int | bool | str]] = [] - self._reported_metrics: dict[str, float] = {} - self._closed = False - self._policy_context = _policy_context_from_env(self._base_env) - self._previous_no_auto_reset = getattr( - self._base_env, - "_demo_no_auto_reset", - _MISSING, - ) - self._base_env._demo_no_auto_reset = True - - @property - def policy_context(self) -> PolicyContext: - """Return the timing used by the original task Environment.""" - return self._policy_context - - @property - def physics_backend(self) -> str: - """Return the backend selected by the original task Environment.""" - return "default" - - @property - def viewer_is_open(self) -> bool: - """Return whether the original task Viewer remains open.""" - world = self._world() - return bool(world is not None and world.is_window_initialized()) - - @property - def current_frame(self) -> EvaluationFrame: - """Return the latest observation and task state.""" - if self._frame is None: - raise RuntimeError("Environment has not been reset") - return self._frame - - @property - def episodes(self) -> tuple[Mapping[str, float | int | bool | str], ...]: - """Return completed episode summaries.""" - return tuple(self._episodes) - - def open_viewer(self, title: str) -> None: - """Apply the evaluation title to the task Viewer.""" - self._world().get_windows().set_window_title(title) - - def reset(self) -> EvaluationFrame: - """Run the task's original reset and return its observation.""" - kwargs = {"seed": self._seed} if self._first_reset else {} - observation, info = self.env.reset(**kwargs) - self._first_reset = False - self._control_step = 0 - self._episode_return = 0.0 - self._episode_length = 0 - self._frame = self._make_frame(observation, {"info": info}) - return self._frame - - def poll(self) -> str | None: - """Report when the native Viewer is closed or Escape is pressed.""" - world = self._world() - if world is None or not world.is_window_initialized(): - return "viewer closed" - from dexsim.types import InputKey - - native = world.get_windows().native() - if native.key_state(InputKey.SCANCODE_ESCAPE): - return "viewer closed" - reset_down = bool(native.key_state(InputKey.SCANCODE_BACKSPACE)) - reset_pressed = reset_down and not self._reset_key_down - self._reset_key_down = reset_down - if reset_pressed: - return "manual reset" - return None - - def step(self, action: object) -> EnvironmentStep: - """Apply one raw Policy action through the task's original action path.""" - if not isinstance(action, torch.Tensor): - raise TypeError("EmbodiChain Policy action must be a torch.Tensor") - started = time.perf_counter() - env_action = convert_policy_action_for_env(self.env, action) - observation, reward, terminated, truncated, info = self.env.step(env_action) - reward_value = _single_float(reward, "reward") - terminated_value = _single_bool(terminated, "terminated") - truncated_value = _single_bool(truncated, "truncated") - self._control_step += 1 - self._episode_return += reward_value - self._episode_length += 1 - task_state = { - "reward": reward, - "terminated": terminated, - "truncated": truncated, - "info": info, - } - self._frame = self._make_frame(observation, task_state) - reason = _termination_reason(info, terminated_value, truncated_value) - metrics = _step_metrics(info, reward_value) - self._reported_metrics.update(metrics) - if reason is not None: - success = _info_bool(info, "success") - self._episodes.append( - { - "index": len(self._episodes), - "reason": reason, - "reward": self._episode_return, - "length": self._episode_length, - "success": success, - } - ) - remaining = self._policy_context.policy_dt - (time.perf_counter() - started) - if remaining > 0.0: - time.sleep(remaining) - return EnvironmentStep( - frame=self._frame, - termination_reason=reason, - metrics=metrics, - ) - - def metrics(self) -> dict[str, float]: - """Return task metrics and completed episode aggregates.""" - result = dict(self._reported_metrics) - if self._episodes: - count = len(self._episodes) - result.update( - { - "eval/avg_reward": sum( - float(episode["reward"]) for episode in self._episodes - ) - / count, - "eval/avg_length": sum( - float(episode["length"]) for episode in self._episodes - ) - / count, - "eval/success_rate": sum( - bool(episode["success"]) for episode in self._episodes - ) - / count, - } - ) - return result - - def wait_for_reset_or_close(self) -> str: - """Keep a paused Viewer responsive until it is closed. - - ``MotionPolicyEvaluator`` calls this method after a task termination - when the selected behavior is ``pause``. - """ - while self.viewer_is_open: - event = self.poll() - if event is not None: - return event - world = self._world() - if world is not None: - world.update(0.0) - time.sleep(0.01) - return "viewer closed" - - def close(self) -> None: - """Close the original task Environment.""" - if self._closed: - return - if self._previous_no_auto_reset is _MISSING: - delattr(self._base_env, "_demo_no_auto_reset") - else: - self._base_env._demo_no_auto_reset = self._previous_no_auto_reset - if getattr(self._base_env, "sim", None) is not None: - self._base_env.close(exit_process=False) - else: - self.env.close() - self._closed = True - - def _make_frame( - self, - observation: object, - task_state: Mapping[str, object], - ) -> EvaluationFrame: - simulation_step = ( - self._control_step * self._policy_context.sim_steps_per_control - ) - return EvaluationFrame( - control_step=self._control_step, - policy_time=self._control_step * self._policy_context.policy_dt, - simulation_step=simulation_step, - simulation_time=simulation_step * self._policy_context.physics_dt, - observation=observation, - task_state=task_state, - ) - - def _world(self) -> Any | None: - sim = getattr(self._base_env, "sim", None) - return None if sim is None else sim.get_world() - - -def evaluate_native_viewer( - runtime: PolicyRuntime, - *, - seed: int, - episodes: int | None, - control_steps: int | None, - duration: float | None, - termination_behavior: str = "auto_reset", -) -> NativeViewerResult: - """Visualize an EmbodiChain Policy in the task used for training.""" - if episodes is not None and episodes <= 0: - raise ValueError("episodes must be positive") - if control_steps is not None and control_steps <= 0: - raise ValueError("control_steps must be positive") - if duration is not None and (duration <= 0.0 or not math.isfinite(duration)): - raise ValueError("duration must be finite and positive") - if control_steps is not None and duration is not None: - raise ValueError("control_steps and duration are mutually exclusive") - if termination_behavior == "continue": - raise ValueError("Native task evaluation supports pause or auto_reset") - - environment = None - adapter = None - evaluator = None - try: - environment = EmbodiChainTaskEnvironment( - runtime.env, - seed=seed, - ) - adapter = EmbodiChainTaskPolicyAdapter( - runtime.policy, - runtime.device, - ) - if duration is not None: - control_steps = math.ceil( - duration / environment.policy_context.policy_dt - 1e-12 - ) - total_steps = 0 - reason = "viewer closed" - options = RunOptions( - headless=False, - termination_behavior=( - "continue" if termination_behavior == "auto_reset" else "pause" - ), - ) - evaluator = create_motion_policy_evaluator( - options=options, - adapter=adapter, - environment=environment, - title=f"{runtime.env_id} - EmbodiChain", - ) - evaluator.reset() - while True: - if control_steps is not None and total_steps >= control_steps: - reason = "control steps reached" - break - if episodes is not None and len(environment.episodes) >= episodes: - reason = "episode target reached" - break - completed_before = len(environment.episodes) - result = evaluator.step() - if result.advanced: - total_steps += 1 - if len(environment.episodes) > completed_before: - if episodes is not None and len(environment.episodes) >= episodes: - reason = "episode target reached" - break - if termination_behavior == "auto_reset": - evaluator.reset() - continue - if result.reason is not None and not result.reset_performed: - reason = result.reason - break - episode_results = environment.episodes - metrics = environment.metrics() - context = environment.policy_context - finally: - if evaluator is not None: - evaluator.close() - elif environment is not None: - if adapter is not None: - adapter.close() - environment.close() - else: - runtime.close() - - simulation_steps = total_steps * context.sim_steps_per_control - return NativeViewerResult( - task_id=runtime.env_id, - reason=reason, - simulation_time=simulation_steps * context.physics_dt, - simulation_steps=simulation_steps, - control_steps=total_steps, - effective_duration=total_steps * context.policy_dt, - requested_duration=duration, - episodes=episode_results, - metrics=metrics, - ) - - -def _single_float(value: object, name: str) -> float: - tensor = torch.as_tensor(value).reshape(-1) - if tensor.numel() != 1: - raise ValueError(f"Native task {name} must contain one value") - return float(tensor.item()) - - -def _single_bool(value: object, name: str) -> bool: - tensor = torch.as_tensor(value, dtype=torch.bool).reshape(-1) - if tensor.numel() != 1: - raise ValueError(f"Native task {name} must contain one value") - return bool(tensor.item()) - - -def _info_bool(info: object, name: str) -> bool: - if not isinstance(info, Mapping) or name not in info: - return False - return _single_bool(info[name], f"info.{name}") - - -def _termination_reason( - info: object, - terminated: bool, - truncated: bool, -) -> str | None: - if _info_bool(info, "success"): - return "success" - if _info_bool(info, "fail"): - return "failure" - if truncated: - return "time limit" - if terminated: - return "terminated" - return None - - -def _step_metrics(info: object, reward: float) -> dict[str, float]: - result = {"reward": reward} - if not isinstance(info, Mapping): - return result - metrics = info.get("metrics") - if not isinstance(metrics, Mapping): - return result - for name, value in metrics.items(): - tensor = torch.as_tensor(value).reshape(-1) - if tensor.numel() == 1: - result[str(name)] = float(tensor.item()) - return result - - -def _policy_context_from_env(env: Any) -> PolicyContext: - """Read timing from the simulator task.""" - return PolicyContext( - robot=None, - physics_dt=float(env.physics_dt), - sim_steps_per_control=int(env.cfg.sim_steps_per_control), - policy_dt=float(env.step_dt), - ) diff --git a/embodichain/learning/rl/runtime.py b/embodichain/learning/rl/runtime.py deleted file mode 100644 index ed8f0c7f7..000000000 --- a/embodichain/learning/rl/runtime.py +++ /dev/null @@ -1,312 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Shared environment and Policy construction for RL training and evaluation.""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import torch - -from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules -from embodichain.lab.gym.utils.profiler import EnvProfilerCfg -from embodichain.lab.gym.utils.registration import build_env -from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.sim.cfg import RenderCfg -from embodichain.learning.rl.env import build_learning_env -from embodichain.learning.rl.models import build_mlp_from_cfg, build_policy -from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation -from embodichain.utils.utility import load_config - -__all__ = [ - "PolicyRuntime", - "build_gym_policy_runtime", - "build_learning_policy_runtime", -] - - -@dataclass(frozen=True) -class _GymEnvironmentRuntime: - """A simulator task reconstructed from one training configuration.""" - - env: Any - env_id: str - env_cfg: Any - gym_config: dict[str, Any] - gym_config_path: Path - - -@dataclass(frozen=True) -class PolicyRuntime: - """An Environment and Policy reconstructed from one training configuration.""" - - env: Any - policy: torch.nn.Module - device: torch.device - env_id: str - env_cfg: Any | None = None - gym_config: dict[str, Any] | None = None - gym_config_path: Path | None = None - - def close(self) -> None: - """Close the Environment without terminating the current process.""" - _close_environment(self.env) - - -def _resolve_config_reference( - value: str | Path, - *, - base_dir: str | Path | None = None, -) -> Path: - """Resolve a referenced config relative to its containing config file.""" - path = Path(value).expanduser() - if path.is_absolute(): - return path - if base_dir is not None: - candidate = Path(base_dir).expanduser().resolve() / path - if candidate.exists(): - return candidate - return path - - -def _build_learning_environment( - config: dict[str, Any], - *, - device: torch.device, - num_envs: int, -) -> tuple[str, Any]: - """Build the lightweight Environment declared by a training config.""" - env_block = config["trainer"]["learning_env"] - if isinstance(env_block, str): - env_name = env_block - env_config: dict[str, Any] = {} - else: - env_name = env_block["name"] - env_config = dict(env_block.get("cfg", {})) - return str(env_name), build_learning_env( - str(env_name), - num_envs=num_envs, - device=device, - **env_config, - ) - - -def build_learning_policy_runtime( - config: dict[str, Any], - *, - device: torch.device, - num_envs: int, -) -> PolicyRuntime: - """Build a lightweight Environment and its configured Policy.""" - env_name, env = _build_learning_environment( - config, - device=device, - num_envs=num_envs, - ) - try: - policy = _build_learning_policy(config["policy"], env, device) - except Exception: - env.close() - raise - return PolicyRuntime(env, policy, device, env_name) - - -def _build_gym_environment( - config: dict[str, Any], - *, - simulation_device: torch.device, - num_envs: int | None, - headless: bool, - renderer: str, - gpu_id: int, - config_dir: str | Path | None = None, - profiler: EnvProfilerCfg | None = None, -) -> _GymEnvironmentRuntime: - """Build the simulator Environment declared by a training config.""" - trainer_cfg = config["trainer"] - gym_config_path = _resolve_config_reference( - trainer_cfg["gym_config"], - base_dir=config_dir, - ) - gym_config = load_config(gym_config_path) - env_cfg = config_to_cfg(gym_config, manager_modules=get_manager_modules()) - if num_envs is not None: - env_cfg.num_envs = int(num_envs) - if env_cfg.sim_cfg is None: - env_cfg.sim_cfg = SimulationManagerCfg() - env_cfg.sim_cfg.sim_device = simulation_device - env_cfg.sim_cfg.headless = headless - env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) - env_cfg.sim_cfg.gpu_id = ( - simulation_device.index - if simulation_device.type == "cuda" and simulation_device.index is not None - else gpu_id - ) - env_cfg.profiler = profiler - env = build_env(gym_config["id"], base_env_cfg=env_cfg) - return _GymEnvironmentRuntime( - env=env, - env_id=str(gym_config["id"]), - env_cfg=env_cfg, - gym_config=gym_config, - gym_config_path=gym_config_path.resolve(), - ) - - -def build_gym_policy_runtime( - config: dict[str, Any], - *, - device: torch.device, - num_envs: int | None, - headless: bool, - renderer: str, - gpu_id: int, - config_dir: str | Path | None = None, - profiler: EnvProfilerCfg | None = None, - simulation_device: torch.device | None = None, -) -> PolicyRuntime: - """Build a simulator task and the Policy declared by its training config.""" - task = _build_gym_environment( - config, - simulation_device=simulation_device or device, - num_envs=num_envs, - headless=headless, - renderer=renderer, - gpu_id=gpu_id, - config_dir=config_dir, - profiler=profiler, - ) - env = task.env - try: - sample_observation, _ = env.reset() - sample_observation_td = dict_to_tensordict(sample_observation, device) - observation_dim = int(flatten_dict_observation(sample_observation_td).shape[-1]) - action_manager = env.get_wrapper_attr("action_manager") - environment_action_dim = ( - action_manager.total_action_dim - if action_manager is not None - else len(env.get_wrapper_attr("active_joint_ids")) - ) - policy = _build_gym_policy( - config["policy"], - env=env, - device=device, - observation_dim=observation_dim, - action_dim=environment_action_dim, - ) - except Exception: - _close_environment(env) - raise - return PolicyRuntime( - env=env, - policy=policy, - device=device, - env_id=task.env_id, - env_cfg=task.env_cfg, - gym_config=task.gym_config, - gym_config_path=task.gym_config_path, - ) - - -def _build_gym_policy( - policy_block: dict[str, Any], - *, - env: Any, - device: torch.device, - observation_dim: int, - action_dim: int, -) -> torch.nn.Module: - configured_action_dim = int(policy_block.get("action_dim", action_dim)) - if configured_action_dim != action_dim: - raise ValueError( - f"Configured policy.action_dim={configured_action_dim} does not match " - f"env action dim {action_dim}." - ) - policy_name = str(policy_block["name"]).lower() - if policy_name == "actor_critic": - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - if actor_cfg is None or critic_cfg is None: - raise ValueError( - "ActorCritic requires policy.actor and policy.critic definitions." - ) - return build_policy( - policy_block, - env.flattened_observation_space, - env.action_space, - device, - actor=build_mlp_from_cfg(actor_cfg, observation_dim, action_dim), - critic=build_mlp_from_cfg(critic_cfg, observation_dim, 1), - ) - if policy_name == "actor_only": - actor_cfg = policy_block.get("actor") - if actor_cfg is None: - raise ValueError("ActorOnly requires a policy.actor definition.") - return build_policy( - policy_block, - env.flattened_observation_space, - env.action_space, - device, - actor=build_mlp_from_cfg(actor_cfg, observation_dim, action_dim), - ) - return build_policy( - policy_block, - env.observation_space, - env.action_space, - device, - ) - - -def _build_learning_policy( - policy_block: dict[str, Any], - env: Any, - device: torch.device, -) -> torch.nn.Module: - observation_dim = int(env.single_observation_space.shape[-1]) - action_dim = int(env.single_action_space.shape[-1]) - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - policy = build_policy( - policy_block, - env.single_observation_space, - env.single_action_space, - device, - actor=( - build_mlp_from_cfg(actor_cfg, observation_dim, action_dim) - if actor_cfg is not None - else None - ), - critic=( - build_mlp_from_cfg(critic_cfg, observation_dim, 1) - if critic_cfg is not None - else None - ), - ) - if "initial_log_std" in policy_block and hasattr(policy, "log_std"): - with torch.no_grad(): - policy.log_std.fill_(float(policy_block["initial_log_std"])) - return policy - - -def _close_environment(env: Any) -> None: - target = getattr(env, "unwrapped", env) - if getattr(target, "sim", None) is not None: - target.close(exit_process=False) - else: - env.close() diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 9556efd31..41c44a3f1 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -20,18 +20,16 @@ import os import time from collections.abc import Sequence -from copy import deepcopy from pathlib import Path import numpy as np import torch import wandb from torch.utils.tensorboard import SummaryWriter +from copy import deepcopy -from embodichain.learning.rl.models import get_registered_policy_names -from embodichain.learning.rl.policy_evaluation.manifest import ( - write_run_manifest, -) +from embodichain.learning.rl.models import build_policy, get_registered_policy_names +from embodichain.learning.rl.models import build_mlp_from_cfg from embodichain.learning.rl.algo import ( RolloutKind, build_algo, @@ -41,12 +39,9 @@ DifferentiableTrainer, DifferentiableTrainerCfg, ) -from embodichain.learning.rl.runtime import ( - _build_learning_environment, - build_gym_policy_runtime, - build_learning_policy_runtime, -) +from embodichain.learning.rl.env import build_learning_env from embodichain.learning.rl.routing import get_trainer_class +from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation from embodichain.learning.rl.utils.trainer import Trainer from embodichain.utils import logger from embodichain.lab.gym.utils.registration import ( @@ -54,13 +49,14 @@ discover_task_packages, execute_init_hooks, ) +from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules from embodichain.lab.gym.utils.profiler import EnvProfilerCfg from embodichain.utils.utility import load_config from embodichain.utils.module_utils import find_function_from_modules +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import RenderCfg from embodichain.lab.gym.envs.managers.cfg import EventCfg -_CAMERA_RECORDERS = {"record_camera_data", "record_camera_data_async"} - def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: """Parse command-line arguments. @@ -118,33 +114,51 @@ def _resolve_profile_output( return str(output.with_name(f"{output.stem}_rank{rank}{output.suffix}")) -def _event_params( - event_info: dict, - *, - run_base: str | Path, - phase: str, -) -> dict: - """Place default camera recordings under the current training run.""" - params = dict(event_info.get("params", {})) - function_name = str(event_info.get("func", "")).rsplit(".", 1)[-1] - if function_name in _CAMERA_RECORDERS: - params.setdefault("save_path", str(Path(run_base) / "videos" / phase)) - return params +def _build_learning_policy( + policy_block: dict, + env, + device: torch.device, +): + obs_dim = int(env.single_observation_space.shape[-1]) + action_dim = int(env.single_action_space.shape[-1]) + policy_name = policy_block["name"].lower() + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + actor = ( + build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) + if actor_cfg is not None + else None + ) + critic = ( + build_mlp_from_cfg(critic_cfg, obs_dim, 1) if critic_cfg is not None else None + ) + policy = build_policy( + policy_block, + env.single_observation_space, + env.single_action_space, + device, + actor=actor, + critic=critic, + ) + if "initial_log_std" in policy_block and hasattr(policy, "log_std"): + with torch.no_grad(): + policy.log_std.fill_(float(policy_block["initial_log_std"])) + return policy def _train_learning_env( cfg_data: dict, *, - config_path: str | Path, distributed: bool | None, profile: bool = False, -) -> dict[str, object]: +): """Train a lightweight registered environment through the unified CLI.""" if profile: raise ValueError( "--profile requires trainer.gym_config; learning_env is unsupported." ) trainer_cfg = cfg_data["trainer"] + policy_block = cfg_data["policy"] algorithm_block = cfg_data["algorithm"] distributed = ( bool(trainer_cfg.get("distributed", False)) @@ -168,25 +182,32 @@ def _train_learning_env( np.random.seed(seed) torch.manual_seed(seed) + env_block = trainer_cfg["learning_env"] + if isinstance(env_block, str): + env_name = env_block + env_cfg = {} + else: + env_name = env_block["name"] + env_cfg = dict(env_block.get("cfg", {})) num_envs = int(trainer_cfg.get("num_envs", 64)) - runtime = build_learning_policy_runtime( - cfg_data, + env = build_learning_env( + env_name, num_envs=num_envs, device=device, + **env_cfg, ) - env = runtime.env - policy = runtime.policy - env_name = runtime.env_id enable_eval = bool(trainer_cfg.get("enable_eval", False)) eval_env = None if enable_eval: - _eval_name, eval_env = _build_learning_environment( - cfg_data, + eval_env = build_learning_env( + env_name, num_envs=int(trainer_cfg.get("num_eval_envs", 16)), device=device, + **env_cfg, ) + policy = _build_learning_policy(policy_block, env, device) algorithm = build_algo( algorithm_block["name"], dict(algorithm_block.get("cfg", {})), @@ -270,7 +291,7 @@ def _train_learning_env( total_timesteps = int(trainer_cfg.get("total_timesteps", default_steps)) trainer.train(total_timesteps) trainer.save_checkpoint() - summary = trainer.get_summary() + return trainer.get_summary() finally: writer.close() if use_wandb: @@ -278,8 +299,6 @@ def _train_learning_env( env.close() if eval_env is not None: eval_env.close() - _write_policy_run_manifest(run_base, config_path, summary) - return summary def train_from_config( @@ -288,7 +307,7 @@ def train_from_config( *, profile: bool = False, profile_output: str | None = None, -) -> dict[str, object] | None: +): """Run training from a config file path. Args: @@ -297,9 +316,6 @@ def train_from_config( If None, use trainer.distributed from config. profile: Enable gym ``EnvProfiler`` on the training environment. profile_output: Optional JSON dump path for the profiling report. - - Returns: - The lightweight trainer summary, or ``None`` for simulator training. """ if profile_output is not None and not profile: raise ValueError("--profile_output requires --profile.") @@ -310,7 +326,6 @@ def train_from_config( if "learning_env" in trainer_cfg: return _train_learning_env( cfg_data, - config_path=config_path, distributed=distributed, profile=profile, ) @@ -426,11 +441,32 @@ def train_from_config( if use_wandb and rank == 0: wandb.init(project=wandb_project_name, name=exp_name, config=cfg_data) + gym_config_path = Path(trainer_cfg["gym_config"]) if rank == 0: logger.log_info(f"Current working directory: {Path.cwd()}") - profiler = ( - EnvProfilerCfg( + gym_config_data = load_config(str(gym_config_path)) + gym_env_cfg = config_to_cfg(gym_config_data, manager_modules=get_manager_modules()) + if num_envs is not None: + gym_env_cfg.num_envs = int(num_envs) + + # Ensure sim configuration mirrors runtime overrides + if gym_env_cfg.sim_cfg is None: + gym_env_cfg.sim_cfg = SimulationManagerCfg() + if device.type == "cuda": + gpu_index = device.index + if gpu_index is None: + gpu_index = torch.cuda.current_device() + gym_env_cfg.sim_cfg.sim_device = torch.device(f"cuda:{gpu_index}") + if hasattr(gym_env_cfg.sim_cfg, "gpu_id"): + gym_env_cfg.sim_cfg.gpu_id = gpu_index + else: + gym_env_cfg.sim_cfg.sim_device = torch.device("cpu") + gym_env_cfg.sim_cfg.headless = headless + gym_env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) + gym_env_cfg.sim_cfg.gpu_id = gpu_id + if profile: + gym_env_cfg.profiler = EnvProfilerCfg( enable_time=True, output_path=_resolve_profile_output( profile_output, @@ -438,37 +474,23 @@ def train_from_config( world_size=world_size, ), ) - if profile - else None - ) - runtime = build_gym_policy_runtime( - cfg_data, - device=device, - num_envs=num_envs, - headless=headless, - renderer=renderer, - gpu_id=gpu_id, - config_dir=Path(config_path).expanduser().resolve().parent, - profiler=profiler, - ) - env = runtime.env - policy = runtime.policy - gym_config_path = runtime.gym_config_path - gym_config_data = runtime.gym_config - gym_env_cfg = runtime.env_cfg - if gym_config_path is None or gym_config_data is None or gym_env_cfg is None: - raise RuntimeError("Simulator Policy runtime is missing task configuration") if rank == 0: logger.log_info( f"Loaded gym_config from {gym_config_path} (env_id={gym_config_data['id']}, num_envs={gym_env_cfg.num_envs}, headless={gym_env_cfg.sim_cfg.headless}, renderer={gym_env_cfg.sim_cfg.render_cfg.renderer}, sim_device={gym_env_cfg.sim_cfg.sim_device})" ) + env = build_env(gym_config_data["id"], base_env_cfg=gym_env_cfg) + sample_obs, _ = env.reset() + sample_obs_td = dict_to_tensordict(sample_obs, device) + obs_dim = flatten_dict_observation(sample_obs_td).shape[-1] + flat_obs_space = env.flattened_observation_space + # Create evaluation environment only if enabled eval_env = None num_eval_envs = trainer_cfg.get("num_eval_envs", 4) if enable_eval and rank == 0: eval_gym_env_cfg = deepcopy(gym_env_cfg) - eval_gym_env_cfg.num_envs = int(num_eval_envs) + eval_gym_env_cfg.num_envs = num_eval_envs eval_gym_env_cfg.sim_cfg.headless = True eval_gym_env_cfg.profiler = None eval_env = build_env(gym_config_data["id"], base_env_cfg=eval_gym_env_cfg) @@ -476,7 +498,59 @@ def train_from_config( f"Evaluation environment created (num_envs={num_eval_envs}, headless=True)" ) + # Build Policy via registry policy_name = policy_block["name"] + env_action_dim = ( + env.get_wrapper_attr("action_manager").total_action_dim + if env.get_wrapper_attr("action_manager") is not None + else len(env.get_wrapper_attr("active_joint_ids")) + ) + action_dim = policy_block.get("action_dim", env_action_dim) + action_dim = int(action_dim) + if action_dim != env_action_dim: + raise ValueError( + f"Configured policy.action_dim={action_dim} does not match env action dim {env_action_dim}." + ) + # Build Policy via registry (actor/critic must be explicitly defined in JSON when using actor_critic/actor_only) + if policy_name.lower() == "actor_critic": + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + if actor_cfg is None or critic_cfg is None: + raise ValueError( + "ActorCritic requires 'actor' and 'critic' definitions in JSON (policy.actor / policy.critic)." + ) + + actor = build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) + critic = build_mlp_from_cfg(critic_cfg, obs_dim, 1) + + policy = build_policy( + policy_block, + flat_obs_space, + env.action_space, + device, + actor=actor, + critic=critic, + ) + elif policy_name.lower() == "actor_only": + actor_cfg = policy_block.get("actor") + if actor_cfg is None: + raise ValueError( + "ActorOnly requires 'actor' definition in JSON (policy.actor)." + ) + + actor = build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) + + policy = build_policy( + policy_block, + flat_obs_space, + env.action_space, + device, + actor=actor, + ) + else: + policy = build_policy( + policy_block, env.observation_space, env.action_space, device + ) # Build Algorithm via factory algo_name = algo_block["name"].lower() @@ -507,7 +581,7 @@ def train_from_config( for event_name, event_info in events_dict.get("train", {}).items(): event_func_str = event_info.get("func") mode = event_info.get("mode", "interval") - params = _event_params(event_info, run_base=run_base, phase="train") + params = event_info.get("params", {}) interval_step = event_info.get("interval_step", 1) event_func = find_function_from_modules( event_func_str, event_modules, raise_if_not_found=True @@ -523,7 +597,7 @@ def train_from_config( for event_name, event_info in events_dict.get("eval", {}).items(): event_func_str = event_info.get("func") mode = event_info.get("mode", "interval") - params = _event_params(event_info, run_base=run_base, phase="eval") + params = event_info.get("params", {}) interval_step = event_info.get("interval_step", 1) event_func = find_function_from_modules( event_func_str, event_modules, raise_if_not_found=True @@ -574,7 +648,6 @@ def train_from_config( f"Total steps: {total_steps} (iterations≈{iterations}, world_size={world_size})" ) - summary = None try: trainer.train(total_steps) except KeyboardInterrupt: @@ -582,8 +655,6 @@ def train_from_config( logger.log_info("Training interrupted by user") finally: trainer.save_checkpoint() - if rank == 0: - summary = trainer.get_summary() if writer is not None: writer.close() if use_wandb and rank == 0: @@ -610,35 +681,8 @@ def train_from_config( if distributed and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() - if summary is not None: - _write_policy_run_manifest( - run_base, - config_path, - summary, - gym_config=gym_config_path, - ) - if rank == 0: - logger.log_info("Training finished") - - -def _write_policy_run_manifest( - run_base: str | Path, - config_path: str | Path, - summary: dict, - *, - gym_config: str | Path | None = None, -) -> Path: - """Write the checkpoint and configuration index for policy evaluation.""" - latest = summary.get("latest_checkpoint_path") - if latest is None: - raise RuntimeError("Training finished without a checkpoint") - return write_run_manifest( - run_base, - train_config=config_path, - gym_config=gym_config, - latest_checkpoint=latest, - best_checkpoint=summary.get("best_checkpoint_path"), - ) + if rank == 0: + logger.log_info("Training finished") def cli(argv: Sequence[str] | None = None) -> None: diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json index f7d951e61..de73223f8 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.json @@ -46,7 +46,8 @@ 600, 320, 240 - ] + ], + "save_path": "./outputs/videos/eval" } } } @@ -86,4 +87,4 @@ "truncate_at_first_done": true } } -} +} \ No newline at end of file diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml index c3c77b95a..d64961242 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/grpo.yaml @@ -40,6 +40,7 @@ trainer: - 600 - 320 - 240 + save_path: ./outputs/videos/eval renderer: hybrid policy: name: actor_only diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json index 62fd4c5a6..a4b04beca 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.json @@ -45,7 +45,8 @@ 600, 320, 240 - ] + ], + "save_path": "./outputs/videos/eval" } } } @@ -92,4 +93,4 @@ "max_grad_norm": 0.5 } } -} +} \ No newline at end of file diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml index 05f10ba96..5b5935b3d 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/agents/ppo.yaml @@ -39,6 +39,7 @@ trainer: - 600 - 320 - 240 + save_path: ./outputs/videos/eval renderer: fast-rt policy: name: actor_critic diff --git a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json index e697ce0d4..21f6cd767 100644 --- a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json +++ b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/grpo.json @@ -28,7 +28,8 @@ "eye": [-1.4, 1.4, 2.0], "target": [0, 0, 0], "up": [0, 0, 1], - "intrinsics": [600, 600, 320, 240] + "intrinsics": [600, 600, 320, 240], + "save_path": "./outputs/videos/eval" } } } diff --git a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json index 25a1928df..399454e6d 100644 --- a/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json +++ b/embodichain_tasks/configs/tasks/manipulation/push_cube/agents/ppo.json @@ -28,7 +28,8 @@ "eye": [-1.4, 1.4, 2.0], "target": [0, 0, 0], "up": [0, 0, 1], - "intrinsics": [600, 600, 320, 240] + "intrinsics": [600, 600, 320, 240], + "save_path": "./outputs/videos_ppo1/eval" } } } @@ -75,4 +76,4 @@ "max_grad_norm": 0.5 } } -} +} \ No newline at end of file diff --git a/examples/learning/policy_evaluation/README.md b/examples/learning/policy_evaluation/README.md deleted file mode 100644 index 8e3155b00..000000000 --- a/examples/learning/policy_evaluation/README.md +++ /dev/null @@ -1,137 +0,0 @@ -# ANYmal-C Velocity Policy Evaluation - -This example connects Newton's public ANYmal-C velocity TorchScript `.pt` to -EmbodiChain and opens it in the DexSim Viewer through Motion Policy Kit. The -model accepts `vx`, `vy`, and `yaw` commands. W/A/S/D and Q/E update these -commands while the Viewer is running. - -The model, configuration, and robot resources come from newton-assets commit -`261cd1f429619d8ef4f546bd788ab9dea906b5e1`. The Policy is distributed under -Apache-2.0, and the robot resources use BSD-3-Clause. The Adapter follows -Newton v1.2.1 -[`example_robot_policy.py`](https://github.com/newton-physics/newton/blob/v1.2.1/newton/examples/robot/example_robot_policy.py) -to reproduce the 48-dimensional observation, TorchScript inference, and joint -target processing. - -## Directory layout - -```text -policy_evaluation/ -├── README.md -├── prepare_resources.py -├── eval_policy.py # Register the local Profile and run the example -└── anymal_c/ - ├── __init__.py # Register newton-anymal-c-velocity - └── profile.py # Policy Spec and AnymalCVelocityAdapter -``` - -Resource preparation creates this local cache: - -```text -~/.cache/embodichain/examples/anymal_c_velocity/ -└── upstream/ - └── anybotics_anymal_c/ - ├── rl_policies/ - │ ├── mjw_anymal.pt - │ ├── anymal.yaml - │ └── LICENSE - ├── urdf/anymal.urdf - ├── meshes/... - └── LICENSE -``` - -## Run the example - -Run these commands from the EmbodiChain repository root. The preparation script -prints the model, asset, checkout, and digest verification progress. Re-running -the command continues an existing Git checkout after an interrupted download. - -```bash -python examples/learning/policy_evaluation/prepare_resources.py -python examples/learning/policy_evaluation/eval_policy.py \ - --viewer \ - --renderer hybrid -``` - -Viewer controls: - -| Key | Command | -|---|---| -| W / S | Increase / decrease `vx` | -| A / D | Increase / decrease `vy` | -| Q / E | Increase / decrease `yaw` | -| M | Set all three commands to zero | -| Backspace | Reset the robot, Policy history, and camera framing | -| T | Switch between tracking and free view | - -The camera follows the robot root in the ground plane. Hold the left mouse -button to change the orbit angle and use the mouse wheel to change the viewing -distance. Right-button panning is locked while tracking is active. Tracking -continues while the orbit angle is being adjusted. Switching back to tracking -centers the camera on the current robot position. - -The terminal prints the path to `evaluation.json` when the Viewer closes. Run a -Headless smoke test with: - -```bash -python examples/learning/policy_evaluation/eval_policy.py \ - --device cpu \ - --sim-device cpu \ - --control-steps 20 -``` - -`eval_policy.py` reads the checkpoint and robot assets from the default cache, -imports the adjacent `anymal_c/profile.py`, and registers the Profile in the -current process. Run the script directly from the repository root. To use -another cache directory: - -```bash -python examples/learning/policy_evaluation/prepare_resources.py \ - --output /tmp/anymal_c_velocity - -ANYMAL_C_EXAMPLE_CACHE=/tmp/anymal_c_velocity \ - python examples/learning/policy_evaluation/eval_policy.py --viewer -``` - -## Execution pipeline - -```mermaid -flowchart LR - CLI[eval-policy] --> Profile[build_profile] - Profile --> Spec[Policy Spec
assets, control parameters, frequency] - Spec --> Setup[Adapter.setup
load TorchScript and joint mapping] - Setup --> State[read RobotState] - Command[WASD + QE command] --> Obs[build 48-dimensional observation] - State --> Obs - Obs --> Actor[TorchScript actor] - Actor --> Action[map to 12 joint targets] - Action --> Sim[advance the Environment] -``` - -`AnymalCVelocityAdapter` restores the upstream data path: - -| Stage | Processing | -|---|---| -| `setup()` | Build a `JointMap` for the 12 ANYmal-C joints, then load and validate the TorchScript inputs and outputs | -| observation | 3 body linear velocity, 3 body angular velocity, 3 projected gravity, 3 command, 12 joint position, 12 joint velocity, and 12 previous action values | -| command | Read `vx`, `vy`, and `yaw` from `frame.controls["command"]`, with ranges ±1.0, ±0.5, and ±1.0 | -| actor | Pass a `[1, 48]` tensor through the model's normalizer and actor to produce `[1, 12]` | -| action | Apply `default_position + 0.5 * action` | -| control | Run simulation at 200 Hz and infer once every four simulation steps for a 50 Hz Policy rate | - -The Adapter clears the previous action during reset. After each inference call, -it stores the current action for the next observation. - -## Integrate another external Policy - -Copy this directory and replace: - -1. the fixed revisions, paths, and digests for the model and robot resources in `prepare_resources.py`; -2. the initial pose, joint control parameters, simulation step, and `sim_steps_per_control` in `build_profile()`; -3. the model format and training joint order in `Adapter.setup()`; -4. observation construction, normalization, network forward pass, action clipping, scale, and offset in `Adapter.infer()`; -5. `PROFILE_ID` and the Profile name used by `eval_policy.py`. - -An Adapter can call an existing project data reader from `__init__()` or -`setup()`. To let Policy Spec resolve a data file path, declare it under -`policy.resources` and read the resolved path from `AdapterRequest.resources`. diff --git a/examples/learning/policy_evaluation/anymal_c/__init__.py b/examples/learning/policy_evaluation/anymal_c/__init__.py deleted file mode 100644 index 0303a1211..000000000 --- a/examples/learning/policy_evaluation/anymal_c/__init__.py +++ /dev/null @@ -1,30 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Register the ANYmal-C velocity Motion Profile.""" - -from __future__ import annotations - -from embodichain.learning.rl.policy_evaluation import register_motion_profile - -from .profile import PROFILE_ID, build_profile - -__all__ = ["register"] - - -def register() -> None: - """Register the ANYmal-C velocity Profile for the example process.""" - register_motion_profile(PROFILE_ID, build_profile) diff --git a/examples/learning/policy_evaluation/anymal_c/profile.py b/examples/learning/policy_evaluation/anymal_c/profile.py deleted file mode 100644 index 2f0e89e73..000000000 --- a/examples/learning/policy_evaluation/anymal_c/profile.py +++ /dev/null @@ -1,278 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""ANYmal-C velocity Profile for a public TorchScript policy.""" - -from __future__ import annotations - -import math -from pathlib import Path - -import numpy as np -import torch -from dexsim.kit.motion_policy import ( - AdapterRequest, - EvaluationFrame, - JointMap, - PolicyContext, - PolicyOutput, - require_finite, -) - -from embodichain.learning.rl.policy_evaluation import ( - MotionProfile, - MotionProfileRequest, -) - -__all__ = ["AnymalCVelocityAdapter", "PROFILE_ID", "build_profile"] - -PROFILE_ID = "newton-anymal-c-velocity" - -_SOURCE_REVISION = "7249270ab41be1c2d4c809aa87536bab3a1a26f4" -_ASSET_REVISION = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" -_ROBOT_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") -_JOINT_NAMES = ( - "LF_HAA", - "LF_HFE", - "LF_KFE", - "LH_HAA", - "LH_HFE", - "LH_KFE", - "RF_HAA", - "RF_HFE", - "RF_KFE", - "RH_HAA", - "RH_HFE", - "RH_KFE", -) -_DEFAULT_POSITION = np.asarray( - (0.0, 0.4, -0.8, 0.0, -0.4, 0.8, 0.0, 0.4, -0.8, 0.0, -0.4, 0.8), - dtype=np.float32, -) - - -def build_profile(request: MotionProfileRequest) -> MotionProfile: - """Build the Policy Spec for the public ANYmal-C checkpoint. - - Args: - request: Checkpoint, resource checkout, device, and renderer. - - Returns: - A Motion Profile ready for DexSim Motion Policy Kit. - """ - if request.resource_root is None: - raise ValueError( - "The ANYmal-C example requires --resource-root from prepare_resources.py" - ) - robot_asset = request.resource_root / _ROBOT_PATH - if not robot_asset.is_file(): - raise FileNotFoundError( - f"ANYmal-C robot asset does not exist: {robot_asset}. " - "Run prepare_resources.py first." - ) - - return MotionProfile( - profile_id=PROFILE_ID, - policy_spec={ - "schema_version": 1, - "kind": "policy", - "id": PROFILE_ID, - "metadata": { - "title": "Newton ANYmal-C velocity policy", - "description": "Public 48-D command locomotion TorchScript policy.", - "status": "example", - "tags": ["external", "quadruped", "velocity", "torchscript"], - }, - "robot": { - "asset": {"path": str(robot_asset)}, - "use_urdf_material": True, - "initial": { - "root_height": 0.76, - "joint_positions": { - "default": 0.0, - "overrides": dict( - zip( - _JOINT_NAMES, - _DEFAULT_POSITION.tolist(), - strict=True, - ) - ), - }, - }, - "control": { - "defaults": { - "stiffness": 300.0, - "damping": 10.0, - "effort_limit": 80.0, - "armature": 0.06, - }, - }, - }, - "policy": { - "models": {"actor": {"path": str(request.checkpoint)}}, - "adapter": { - "type": "python", - "entrypoint": ("anymal_c.profile:AnymalCVelocityAdapter"), - "config": { - "inference_device": str(request.device), - "joint_names": list(_JOINT_NAMES), - }, - }, - }, - "evaluation": { - "initial_command": [0.0, 0.0, 0.0], - "termination": {"behavior": "pause"}, - }, - "runtime": { - "physics_dt": 0.005, - "sim_steps_per_control": 4, - "physics_backend": "default", - "simulation_device": "cpu", - "inference_provider": ( - "cuda" if request.device.type == "cuda" else "cpu" - ), - "renderer": request.renderer, - }, - }, - provenance={ - "source": "Newton ANYmal-C keyboard policy example", - "source_revision": _SOURCE_REVISION, - "source_example": "newton/examples/robot/example_robot_policy.py", - "asset_revision": _ASSET_REVISION, - "model_format": "torchscript", - "observation_size": 48, - "action_size": 12, - }, - ) - - -class AnymalCVelocityAdapter: - """Reproduce the upstream command locomotion observation and action path.""" - - command_enabled = True - command_step = (0.1, 0.05, 0.1) - command_limits = (1.0, 0.5, 1.0) - - def __init__(self, request: AdapterRequest) -> None: - config = dict(request.config) - self.device = torch.device(str(config["inference_device"])) - self.joint_names = tuple(config["joint_names"]) - self.checkpoint = request.models["actor"] - self.previous_action = np.zeros(12, dtype=np.float32) - self.joints: JointMap | None = None - self.model: torch.jit.ScriptModule | None = None - - def setup(self, context: PolicyContext) -> None: - """Load the model and bind the runtime joint order.""" - robot = context.robot - if robot is None: - raise RuntimeError("ANYmal-C robot description is required") - self.joints = JointMap.from_joint_names( - robot.joint_names, - self.joint_names, - ) - self.model = torch.jit.load( - str(self.checkpoint), - map_location=self.device, - ).eval() - with torch.inference_mode(): - output = self.model(torch.zeros((1, 48), device=self.device)) - if not isinstance(output, torch.Tensor) or tuple(output.shape) != (1, 12): - shape = ( - None if not isinstance(output, torch.Tensor) else tuple(output.shape) - ) - raise ValueError(f"ANYmal-C policy output must be (1, 12), got {shape}") - - def reset(self, frame: EvaluationFrame) -> None: - """Reset the previous action used by the policy observation.""" - self.previous_action.fill(0.0) - - def infer(self, frame: EvaluationFrame) -> PolicyOutput: - """Build one 48-D observation and return 12 joint targets.""" - observation = self._build_observation(frame) - tensor = torch.from_numpy(observation).to(self.device).unsqueeze(0) - with torch.inference_mode(): - output = self._model()(tensor) - action = require_finite( - "ANYmal-C action", - output[0].detach().cpu().numpy(), - ) - self.previous_action = action.copy() - return PolicyOutput( - action=self._joints().command( - position=_DEFAULT_POSITION + 0.5 * action, - ), - termination_reason=_fall_reason(frame.robot_state.root_pose), - ) - - def metrics(self) -> dict[str, float]: - """Return the metrics produced by this velocity example.""" - return {} - - def close(self) -> None: - """Release the loaded TorchScript model.""" - self.model = None - - def _build_observation(self, frame: EvaluationFrame) -> np.ndarray: - state = frame.robot_state - if state is None: - raise RuntimeError("ANYmal-C robot state is required") - pose = np.asarray(state.root_pose, dtype=np.float32) - velocity = np.asarray(state.root_velocity, dtype=np.float32) - rotation = pose[:3, :3] - qpos = self._joints().to_model(state.qpos) - qvel = self._joints().to_model(state.qvel) - command = require_finite( - "ANYmal-C command", - frame.controls["command"], - ) - if command.shape != (3,): - raise ValueError("ANYmal-C command must contain vx, vy, and yaw rate") - observation = np.concatenate( - ( - rotation.T @ velocity[:3], - rotation.T @ velocity[3:], - rotation.T @ np.asarray((0.0, 0.0, -1.0), dtype=np.float32), - command, - qpos - _DEFAULT_POSITION, - qvel, - self.previous_action, - ), - dtype=np.float32, - ) - return require_finite("ANYmal-C observation", observation) - - def _joints(self) -> JointMap: - if self.joints is None: - raise RuntimeError("ANYmal-C Adapter is not set up") - return self.joints - - def _model(self) -> torch.jit.ScriptModule: - if self.model is None: - raise RuntimeError("ANYmal-C Adapter is not set up") - return self.model - - -def _fall_reason(root_pose: np.ndarray) -> str | None: - pose = np.asarray(root_pose, dtype=np.float64) - height = float(pose[2, 3]) - tilt = math.acos(float(np.clip(pose[2, 2], -1.0, 1.0))) - reasons = [] - if height < 0.25: - reasons.append(f"base_height_below_minimum: {height:.3f} m") - if tilt > math.pi * 0.4: - reasons.append(f"bad_orientation: {tilt:.3f} rad") - return "; ".join(reasons) or None diff --git a/examples/learning/policy_evaluation/eval_policy.py b/examples/learning/policy_evaluation/eval_policy.py deleted file mode 100644 index b01048b30..000000000 --- a/examples/learning/policy_evaluation/eval_policy.py +++ /dev/null @@ -1,72 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Evaluate the public ANYmal-C checkpoint from the example directory.""" - -from __future__ import annotations - -import os -import sys -from pathlib import Path - -__all__ = ["example_arguments", "main"] - -EXAMPLE_ROOT = Path(__file__).resolve().parent -REPOSITORY_ROOT = EXAMPLE_ROOT.parents[2] -DEFAULT_CACHE = Path.home() / ".cache/embodichain/examples/anymal_c_velocity" - - -def example_arguments(argv: list[str]) -> list[str]: - """Add the example Profile, checkpoint, and resource paths. - - Args: - argv: Evaluation options accepted by ``eval-policy``. - - Returns: - Arguments ready for the EmbodiChain evaluation CLI. - """ - cache = Path(os.environ.get("ANYMAL_C_EXAMPLE_CACHE", DEFAULT_CACHE)) - resource_root = cache / "upstream" - checkpoint = resource_root / "anybotics_anymal_c/rl_policies/mjw_anymal.pt" - return [ - "--profile", - "newton-anymal-c-velocity", - "--checkpoint", - str(checkpoint), - "--resource-root", - str(resource_root), - *argv, - ] - - -def main(argv: list[str] | None = None) -> None: - """Register the local Profile and run visual policy evaluation. - - Args: - argv: Evaluation options. Uses command-line arguments when omitted. - """ - if str(REPOSITORY_ROOT) not in sys.path: - sys.path.insert(0, str(REPOSITORY_ROOT)) - - from anymal_c import register - from embodichain.learning.rl.policy_evaluation.cli import cli - - register() - cli(example_arguments(sys.argv[1:] if argv is None else argv)) - - -if __name__ == "__main__": - main() diff --git a/examples/learning/policy_evaluation/prepare_resources.py b/examples/learning/policy_evaluation/prepare_resources.py deleted file mode 100644 index 7e000e2a4..000000000 --- a/examples/learning/policy_evaluation/prepare_resources.py +++ /dev/null @@ -1,209 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Prepare the pinned public policy and assets for the ANYmal-C example.""" - -from __future__ import annotations - -import argparse -import hashlib -import os -import subprocess -from pathlib import Path - -__all__ = ["main", "prepare_resources"] - -UPSTREAM_URL = "https://github.com/newton-physics/newton-assets.git" -UPSTREAM_REVISION = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" -MODEL_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/mjw_anymal.pt") -POLICY_CONFIG_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/anymal.yaml") -POLICY_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/LICENSE") -ROBOT_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/LICENSE") -ROBOT_RELATIVE_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") -MESH_RELATIVE_PATH = Path("anybotics_anymal_c/meshes/base.dae") -SHA256 = { - MODEL_RELATIVE_PATH: "00765c1c07e497be3825672b05f9cefff9238f2df72fb0bcb5ac9541155b945f", - POLICY_CONFIG_RELATIVE_PATH: "b5a463ac418c7f40ebe494c7bcf0d8031f021db70a0625dbcc28a718de8ee817", - POLICY_LICENSE_RELATIVE_PATH: "59899c6091b540582ed617e8eeaac4919dc985ccfc35459ee9752b699be5205b", - ROBOT_LICENSE_RELATIVE_PATH: "cef384faae108293b03b5e16a00bc3db8212d44575f69df6296438a3f901700b", - ROBOT_RELATIVE_PATH: "d6bd20292cdd4873ffdeeb6f8ca3f96c4a0096565d78d8b6204f6edf0d19fb83", - MESH_RELATIVE_PATH: "785bea9b33831f8c741fc0ca070162e73cbf560ea9b03c53abf8978be877fc48", -} - - -def prepare_resources(output: Path) -> tuple[Path, Path]: - """Fetch and verify the pinned upstream files. - - Args: - output: Cache directory that will contain the sparse Git checkout. - - Returns: - The local checkpoint and resource-root paths. - """ - output = output.expanduser().resolve() - checkout = output / "upstream" - _status("Preparing the ANYmal-C command policy and robot assets") - _prepare_checkout( - checkout, - UPSTREAM_URL, - UPSTREAM_REVISION, - ( - f"/{MODEL_RELATIVE_PATH}", - f"/{POLICY_CONFIG_RELATIVE_PATH}", - f"/{POLICY_LICENSE_RELATIVE_PATH}", - f"/{ROBOT_LICENSE_RELATIVE_PATH}", - "/anybotics_anymal_c/urdf/**", - "/anybotics_anymal_c/meshes/**", - ), - ) - - checkpoint = checkout / MODEL_RELATIVE_PATH - for relative, digest in SHA256.items(): - _verify_sha256(checkout / relative, digest) - _status("Resource verification completed") - return checkpoint, checkout - - -def main() -> None: - """Prepare resources and print the paths used by the evaluation command.""" - parser = argparse.ArgumentParser() - parser.add_argument( - "--output", - type=Path, - default=Path.home() / ".cache/embodichain/examples/anymal_c_velocity", - help="Directory used for the pinned upstream checkout", - ) - args = parser.parse_args() - checkpoint, resource_root = prepare_resources(args.output) - print(f"Checkpoint: {checkpoint}") - print(f"Resource root: {resource_root}") - - -def _git( - checkout: Path, - *args: str, - capture_output: bool = False, -) -> subprocess.CompletedProcess[str]: - command = ["git", "-C", str(checkout), *args] - environment = os.environ.copy() - environment["GIT_TERMINAL_PROMPT"] = "0" - try: - return subprocess.run( - command, - check=True, - text=True, - capture_output=capture_output, - env=environment, - timeout=600, - ) - except subprocess.TimeoutExpired as error: - raise RuntimeError( - f"Git command did not finish within 10 minutes: {' '.join(command)}" - ) from error - - -def _git_output(checkout: Path, *args: str) -> str | None: - try: - return _git(checkout, *args, capture_output=True).stdout.strip() - except subprocess.CalledProcessError: - return None - - -def _prepare_checkout( - checkout: Path, - url: str, - revision: str, - includes: tuple[str, ...], -) -> None: - if checkout.exists() and not (checkout / ".git").is_dir(): - raise RuntimeError( - f"Resource path exists but is not a Git checkout: {checkout}" - ) - - if checkout.exists(): - remote_url = _git_output(checkout, "remote", "get-url", "origin") - if remote_url != url: - raise RuntimeError( - f"Resource checkout uses an unexpected remote: {remote_url}" - ) - if _git_output(checkout, "rev-parse", "HEAD") == revision: - tracked_changes = _git_output( - checkout, "status", "--porcelain", "--untracked-files=no" - ) - if tracked_changes == "": - _status(f"Using cached revision {revision[:8]} from {checkout}") - return - else: - checkout.parent.mkdir(parents=True, exist_ok=True) - checkout.mkdir() - _git(checkout, "init", "--quiet") - _git(checkout, "remote", "add", "origin", url) - - _git(checkout, "sparse-checkout", "init", "--no-cone") - _git(checkout, "sparse-checkout", "set", *includes) - - _status(f"Fetching revision {revision[:8]} from {url}") - _git( - checkout, - "fetch", - "--progress", - "--filter=blob:none", - "--depth", - "1", - "origin", - revision, - ) - - _status(f"Checking out required files in {checkout}") - _git( - checkout, - "-c", - "advice.detachedHead=false", - "checkout", - "--progress", - "--force", - "--detach", - "FETCH_HEAD", - ) - actual_revision = _git_output(checkout, "rev-parse", "HEAD") - if actual_revision != revision: - raise RuntimeError( - f"Checkout revision mismatch: expected {revision}, got {actual_revision}" - ) - - -def _status(message: str) -> None: - print(f"[resources] {message}", flush=True) - - -def _verify_sha256(path: Path, expected: str) -> None: - actual = _sha256(path) - if actual != expected: - raise RuntimeError( - f"SHA256 mismatch for {path}: expected {expected}, got {actual}" - ) - - -def _sha256(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() - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py index 6fb13a58f..a7b7e6db9 100644 --- a/scripts/benchmark/atomic_action/common.py +++ b/scripts/benchmark/atomic_action/common.py @@ -240,7 +240,7 @@ def add_common_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "offline-rt"), + choices=("auto", "hybrid", "fast-rt", "rt"), default="auto", help="Renderer backend used by SimulationManager.", ) diff --git a/scripts/benchmark/atomic_action/run_benchmark.py b/scripts/benchmark/atomic_action/run_benchmark.py index d6cbc462e..d50107756 100644 --- a/scripts/benchmark/atomic_action/run_benchmark.py +++ b/scripts/benchmark/atomic_action/run_benchmark.py @@ -102,7 +102,7 @@ def add_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "offline-rt"), + choices=("auto", "hybrid", "fast-rt", "rt"), default="auto", help="Renderer backend forwarded to each selected benchmark.", ) diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 763cf5e08..ae0b855b3 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -385,12 +385,12 @@ def test_launcher_preserves_gym_renderer_when_cli_omits_override(): add_env_launcher_args_to_parser(parser, require_gym_config=True) args = parser.parse_args(["--gym_config", "gym_config.yaml"]) - gym_config = {"id": "Dummy-v0", "render_cfg": {"renderer": "offline-rt"}} + gym_config = {"id": "Dummy-v0", "render_cfg": {"renderer": "rt"}} merged_config = merge_args_with_gym_config(args, gym_config) assert args.renderer is None assert "renderer" not in merged_config - assert merged_config["render_cfg"]["renderer"] == "offline-rt" + assert merged_config["render_cfg"]["renderer"] == "rt" def test_env_launcher_includes_viser_arguments(): @@ -1155,7 +1155,7 @@ def test_yaml_gym_config_parses_to_cfg(self, tmp_path): "speed_tolerance": 0.1, }, "render_cfg": { - "renderer": "offline-rt", + "renderer": "rt", "spp": 4, "tone_mapping_enabled": True, "tone_mapping_exposure": 1.25, @@ -1208,7 +1208,7 @@ def test_yaml_gym_config_parses_to_cfg(self, tmp_path): assert cfg.sim_cfg.physics_config.enable_ccd is True assert cfg.sim_cfg.physics_config.length_tolerance == 0.02 assert cfg.sim_cfg.physics_config.speed_tolerance == 0.1 - assert cfg.sim_cfg.render_cfg.renderer == "offline-rt" + assert cfg.sim_cfg.render_cfg.renderer == "rt" assert cfg.sim_cfg.render_cfg.spp == 4 assert cfg.sim_cfg.render_cfg.tone_mapping_enabled is True assert cfg.sim_cfg.render_cfg.tone_mapping_exposure == 1.25 @@ -1267,7 +1267,7 @@ def test_build_env_cfg_applies_modifier_before_parsing(self, tmp_path): "enable_ccd": True, }, "render_cfg": { - "renderer": "offline-rt", + "renderer": "rt", "spp": 8, "tone_mapping_enabled": True, }, diff --git a/tests/learning/rl/policy_evaluation/test_anymal_c_example.py b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py deleted file mode 100644 index c07d6ef9f..000000000 --- a/tests/learning/rl/policy_evaluation/test_anymal_c_example.py +++ /dev/null @@ -1,181 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest -import torch - -pytest.importorskip("dexsim.kit.motion_policy.evaluator") - -from dexsim.kit.motion_policy import ( - AdapterRequest, - EvaluationFrame, - PolicyContext, - RobotDescription, - RobotState, - parse_policy_spec, -) - -from embodichain.learning.rl.policy_evaluation import ( - MotionProfileRequest, - build_motion_profile, -) - -_JOINT_NAMES = ( - "LF_HAA", - "LF_HFE", - "LF_KFE", - "LH_HAA", - "LH_HFE", - "LH_KFE", - "RF_HAA", - "RF_HFE", - "RF_KFE", - "RH_HAA", - "RH_HFE", - "RH_KFE", -) -_DEFAULT_POSITION = np.asarray( - (0.0, 0.4, -0.8, 0.0, -0.4, 0.8, 0.0, 0.4, -0.8, 0.0, -0.4, 0.8), - dtype=np.float32, -) -_COMMAND = np.asarray((0.4, -0.2, 0.6), dtype=np.float32) - - -class _CommandPolicy(torch.nn.Module): - def forward(self, observation: torch.Tensor) -> torch.Tensor: - padding = torch.zeros( - (observation.shape[0], 9), - dtype=observation.dtype, - device=observation.device, - ) - return torch.cat((observation[:, 9:12], padding), dim=1) - - -def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): - example_root = ( - Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" - ) - monkeypatch.syspath_prepend(str(example_root)) - from anymal_c.profile import ( - AnymalCVelocityAdapter, - ) - from anymal_c import register - - checkpoint = tmp_path / "mjw_anymal.pt" - traced = torch.jit.trace(_CommandPolicy().eval(), torch.zeros((1, 48))) - torch.jit.save(traced, checkpoint) - - robot_asset = tmp_path / "anybotics_anymal_c/urdf/anymal.urdf" - robot_asset.parent.mkdir(parents=True) - robot_asset.write_text("\n", encoding="utf-8") - - request = MotionProfileRequest( - checkpoint=checkpoint, - device=torch.device("cpu"), - resource_root=tmp_path, - ) - register() - profile = build_motion_profile("newton-anymal-c-velocity", request) - spec = parse_policy_spec(profile.policy_spec) - assert spec.environment.entrypoint is None - config = profile.policy_spec["policy"]["adapter"]["config"] - adapter = AnymalCVelocityAdapter( - AdapterRequest( - asset_path=robot_asset, - models={"actor": checkpoint}, - resources={}, - config=config, - ) - ) - context = PolicyContext( - robot=RobotDescription( - _JOINT_NAMES, - ("base",), - "base", - ), - physics_dt=0.005, - sim_steps_per_control=4, - policy_dt=0.02, - ) - pose = np.eye(4, dtype=np.float32) - pose[2, 3] = 0.76 - frame = EvaluationFrame( - control_step=0, - policy_time=0.0, - simulation_time=0.0, - simulation_step=0, - robot_state=RobotState( - joint_names=_JOINT_NAMES, - qpos=_DEFAULT_POSITION.copy(), - qvel=np.zeros(12, dtype=np.float32), - target_qpos=_DEFAULT_POSITION.copy(), - target_qvel=np.zeros(12, dtype=np.float32), - joint_effort=np.zeros(12, dtype=np.float32), - root_name="base", - root_pose=pose, - root_velocity=np.zeros(6, dtype=np.float32), - link_names=("base",), - link_poses=pose[None, ...], - link_velocities=np.zeros((1, 6), dtype=np.float32), - ), - controls={"command": _COMMAND}, - ) - - adapter.setup(context) - adapter.reset(frame) - output = adapter.infer(frame) - - assert output.action.joint_names == _JOINT_NAMES - expected_position = _DEFAULT_POSITION.copy() - expected_position[:3] += 0.5 * _COMMAND - np.testing.assert_allclose(output.action.position, expected_position) - np.testing.assert_allclose( - adapter.previous_action, - np.concatenate((_COMMAND, np.zeros(9, dtype=np.float32))), - ) - assert adapter.command_enabled - assert adapter.command_limits == (1.0, 0.5, 1.0) - assert output.termination_reason is None - adapter.reset(frame) - np.testing.assert_array_equal(adapter.previous_action, np.zeros(12)) - adapter.close() - - -def test_example_script_supplies_default_resource_paths(tmp_path, monkeypatch): - example_root = ( - Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" - ) - monkeypatch.syspath_prepend(str(example_root)) - monkeypatch.setenv("ANYMAL_C_EXAMPLE_CACHE", str(tmp_path)) - from eval_policy import example_arguments - - arguments = example_arguments(["--control-steps", "5"]) - - assert arguments == [ - "--profile", - "newton-anymal-c-velocity", - "--checkpoint", - str(tmp_path / "upstream/anybotics_anymal_c/rl_policies/mjw_anymal.pt"), - "--resource-root", - str(tmp_path / "upstream"), - "--control-steps", - "5", - ] diff --git a/tests/learning/rl/policy_evaluation/test_bridge.py b/tests/learning/rl/policy_evaluation/test_bridge.py deleted file mode 100644 index eb3f120f5..000000000 --- a/tests/learning/rl/policy_evaluation/test_bridge.py +++ /dev/null @@ -1,84 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -pytest.importorskip("dexsim.kit.motion_policy.evaluator") - -from embodichain.learning.rl.policy_evaluation.bridge import ( - evaluate_motion_profile, -) -from embodichain.learning.rl.policy_evaluation.profile import MotionProfile - - -def test_bridge_forwards_physics_backend_and_exact_control_steps( - monkeypatch, -): - profile = MotionProfile( - profile_id="example", - policy_spec={"schema_version": 1}, - ) - options = [] - monkeypatch.setattr( - "embodichain.learning.rl.policy_evaluation.bridge.parse_policy_spec", - lambda value: "parsed", - ) - monkeypatch.setattr( - "embodichain.learning.rl.policy_evaluation.bridge.resolve_policy_spec", - lambda spec, resolver: "resolved", - ) - monkeypatch.setattr( - "embodichain.learning.rl.policy_evaluation.bridge.policy_spec_to_dict", - lambda value: {"policy_id": "example"}, - ) - monkeypatch.setattr( - "embodichain.learning.rl.policy_evaluation.bridge.scene_config_to_dict", - lambda value: {"style": "standard"}, - ) - - def run_policy(resolved, run_options): - options.append(run_options) - return SimpleNamespace( - reason="control steps reached", - simulation_time=0.2, - simulation_steps=40, - control_steps=10, - physics_backend="default", - requested_duration=None, - effective_duration=0.2, - metrics={"tracking/error": 0.25}, - ) - - monkeypatch.setattr( - "embodichain.learning.rl.policy_evaluation.bridge.run_motion_policy", - run_policy, - ) - - result = evaluate_motion_profile( - profile, - control_steps=10, - physics_backend="default", - ) - - assert options[0].control_steps == 10 - assert options[0].physics_backend == "default" - assert result.episodes[0]["control_steps"] == 10 - assert result.episodes[0]["effective_duration"] == 0.2 - assert result.summary["metrics"]["tracking/error"] == 0.25 diff --git a/tests/learning/rl/policy_evaluation/test_cli.py b/tests/learning/rl/policy_evaluation/test_cli.py deleted file mode 100644 index d36059fee..000000000 --- a/tests/learning/rl/policy_evaluation/test_cli.py +++ /dev/null @@ -1,130 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import importlib - -import pytest - -from embodichain.learning.rl.policy_evaluation.cli import ( - _resolve_input, - _validate_native_options, - parse_args, -) -from embodichain.learning.rl.policy_evaluation.manifest import write_run_manifest - - -def _run(tmp_path): - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - checkpoint.write_bytes(b"checkpoint") - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - write_run_manifest( - run, - train_config=train, - latest_checkpoint=checkpoint, - ) - return run, checkpoint - - -def test_run_defaults_to_latest_checkpoint(tmp_path): - run, checkpoint = _run(tmp_path) - - resolved = _resolve_input(parse_args((str(run),))) - - assert resolved.checkpoint == checkpoint - assert resolved.requested_checkpoint == "latest" - assert resolved.selected_checkpoint == "latest" - - -@pytest.mark.parametrize( - "arguments, handler", - [ - ((), "_run_native_headless"), - (("--viewer",), "_run_native_viewer"), - (("--profile", "example"), "_run_profile"), - ], -) -def test_cli_routes_one_command_to_the_selected_evaluation( - tmp_path, - monkeypatch, - arguments, - handler, -): - module = importlib.import_module("embodichain.learning.rl.policy_evaluation.cli") - run, _checkpoint = _run(tmp_path) - expected = tmp_path / "evaluation.json" - calls = [] - - monkeypatch.setattr(module, "discover_task_packages", lambda: None) - monkeypatch.setattr(module, "execute_init_hooks", lambda: None) - for name in ("_run_native_headless", "_run_native_viewer", "_run_profile"): - monkeypatch.setattr( - module, - name, - lambda args, resolved, name=name: calls.append(name) or expected, - ) - - report = module.run(parse_args((str(run), *arguments))) - - assert report == expected - assert calls == [handler] - - -def test_explicit_checkpoint_requires_training_config(tmp_path): - checkpoint = tmp_path / "policy.pt" - checkpoint.write_bytes(b"checkpoint") - - with pytest.raises(ValueError, match="--config is required"): - _resolve_input(parse_args(("--checkpoint", str(checkpoint)))) - - -@pytest.mark.parametrize("renderer", ("hybrid", "fast-rt", "offline-rt")) -def test_cli_accepts_dexsim_renderer_names(renderer): - args = parse_args(("--renderer", renderer)) - - assert args.renderer == renderer - - -def test_native_options_keep_profile_and_viewer_inputs_explicit(): - profile_args = parse_args( - ( - "--checkpoint", - "policy.pt", - "--config", - "train.yaml", - "--command", - "0.5", - ) - ) - viewer_args = parse_args( - ( - "--checkpoint", - "policy.pt", - "--config", - "train.yaml", - "--control-steps", - "10", - ) - ) - - with pytest.raises(ValueError, match="--command requires --profile"): - _validate_native_options(profile_args) - with pytest.raises(ValueError, match="--control-steps.*require --viewer"): - _validate_native_options(viewer_args) diff --git a/tests/learning/rl/policy_evaluation/test_viewer.py b/tests/learning/rl/policy_evaluation/test_viewer.py deleted file mode 100644 index bde501fe0..000000000 --- a/tests/learning/rl/policy_evaluation/test_viewer.py +++ /dev/null @@ -1,221 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch -from tensordict import TensorDict - -pytest.importorskip("dexsim.kit.motion_policy.evaluator") - -from dexsim.kit.motion_policy import EvaluationFrame, PolicyContext - -from embodichain.learning.rl.evaluation import infer_policy_action -from embodichain.learning.rl.policy_evaluation.viewer import ( - EmbodiChainTaskEnvironment, - EmbodiChainTaskPolicyAdapter, - evaluate_native_viewer, -) -from embodichain.learning.rl.runtime import PolicyRuntime - - -class Policy(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.weight = torch.nn.Parameter(torch.tensor([[2.0], [-1.0]])) - - def get_action( - self, - tensordict: TensorDict, - deterministic: bool = False, - ) -> TensorDict: - assert deterministic - tensordict["action"] = tensordict["obs"] @ self.weight - return tensordict - - -class Window: - def __init__(self) -> None: - self.titles = [] - self.keys = set() - - def set_window_title(self, title): - self.titles.append(title) - - def native(self): - return self - - def key_state(self, key): - return key in self.keys - - -class World: - def __init__(self) -> None: - self.window = Window() - self.open = True - - def is_window_initialized(self): - return self.open - - def get_windows(self): - return self.window - - -class ActionManager: - def __init__(self) -> None: - self.calls = 0 - - def convert_policy_action_to_env_action(self, action: torch.Tensor): - self.calls += 1 - return action + 0.5 - - -class Environment: - num_envs = 1 - physics_dt = 0.005 - step_dt = 0.02 - - def __init__(self) -> None: - self.unwrapped = self - self.cfg = SimpleNamespace(sim_steps_per_control=4) - self.action_manager = ActionManager() - self.world = World() - self.sim = SimpleNamespace(get_world=lambda: self.world) - self.actions = [] - self.episode_step = 0 - self.reset_seeds = [] - self.exit_process_values = [] - - def reset(self, seed=None): - self.reset_seeds.append(seed) - self.episode_step = 0 - return self._observation(), {} - - def step(self, action): - self.actions.append(action.clone()) - self.episode_step += 1 - done = self.episode_step == 2 - return ( - self._observation(), - torch.tensor([1.25]), - torch.tensor([done]), - torch.tensor([False]), - { - "success": torch.tensor([done]), - "metrics": {"task_progress": torch.tensor([self.episode_step])}, - }, - ) - - def close(self, *, exit_process=None): - self.exit_process_values.append(exit_process) - - def _observation(self): - return { - "policy": torch.tensor( - [[float(self.episode_step), 1.0]], - dtype=torch.float32, - ) - } - - -def _runtime(env: Environment, policy: Policy | None = None) -> PolicyRuntime: - return PolicyRuntime( - env=env, - policy=policy or Policy(), - device=torch.device("cpu"), - env_id="ExampleTask", - ) - - -def test_viewer_adapter_uses_the_shared_deterministic_inference_chain(): - policy = Policy() - observation = {"policy": torch.tensor([[0.25, 0.75]])} - expected = infer_policy_action(policy, observation, device="cpu", num_envs=1) - adapter = EmbodiChainTaskPolicyAdapter(policy, torch.device("cpu")) - adapter.setup(PolicyContext(None, 0.005, 4, 0.02)) - - output = adapter.infer(EvaluationFrame(0, 0.0, 0.0, 0, observation=observation)) - - assert torch.equal(output.action, expected) - adapter.close() - - -def test_viewer_reuses_task_actions_resets_and_metrics(): - env = Environment() - - result = evaluate_native_viewer( - _runtime(env), - seed=17, - episodes=2, - control_steps=None, - duration=None, - ) - - assert result.control_steps == 4 - assert result.simulation_steps == 16 - assert len(result.episodes) == 2 - assert result.metrics == pytest.approx( - { - "reward": 1.25, - "task_progress": 2.0, - "eval/avg_reward": 2.5, - "eval/avg_length": 2.0, - "eval/success_rate": 1.0, - } - ) - assert env.action_manager.calls == 4 - assert env.reset_seeds == [17, None] - assert env.exit_process_values == [False] - - -def test_backspace_requests_one_reset_per_key_press(): - from dexsim.types import InputKey - - env = Environment() - task = EmbodiChainTaskEnvironment(env, seed=1) - env.world.window.keys.add(InputKey.SCANCODE_BACKSPACE) - - assert task.poll() == "manual reset" - assert task.poll() is None - env.world.window.keys.clear() - assert task.poll() is None - env.world.window.keys.add(InputKey.SCANCODE_BACKSPACE) - assert task.poll() == "manual reset" - task.close() - - -def test_viewer_closes_resources_when_evaluator_creation_fails(monkeypatch): - env = Environment() - policy = Policy() - monkeypatch.setattr( - "embodichain.learning.rl.policy_evaluation.viewer.create_motion_policy_evaluator", - lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("setup failed")), - ) - - with pytest.raises(RuntimeError, match="setup failed"): - evaluate_native_viewer( - _runtime(env, policy), - seed=1, - episodes=1, - control_steps=None, - duration=None, - ) - - assert env.exit_process_values == [False] - assert policy.training is True diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py index 43ceb9da9..21012c4e6 100644 --- a/tests/learning/test_point_mass.py +++ b/tests/learning/test_point_mass.py @@ -31,7 +31,6 @@ build_learning_env, ) from embodichain.learning.rl.models import ActorCritic -from embodichain.learning.rl.policy_evaluation.manifest import RunManifest from embodichain.learning.rl.train import train_from_config from embodichain.learning.rl.utils import OptimizerCfg from embodichain.learning.rl.utils.trainer import Trainer @@ -184,13 +183,6 @@ def test_unified_train_entry_runs_apg_and_ppo( assert summary["global_step"] == 8 assert summary["latest_checkpoint_path"] is not None - checkpoint = Path(summary["latest_checkpoint_path"]).resolve() - run = checkpoint.parents[1] - manifest = RunManifest.load(run) - assert ( - manifest.configs["train"] == run / "configs" / f"train.{config_path.suffix[1:]}" - ) - assert manifest.select_checkpoint("latest")[1] == checkpoint def test_sync_collector_accepts_tensor_point_mass_observations() -> None: diff --git a/tests/learning/test_runtime.py b/tests/learning/test_runtime.py deleted file mode 100644 index ca35bab74..000000000 --- a/tests/learning/test_runtime.py +++ /dev/null @@ -1,125 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from types import SimpleNamespace - -import gymnasium as gym -import torch - -from embodichain.learning.rl.runtime import ( - _GymEnvironmentRuntime, - _build_gym_environment, - build_gym_policy_runtime, -) - - -def _policy_config() -> dict: - network = { - "type": "mlp", - "network_cfg": {"hidden_sizes": [8], "activation": "relu"}, - } - return { - "name": "actor_critic", - "actor": network, - "critic": network, - } - - -class GymEnvironment: - flattened_observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,)) - observation_space = gym.spaces.Dict( - {"policy": gym.spaces.Box(-1.0, 1.0, shape=(1, 3))} - ) - action_space = gym.spaces.Box(-1.0, 1.0, shape=(1, 2)) - - def __init__(self) -> None: - self.closed = 0 - self.action_manager = SimpleNamespace(total_action_dim=2) - - def reset(self): - return {"policy": torch.zeros(1, 3)}, {} - - def get_wrapper_attr(self, name): - return getattr(self, name) - - def close(self) -> None: - self.closed += 1 - - -def test_gym_environment_applies_runtime_overrides(tmp_path, monkeypatch): - gym_config = tmp_path / "gym.yaml" - gym_config.write_text("id: Example\n", encoding="utf-8") - env_cfg = SimpleNamespace( - num_envs=8, - sim_cfg=None, - profiler=None, - ) - built = GymEnvironment() - monkeypatch.setattr( - "embodichain.learning.rl.runtime.config_to_cfg", - lambda value, manager_modules: env_cfg, - ) - monkeypatch.setattr( - "embodichain.learning.rl.runtime.build_env", - lambda env_id, base_env_cfg: built, - ) - - runtime = _build_gym_environment( - {"trainer": {"gym_config": "gym.yaml"}}, - simulation_device=torch.device("cpu"), - num_envs=1, - headless=True, - renderer="hybrid", - gpu_id=0, - config_dir=tmp_path, - ) - - assert runtime.env is built - assert runtime.env_id == "Example" - assert runtime.env_cfg.num_envs == 1 - assert runtime.env_cfg.sim_cfg.sim_device == torch.device("cpu") - assert runtime.env_cfg.sim_cfg.headless is True - assert runtime.env_cfg.sim_cfg.render_cfg.renderer == "hybrid" - - -def test_gym_runtime_uses_the_same_task_spaces_for_policy_build(monkeypatch): - env = GymEnvironment() - task = _GymEnvironmentRuntime( - env=env, - env_id="Example", - env_cfg=SimpleNamespace(), - gym_config={"id": "Example"}, - gym_config_path=SimpleNamespace(resolve=lambda: None), - ) - monkeypatch.setattr( - "embodichain.learning.rl.runtime._build_gym_environment", - lambda *args, **kwargs: task, - ) - - runtime = build_gym_policy_runtime( - {"trainer": {"gym_config": "gym.yaml"}, "policy": _policy_config()}, - device=torch.device("cpu"), - num_envs=1, - headless=True, - renderer="hybrid", - gpu_id=0, - ) - - assert runtime.env is env - assert runtime.policy.actor[0].in_features == 3 - assert runtime.policy.actor[-1].out_features == 2 diff --git a/tests/learning/test_train_profile.py b/tests/learning/test_train_profile.py index 7934fb7e2..8b30b74f2 100644 --- a/tests/learning/test_train_profile.py +++ b/tests/learning/test_train_profile.py @@ -21,7 +21,6 @@ import pytest from embodichain.learning.rl.train import ( - _event_params, _resolve_profile_output, parse_args, train_from_config, @@ -66,13 +65,3 @@ def test_learning_env_rejects_profile(tmp_path): with pytest.raises(ValueError, match="--profile_output requires --profile"): train_from_config(str(config_path), profile_output="prof.json") - - -def test_camera_recording_defaults_to_the_run_directory(tmp_path): - params = _event_params( - {"func": "record_camera_data_async", "params": {"name": "main"}}, - run_base=tmp_path / "run", - phase="eval", - ) - - assert params["save_path"] == str(tmp_path / "run" / "videos" / "eval") diff --git a/tests/sim/test_cfg.py b/tests/sim/test_cfg.py index aef03c050..c9cfc28fe 100644 --- a/tests/sim/test_cfg.py +++ b/tests/sim/test_cfg.py @@ -94,14 +94,13 @@ def test_render_cfg_applies_tone_mapping_and_fixed_exposure() -> None: expected_exposure = 1.25 world_config = dexsim.WorldConfig() render_cfg = RenderCfg( - renderer="offline-rt", + renderer="rt", tone_mapping_enabled=True, tone_mapping_exposure=expected_exposure, ) render_cfg.apply_to_dexsim_config(world_config) - assert world_config.renderer == Renderer.OFFLINERT assert world_config.postprocess_config.tone_mapping_enabled is True assert ( world_config.postprocess_config.tone_mapping_type From 2953d0bf49c5a15286dc3fd5498baf4568e10e78 Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 19 Aug 2026 16:46:26 +0800 Subject: [PATCH 57/85] stash --- .../overview/sim/atomic_actions/index.md | 2 +- tests/sim/atomic_actions/test_actions.py | 159 ++++++++++++++++++ tests/sim/atomic_actions/test_affordance.py | 23 +++ 3 files changed, 183 insertions(+), 1 deletion(-) diff --git a/docs/source/overview/sim/atomic_actions/index.md b/docs/source/overview/sim/atomic_actions/index.md index e7a21f3dc..b2be92017 100644 --- a/docs/source/overview/sim/atomic_actions/index.md +++ b/docs/source/overview/sim/atomic_actions/index.md @@ -352,7 +352,7 @@ instances to the engine's planning services: ```python engine = AtomicActionEngine(motion_generator, control_profiles=profiles) -# All eleven built-ins are immediately usable by stable skill ID. +# All twelve built-ins are immediately usable by stable skill ID. assert "move_end_effector" in engine.actions assert "pick_up" in engine.actions ``` diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index ad0619b02..8f8952283 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -1957,6 +1957,165 @@ def test_axis_align_handles_opposite_axes_without_nan() -> None: ) +def test_axis_align_plans_seven_segments_and_aligns_the_object_axis() -> None: + generator = _motion_generator() + solved_poses: list[torch.Tensor] = [] + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + solved_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + action = _bind_action(generator, AxisAlign()) + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(internal_axis=torch.tensor([1.0, 0.0, 0.0])), + geometry={}, + label="axis-object", + entity_id="target", + ) + context = _context(scene=_target_scene(object_pose, timestamp=0.0, version=0)) + original_task = context.task + + plan = _plan_action( + action, + ActionInvocation( + skill_id="axis_align", + goal=AxisAlignGoal(semantics=semantics, grasp_xpos=torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=20), + skill_options=AxisAlignOptions( + target_axis=torch.tensor([0.0, 0.0, 1.0]), + lift_height=0.1, + lower_distance=0.03, + ), + ), + context, + ) + + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 20, ROBOT_DOF) + assert torch.equal(plan.trajectory.env_ids, context.env_ids) + assert plan.trajectory.duration.tolist() == pytest.approx([19.0 / 60.0] * NUM_ENVS) + assert [segment.name for segment in plan.segments] == [ + "approach", + "reach", + "close", + "lift", + "align", + "lower", + "open", + ] + assert plan.expected_effects.is_empty + assert context.task is original_task + assert plan.scene_dependencies == ("target",) + final_object_rotation = solved_poses[-1][:, :3, :3] + final_world_axis = torch.matmul( + final_object_rotation, + torch.tensor([1.0, 0.0, 0.0]), + ) + assert torch.allclose( + final_world_axis, + torch.tensor([0.0, 0.0, 1.0]).expand(NUM_ENVS, -1), + atol=1.0e-6, + ) + assert solved_poses[-1][:, 2, 3].tolist() == pytest.approx([0.07, 0.07]) + + +def test_axis_align_holds_only_failed_environment_rows() -> None: + generator = _motion_generator() + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.tensor([True, False]), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + action = _bind_action(generator, AxisAlign()) + object_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(), + geometry={}, + label="partially-alignable-object", + entity_id="target", + ) + context = _context(scene=_target_scene(object_pose, timestamp=0.0, version=0)) + + plan = _plan_action( + action, + _invocation( + "axis_align", + AxisAlignGoal(semantics=semantics, grasp_xpos=torch.eye(4)), + sample_count=20, + ), + context, + ) + + assert plan.plan_success.tolist() == [True, False] + assert not torch.allclose(plan.trajectory.positions[0], context.robot.qpos[0]) + assert torch.allclose( + plan.trajectory.positions[1], + context.robot.qpos[1].unsqueeze(0).expand(plan.trajectory.waypoint_count, -1), + ) + + +def test_axis_align_validates_goal_and_binding_contract() -> None: + action = _bind_action(_motion_generator(), AxisAlign()) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(), + geometry={}, + label="axis-object", + ) + + with pytest.raises(TypeError, match="expects goal AxisAlignGoal"): + action.resolve_request( + ActionInvocation( + skill_id="axis_align", + goal=object(), + binding=_binding(), + ) + ) + with pytest.raises(KeyError, match="end effector"): + action.resolve_request( + ActionInvocation( + skill_id="axis_align", + goal=AxisAlignGoal(semantics), + binding=ActionBinding(manipulators={"primary": "arm"}), + ) + ) + + +def test_axis_align_handles_opposite_axes_without_nan() -> None: + action = _bind_action(_motion_generator(), AxisAlign()) + identity = torch.eye(4).repeat(NUM_ENVS, 1, 1) + + eef_keyframes = action._axis_alignment_eef_keyframes( + identity, + identity, + torch.tensor([1.0, 0.0, 0.0]), + torch.tensor([-1.0, 0.0, 0.0]), + waypoint_count=3, + ) + + final_axis = torch.matmul( + eef_keyframes[:, -1, :3, :3], torch.tensor([1.0, 0.0, 0.0]) + ) + assert torch.isfinite(eef_keyframes).all() + assert torch.allclose( + final_axis, + torch.tensor([-1.0, 0.0, 0.0]).expand(NUM_ENVS, -1), + atol=1.0e-6, + ) + + def test_pick_holds_only_environment_without_a_feasible_grasp() -> None: generator = _motion_generator() semantics = ObjectSemantics( diff --git a/tests/sim/atomic_actions/test_affordance.py b/tests/sim/atomic_actions/test_affordance.py index 1fdd5fc87..79a2e8e40 100644 --- a/tests/sim/atomic_actions/test_affordance.py +++ b/tests/sim/atomic_actions/test_affordance.py @@ -155,6 +155,29 @@ def test_rejects_invalid_internal_axis(self, internal_axis): AxisAlignAffordance(internal_axis=internal_axis) +class TestAxisAlignAffordance: + def test_extends_antipodal_affordance_with_owned_internal_axis(self): + internal_axis = torch.tensor([1.0, 0.0, 0.0]) + + affordance = AxisAlignAffordance(internal_axis=internal_axis) + internal_axis[0] = 0.0 + + assert isinstance(affordance, AntipodalAffordance) + assert torch.equal(affordance.internal_axis, torch.tensor([1.0, 0.0, 0.0])) + + @pytest.mark.parametrize( + "internal_axis", + ( + torch.zeros(3), + torch.tensor([float("nan"), 0.0, 0.0]), + torch.zeros(2), + ), + ) + def test_rejects_invalid_internal_axis(self, internal_axis): + with pytest.raises(ValueError, match="internal_axis"): + AxisAlignAffordance(internal_axis=internal_axis) + + class TestTwistAffordance: def test_requires_explicit_grasp_position_and_axis_origin(self): with pytest.raises(TypeError, match="grasp_position"): From ccfb7a74ee44a77d07790cbdffa29da434480bbb Mon Sep 17 00:00:00 2001 From: matafela Date: Wed, 19 Aug 2026 17:12:37 +0800 Subject: [PATCH 58/85] fix --- embodichain/data/assets/obj_assets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/embodichain/data/assets/obj_assets.py b/embodichain/data/assets/obj_assets.py index 403549138..169f0d946 100644 --- a/embodichain/data/assets/obj_assets.py +++ b/embodichain/data/assets/obj_assets.py @@ -291,7 +291,7 @@ class Drawer(EmbodiChainDataset): def __init__(self, data_root: str = None): data_descriptor = o3d.data.DataDescriptor( os.path.join(EMBODICHAIN_DOWNLOAD_PREFIX, obj_assets, "Drawer.zip"), - "3981636db1f4188146fce25d54084612", + "eba30c852074388c2e5b634b1ae37572", ) prefix = type(self).__name__ path = EMBODICHAIN_DEFAULT_DATA_ROOT if data_root is None else data_root From 34178cae785c3ea5077998f6562b5a321c5660b5 Mon Sep 17 00:00:00 2001 From: matafela Date: Thu, 20 Aug 2026 15:46:01 +0800 Subject: [PATCH 59/85] add pour action --- tests/sim/atomic_actions/test_actions.py | 172 +++++++++++++++++++++++ 1 file changed, 172 insertions(+) diff --git a/tests/sim/atomic_actions/test_actions.py b/tests/sim/atomic_actions/test_actions.py index 8f8952283..00da11933 100644 --- a/tests/sim/atomic_actions/test_actions.py +++ b/tests/sim/atomic_actions/test_actions.py @@ -1464,6 +1464,178 @@ def move_ik( assert not torch.allclose(trajectory.positions[1], context.robot.qpos[1]) +def test_pour_rotates_held_object_about_internal_axis_and_returns() -> None: + generator = _motion_generator() + solved_poses: list[torch.Tensor] = [] + + def compute_ik( + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + solved_poses.append(pose.clone()) + return torch.ones(NUM_ENVS, dtype=torch.bool), joint_seed + 0.1 + + generator.robot.compute_ik.side_effect = compute_ik + current_eef_pose = torch.eye(4).repeat(NUM_ENVS, 1, 1) + current_eef_pose[:, 0, 3] = torch.tensor([0.4, 0.7]) + generator.robot.compute_fk.return_value = current_eef_pose + generator.robot.compute_fk.side_effect = None + action = _bind_action(generator, Pour()) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(internal_axis=torch.tensor([1.0, 0.0, 0.0])), + geometry={}, + label="pourable-object", + ) + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": _held(semantics)}, + ) + context = _context(task) + + plan = _plan_action( + action, + ActionInvocation( + skill_id="pour", + goal=PourGoal(), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=10), + skill_options=PourOptions(rotate_angle=math.pi / 2.0), + ), + context, + ) + + expected_rotation = axis_angle_to_rotation_matrix( + torch.tensor([math.pi / 2.0, 0.0, 0.0]) + ) + assert plan.plan_success.tolist() == [True, True] + assert plan.trajectory.positions.shape == (NUM_ENVS, 10, ROBOT_DOF) + assert plan.trajectory.duration.tolist() == pytest.approx([9.0 / 60.0] * NUM_ENVS) + assert [segment.name for segment in plan.segments] == ["pour"] + assert torch.allclose( + solved_poses[0][:, :3, :3], + expected_rotation.expand(NUM_ENVS, -1, -1), + atol=1.0e-6, + ) + assert torch.allclose( + solved_poses[0][:, :3, 3], + current_eef_pose[:, :3, 3], + ) + assert len(solved_poses) == 2 + assert torch.allclose(solved_poses[1], current_eef_pose, atol=1.0e-6) + assert torch.all(plan.trajectory.positions[:, :, ARM_DOF:] == 1.0) + assert plan.expected_effects.is_empty + assert context.task is task + + +def test_engine_compiles_pickup_followed_by_pour() -> None: + generator = _motion_generator() + engine = AtomicActionEngine( + generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(HAND_DOF), + grasp=torch.ones(HAND_DOF), + ) + }, + ) + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(internal_axis=torch.tensor([1.0, 0.0, 0.0])), + geometry={}, + label="pourable-object", + entity_id="target", + ) + context = _context( + scene=_target_scene( + torch.eye(4).repeat(NUM_ENVS, 1, 1), + timestamp=0.0, + version=0, + ) + ) + + compiled = engine.compile( + ( + ActionInvocation( + skill_id="pick_up", + goal=GraspGoal(semantics, grasp_xpos=torch.eye(4)), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=20), + ), + ActionInvocation( + skill_id="pour", + goal=PourGoal(), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=10), + skill_options=PourOptions(rotate_angle=math.pi / 2.0), + ), + ), + context, + ) + + assert compiled.plan_success.tolist() == [True, True] + assert [plan.skill_id for plan in compiled.action_plans] == ["pick_up", "pour"] + assert compiled.projected_context.get_held_object("arm") is not None + + +def test_pour_requires_exclusively_held_axis_align_affordance() -> None: + generator = _motion_generator() + action = _bind_action(generator, Pour()) + invocation = ActionInvocation( + skill_id="pour", + goal=PourGoal(), + binding=_binding(), + motion_policy=MotionPolicy(sample_count=10), + ) + + with pytest.raises(ValueError, match="run PickUp first"): + _plan_action(action, invocation, _context()) + + invalid_task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={"arm": _held(_semantics())}, + ) + with pytest.raises(ValueError, match="AxisAlignAffordance"): + _plan_action(action, invocation, _context(invalid_task)) + + semantics = ObjectSemantics( + affordance=AxisAlignAffordance(), + geometry={}, + label="shared-pourable-object", + ) + task = TaskState( + batch_size=NUM_ENVS, + device="cpu", + held_objects={ + "arm": _held(semantics), + "alternate_arm": _held( + semantics, + env_mask=torch.tensor([True, False]), + ), + }, + ) + + plan = _plan_action(action, invocation, _context(task)) + + assert plan.plan_success.tolist() == [False, True] + assert torch.allclose( + plan.trajectory.positions[0], + _context(task) + .robot.qpos[0] + .unsqueeze(0) + .expand(plan.trajectory.waypoint_count, -1), + ) + + +def test_pour_options_only_contain_rotate_angle_and_require_finite_value() -> None: + assert set(PourOptions.__dataclass_fields__) == {"rotate_angle"} + assert PourOptions().rotate_angle == pytest.approx(math.pi / 4.0) + with pytest.raises(ValueError, match="rotate_angle must be finite"): + PourOptions(rotate_angle=float("nan")) + + def test_strategy_and_sample_count_are_not_action_config_fields() -> None: with pytest.raises(TypeError): MoveEndEffectorOptions(strategy="motion_gen") # type: ignore[call-arg] From bd3c0dedba744d2598ef230131c6c5158fb88310 Mon Sep 17 00:00:00 2001 From: matafela Date: Mon, 24 Aug 2026 11:46:37 +0800 Subject: [PATCH 60/85] fix docs --- .../scene_engine/core/scene_edit_plan.py | 36 +++++++++++++++++-- .../editing/scene_edit_asset_preparation.py | 33 ++++++++++++++++- 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index 0f78b9699..7dffaae2c 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -35,7 +35,24 @@ @dataclass(frozen=True) class SceneEditOperation: - """One normalized edit operation produced from an LLM edit draft.""" + """Describe one normalized add, move, or delete operation. + + Attributes: + op: Operation kind. Add creates a new object, move repositions an + existing object, and delete removes an existing object. + object_id: Existing object ID for move/delete, or the generated ID for + an added object. + target_id: Optional existing scene object used as the spatial target. + relation: Spatial relation between the edited object and ``target_id``. + table_region: Optional named tabletop region. It is valid only when the + target is the table and the relation is ``"on"``. + category: Semantic category required for an added object. + name: Human-readable name required for an added object. + description: Generation prompt and semantic description required for + an added object. + orientation_state: Optional standing or lying intent for an added + object. Move operations may only preserve the existing state. + """ op: SceneEditOperationType object_id: str | None = None @@ -64,7 +81,22 @@ def to_dict(self) -> dict[str, object]: @dataclass class SceneEditPlan: - """Validated operations against one immutable pre-edit scene state.""" + """Validate edit operations against one pre-edit scene state. + + Construction validates every object reference and rejects conflicting + operations without mutating the supplied scene or scene graph. + + Attributes: + scene: Scene state that exists before the edit is applied. + scene_graph: Pre-edit support and spatial-relation graph. Its node IDs + must match the scene object IDs. + operations: Normalized operations in application order. + + Raises: + ValueError: If scene IDs are inconsistent, an operation has invalid + fields or references, edits conflict, or a deletion would orphan a + support descendant. + """ scene: Scene scene_graph: SceneGraph diff --git a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py index 220a91014..616d647e3 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py +++ b/embodichain/gen_sim/scene_engine/pipeline/editing/scene_edit_asset_preparation.py @@ -72,7 +72,38 @@ def prepare_scene_edit_assets( image_segmentation_client: ImageSegmentationClient, vlm_client: OpenAICompatibleVLM | None = None, ) -> list[SceneObject]: - """Prepare and return SimReady assets required by add operations.""" + """Generate canonical SimReady assets for a scene edit's add operations. + + Move-only and delete-only plans return immediately without modifying an + existing asset-preparation directory. For add operations, the function + generates and segments one image per object, creates coarse geometry, + processes it into SimReady geometry, and resets the returned objects to + identity edit-time poses. + + Args: + scene_edit_plan: Validated edit plan whose add operations define the + objects to generate. + output_root: Scene Engine output root. Intermediate artifacts are + written below ``scene_editing/asset_preparation``. + image_generation_client: Client used to render object images from the + operation descriptions. + geometry_generation_client: Client used to create coarse GLB geometry + from each generated image and mask. + image_segmentation_client: Client used to isolate the generated object + in each image. + vlm_client: Optional VLM used by SimReady processing to estimate asset + scale and orientation. + + Returns: + Added ``SceneObject`` assets in edit-plan order, or an empty list when + the plan contains no add operations. + + Raises: + ValueError: If add metadata or generated image, mask, and geometry + mappings are incomplete or inconsistent. + FileNotFoundError: If geometry generation does not produce an expected + GLB file. + """ # Prepare descriptions for all newly added objects. added_asset_descriptions = _collect_added_asset_descriptions(scene_edit_plan) # Skip asset generation when the edit plan only moves or deletes existing objects. From adfe065b1be0f6bdf60c763bad0c73299c34d73f Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:01:05 +0800 Subject: [PATCH 61/85] feat(scene-engine): add versioned authoring API --- docs/source/api_reference/public_api.rst | 52 +++ embodichain/gen_sim/scene_engine/errors.py | 23 ++ .../gen_sim/scene_engine/pipeline/__init__.py | 24 +- .../gen_sim/scene_engine/pipeline/api.py | 351 ++++++++++++++++++ .../gen_sim/scene_engine/test_pipeline_api.py | 244 ++++++++++++ 5 files changed, 693 insertions(+), 1 deletion(-) create mode 100644 embodichain/gen_sim/scene_engine/errors.py create mode 100644 embodichain/gen_sim/scene_engine/pipeline/api.py create mode 100644 tests/gen_sim/scene_engine/test_pipeline_api.py diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index cdf2f0611..92103798d 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -238,6 +238,58 @@ embodichain.gen_sim.scene_engine.core.scene_edit_plan SceneEditOperation SceneEditPlan +embodichain.gen_sim.scene_engine.errors +--------------------------------------- + +.. currentmodule:: embodichain.gen_sim.scene_engine.errors + +Scene service failures preserve typed preparation and materialization errors +across the Task Engine boundary. + +.. autosummary:: + + SceneServiceError + +embodichain.gen_sim.scene_engine.pipeline +----------------------------------------- + +.. currentmodule:: embodichain.gen_sim.scene_engine.pipeline + +The public authoring boundary separates deterministic scene analysis from +side-effecting materialization for generated and edited scenes. + +.. autosummary:: + + SCENE_BLUEPRINT_SCHEMA + SCENE_EDIT_BLUEPRINT_SCHEMA + SceneBlueprintPackage + SceneEditBlueprintPackage + SceneMaterialization + analyze_edit + analyze_image + materialize_blueprint + materialize_edit + +embodichain.gen_sim.scene_engine.pipeline.api +--------------------------------------------- + +.. currentmodule:: embodichain.gen_sim.scene_engine.pipeline.api + +Versioned blueprint artifacts and analyze/materialize operations provide the +implementation-level Scene Engine authoring contract. + +.. autosummary:: + + SCENE_BLUEPRINT_SCHEMA + SCENE_EDIT_BLUEPRINT_SCHEMA + SceneBlueprintPackage + SceneEditBlueprintPackage + SceneMaterialization + analyze_edit + analyze_image + materialize_blueprint + materialize_edit + embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation ------------------------------------------------------------------------------- diff --git a/embodichain/gen_sim/scene_engine/errors.py b/embodichain/gen_sim/scene_engine/errors.py new file mode 100644 index 000000000..dd23d2aa9 --- /dev/null +++ b/embodichain/gen_sim/scene_engine/errors.py @@ -0,0 +1,23 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +__all__ = ["SceneServiceError"] + + +class SceneServiceError(RuntimeError): + """A transient or remote Scene Engine service failure.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/__init__.py b/embodichain/gen_sim/scene_engine/pipeline/__init__.py index 015c41510..ecf448d22 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/__init__.py +++ b/embodichain/gen_sim/scene_engine/pipeline/__init__.py @@ -16,4 +16,26 @@ from __future__ import annotations -__all__: list[str] = [] +from .api import ( + SCENE_BLUEPRINT_SCHEMA, + SCENE_EDIT_BLUEPRINT_SCHEMA, + SceneBlueprintPackage, + SceneEditBlueprintPackage, + SceneMaterialization, + analyze_edit, + analyze_image, + materialize_blueprint, + materialize_edit, +) + +__all__ = [ + "SCENE_BLUEPRINT_SCHEMA", + "SCENE_EDIT_BLUEPRINT_SCHEMA", + "SceneBlueprintPackage", + "SceneEditBlueprintPackage", + "SceneMaterialization", + "analyze_edit", + "analyze_image", + "materialize_blueprint", + "materialize_edit", +] diff --git a/embodichain/gen_sim/scene_engine/pipeline/api.py b/embodichain/gen_sim/scene_engine/pipeline/api.py new file mode 100644 index 000000000..6f219615c --- /dev/null +++ b/embodichain/gen_sim/scene_engine/pipeline/api.py @@ -0,0 +1,351 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Auditable stage boundaries for Scene Engine generation and editing.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +import hashlib +import json +from pathlib import Path +from typing import Any, Final + +from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( + GeometryGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_generation import ( + ImageGenerationClient, +) +from embodichain.gen_sim.scene_engine.clients.image_segmentation import ( + ImageSegmentationClient, +) +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import SceneGraph +from embodichain.gen_sim.scene_engine.llms.openai_compatible_client import ( + OpenAICompatibleVLM, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_asset_preparation import ( + prepare_scene_edit_assets, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_layout_generation import ( + edit_layout, +) +from embodichain.gen_sim.scene_engine.pipeline.editing.scene_edit_understanding import ( + understand_scene_edit, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_generation import ( + generate_scene_and_refine, +) +from embodichain.gen_sim.scene_engine.pipeline.generation.scene_understanding import ( + understand_scene, +) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_importer import ( + SceneExportImporter, +) +from embodichain.utils.logger import log_info + +__all__ = [ + "SCENE_BLUEPRINT_SCHEMA", + "SCENE_EDIT_BLUEPRINT_SCHEMA", + "SceneBlueprintPackage", + "SceneEditBlueprintPackage", + "SceneMaterialization", + "analyze_edit", + "analyze_image", + "materialize_blueprint", + "materialize_edit", +] + +SCENE_BLUEPRINT_SCHEMA: Final = "embodichain.scene-blueprint/v1" +SCENE_EDIT_BLUEPRINT_SCHEMA: Final = "embodichain.scene-edit-blueprint/v1" + + +@dataclass(frozen=True) +class SceneBlueprintPackage: + """In-process scene semantics plus their persisted audit document.""" + + blueprint_id: str + image_path: Path + output_root: Path + manifest_path: Path + scene: Scene + scene_graph: SceneGraph + + +@dataclass(frozen=True) +class SceneEditBlueprintPackage: + """Validated edit intent before added assets and layout are materialized.""" + + blueprint_id: str + edit_prompt: str + output_root: Path + manifest_path: Path + scene_edit_plan: SceneEditPlan + updated_scene_graph: SceneGraph + + +@dataclass(frozen=True) +class SceneMaterialization: + """One exported materialized scene revision.""" + + scene: Scene + scene_graph: SceneGraph + output_root: Path + scene_config_path: Path + + +def analyze_image( + image_path: str | Path, + output_root: str | Path, + *, + vlm_client: OpenAICompatibleVLM | None = None, + image_segmentation_client: ImageSegmentationClient | None = None, +) -> SceneBlueprintPackage: + """Understand an image and persist the pre-generation semantic blueprint.""" + resolved_image = Path(image_path).expanduser().resolve() + resolved_output = Path(output_root).expanduser().resolve() + resolved_output.mkdir(parents=True, exist_ok=True) + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + segmentation = image_segmentation_client or ImageSegmentationClient.from_dotenv() + owns_segmentation = image_segmentation_client is None + log_info("Starting Scene Understanding") + try: + segmentation.check_health() + scene, scene_graph = understand_scene( + scene=Scene(), + image_path=resolved_image, + output_root=resolved_output, + vlm_client=effective_vlm, + image_segmentation_client=segmentation, + ) + finally: + if owns_segmentation: + segmentation.close() + log_info("Completed Scene Understanding") + + payload = { + "schema_version": SCENE_BLUEPRINT_SCHEMA, + "image_path": resolved_image.as_posix(), + "scene": scene.to_dict(), + "scene_graph": scene_graph.to_dict(), + "artifacts": _artifact_records(resolved_output / "scene_understanding"), + } + blueprint_id = _canonical_hash(payload) + document = {**payload, "blueprint_id": blueprint_id} + manifest_path = resolved_output / "scene_blueprint.json" + _write_json(manifest_path, document) + return SceneBlueprintPackage( + blueprint_id=blueprint_id, + image_path=resolved_image, + output_root=resolved_output, + manifest_path=manifest_path, + scene=scene, + scene_graph=scene_graph, + ) + + +def materialize_blueprint( + blueprint: SceneBlueprintPackage, + *, + vlm_client: OpenAICompatibleVLM | None = None, + geometry_generation_client: GeometryGenerationClient | None = None, + seed: int | None = None, +) -> SceneMaterialization: + """Generate assets and layout for one image-derived blueprint.""" + scene = deepcopy(blueprint.scene) + scene_graph = deepcopy(blueprint.scene_graph) + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + geometry = geometry_generation_client or GeometryGenerationClient.from_dotenv() + owns_geometry = geometry_generation_client is None + log_info("Starting Objects + Coarse Layout Generation") + try: + geometry.check_health() + scene = generate_scene_and_refine( + image_path=blueprint.image_path, + output_root=blueprint.output_root, + scene=scene, + scene_graph=scene_graph, + geometry_generation_client=geometry, + vlm_client=effective_vlm, + seed=seed, + ) + finally: + if owns_geometry: + geometry.close() + log_info("Completed Objects + Coarse Layout Generation") + return _export_materialization( + scene=scene, + scene_graph=scene_graph, + output_root=blueprint.output_root, + ) + + +def analyze_edit( + *, + output_root: str | Path, + edit_prompt: str, + vlm_client: OpenAICompatibleVLM | None = None, +) -> SceneEditBlueprintPackage: + """Interpret and persist one edit against an already generated scene.""" + resolved_output = Path(output_root).expanduser().resolve() + normalized_prompt = str(edit_prompt).strip() + if not normalized_prompt: + raise ValueError("Edit prompt must not be empty.") + scene, scene_graph = SceneExportImporter( + output_root=resolved_output + ).import_scene_and_graph() + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + log_info("Starting Edit Understanding") + scene_edit_plan, updated_scene_graph = understand_scene_edit( + scene=scene, + scene_graph=scene_graph, + edit_prompt=normalized_prompt, + vlm_client=effective_vlm, + ) + log_info("Completed Edit Understanding") + payload = { + "schema_version": SCENE_EDIT_BLUEPRINT_SCHEMA, + "edit_prompt": normalized_prompt, + "scene_edit_plan": scene_edit_plan.to_dict(), + "updated_scene_graph": updated_scene_graph.to_dict(), + } + blueprint_id = _canonical_hash(payload) + manifest_path = resolved_output / "scene_edit" / "scene_edit_blueprint.json" + _write_json(manifest_path, {**payload, "blueprint_id": blueprint_id}) + return SceneEditBlueprintPackage( + blueprint_id=blueprint_id, + edit_prompt=normalized_prompt, + output_root=resolved_output, + manifest_path=manifest_path, + scene_edit_plan=scene_edit_plan, + updated_scene_graph=updated_scene_graph, + ) + + +def materialize_edit( + blueprint: SceneEditBlueprintPackage, + *, + vlm_client: OpenAICompatibleVLM | None = None, + image_generation_client: ImageGenerationClient | None = None, + geometry_generation_client: GeometryGenerationClient | None = None, + image_segmentation_client: ImageSegmentationClient | None = None, + seed: int | None = None, +) -> SceneMaterialization: + """Generate added assets, apply layout edits, and export the new revision.""" + scene_edit_plan = deepcopy(blueprint.scene_edit_plan) + updated_scene_graph = deepcopy(blueprint.updated_scene_graph) + effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() + image_generation = image_generation_client or ImageGenerationClient.from_dotenv() + geometry = geometry_generation_client or GeometryGenerationClient.from_dotenv() + segmentation = image_segmentation_client or ImageSegmentationClient.from_dotenv() + owned_clients = ( + (image_generation, image_generation_client is None), + (geometry, geometry_generation_client is None), + (segmentation, image_segmentation_client is None), + ) + log_info("Starting Objects Preparation") + try: + for client, _ in owned_clients: + client.check_health() + added_assets = prepare_scene_edit_assets( + scene_edit_plan=scene_edit_plan, + output_root=blueprint.output_root, + image_generation_client=image_generation, + geometry_generation_client=geometry, + image_segmentation_client=segmentation, + vlm_client=effective_vlm, + seed=seed, + ) + finally: + for client, owned in owned_clients: + if owned: + client.close() + log_info("Completed Objects Preparation") + log_info("Starting Layout Generation") + scene = edit_layout( + scene=scene_edit_plan.scene, + scene_edit_plan=scene_edit_plan, + updated_scene_graph=updated_scene_graph, + added_assets=added_assets, + output_root=blueprint.output_root, + ) + log_info("Completed Layout Generation") + return _export_materialization( + scene=scene, + scene_graph=updated_scene_graph, + output_root=blueprint.output_root, + ) + + +def _export_materialization( + *, + scene: Scene, + scene_graph: SceneGraph, + output_root: Path, +) -> SceneMaterialization: + log_info("Starting Scene Export") + scene_config_path = SceneExporter( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ).export() + log_info("Completed Scene Export") + return SceneMaterialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + scene_config_path=scene_config_path, + ) + + +def _artifact_records(root: Path) -> list[dict[str, Any]]: + if not root.is_dir(): + return [] + records = [] + for path in sorted(item for item in root.rglob("*") if item.is_file()): + records.append( + { + "path": path.resolve().as_posix(), + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + "size": path.stat().st_size, + } + ) + return records + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text( + json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False) + "\n", + encoding="utf-8", + ) + temporary.replace(path) diff --git a/tests/gen_sim/scene_engine/test_pipeline_api.py b/tests/gen_sim/scene_engine/test_pipeline_api.py new file mode 100644 index 000000000..9b10361f7 --- /dev/null +++ b/tests/gen_sim/scene_engine/test_pipeline_api.py @@ -0,0 +1,244 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +import json +from pathlib import Path + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline import api + + +class _HealthyClient: + def __init__(self) -> None: + self.health_checks = 0 + + def check_health(self) -> None: + self.health_checks += 1 + + +def _materialization( + *, + scene: Scene, + scene_graph: SceneGraph, + output_root: Path, +) -> api.SceneMaterialization: + return api.SceneMaterialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + scene_config_path=output_root / "scene_export" / "scene_config.json", + ) + + +def _table_scene() -> tuple[Scene, SceneGraph]: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="A work table.", + ) + ] + ) + graph = SceneGraph(nodes=[SceneGraphNode(object_id="table", parent_id=None)]) + return scene, graph + + +def test_analyze_image_persists_blueprint_and_artifact_hashes( + tmp_path: Path, + monkeypatch, +) -> None: + image_path = tmp_path / "input.png" + image_path.write_bytes(b"image") + scene, graph = _table_scene() + + def fake_understand_scene(**kwargs): + stage_root = Path(kwargs["output_root"]) / "scene_understanding" + stage_root.mkdir(parents=True) + (stage_root / "table-mask.png").write_bytes(b"mask") + return scene, graph + + monkeypatch.setattr(api, "understand_scene", fake_understand_scene) + segmentation = _HealthyClient() + package = api.analyze_image( + image_path, + tmp_path / "output", + vlm_client=object(), + image_segmentation_client=segmentation, + ) + document = json.loads(package.manifest_path.read_text(encoding="utf-8")) + + assert segmentation.health_checks == 1 + assert document["blueprint_id"] == package.blueprint_id + assert document["scene_graph"] == graph.to_dict() + assert document["artifacts"][0]["path"].endswith("table-mask.png") + assert len(document["artifacts"][0]["sha256"]) == 64 + + +def test_analyze_edit_persists_post_edit_blueprint( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + plan = SceneEditPlan(scene=scene, scene_graph=graph, operations=[]) + + class FakeImporter: + def __init__(self, *, output_root: Path) -> None: + self.output_root = output_root + + def import_scene_and_graph(self): + return scene, graph + + monkeypatch.setattr(api, "SceneExportImporter", FakeImporter) + monkeypatch.setattr( + api, + "understand_scene_edit", + lambda **_: (plan, graph), + ) + package = api.analyze_edit( + output_root=tmp_path, + edit_prompt="Keep the scene unchanged.", + vlm_client=object(), + ) + document = json.loads(package.manifest_path.read_text(encoding="utf-8")) + + assert document["blueprint_id"] == package.blueprint_id + assert document["scene_edit_plan"] == plan.to_dict() + assert document["updated_scene_graph"] == graph.to_dict() + + +def test_materialize_blueprint_does_not_mutate_audited_snapshot( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + manifest_path = tmp_path / "scene_blueprint.json" + manifest_path.write_text("audited blueprint\n", encoding="utf-8") + package = api.SceneBlueprintPackage( + blueprint_id="blueprint", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=manifest_path, + scene=scene, + scene_graph=graph, + ) + original_scene = deepcopy(scene.to_dict()) + original_graph = deepcopy(graph.to_dict()) + + def fake_generate_scene_and_refine(**kwargs): + assert kwargs["seed"] == 31 + assert kwargs["scene"] is not package.scene + assert kwargs["scene_graph"] is not package.scene_graph + kwargs["scene"].objects[0].name = "materialized table" + return kwargs["scene"] + + monkeypatch.setattr( + api, + "generate_scene_and_refine", + fake_generate_scene_and_refine, + ) + monkeypatch.setattr( + api, + "_export_materialization", + lambda *, scene, scene_graph, output_root: _materialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ), + ) + + result = api.materialize_blueprint( + package, + vlm_client=object(), + geometry_generation_client=_HealthyClient(), + seed=31, + ) + + assert result.scene.objects[0].name == "materialized table" + assert package.scene.to_dict() == original_scene + assert package.scene_graph.to_dict() == original_graph + assert manifest_path.read_text(encoding="utf-8") == "audited blueprint\n" + + +def test_materialize_edit_does_not_mutate_audited_snapshot( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + plan = SceneEditPlan(scene=scene, scene_graph=graph, operations=[]) + manifest_path = tmp_path / "scene_edit_blueprint.json" + manifest_path.write_text("audited edit blueprint\n", encoding="utf-8") + package = api.SceneEditBlueprintPackage( + blueprint_id="edit-blueprint", + edit_prompt="Keep the scene unchanged.", + output_root=tmp_path, + manifest_path=manifest_path, + scene_edit_plan=plan, + updated_scene_graph=graph, + ) + original_plan = deepcopy(plan.to_dict()) + original_graph = deepcopy(graph.to_dict()) + + def fake_prepare_scene_edit_assets(**kwargs): + assert kwargs["seed"] == 32 + return [] + + monkeypatch.setattr( + api, "prepare_scene_edit_assets", fake_prepare_scene_edit_assets + ) + + def fake_edit_layout(**kwargs): + assert kwargs["scene_edit_plan"] is not package.scene_edit_plan + assert kwargs["updated_scene_graph"] is not package.updated_scene_graph + kwargs["scene"].objects[0].name = "edited table" + return kwargs["scene"] + + monkeypatch.setattr(api, "edit_layout", fake_edit_layout) + monkeypatch.setattr( + api, + "_export_materialization", + lambda *, scene, scene_graph, output_root: _materialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ), + ) + clients = [_HealthyClient(), _HealthyClient(), _HealthyClient()] + + result = api.materialize_edit( + package, + vlm_client=object(), + image_generation_client=clients[0], + geometry_generation_client=clients[1], + image_segmentation_client=clients[2], + seed=32, + ) + + assert result.scene.objects[0].name == "edited table" + assert package.scene_edit_plan.to_dict() == original_plan + assert package.updated_scene_graph.to_dict() == original_graph + assert manifest_path.read_text(encoding="utf-8") == "audited edit blueprint\n" From 838b167644717ef06c8eeb3d3d470a2cfcbe92fd Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:01:17 +0800 Subject: [PATCH 62/85] feat(atomic-actions): preserve GenSim execution contracts --- embodichain/lab/sim/atomic_actions/__init__.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/__init__.py b/embodichain/lab/sim/atomic_actions/__init__.py index 5ce887e8a..7056efc09 100644 --- a/embodichain/lab/sim/atomic_actions/__init__.py +++ b/embodichain/lab/sim/atomic_actions/__init__.py @@ -31,8 +31,8 @@ from .affordance import ( Affordance, AntipodalAffordance, - AssembleAffordance, AxisAlignAffordance, + AssembleAffordance, InteractionPoints, OpenDoorAffordance, PressAffordance, @@ -171,7 +171,6 @@ EndEffectorPoseGoal, GraspGoal, HandOver, - HandOverGoal, HandOverOptions, HeldObjectPoseGoal, JointPositionGoal, @@ -253,11 +252,11 @@ "ActionPlanningServices", "Affordance", "AntipodalAffordance", + "AxisAlignAffordance", "ArticulationJointState", "AssembleAffordance", "AssembleGoal", "AxisAlign", - "AxisAlignAffordance", "AxisAlignGoal", "AxisAlignOptions", "AtomicAction", @@ -318,7 +317,6 @@ "GRASP_CAPABILITY", "GraspGoal", "HandOver", - "HandOverGoal", "HandOverOptions", "HeldObjectPoseGoal", "HeldObjectState", From 168ed64f3124f4a860db62ae9a4b7eac0bf9d6a8 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:01:28 +0800 Subject: [PATCH 63/85] feat(gen-sim): migrate action engine onto axis baseline --- .../gen_sim/action_engine/ARCHITECTURE.md | 398 ++ embodichain/gen_sim/action_engine/__init__.py | 50 + embodichain/gen_sim/action_engine/agent.py | 762 ++ .../action_engine/capabilities/__init__.py | 57 + .../action_engine/capabilities/atomic.py | 1040 +++ .../action_engine/capabilities/builtins.py | 1018 +++ .../action_engine/capabilities/registry.py | 183 + .../gen_sim/action_engine/cli/__init__.py | 21 + .../cli/generate_action_agent_config.py | 240 + .../gen_sim/action_engine/cli/run_agent.py | 1620 +++++ .../action_engine/compiler/__init__.py | 33 + .../gen_sim/action_engine/compiler/core.py | 594 ++ .../gen_sim/action_engine/compiler/v2.py | 424 ++ .../gen_sim/action_engine/config/__init__.py | 41 + .../action_engine/config/defaults.yaml | 344 + .../action_engine/config/runtime_policy.py | 775 ++ .../gen_sim/action_engine/domain/__init__.py | 93 + .../gen_sim/action_engine/domain/motion.py | 111 + .../gen_sim/action_engine/domain/programs.py | 848 +++ .../action_engine/domain/task_contracts.py | 125 + .../gen_sim/action_engine/domain/v2.py | 1297 ++++ .../action_engine/domain/visual_contracts.py | 60 + .../action_engine/environment/__init__.py | 23 + .../action_engine/environment/agent_env.py | 535 ++ .../action_engine/evaluation/__init__.py | 27 + .../gen_sim/action_engine/evaluation/ab.py | 831 +++ .../evaluation/e1_e2_scene_action.py | 448 ++ .../action_engine/generation/__init__.py | 33 + .../action_engine/generation/artifacts.py | 141 + .../action_engine/generation/assets.py | 158 + .../generation/config_builder.py | 852 +++ .../action_engine/generation/generator.py | 877 +++ .../action_engine/generation/models.py | 83 + .../action_engine/generation/source_scene.py | 644 ++ .../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/robot_profiles.json | 45 + .../generation/templates/vlm_sensors.json | 58 + .../action_engine/graph_visualization.py | 938 +++ .../gen_sim/action_engine/orientation.py | 226 + .../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 +++ embodichain/gen_sim/action_engine/protocol.py | 60 + .../gen_sim/action_engine/runtime/__init__.py | 66 + .../gen_sim/action_engine/runtime/actions.py | 1534 ++++ .../action_engine/runtime/atomic_compat.py | 85 + .../gen_sim/action_engine/runtime/dynamic.py | 164 + .../gen_sim/action_engine/runtime/executor.py | 4359 ++++++++++++ .../gen_sim/action_engine/runtime/frames.py | 165 + .../runtime/grasp_collision_cache.py | 330 + .../action_engine/runtime/grounding.py | 3155 +++++++++ .../gen_sim/action_engine/runtime/loader.py | 296 + .../gen_sim/action_engine/runtime/models.py | 314 + .../action_engine/runtime/motion_policy.py | 106 + .../action_engine/runtime/predicates.py | 844 +++ .../action_engine/runtime/recording.py | 391 ++ .../gen_sim/action_engine/runtime/recovery.py | 649 ++ .../action_engine/runtime/reporting.py | 348 + .../action_engine/runtime/robot_parts.py | 34 + .../action_engine/runtime/solver_compat.py | 234 + .../gen_sim/action_engine/runtime/state.py | 84 + .../gen_sim/action_engine/tasks/__init__.py | 50 + .../gen_sim/action_engine/tasks/assembly.py | 425 ++ .../gen_sim/action_engine/tasks/grounding.py | 513 ++ .../action_engine/tasks/interpretation.py | 404 ++ .../gen_sim/action_engine/tasks/recipes.py | 1061 +++ .../gen_sim/action_engine/tasks/scene.py | 177 + embodichain/gen_sim/action_engine/unbound.py | 205 + tests/gen_sim/__init__.py | 4 + tests/gen_sim/action_engine/__init__.py | 19 + .../action_engine/capabilities/__init__.py | 19 + .../capabilities/test_atomic_v2.py | 233 + .../action_engine/cli/test_run_agent.py | 219 + .../action_engine/compiler/__init__.py | 19 + .../action_engine/compiler/test_compiler.py | 572 ++ .../gen_sim/action_engine/compiler/test_v2.py | 172 + .../gen_sim/action_engine/config/__init__.py | 17 + .../config/test_runtime_policy.py | 308 + .../gen_sim/action_engine/domain/__init__.py | 19 + .../action_engine/domain/test_programs.py | 150 + .../domain/test_task_contracts.py | 71 + tests/gen_sim/action_engine/domain/test_v2.py | 280 + .../action_engine/evaluation/__init__.py | 19 + .../action_engine/evaluation/test_ab.py | 445 ++ .../generation/test_generation.py | 1590 +++++ .../action_engine/planning/__init__.py | 21 + .../action_engine/planning/test_linker.py | 407 ++ .../action_engine/planning/test_online_v2.py | 480 ++ .../action_engine/planning/test_planner.py | 730 ++ .../gen_sim/action_engine/runtime/__init__.py | 19 + .../action_engine/runtime/test_actions.py | 835 +++ .../runtime/test_atomic_compat.py | 88 + .../runtime/test_grasp_collision_cache.py | 354 + .../action_engine/runtime/test_recovery_v2.py | 720 ++ .../runtime/test_runtime_contracts.py | 6253 +++++++++++++++++ tests/gen_sim/action_engine/task_fixtures.py | 229 + tests/gen_sim/action_engine/tasks/__init__.py | 19 + .../action_engine/tasks/test_factory.py | 522 ++ .../action_engine/tasks/test_grounding.py | 339 + .../tasks/test_interpretation.py | 1939 +++++ .../tasks/test_language_decoupling.py | 420 ++ tests/gen_sim/action_engine/test_agent.py | 219 + .../action_engine/test_architecture.py | 155 + .../action_engine/test_graph_visualization.py | 362 + .../action_engine/test_motion_policy.py | 72 + .../gen_sim/action_engine/test_orientation.py | 145 + tests/gen_sim/action_engine/test_unbound.py | 111 + 115 files changed, 54850 insertions(+) create mode 100644 embodichain/gen_sim/action_engine/ARCHITECTURE.md create mode 100644 embodichain/gen_sim/action_engine/__init__.py create mode 100644 embodichain/gen_sim/action_engine/agent.py create mode 100644 embodichain/gen_sim/action_engine/capabilities/__init__.py create mode 100644 embodichain/gen_sim/action_engine/capabilities/atomic.py create mode 100644 embodichain/gen_sim/action_engine/capabilities/builtins.py create mode 100644 embodichain/gen_sim/action_engine/capabilities/registry.py 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/cli/run_agent.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/config/__init__.py create mode 100644 embodichain/gen_sim/action_engine/config/defaults.yaml create mode 100644 embodichain/gen_sim/action_engine/config/runtime_policy.py create mode 100644 embodichain/gen_sim/action_engine/domain/__init__.py create mode 100644 embodichain/gen_sim/action_engine/domain/motion.py create mode 100644 embodichain/gen_sim/action_engine/domain/programs.py create mode 100644 embodichain/gen_sim/action_engine/domain/task_contracts.py create mode 100644 embodichain/gen_sim/action_engine/domain/v2.py create mode 100644 embodichain/gen_sim/action_engine/domain/visual_contracts.py create mode 100644 embodichain/gen_sim/action_engine/environment/__init__.py create mode 100644 embodichain/gen_sim/action_engine/environment/agent_env.py create mode 100644 embodichain/gen_sim/action_engine/evaluation/__init__.py create mode 100644 embodichain/gen_sim/action_engine/evaluation/ab.py create mode 100644 embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.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/config_builder.py create mode 100644 embodichain/gen_sim/action_engine/generation/generator.py create mode 100644 embodichain/gen_sim/action_engine/generation/models.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/robot_profiles.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/orientation.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/protocol.py create mode 100644 embodichain/gen_sim/action_engine/runtime/__init__.py create mode 100644 embodichain/gen_sim/action_engine/runtime/actions.py create mode 100644 embodichain/gen_sim/action_engine/runtime/atomic_compat.py create mode 100644 embodichain/gen_sim/action_engine/runtime/dynamic.py create mode 100644 embodichain/gen_sim/action_engine/runtime/executor.py create mode 100644 embodichain/gen_sim/action_engine/runtime/frames.py create mode 100644 embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py create mode 100644 embodichain/gen_sim/action_engine/runtime/grounding.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/motion_policy.py create mode 100644 embodichain/gen_sim/action_engine/runtime/predicates.py create mode 100644 embodichain/gen_sim/action_engine/runtime/recording.py create mode 100644 embodichain/gen_sim/action_engine/runtime/recovery.py create mode 100644 embodichain/gen_sim/action_engine/runtime/reporting.py create mode 100644 embodichain/gen_sim/action_engine/runtime/robot_parts.py create mode 100644 embodichain/gen_sim/action_engine/runtime/solver_compat.py create mode 100644 embodichain/gen_sim/action_engine/runtime/state.py create mode 100644 embodichain/gen_sim/action_engine/tasks/__init__.py create mode 100644 embodichain/gen_sim/action_engine/tasks/assembly.py create mode 100644 embodichain/gen_sim/action_engine/tasks/grounding.py create mode 100644 embodichain/gen_sim/action_engine/tasks/interpretation.py create mode 100644 embodichain/gen_sim/action_engine/tasks/recipes.py create mode 100644 embodichain/gen_sim/action_engine/tasks/scene.py create mode 100644 embodichain/gen_sim/action_engine/unbound.py create mode 100644 tests/gen_sim/action_engine/__init__.py create mode 100644 tests/gen_sim/action_engine/capabilities/__init__.py create mode 100644 tests/gen_sim/action_engine/capabilities/test_atomic_v2.py create mode 100644 tests/gen_sim/action_engine/cli/test_run_agent.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/config/__init__.py create mode 100644 tests/gen_sim/action_engine/config/test_runtime_policy.py create mode 100644 tests/gen_sim/action_engine/domain/__init__.py create mode 100644 tests/gen_sim/action_engine/domain/test_programs.py create mode 100644 tests/gen_sim/action_engine/domain/test_task_contracts.py create mode 100644 tests/gen_sim/action_engine/domain/test_v2.py create mode 100644 tests/gen_sim/action_engine/evaluation/__init__.py create mode 100644 tests/gen_sim/action_engine/evaluation/test_ab.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_linker.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/runtime/__init__.py create mode 100644 tests/gen_sim/action_engine/runtime/test_actions.py create mode 100644 tests/gen_sim/action_engine/runtime/test_atomic_compat.py create mode 100644 tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py create mode 100644 tests/gen_sim/action_engine/runtime/test_recovery_v2.py create mode 100644 tests/gen_sim/action_engine/runtime/test_runtime_contracts.py create mode 100644 tests/gen_sim/action_engine/task_fixtures.py create mode 100644 tests/gen_sim/action_engine/tasks/__init__.py create mode 100644 tests/gen_sim/action_engine/tasks/test_factory.py create mode 100644 tests/gen_sim/action_engine/tasks/test_grounding.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_agent.py create mode 100644 tests/gen_sim/action_engine/test_architecture.py create mode 100644 tests/gen_sim/action_engine/test_graph_visualization.py create mode 100644 tests/gen_sim/action_engine/test_motion_policy.py create mode 100644 tests/gen_sim/action_engine/test_orientation.py create mode 100644 tests/gen_sim/action_engine/test_unbound.py diff --git a/embodichain/gen_sim/action_engine/ARCHITECTURE.md b/embodichain/gen_sim/action_engine/ARCHITECTURE.md new file mode 100644 index 000000000..6ead4afa6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/ARCHITECTURE.md @@ -0,0 +1,398 @@ +# Action Engine v2 Architecture + +Action Engine v2 uses a task-first protocol and executes a direct +`AtomicAction` graph. The persisted graph is symbolic and coordinate-free; +simulator geometry is resolved immediately before each action executes. + +The Task Engine entry point wraps that existing pipeline with three narrow +owners: + +1. `TaskAgent` produces three scene-independent `TaskDraft` candidates and + deterministically derives each `SceneRequest` and `SuccessSpec`. +2. `SceneAdapter` binds one verified candidate to a read-only existing scene, + producing `SceneManifest`, `RoleBindings`, and a complete `BindingReport`. +3. `ActionAgent` lowers the selected `GroundedTaskPlan` to the existing + `action_engine_seed_graph_v3`, performs executable capability preflight, + runs it through `ProgramExecutor`, and emits a tensor-free + `ExecutionReport`. Report v2 records the episode seed, package and Python + versions, Git commit/dirty state when available, and structured runtime + arguments alongside the existing plan and graph hashes. + +The public CLI is `python -m embodichain.gen_sim.task_engine --mode ...` with +strict `image`, `image-edit`, `scene`, and `scene-edit` input modes. Every +invocation runs the complete workflow through real trajectory acceptance and +publishes an isolated, timestamped child under `--output-root`. Source projects +are referenced in place and integrity-hashed; Task Engine does not modify them +or copy them into a scene package store. + +## Package Ownership + +The cross-engine workflow is owned by Task Engine rather than nested under +Action Engine: + +- `embodichain.gen_sim.task_engine` owns scene-independent interpretation, + E1-E9 semantic ontology, `TaskDraft`, `SceneRequest`, `SuccessSpec`, and + `TaskAgent`. +- `embodichain.gen_sim.scene_engine` remains the scene generation subsystem and + exposes auditable image-understanding, materialization, edit-understanding, + and edit-materialization stages. +- `embodichain.gen_sim.task_engine.scene` owns Scene Engine adaptation, the + richer static manifest, and deterministic scene/action feasibility reports. +- `embodichain.gen_sim.action_engine.agent` owns `ActionAgent`; Action Engine's + existing `domain`, `planning`, and `runtime` packages remain authoritative + for graph compilation and execution. +- `embodichain.gen_sim.task_engine.orchestration` owns cross-engine contracts, + read-only source references, scene adaptation, orchestration, and artifacts. + `embodichain.gen_sim.task_engine.cli` owns the unified CLI. + +The former `scene_bridge`, `collaboration`, and +`action_engine.collaboration` namespaces were removed; there are no import +bridges or fallback entry points. + +## Data Flow + +1. `TaskAgent` or a caller creates a validated `TaskSpec`. +2. Action Engine emits `SceneRequirements` for the external Scene Engine. +3. After scene generation, Task Engine's scene adapter preserves geometry, physics, + articulation, affordance evidence, and provenance in a versioned + `StaticSceneManifest` while the existing redacted manifest remains compatible. +4. `FeasibilityBroker` intersects the selected task, role bindings, static scene, + robot profile, and executable capability catalog without repairing unknowns. +5. Offline recipes and the online planner independently create complete + `SeedGraph` candidates whose nodes are `AtomicAction` calls. +6. Runtime preflight checks the capability catalog and rejects planning-only + actions before simulator motion starts. +7. `ActionGrounder` reads live robot, object, articulation, and camera state and + materializes the typed goal and immutable action options just in time. +8. `ProgramExecutor` schedules the DAG, executes vectorized action masks, and + verifies semantic postconditions from live state. + +There is no persisted semantic task graph between `TaskSpec` and `SeedGraph`. +The standalone TaskAgent v1 planner and compiler remain available to their +existing callers, but the generation pipeline neither accepts nor publishes +TaskAgent v1 JSON. + +## Protocols + +### TaskSpec + +`TaskSpec` owns the level, public instruction, E1-E9 task instances, +dependencies, path-independent success conditions, and a private oracle. The +online planner receives `public_task_spec(...)`, which removes the oracle and, +for L4, the hidden reference task instances. Online L4 TaskGroups are inferred +from the instruction and observations rather than matched to an oracle path. + +Levels classify how the task is specified, not action count: + +- L1: one E instance. +- L2: two or more instances of the same E type. +- L3: two or more different E types explicitly composed. +- L4: an abstract instruction that requires memory, visual semantics, pattern, + logic, common-sense, or constraint reasoning. + +Free-language L1-L3 generation uses two structured model calls. The first sees +only the instruction and E1-E9 catalog and emits typed steps whose scene +selectors are open natural-language references. The second sees those +references plus a coordinate-free semantic inventory and may return only +existing scene UIDs, status, and confidence. Local validation enforces complete +request coverage, candidate roles, cardinality, confidence, and non-self +targets; unresolved or ambiguous references fail instead of being guessed. +Each structured model stage gets at most one bounded repair attempt after an +invalid response. If repair still fails, or the model call itself fails, +generation stops before recipe expansion and artifact publication. It never +switches to keyword or rule-based instruction parsing. + +The older public `planning.plan_task` adapter still produces a standalone +`TaskAgent` from structured LLM output, but it does not reinterpret the +instruction after that output exists. Axis, orientation, and arm-allocation +fields come only from the validated model result. It has no keyword fallback. + +Generator callers without a configured LLM must provide a validated `TaskSpec` +with explicit role bindings (or a matching `SceneRequirements` sidecar). This +path is fully offline and never imports or calls an LLM client. + +### SceneRequirements + +`SceneRequirements` is the JSON hand-off to the external Scene Engine. It +declares object roles, counts, categories, affordances, initial states, spatial +constraints, camera requirements, and distractors. Scene results are never +silently repaired. Structural contradictions and explicit affordance +contradictions invalidate the task instance; an absent affordance declaration +remains unknown and is deferred to runtime physical validation. + +The current tabletop importer has one deliberately narrow structural contract: +exactly one `background` object is the support surface and receives runtime UID +`table`; every movable object is assumed to begin on that surface. Zero or +multiple backgrounds are rejected rather than resolved from position, UID, or +description text. Semantic `category`, `color`, and `attributes` come only from +their explicit scene fields. Physics `attrs` are not semantic metadata. + +For task-first inputs, explicit role bindings are authoritative unless they +contradict metadata that the scene actually declares. Automatic role binding +requires a unique match with complete structured category, attribute, state, +and affordance evidence. Object names and descriptions remain available to the +LLM grounding call, but deterministic validation never searches them for +semantic substrings. + +### StaticSceneManifest And FeasibilityReport + +`StaticSceneManifest` is an additive Scene Bridge artifact. It keeps the legacy +scene manifest stable while recording initial poses, geometry hashes, physics, +articulation payloads, structured affordance evidence, and provenance. Legacy +affordance strings become `declared` evidence; only structural facts derived by +the adapter are marked `verified`. + +`FeasibilityReport` classifies each check as `proven`, `runtime_probe`, `unknown`, +or `contradicted`. Missing evidence remains unknown. Declared geometric +affordances require a runtime probe, and unavailable AtomicActions are explicit +contradictions. A contradicted report publishes an `infeasible` audit result and +does not invoke graph or bundle generation. Executable preflight remains a +second authoritative gate before bundle generation. + +Scene Bridge reports arm-layout and whole-task pickup, handover, +target-interaction, and safety-clearance phases as `runtime_probe` evidence. +Arm-side compatibility is not claimed without live arm-base poses and workspace +geometry. + +### SeedGraph + +Every node directly names an `atomic_action`, scene `object_uid`, symbolic +`target_binding`, actor, control, dependencies, resources, pre/postconditions, +motion policy, E type, `task_instance_id`, and an Action Contract v2 +`failure_policy`. `task_required` and `safety_required` failures invalidate the +candidate; `best_effort` failures remain observable but do not erase an already +verified task and safety result. `TaskGroup` groups all nodes of one E instance +with `role=primary|recovery`; it is metadata over the same DAG, not a second +graph. + +Validation guarantees: + +- node and TaskGroup dependencies are DAGs; +- every node belongs to exactly one TaskGroup; +- E groups contain their required core actions; +- concurrent nodes do not claim the same exclusive arm/object resource; +- object references resolve to scene UIDs; +- world poses, qpos, trajectories, grasp poses, and waypoints are rejected + recursively; +- hashes use canonical strict JSON and are stable across processes. + +The production loader accepts v3 graphs only. An older graph, whether supplied as +JSON or an in-memory mapping, receives an explicit regeneration error rather +than an implicit migration. + +## Capability Boundary + +`AtomicCapabilityRegistry` is the single runtime catalog. A descriptor declares +the action/option types, accepted symbolic bindings and controls, resource mode, +held-object state effect, target and config materializers, verifier, failure +classifier, retry mode, and runtime availability. + +The executable catalog currently contains: + +- `PickUp`, `MoveHeldObject`, `MoveEndEffector`, `MoveJoints`, and `Place` +- `Press` +- `CoordinatedPickment` and `CoordinatedPlacement` +- `HandOver` + +`Pour`, `PullArticulatedPart`, `PushArticulatedPart`, and `TurnKnob` are +planning-only until matching lower-level implementations exist. They can be +generated and statically checked, but preflight fails before any motion with +the descriptor's unavailable reason. + +Adding an executable skill consists of registering its descriptor and reusable +materializer/verifier hooks plus focused tests. Planner and executor dispatch +do not maintain a parallel action-class table. + +## Offline And Online Planning + +Offline recipes deterministically instantiate E1-E9 task instances. Current +task mappings are: + +- `place_relative -> E1` +- `orient_object -> E2` +- `coordinated_transport -> E5` +- every member of `build_stack` and `arrange_line` -> one E1 instance + +E5 uses `coordinated_transport` only as the semantic task-group operator. Its +motion graph contains one `CoordinatedPickment`; a `place` terminal behavior +adds synchronized left/right `MoveJoints(gripper_open)` nodes. The executor +clears coordinated hold state only after both grippers are observed open. + +The online path first extracts auditable visual facts from multi-view RGB and, +when available, depth and camera calibration. Facts contain only known UIDs, +normalized bboxes/keypoints, canonical spatial relations, task predicates, and +confidence. Spatial relations use a shared ontology and fixed participant +order. Task-level judgments such as visual or pattern completion are accepted +only when the current `TaskSpec.success` explicitly requests them. A second +structured call produces a complete direct `AtomicAction` graph. Prompts +request facts and graph JSON only; hidden chain-of-thought is neither requested +nor stored. + +Image-space constraints may use normalized keypoints, masks, bboxes, and +relative relations. The Grounder uses live depth and camera calibration to +convert them to world targets. The SeedGraph never stores that result. + +## JIT Grounding + +Each action is grounded again immediately before planning/execution. Grounding +therefore observes object displacement, current qpos, current held-object +ownership, articulation state, and fresh camera measurements. Coordinated and +handover actions are grounded as synchronized execution units. Automatic arm +selection, collision checks, live arrangement slots, and current predicate +semantics remain deterministic runtime responsibilities. + +Arm allocation uses the live right-to-left arm-base axis and the live table +center. The preference therefore follows translated and rotated robot +workspaces; live motion planning remains authoritative for reachability. + +Placement support is a relation, not an entity category. Static adaptation +accepts rigid `physical_entity` targets without requiring a `support_surface` +affordance. Articulations require a link-level runtime target interface and are +rejected until that interface is available. Runtime evaluates +`object_supported_by(payload, support, pose)` from live geometry and center of +mass, applies the requested `orientation_goal`, and requires low motion across a +bounded stability window. Successful relations form a per-environment support +graph that is checked for cycles and revalidated at task completion. + +Orientation is compiled into hard `align_axis` or `match_rotation` terms plus a +separate minimum-rotation planning preference. An omitted orientation request +adds no hard acceptance term; `preserve` remains an explicit full-rotation +contract for persisted bundles, while `upright` constrains only the requested +local axis and declares whether that axis is directed. Grounding and runtime +verification consume the same compiled contract so reachability search cannot +silently relax a required terminal orientation. With no hard term, a live +upright state may still select upright-preserving yaw candidates as a planning +preference; this follows current state and automatically stops after that state +is invalidated rather than becoming a sticky success requirement. + +Grasp generation keeps support-plane collision filtering as its strict first +pass. If diagnostics show that this heuristic alone exhausted otherwise +object-collision-free candidates, Action Engine retries without the heuristic; +the relaxed candidates still pass through the live robot and scene collision +planner before execution. This avoids treating a local support-plane proxy as +a proof of scene-level infeasibility, including for objects already held above +the support surface. + +Grounding samples bounded support-relative placement poses. Planning failures +try the next pose before release; instability after release requires a fresh +grasp and an unused pose. The recovery keeps the original actor contract, and +its edges and failure provenance are recorded separately from the primary +attempt. + +## Mainline Planning Contract + +The runtime keeps only an Action Engine-local `ExecutionState` for full-robot +qpos and held-object relations. Each plan converts that state to the mainline +`PlanningContext` (`RobotObservation`, `TaskState`, and `SceneSnapshot`) and +submits an `ActionInvocation` to `AtomicActionEngine`. The returned +`StateDelta` remains speculative until physical and semantic verification; only +verified vectorized rows are committed. + +`AtomicActionAdapter` accepts a shared `SceneProvider` and otherwise builds a +`RigidObjectSceneProvider` from live simulation entities. Planning snapshots now +carry monotonic timestamps and material-change scene/collision revisions. The +adapter also exposes `start_session(...)` for callers adopting +`ExecutionSession`; the existing compound and per-arm merged trajectory +scheduler remains as the compatibility execution path. + +Single-arm arm motion uses cuRobo `motion_gen` by default. Hand-only and +coordinated dual-arm actions use `ik_interp`, because mainline coordinated +primitives do not support cuRobo motion generation. A failed single-arm cuRobo +row may fall back to `ik_interp` without replacing successful rows. Generated +background objects form the static cuRobo collision world; dynamic obstacles +are an explicit runtime-policy opt-in. + +Generated mesh objects carry V-HACD settings in both the current shape-level +schema and legacy top-level fields. Before antipodal grasp construction, the +runtime prepares a checksummed V-HACD payload at the shared collision-checker +cache path so the unchanged mainline checker does not silently recompute CoACD. +Grasp generation samples multiple deviated approach directions and filters them +through the existing gripper collision model. Safety retreat planning searches +a bounded set of live height and baseward targets instead of treating one exact +height as a geometric reachability certificate. + +## A/B Evaluation + +Test mode retains both candidates. `run_strict_ab` creates distinct offline and +online environments with the same task, scene configuration, seed, Grounder, +verifiers, and retry policy. Both environments reset before execution and a +digest over robot qpos and object state must match exactly; a mismatch aborts +before either branch executes. + +Artifacts are written under `offline/` and `online/`, with a shared +`comparison.json`. The comparison records graph hashes/differences, action and +path lengths, success, retries, recoveries, revisions, latency, record paths, +and planner/VLM metadata supplied by each candidate. + +L4 A/B runs must supply a private-oracle evaluator. The built-in evaluator +checks memory reconstruction, visual completion, pattern completion, numeric +selection, functional placement, and stable/unobstructed goals from the final +state only. The comparison labels whether success came from runtime step +postconditions or the private oracle. + +## Dynamic Recovery + +The persisted `SeedGraph` is immutable. `RuntimeGraph` keeps a detached working +copy and an ordered revision log. One failed `AtomicAction` can be freshly +grounded and retried twice, for three total attempts, and only while its live +precondition remains true. + +Failures use the bounded taxonomy `search_exhausted`, `plan_failed`, +`grasp_missed`, `object_fallen`, `object_dropped`, and +`postcondition_failed`. `search_exhausted` records the blocking edge, planning +stage, strategy, finite budget, and observed evidence; it does not claim that a +target is geometrically unreachable. Known recoverable states can insert a +complete `role=recovery` TaskGroup, such as an E2 upright group. Recovery keeps +the failed TaskGroup's actor contract, and primary, recovery, and replay events +are recorded separately. After recovery, the selected route replans only the +unfinished suffix. Offline and online dynamic replanners are explicit, +separate modes. Revision, recovery-action, transition, and retry budgets bound +every loop. + +## Selection And Fusion + +Product mode statically scores offline and online candidates using schema +validity, capability availability, UID validity, task coverage, visual +confidence, and estimated action cost. Exact mature-template matches favor the +offline route; L4 visual tasks favor sufficiently confident online results. + +Fusion is conservative. It may choose only complete `TaskGroup` units, rewires +dependencies at group boundaries, and rejects unordered state changes to the +same object. It never splits one E instance across candidates. + +## Artifacts + +A normal generated bundle contains: + +- `task_spec.json` +- `scene_requirements.json` +- `seed_task_graph.json` +- `seed_task_graph.png` +- `agent_config.json` +- `fast_gym_config.json` + +Strict A/B adds branch-local graph/result artifacts and `comparison.json`. +Review graphs, runtime records, and videos never become execution inputs. + +`prepare` lowers and preflights resolved semantic candidates in selection order. +A candidate-local lowering, symbolic planning, or preflight error rejects only +that candidate. If no resolved candidate is executable, the transaction +publishes `preparation_failure.json` with each attempted draft, verified +bindings, available grounded plan, failure stage, and exception instead of +leaving an older successful bundle in place. + +## Invariants + +- SeedGraph nodes are direct AtomicActions, not E-level operators. +- Lowering uses original instruction-step order as the stable tie-break among + dependency-ready steps. Independent steps remain independent; the contract + linker serializes only actual resource conflicts. +- E labels are subgraph grouping semantics only. +- Planning artifacts contain no grounded motion coordinates. +- Online planning never receives the private oracle. +- Runtime uses one capability registry for preflight, Grounding, config + construction, execution, verification policy, and recovery policy. +- Required arms are never silently replaced. +- Failed or inactive vectorized rows preserve their last valid state. +- Current five task families preserve their v1 AtomicAction topology and live + Grounding behavior after regeneration. diff --git a/embodichain/gen_sim/action_engine/__init__.py b/embodichain/gen_sim/action_engine/__init__.py new file mode 100644 index 000000000..cd941d35b --- /dev/null +++ b/embodichain/gen_sim/action_engine/__init__.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Capability-driven planning and live execution for generated simulations. + +Action Engine deliberately exposes a small public surface. Natural-language +goals become a typed TaskSpec, deterministic planning lowers them into a +coordinate-free SeedGraph, and the runtime grounds that graph only against live +simulator state. +""" + +from __future__ import annotations + +from .unbound import ( + UNBOUND_ACTION_PLAN_SCHEMA, + UnboundActionPlan, + build_unbound_action_plan, + validate_unbound_action_plan, +) + +from .protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + ACTION_ENGINE_ENV_ID, + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, +) + +__all__ = [ + "ACTION_ENGINE_CONFIG_SCHEMA", + "ACTION_ENGINE_ENV_ID", + "EXECUTION_PROGRAM_SCHEMA", + "TASK_AGENT_SCHEMA", + "UNBOUND_ACTION_PLAN_SCHEMA", + "UnboundActionPlan", + "build_unbound_action_plan", + "validate_unbound_action_plan", +] diff --git a/embodichain/gen_sim/action_engine/agent.py b/embodichain/gen_sim/action_engine/agent.py new file mode 100644 index 000000000..55def0bec --- /dev/null +++ b/embodichain/gen_sim/action_engine/agent.py @@ -0,0 +1,762 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Grounded-plan compilation and compact execution reporting.""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from dataclasses import asdict, is_dataclass +from datetime import datetime, timezone +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Callable, TypeAlias + +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + seed_graph_hash, + validate_seed_graph, +) +from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionProgram, + ExecutionReport, + ExecutionResult, + ProgramExecutor, + build_execution_provenance, + load_execution_program, + validate_execution_report, + write_execution_report, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph +from embodichain.gen_sim.action_engine.unbound import ( + ActionCapabilityError, + UnboundActionPlan, + build_unbound_action_plan, + validate_unbound_action_plan, +) + +__all__ = ["ActionAgent", "ActionGraph"] + +ExecutorFactory = Callable[..., ProgramExecutor] +ActionGraph: TypeAlias = dict[str, Any] + + +class ActionAgent: + """Compile, preflight, execute, and report one grounded task plan.""" + + def __init__( + self, + *, + registry: AtomicCapabilityRegistry | None = None, + executor_factory: ExecutorFactory = ProgramExecutor, + ) -> None: + self.registry = registry or build_atomic_capability_registry() + self.executor_factory = executor_factory + + def plan(self, grounded_plan: Mapping[str, Any]) -> ActionGraph: + """Compile a validated GroundedTaskPlan to the public SeedGraph v3.""" + plan = _validate_grounded_plan(grounded_plan) + task_spec = _mapping(plan.get("task_spec"), "GroundedTaskPlan.task_spec") + bindings = _role_binding_map(plan.get("role_bindings")) + graph = instantiate_seed_graph( + task_spec, + bindings, + registry=self.registry, + ) + known_uids = _known_uids( + plan.get("scene_manifest"), + bindings=bindings, + ) + known_uids.add("table") + graph = validate_seed_graph( + graph, + known_objects=known_uids or None, + known_actions=self.registry.names(), + executable_actions=self.registry.executable_names(), + require_executable=False, + ) + validate_persisted_contracts(graph, self.registry) + return graph + + def draft(self, candidate: Mapping[str, Any]) -> UnboundActionPlan: + """Create an Action-owned draft before final scene UID binding. + + Args: + candidate: One validated Task Engine candidate. + + Returns: + A scene-independent action plan whose selectors contain no UIDs. + """ + plan = build_unbound_action_plan(candidate) + names = getattr(self.registry, "names", None) + executable_names = getattr(self.registry, "executable_names", None) + if callable(names): + missing = sorted(set(plan["required_actions"]) - set(names())) + if missing: + raise ActionCapabilityError( + "Required AtomicAction is not registered: " + ", ".join(missing) + ) + if callable(executable_names): + unavailable = sorted( + set(plan["required_actions"]) - set(executable_names()) + ) + if unavailable: + raise ActionCapabilityError( + "Required AtomicAction is not executable: " + ", ".join(unavailable) + ) + return plan + + def bind_and_plan( + self, + unbound_plan: Mapping[str, Any], + grounded_plan: Mapping[str, Any], + ) -> ActionGraph: + """Bind one audited unbound plan through a final GroundedTaskPlan. + + The grounded task draft must reproduce the exact unbound IR. This + prevents the final planner from silently reinterpreting a candidate + after Scene Engine work has run concurrently. + """ + unbound = validate_unbound_action_plan(unbound_plan) + grounded = _validate_grounded_plan(grounded_plan) + expected = build_unbound_action_plan( + { + "candidate_id": grounded["selected_candidate_id"], + "draft": grounded["task_draft"], + } + ) + if unbound != expected: + raise ValueError( + "UnboundActionPlan does not match the final GroundedTaskPlan." + ) + return self.plan(grounded) + + def preflight( + self, + action_graph: Mapping[str, Any] | str | Path, + *, + scene_manifest: Mapping[str, Any] | None = None, + known_uids: Collection[str] | None = None, + ) -> ExecutionProgram: + """Reject invalid and planning-only graphs before simulator motion.""" + known = set(str(uid) for uid in (known_uids or ()) if str(uid)) + known.update(_known_uids(scene_manifest)) + if isinstance(action_graph, Mapping): + metadata = action_graph.get("metadata", {}) + if isinstance(metadata, Mapping): + bindings = metadata.get("role_bindings", {}) + if isinstance(bindings, Mapping): + known.update(str(uid) for uid in bindings.values() if str(uid)) + if known: + known.add("table") + return load_execution_program( + action_graph, + known_objects=known or None, + registry=self.registry, + require_executable=True, + ) + + def execute( + self, + action_graph: Mapping[str, Any] | str | Path, + env: Any, + *, + grounded_plan: Mapping[str, Any] | None = None, + scene_manifest: Mapping[str, Any] | None = None, + known_uids: Collection[str] | None = None, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + executor_kwargs: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Preflight and execute a graph, converting all outcomes to a report.""" + task_id = _task_id(grounded_plan, action_graph) + plan_hash = _plan_hash(grounded_plan) + graph_hash = _action_graph_hash(action_graph) + effective_run_id = run_id or _new_run_id() + provenance = build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + effective_manifest = scene_manifest + if effective_manifest is None and grounded_plan is not None: + value = grounded_plan.get("scene_manifest") + if isinstance(value, Mapping): + effective_manifest = value + + try: + program = self.preflight( + action_graph, + scene_manifest=effective_manifest, + known_uids=known_uids, + ) + except (TypeError, ValueError, OSError) as exc: + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="rejected", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error=_error_message(exc), + ) + + kwargs = dict(executor_kwargs or {}) + kwargs.setdefault("capability_registry", self.registry) + try: + executor = self.executor_factory(program, env, **kwargs) + result = executor.run( + run_id=effective_run_id, + episode_index=episode_index, + ) + except Exception as exc: + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="aborted", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error=_error_message(exc), + ) + if not isinstance(result, ExecutionResult): + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="aborted", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error="TypeError: ProgramExecutor.run must return ExecutionResult.", + ) + try: + return self.report_execution_result( + result, + action_graph=action_graph, + grounded_plan=grounded_plan, + run_id=effective_run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + except (TypeError, ValueError, OverflowError) as exc: + return self._empty_report( + env, + task_id=task_id, + plan_hash=plan_hash, + graph_hash=graph_hash, + status="aborted", + run_id=effective_run_id, + episode_index=episode_index, + provenance=provenance, + error=_error_message(exc), + ) + + def run( + self, + grounded_plan: Mapping[str, Any], + env: Any, + *, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + executor_kwargs: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Compile and execute one GroundedTaskPlan through the full pipeline.""" + effective_run_id = run_id or _new_run_id() + try: + graph = self.plan(grounded_plan) + except (TypeError, ValueError, OSError) as exc: + return self._empty_report( + env, + task_id=_task_id(grounded_plan, {}), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_document_hash({}), + status="rejected", + run_id=effective_run_id, + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + error=_error_message(exc), + ) + return self.execute( + graph, + env, + grounded_plan=grounded_plan, + run_id=effective_run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + executor_kwargs=executor_kwargs, + ) + + def report_execution_result( + self, + result: ExecutionResult, + *, + action_graph: Mapping[str, Any] | str | Path, + grounded_plan: Mapping[str, Any] | None = None, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Convert a result already executed by the legacy runner to a report.""" + if not isinstance(result, ExecutionResult): + raise TypeError("result must be an ExecutionResult.") + return self._result_report( + result, + task_id=_task_id(grounded_plan, action_graph), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_action_graph_hash(action_graph), + run_id=run_id or _new_run_id(), + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + ) + + def rejection_report( + self, + action_graph: Mapping[str, Any] | str | Path, + error: BaseException | str, + *, + grounded_plan: Mapping[str, Any] | None = None, + environment_count: int = 1, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Build a zero-action report for a preflight rejection.""" + message = error if isinstance(error, str) else _error_message(error) + return self._empty_report( + SimpleNamespace(num_envs=max(1, int(environment_count))), + task_id=_task_id(grounded_plan, action_graph), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_action_graph_hash(action_graph), + status="rejected", + run_id=run_id or _new_run_id(), + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + error=str(message), + ) + + def abortion_report( + self, + action_graph: Mapping[str, Any] | str | Path, + error: BaseException | str, + *, + grounded_plan: Mapping[str, Any] | None = None, + environment_count: int = 1, + run_id: str | None = None, + episode_index: int = 0, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, + ) -> ExecutionReport: + """Build a zero-action report for an unexpected runtime exception.""" + message = error if isinstance(error, str) else _error_message(error) + return self._empty_report( + SimpleNamespace(num_envs=max(1, int(environment_count))), + task_id=_task_id(grounded_plan, action_graph), + plan_hash=_plan_hash(grounded_plan), + graph_hash=_action_graph_hash(action_graph), + status="aborted", + run_id=run_id or _new_run_id(), + episode_index=episode_index, + provenance=build_execution_provenance( + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ), + error=str(message), + ) + + def _result_report( + self, + result: ExecutionResult, + *, + task_id: str, + plan_hash: str, + graph_hash: str, + run_id: str, + episode_index: int, + provenance: Mapping[str, Any], + ) -> ExecutionReport: + _persist_executed_trajectory(result) + success = _bool_vector(result.success) + semantics = { + str(step_id): _bool_vector(mask) + for step_id, mask in result.semantic_success.items() + } + failures = tuple(_json_safe(item) for item in result.failure_events) + revisions = tuple(_json_safe(item) for item in result.runtime_revisions) + action_count = len(result.actions) + environments = tuple( + { + "env_id": str(env_id), + "success": value, + "semantic_success": { + step_id: values[env_id] + for step_id, values in semantics.items() + if env_id < len(values) + }, + "action_count": action_count, + "retry_count": _retry_count_for_env(result, env_id), + "recovery_count": _revision_count_for_env( + revisions, env_id, kind="insert_recovery" + ), + "revision_count": _revision_count_for_env(revisions, env_id), + "failures": _events_for_env(failures, env_id), + } + for env_id, value in enumerate(success) + ) + report = ExecutionReport( + task_id=task_id, + plan_hash=plan_hash, + action_graph_hash=graph_hash, + status="succeeded" if all(success) else "failed", + run_id=run_id, + episode_id=str(episode_index), + provenance=deepcopy(dict(provenance)), + environments=environments, + action_count=action_count, + retry_count=int(result.retry_count), + recovery_count=int(result.recovery_count), + revision_count=int(result.revision_count), + failure_events=failures, + graph_revisions=revisions, + record_dir=result.record_dir, + error=None, + ) + validated = _validated_report(report) + _publish_execution_report(validated) + return validated + + def _empty_report( + self, + env: Any, + *, + task_id: str, + plan_hash: str, + graph_hash: str, + status: str, + run_id: str, + episode_index: int, + provenance: Mapping[str, Any], + error: str, + ) -> ExecutionReport: + count = _environment_count(env) + report = ExecutionReport( + task_id=task_id, + plan_hash=plan_hash, + action_graph_hash=graph_hash, + status=status, + run_id=run_id, + episode_id=str(episode_index), + provenance=deepcopy(dict(provenance)), + environments=tuple( + { + "env_id": str(env_id), + "success": False, + "semantic_success": {}, + "action_count": 0, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + } + for env_id in range(count) + ), + error=error, + ) + return _validated_report(report) + + +def _validate_grounded_plan(value: Mapping[str, Any]) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError("GroundedTaskPlan must be a mapping.") + # GroundedTaskPlan is a cross-engine protocol owned by Task Engine. + # Import lazily so Action Engine remains importable without initializing + # the coordinator or Scene Adapter. + try: + from embodichain.gen_sim.task_engine.orchestration.contracts import ( + validate_grounded_task_plan, + ) + except (ImportError, AttributeError): + return deepcopy(dict(value)) + return validate_grounded_task_plan(value) + + +def _validated_report(report: ExecutionReport) -> ExecutionReport: + payload = report.as_mapping() + validate_execution_report(payload) + return report + + +def _persist_executed_trajectory(result: ExecutionResult) -> None: + """Persist every emitted control tensor beside the runtime graph audit.""" + if not result.record_dir: + return + root = Path(result.record_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + actions = [action.detach().cpu() for action in result.actions] + temporary = root / ".executed_trajectory.pt.tmp" + destination = root / "executed_trajectory.pt" + torch.save({"actions": actions}, temporary) + temporary.replace(destination) + manifest = { + "schema_version": "action_engine_executed_trajectory/v1", + "path": destination.name, + "action_count": len(actions), + "actions": [ + { + "index": index, + "shape": list(action.shape), + "dtype": str(action.dtype), + } + for index, action in enumerate(actions) + ], + } + manifest_path = root / "executed_trajectory.json" + manifest_path.write_text( + json.dumps(manifest, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def _mapping(value: Any, label: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{label} must be a mapping.") + return deepcopy(dict(value)) + + +def _role_binding_map(value: Any) -> dict[str, str]: + source = _mapping(value, "GroundedTaskPlan.role_bindings") + nested = source.get("role_bindings") + if isinstance(nested, Mapping): + source = dict(nested) + result = {str(role): str(uid) for role, uid in source.items()} + if not result or any(not role or not uid for role, uid in result.items()): + raise ValueError("GroundedTaskPlan role bindings must not be empty.") + return result + + +def _known_uids( + manifest: Any, + *, + bindings: Mapping[str, str] | None = None, +) -> set[str]: + result = {str(uid) for uid in (bindings or {}).values() if str(uid)} + if not isinstance(manifest, Mapping): + return result + objects = manifest.get("objects", ()) + if isinstance(objects, Sequence) and not isinstance( + objects, (str, bytes, bytearray) + ): + for item in objects: + if isinstance(item, Mapping): + uid = item.get("uid", item.get("runtime_uid")) + if isinstance(uid, str) and uid: + result.add(uid) + return result + + +def _task_id( + plan: Mapping[str, Any] | None, + graph: Mapping[str, Any] | str | Path, +) -> str: + if isinstance(plan, Mapping): + value = plan.get("task_id") + if isinstance(value, str) and value: + return value + if isinstance(graph, Mapping): + value = graph.get("task_id") + if isinstance(value, str) and value: + return value + return "unknown_task" + + +def _plan_hash(plan: Mapping[str, Any] | None) -> str: + if isinstance(plan, Mapping): + hashes = plan.get("hashes", {}) + if isinstance(hashes, Mapping): + value = hashes.get("plan") + if isinstance(value, str) and value: + return value + return _safe_document_hash(plan) + return _document_hash({}) + + +def _action_graph_hash(value: Mapping[str, Any] | str | Path) -> str: + if isinstance(value, Mapping): + try: + return seed_graph_hash(value) + except (TypeError, ValueError): + return _safe_document_hash(value) + path = Path(value).expanduser() + try: + loaded = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return hashlib.sha256(str(path).encode("utf-8")).hexdigest() + return ( + _action_graph_hash(loaded) + if isinstance(loaded, Mapping) + else _safe_document_hash(loaded) + ) + + +def _safe_document_hash(value: Any) -> str: + try: + return _document_hash(value) + except (TypeError, ValueError, OverflowError): + return hashlib.sha256(repr(value).encode("utf-8")).hexdigest() + + +def _document_hash(value: Any) -> str: + payload = json.dumps( + _json_safe(value), + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _bool_vector(value: Any) -> list[bool]: + if isinstance(value, torch.Tensor): + return [bool(item) for item in value.detach().cpu().reshape(-1).tolist()] + if isinstance(value, np.ndarray): + return [bool(item) for item in value.reshape(-1).tolist()] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + return [bool(item) for item in value] + return [bool(value)] + + +def _events_for_env( + events: Sequence[Mapping[str, Any]], env_id: int +) -> list[dict[str, Any]]: + result = [] + for event in events: + env_ids = event.get("env_ids") + if isinstance(env_ids, Sequence) and not isinstance( + env_ids, (str, bytes, bytearray) + ): + if env_id not in env_ids: + continue + item = deepcopy(dict(event)) + item["env_ids"] = [env_id] + result.append(item) + else: + result.append(deepcopy(dict(event))) + return result + + +def _revision_count_for_env( + revisions: Sequence[Mapping[str, Any]], + env_id: int, + *, + kind: str | None = None, +) -> int: + count = 0 + for revision in revisions: + if kind is not None and revision.get("kind") != kind: + continue + active = revision.get("active_env_ids") + if ( + isinstance(active, Sequence) + and not isinstance(active, (str, bytes, bytearray)) + and env_id not in active + ): + continue + count += 1 + return count + + +def _retry_count_for_env(result: ExecutionResult, env_id: int) -> int: + counts = result.retry_counts + if env_id < len(counts): + return int(counts[env_id]) + return int(result.retry_count) + + +def _environment_count(env: Any) -> int: + value = getattr(env, "num_envs", 1) + try: + return max(1, int(value)) + except (TypeError, ValueError): + return 1 + + +def _json_safe(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Path): + return value.as_posix() + if is_dataclass(value): + return _json_safe(asdict(value)) + if isinstance(value, Mapping): + return {str(key): _json_safe(item) for key, item in value.items()} + if isinstance(value, (list, tuple, set, frozenset)): + return [_json_safe(item) for item in value] + if value is None or isinstance(value, (str, int, float, bool)): + return value + return str(value) + + +def _new_run_id() -> str: + return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + + +def _error_message(exc: BaseException) -> str: + return f"{type(exc).__name__}: {exc}" + + +def _publish_execution_report(report: ExecutionReport) -> None: + """Atomically publish the compact report beside runtime episode records.""" + if not report.record_dir: + return + write_execution_report(report.record_dir, report) diff --git a/embodichain/gen_sim/action_engine/capabilities/__init__.py b/embodichain/gen_sim/action_engine/capabilities/__init__.py new file mode 100644 index 000000000..5e7b6935e --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/__init__.py @@ -0,0 +1,57 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Semantic operator and atomic-action capability registry.""" + +from __future__ import annotations + +from .atomic import ( + ACTION_CONTRACT_VERSION, + AtomicCapability, + AtomicCapabilityRegistry, + ResolvedActionContract, + ResourceClaim, + StateAtom, + StateEffect, + build_atomic_capability_registry, + capability_precondition, +) +from .builtins import build_default_registry +from .registry import ( + ActionCapability, + ActionTemplate, + CapabilityRegistry, + OperatorCapability, + PhaseTemplate, +) + +__all__ = [ + "ACTION_CONTRACT_VERSION", + "ActionCapability", + "ActionTemplate", + "AtomicCapability", + "AtomicCapabilityRegistry", + "CapabilityRegistry", + "OperatorCapability", + "PhaseTemplate", + "ResolvedActionContract", + "ResourceClaim", + "StateAtom", + "StateEffect", + "build_atomic_capability_registry", + "build_default_registry", + "capability_precondition", +] diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py new file mode 100644 index 000000000..4731ed7e5 --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -0,0 +1,1040 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Single-source AtomicAction capability descriptors.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +import hashlib +import json +from typing import Any + +import torch + +__all__ = [ + "ACTION_CONTRACT_VERSION", + "AtomicCapability", + "AtomicCapabilityRegistry", + "ResolvedActionContract", + "ResourceClaim", + "StateAtom", + "StateEffect", + "build_atomic_capability_registry", + "capability_precondition", +] + +_RETRY_MODES = frozenset({"direct", "recover_then_retry", "non_retryable"}) +ACTION_CONTRACT_VERSION = "action_contract_v2" +_PREDICATES = frozenset( + { + "arm_free", + "object_free", + "object_held", + "object_coordinated_held", + "handover_complete", + "arm_clear", + "arm_home", + } +) +_EFFECT_OPERATIONS = frozenset({"add", "delete"}) +_RESOURCE_ACCESS = frozenset({"shared_read", "exclusive"}) +_RESOURCE_LIFETIMES = frozenset({"action", "until_release"}) +_COMPLETION_MODES = frozenset({"ordinary", "cleanup", "terminal_barrier"}) +_FAILURE_POLICIES = frozenset({"task_required", "safety_required", "best_effort"}) + + +@dataclass(frozen=True) +class StateAtom: + """One symbolic state fact used by an Action Contract.""" + + predicate: str + object_uid: str | None = None + arm: str | None = None + + def __post_init__(self) -> None: + if self.predicate not in _PREDICATES: + raise ValueError(f"Unknown Action Contract predicate {self.predicate!r}.") + if self.object_uid is not None and not self.object_uid: + raise ValueError("StateAtom.object_uid must not be empty.") + if self.arm is not None and not self.arm: + raise ValueError("StateAtom.arm must not be empty.") + + def as_mapping(self) -> dict[str, str]: + """Return the stable JSON representation of this fact.""" + result = {"predicate": self.predicate} + if self.object_uid is not None: + result["object_uid"] = self.object_uid + if self.arm is not None: + result["arm"] = self.arm + return result + + +@dataclass(frozen=True) +class StateEffect: + """Add or delete one symbolic state fact.""" + + op: str + atom: StateAtom + + def __post_init__(self) -> None: + if self.op not in _EFFECT_OPERATIONS: + raise ValueError(f"Unknown Action Contract effect operation {self.op!r}.") + + def as_mapping(self) -> dict[str, Any]: + """Return the stable JSON representation of this effect.""" + return {"op": self.op, "atom": self.atom.as_mapping()} + + +@dataclass(frozen=True) +class ResourceClaim: + """One resource access claim made by an AtomicAction.""" + + resource: str + access: str = "exclusive" + lifetime: str = "action" + + def __post_init__(self) -> None: + if not self.resource: + raise ValueError("ResourceClaim.resource must not be empty.") + if self.access not in _RESOURCE_ACCESS: + raise ValueError(f"Unknown resource access mode {self.access!r}.") + if self.lifetime not in _RESOURCE_LIFETIMES: + raise ValueError(f"Unknown resource lifetime {self.lifetime!r}.") + + def as_mapping(self) -> dict[str, str]: + """Return the stable JSON representation of this claim.""" + return { + "resource": self.resource, + "access": self.access, + "lifetime": self.lifetime, + } + + +@dataclass(frozen=True) +class ResolvedActionContract: + """Fully resolved, serializable contract for one action node.""" + + requires: tuple[StateAtom, ...] = () + effects: tuple[StateEffect, ...] = () + claims: tuple[ResourceClaim, ...] = () + completion: str = "ordinary" + failure_policy: str = "task_required" + version: str = ACTION_CONTRACT_VERSION + + def __post_init__(self) -> None: + if self.version != ACTION_CONTRACT_VERSION: + raise ValueError( + f"Unsupported Action Contract version {self.version!r}; " + f"expected {ACTION_CONTRACT_VERSION!r}." + ) + if self.completion not in _COMPLETION_MODES: + raise ValueError(f"Unknown Action Contract completion {self.completion!r}.") + if self.failure_policy not in _FAILURE_POLICIES: + raise ValueError( + f"Unknown Action Contract failure policy {self.failure_policy!r}." + ) + + def as_mapping(self) -> dict[str, Any]: + """Return the stable JSON representation persisted in SeedGraph v3.""" + return { + "version": self.version, + "requires": [atom.as_mapping() for atom in self.requires], + "effects": [effect.as_mapping() for effect in self.effects], + "claims": [claim.as_mapping() for claim in self.claims], + "completion": self.completion, + "failure_policy": self.failure_policy, + } + + +@dataclass(frozen=True) +class AtomicCapability: + """Describe planning, grounding, execution, and recovery for one skill.""" + + name: str + action_type: type | None + config_type: type | None + binding_kinds: frozenset[str] + controls: frozenset[str] + resource_mode: str + state_effect: str + target_materializer: str + motion_base: str | None = None + config_materializer: str = "single_arm" + verifier: str = "postcondition" + failure_classifier: str = "default" + retry_mode: str = "direct" + runtime_available: bool = True + unavailable_reason: str | None = None + target_materializer_hook: Callable[..., Any] | None = None + config_materializer_hook: Callable[..., Any] | None = None + verifier_hook: Callable[..., Any] | None = None + failure_classifier_hook: Callable[..., str] | None = None + contract_resolver_hook: ( + Callable[[Mapping[str, Any]], ResolvedActionContract] | None + ) = None + allows_target_contact: bool = False + """Whether motion planning may temporarily exclude the action target.""" + + def __post_init__(self) -> None: + if not self.name: + raise ValueError("AtomicCapability.name must not be empty.") + if self.motion_base is not None and not self.motion_base: + raise ValueError( + f"AtomicCapability {self.name!r} motion_base must not be empty." + ) + if not self.binding_kinds or not self.controls: + raise ValueError( + f"AtomicCapability {self.name!r} requires bindings and controls." + ) + if self.retry_mode not in _RETRY_MODES: + raise ValueError( + f"AtomicCapability {self.name!r} has invalid retry_mode {self.retry_mode!r}." + ) + if not isinstance(self.allows_target_contact, bool): + raise TypeError("allows_target_contact must be a boolean.") + if self.runtime_available: + if self.action_type is None or self.config_type is None: + raise ValueError( + f"Executable AtomicCapability {self.name!r} requires action/config types." + ) + if self.unavailable_reason is not None: + raise ValueError( + f"Executable AtomicCapability {self.name!r} cannot have an unavailable reason." + ) + elif not self.unavailable_reason: + raise ValueError( + f"Planning-only AtomicCapability {self.name!r} requires unavailable_reason." + ) + for field_name in ( + "target_materializer_hook", + "config_materializer_hook", + "verifier_hook", + "failure_classifier_hook", + "contract_resolver_hook", + ): + value = getattr(self, field_name) + if value is not None and not callable(value): + raise TypeError( + f"AtomicCapability {self.name!r} {field_name} must be callable." + ) + + def resolve_contract(self, node: Mapping[str, Any]) -> ResolvedActionContract: + """Resolve the deterministic Action Contract for one bound node.""" + if self.contract_resolver_hook is not None: + contract = self.contract_resolver_hook(node) + if not isinstance(contract, ResolvedActionContract): + raise TypeError( + f"AtomicCapability {self.name!r} contract resolver must return " + "ResolvedActionContract." + ) + return contract + return _resolve_default_contract(self, node) + + def as_catalog_entry(self) -> dict[str, Any]: + """Return the stable, JSON-safe planning view of this capability.""" + return { + "name": self.name, + "binding_kinds": sorted(self.binding_kinds), + "controls": sorted(self.controls), + "resource_mode": self.resource_mode, + "state_effect": self.state_effect, + "target_materializer": self.target_materializer, + "motion_base": self.motion_base or self.name, + "config_materializer": self.config_materializer, + "verifier": self.verifier, + "failure_classifier": self.failure_classifier, + "retry_mode": self.retry_mode, + "runtime_available": self.runtime_available, + "unavailable_reason": self.unavailable_reason, + "allows_target_contact": self.allows_target_contact, + "custom_target_materializer": _callable_name(self.target_materializer_hook), + "custom_config_materializer": _callable_name(self.config_materializer_hook), + "custom_verifier": _callable_name(self.verifier_hook), + "custom_failure_classifier": _callable_name(self.failure_classifier_hook), + "contract_version": ACTION_CONTRACT_VERSION, + "contract_resolver": _callable_name(self.contract_resolver_hook) + or f"{__name__}._resolve_default_contract", + } + + +class AtomicCapabilityRegistry: + """Strict registry shared by planners, validators, grounders, and runtime.""" + + def __init__(self) -> None: + self._capabilities: dict[str, AtomicCapability] = {} + + def register(self, capability: AtomicCapability) -> None: + if capability.name in self._capabilities: + raise ValueError( + f"AtomicCapability {capability.name!r} is already registered." + ) + self._capabilities[capability.name] = capability + + def get(self, name: str) -> AtomicCapability: + try: + return self._capabilities[name] + except KeyError as exc: + raise ValueError( + f"Unknown AtomicAction {name!r}; available actions are {list(self.names())}." + ) from exc + + def require_executable(self, name: str) -> AtomicCapability: + capability = self.get(name) + if not capability.runtime_available: + raise ValueError( + f"AtomicAction {name!r} is planning-only and cannot be executed: " + f"{capability.unavailable_reason}" + ) + return capability + + def validate_binding(self, action: Mapping[str, Any]) -> None: + name = str(action.get("atomic_action", action.get("atomic_action_class", ""))) + capability = self.get(name) + binding = action.get("target_binding") + if not isinstance(binding, Mapping): + raise ValueError( + f"AtomicAction {name!r} requires a target_binding mapping." + ) + kind = str(binding.get("kind", "")) + if kind not in capability.binding_kinds: + raise ValueError( + f"AtomicAction {name!r} does not accept binding kind {kind!r}; " + f"expected one of {sorted(capability.binding_kinds)}." + ) + control = str(action.get("control", "arm")) + if control not in capability.controls: + raise ValueError( + f"AtomicAction {name!r} does not support control {control!r}; " + f"expected one of {sorted(capability.controls)}." + ) + + def names(self) -> tuple[str, ...]: + return tuple(sorted(self._capabilities)) + + def executable_names(self) -> tuple[str, ...]: + return tuple( + name for name in self.names() if self._capabilities[name].runtime_available + ) + + def catalog(self) -> dict[str, dict[str, Any]]: + return { + name: self._capabilities[name].as_catalog_entry() for name in self.names() + } + + def catalog_hash(self) -> str: + payload = json.dumps( + self.catalog(), sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def build_atomic_capability_registry() -> AtomicCapabilityRegistry: + """Build the default catalog, including explicit planning-only skills.""" + from embodichain.lab.sim.atomic_actions import ( + CoordinatedPickment, + CoordinatedPickmentOptions, + CoordinatedPlacement, + CoordinatedPlacementOptions, + AxisAlign, + AxisAlignOptions, + HandOver, + HandOverOptions, + MoveEndEffector, + MoveEndEffectorOptions, + MoveHeldObject, + MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + PickUp, + PickUpOptions, + Place, + PlaceOptions, + Pour, + PourOptions, + Press, + PressOptions, + Slide, + SlideOptions, + Twist, + TwistOptions, + ) + + registry = AtomicCapabilityRegistry() + definitions = ( + AtomicCapability( + "AxisAlign", + AxisAlign, + AxisAlignOptions, + frozenset({"object"}), + frozenset({"arm"}), + "single_arm_object", + "preserve", + "axis_align", + motion_base="PickUp", + verifier="postcondition", + failure_classifier="grasp", + contract_resolver_hook=_resolve_axis_align_contract, + allows_target_contact=True, + ), + AtomicCapability( + "PickUp", + PickUp, + PickUpOptions, + frozenset({"object"}), + frozenset({"arm"}), + "single_arm_object", + "hold", + "object_grasp", + verifier="held_object", + failure_classifier="grasp", + allows_target_contact=True, + ), + AtomicCapability( + "MoveHeldObject", + MoveHeldObject, + MoveHeldObjectOptions, + frozenset({"semantic_goal", "visual_constraint", "handover_staging"}), + frozenset({"arm"}), + "single_arm_object", + "preserve_hold", + "semantic_held_object", + ), + AtomicCapability( + "MoveEndEffector", + MoveEndEffector, + MoveEndEffectorOptions, + frozenset({"policy_pose", "visual_constraint"}), + frozenset({"arm"}), + "single_arm", + "preserve", + "eef_pose", + verifier_hook=_verify_arm_clearance, + contract_resolver_hook=_resolve_end_effector_contract, + ), + AtomicCapability( + "MoveJoints", + MoveJoints, + MoveJointsOptions, + frozenset({"joint_state"}), + frozenset({"arm", "hand"}), + "control_part", + "preserve", + "joint_state", + contract_resolver_hook=_resolve_joints_contract, + ), + AtomicCapability( + "Place", + Place, + PlaceOptions, + frozenset({"current_held_pose"}), + frozenset({"arm"}), + "single_arm_object", + "release", + "current_held_pose", + ), + AtomicCapability( + "Pour", + Pour, + PourOptions, + frozenset({"pour_goal"}), + frozenset({"arm"}), + "single_arm_object", + "preserve_hold", + "pour", + motion_base="MoveHeldObject", + verifier="postcondition", + retry_mode="non_retryable", + contract_resolver_hook=_resolve_pour_contract, + ), + AtomicCapability( + "PullArticulatedPart", + Slide, + SlideOptions, + frozenset({"articulation_goal"}), + frozenset({"arm"}), + "single_arm_object", + "articulation_change", + "slide", + motion_base="Press", + verifier="postcondition", + retry_mode="non_retryable", + contract_resolver_hook=_resolve_articulation_contract, + allows_target_contact=True, + ), + AtomicCapability( + "PushArticulatedPart", + Slide, + SlideOptions, + frozenset({"articulation_goal"}), + frozenset({"arm"}), + "single_arm_object", + "articulation_change", + "slide", + motion_base="Press", + verifier="postcondition", + retry_mode="non_retryable", + contract_resolver_hook=_resolve_articulation_contract, + allows_target_contact=True, + ), + AtomicCapability( + "TurnKnob", + Twist, + TwistOptions, + frozenset({"articulation_goal"}), + frozenset({"arm"}), + "single_arm_object", + "articulation_change", + "twist", + motion_base="Press", + verifier="postcondition", + retry_mode="non_retryable", + contract_resolver_hook=_resolve_articulation_contract, + allows_target_contact=True, + ), + AtomicCapability( + "Press", + Press, + PressOptions, + frozenset({"object", "semantic_goal"}), + frozenset({"arm"}), + "single_arm_object", + "preserve", + "press", + verifier="pressed", + allows_target_contact=True, + ), + AtomicCapability( + "CoordinatedPickment", + CoordinatedPickment, + CoordinatedPickmentOptions, + frozenset({"object", "coordinated_goal"}), + frozenset({"coordinated"}), + "coordinated_object", + "coordinated_hold", + "coordinated_pickment", + config_materializer="coordinated_pickment", + verifier="coordinated_hold", + failure_classifier="grasp", + ), + AtomicCapability( + "CoordinatedPlacement", + CoordinatedPlacement, + CoordinatedPlacementOptions, + frozenset({"coordinated_placement_goal"}), + frozenset({"coordinated"}), + "coordinated_object", + "coordinated_release", + "coordinated_placement", + config_materializer="coordinated_placement", + ), + AtomicCapability( + "HandOver", + HandOver, + HandOverOptions, + frozenset({"handover_goal"}), + frozenset({"coordinated"}), + "coordinated_object", + "transfer_hold", + "handover", + config_materializer="handover", + verifier="receiver_holds", + failure_classifier="handover", + retry_mode="recover_then_retry", + ), + ) + for capability in definitions: + registry.register(capability) + + return registry + + +def capability_precondition( + capability: AtomicCapability, + *, + object_uid: str, + actor: Mapping[str, Any], + target_binding: Mapping[str, Any], +) -> dict[str, Any]: + """Build the generic live precondition used to authorize a retry.""" + if target_binding.get("coordinated_release_role") is not None: + # Opening a gripper is idempotent. A retry must remain legal when one + # hand opened on the first attempt and the physical dual-hold predicate + # therefore no longer holds. + return {} + if capability.state_effect == "coordinated_release": + return {"type": "held_by_both_grippers", "object": object_uid} + if capability.state_effect in {"preserve_hold", "release", "transfer_hold"}: + result = {"type": "object_held", "object": object_uid} + arm = target_binding.get("transfer_arm") + if arm is None and actor.get("mode") in {"required", "preferred"}: + arm = actor.get("arm") + if isinstance(arm, str) and arm: + result["arm"] = arm + return result + return {} + + +def _resolve_default_contract( + capability: AtomicCapability, node: Mapping[str, Any] +) -> ResolvedActionContract: + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + if not isinstance(actor, Mapping): + raise ValueError("Action Contract resolution requires a mapping actor.") + binding = node.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("Action Contract resolution requires a target_binding.") + arms = _actor_arms(actor) + arm = arms[0] if len(arms) == 1 else None + arm_claims = tuple(ResourceClaim(f"arm:{item}") for item in arms) + object_claim = ResourceClaim(f"object:{object_uid}") + payload_claims = _payload_resource_claims(binding, object_uid) + + if capability.name == "PickUp": + required_arm = _required_arm(arm, capability.name) + return ResolvedActionContract( + requires=( + StateAtom("arm_free", arm=required_arm), + StateAtom("object_free", object_uid=object_uid), + ), + effects=( + StateEffect("delete", StateAtom("arm_free", arm=required_arm)), + StateEffect("delete", StateAtom("object_free", object_uid=object_uid)), + StateEffect( + "add", + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + ), + claims=( + ResourceClaim(f"arm:{required_arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + + payload_claims, + ) + if capability.name == "MoveHeldObject": + required_arm = _required_arm(arm, capability.name) + return ResolvedActionContract( + requires=( + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + claims=( + ResourceClaim(f"arm:{required_arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + + payload_claims, + ) + if capability.name == "Place": + required_arm = _required_arm(arm, capability.name) + return ResolvedActionContract( + requires=( + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + effects=( + StateEffect( + "delete", + StateAtom("object_held", object_uid=object_uid, arm=required_arm), + ), + StateEffect("add", StateAtom("arm_free", arm=required_arm)), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + ), + claims=arm_claims + (object_claim,) + payload_claims, + ) + if capability.name == "HandOver": + transfer = _required_string( + binding.get("transfer_arm"), "target_binding.transfer_arm" + ) + receive = _required_string( + binding.get("receive_arm"), "target_binding.receive_arm" + ) + if transfer == receive: + raise ValueError("HandOver requires distinct transfer and receive arms.") + return ResolvedActionContract( + requires=( + StateAtom("object_held", object_uid=object_uid, arm=transfer), + StateAtom("arm_free", arm=receive), + ), + effects=( + StateEffect( + "delete", + StateAtom("object_held", object_uid=object_uid, arm=transfer), + ), + StateEffect("delete", StateAtom("arm_free", arm=receive)), + StateEffect("add", StateAtom("arm_free", arm=transfer)), + StateEffect( + "add", + StateAtom("object_held", object_uid=object_uid, arm=receive), + ), + StateEffect( + "add", StateAtom("handover_complete", object_uid=object_uid) + ), + ), + claims=( + ResourceClaim(f"arm:{transfer}"), + ResourceClaim(f"arm:{receive}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ), + ) + if capability.name == "CoordinatedPickment": + coordinated_arms = _coordinated_arms(arms, capability.name) + requires = tuple(StateAtom("arm_free", arm=item) for item in coordinated_arms) + effects = tuple( + StateEffect("delete", StateAtom("arm_free", arm=item)) + for item in coordinated_arms + ) + ( + StateEffect("delete", StateAtom("object_free", object_uid=object_uid)), + StateEffect( + "add", StateAtom("object_coordinated_held", object_uid=object_uid) + ), + ) + claims = tuple( + ResourceClaim(f"arm:{item}", lifetime="until_release") + for item in coordinated_arms + ) + ( + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + return ResolvedActionContract( + requires=requires + (StateAtom("object_free", object_uid=object_uid),), + effects=effects, + claims=claims + payload_claims, + ) + if capability.name == "CoordinatedPlacement": + coordinated_arms = _coordinated_arms(arms, capability.name) + effects = ( + StateEffect( + "delete", StateAtom("object_coordinated_held", object_uid=object_uid) + ), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + ) + tuple( + StateEffect("add", StateAtom("arm_free", arm=item)) + for item in coordinated_arms + ) + claims = tuple(ResourceClaim(f"arm:{item}") for item in coordinated_arms) + ( + object_claim, + ) + return ResolvedActionContract( + requires=(StateAtom("object_coordinated_held", object_uid=object_uid),), + effects=effects, + claims=claims + payload_claims, + ) + + requirements: tuple[StateAtom, ...] = () + if capability.state_effect == "coordinated_hold": + requirements = (StateAtom("object_free", object_uid=object_uid),) + elif capability.state_effect == "coordinated_release": + requirements = (StateAtom("object_coordinated_held", object_uid=object_uid),) + elif capability.resource_mode in {"single_arm", "single_arm_object"}: + required_arm = _required_arm(arm, capability.name) + requirements = (StateAtom("arm_free", arm=required_arm),) + claims = arm_claims + if "object" in capability.resource_mode: + claims += (object_claim,) + return ResolvedActionContract( + requires=requirements, + claims=claims + payload_claims, + ) + + +def _payload_resource_claims( + binding: Mapping[str, Any], object_uid: str +) -> tuple[ResourceClaim, ...]: + """Resolve exclusive claims for objects physically carried by a carrier.""" + raw_payloads = binding.get("payloads", ()) + if not isinstance(raw_payloads, Sequence) or isinstance( + raw_payloads, (str, bytes, bytearray) + ): + raise ValueError("target_binding.payloads must be a list.") + payload_uids: list[str] = [] + for index, raw_payload in enumerate(raw_payloads): + value = ( + raw_payload.get("object") + if isinstance(raw_payload, Mapping) + else raw_payload + ) + if not isinstance(value, str) or not value: + raise ValueError( + f"target_binding.payloads[{index}] requires an object UID." + ) + if value == object_uid: + raise ValueError("An AtomicAction carrier cannot be its own payload.") + payload_uids.append(value) + if len(payload_uids) != len(set(payload_uids)): + raise ValueError("target_binding payload objects must be unique.") + return tuple(ResourceClaim(f"object:{uid}") for uid in payload_uids) + + +def _resolve_end_effector_contract( + node: Mapping[str, Any], +) -> ResolvedActionContract: + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + binding = node.get("target_binding", {}) + if not isinstance(actor, Mapping) or not isinstance(binding, Mapping): + raise ValueError("MoveEndEffector contract requires actor and target_binding.") + arm = _required_arm(_actor_arms(actor)[0], "MoveEndEffector") + if binding.get("operation") == "retreat" or node.get("role") == "cleanup": + requires = [StateAtom("arm_free", arm=arm)] + if binding.get("source") == "handover": + requires.append(StateAtom("handover_complete", object_uid=object_uid)) + return ResolvedActionContract( + requires=tuple(requires), + effects=(StateEffect("add", StateAtom("arm_clear", arm=arm)),), + claims=(ResourceClaim(f"arm:{arm}"),), + completion="cleanup", + failure_policy="safety_required", + ) + return ResolvedActionContract( + requires=(StateAtom("arm_free", arm=arm),), + claims=(ResourceClaim(f"arm:{arm}"),), + ) + + +def _resolve_axis_align_contract(node: Mapping[str, Any]) -> ResolvedActionContract: + """Keep the object free while enforcing one verified E2 terminal barrier.""" + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + if not isinstance(actor, Mapping): + raise ValueError("AxisAlign contract requires an actor mapping.") + arm = _required_arm(_actor_arms(actor)[0], "AxisAlign") + return ResolvedActionContract( + requires=( + StateAtom("arm_free", arm=arm), + StateAtom("object_free", object_uid=object_uid), + ), + effects=( + StateEffect("add", StateAtom("arm_free", arm=arm)), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + ), + claims=( + ResourceClaim(f"arm:{arm}"), + ResourceClaim(f"object:{object_uid}"), + ), + completion="terminal_barrier", + failure_policy="task_required", + ) + + +def _resolve_pour_contract(node: Mapping[str, Any]) -> ResolvedActionContract: + """Retain one verified holder until observable content transfer succeeds.""" + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + binding = node.get("target_binding", {}) + if not isinstance(actor, Mapping) or not isinstance(binding, Mapping): + raise ValueError("Pour contract requires actor and target_binding mappings.") + arm = _required_arm(_actor_arms(actor)[0], "Pour") + return ResolvedActionContract( + requires=(StateAtom("object_held", object_uid=object_uid, arm=arm),), + claims=( + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + + _payload_resource_claims(binding, object_uid), + completion="terminal_barrier", + failure_policy="task_required", + ) + + +def _resolve_articulation_contract( + node: Mapping[str, Any], +) -> ResolvedActionContract: + """Require one free arm and verify the observed articulation terminal state.""" + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + if not isinstance(actor, Mapping): + raise ValueError("Articulation action contract requires an actor mapping.") + arm = _required_arm(_actor_arms(actor)[0], "articulation action") + return ResolvedActionContract( + requires=(StateAtom("arm_free", arm=arm),), + claims=( + ResourceClaim(f"arm:{arm}"), + ResourceClaim(f"object:{object_uid}"), + ), + completion="terminal_barrier", + failure_policy="task_required", + ) + + +def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: + object_uid = _required_string(node.get("object_uid"), "node.object_uid") + actor = node.get("actor", {}) + if not isinstance(actor, Mapping): + raise ValueError("MoveJoints contract requires an actor mapping.") + arm = _required_arm(_actor_arms(actor)[0], "MoveJoints") + binding = node.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("MoveJoints contract requires a target_binding mapping.") + release_role = binding.get("coordinated_release_role") + if release_role is not None: + if ( + node.get("task_type") != "E5" + or node.get("control") != "hand" + or binding.get("source") != "gripper_open" + or not node.get("sync_group") + ): + raise ValueError( + "Coordinated MoveJoints release requires an E5 synchronized " + "hand action targeting gripper_open." + ) + if release_role not in {"participant", "commit"}: + raise ValueError( + "coordinated_release_role must be 'participant' or 'commit'." + ) + claims = ( + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ) + if release_role == "participant": + return ResolvedActionContract( + requires=(StateAtom("object_coordinated_held", object_uid=object_uid),), + claims=claims, + ) + return ResolvedActionContract( + requires=(StateAtom("object_coordinated_held", object_uid=object_uid),), + effects=( + StateEffect( + "delete", + StateAtom("object_coordinated_held", object_uid=object_uid), + ), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + StateEffect("add", StateAtom("arm_free", arm="left_arm")), + StateEffect("add", StateAtom("arm_free", arm="right_arm")), + ), + claims=claims, + ) + if node.get("control") == "hand": + return ResolvedActionContract( + requires=(StateAtom("object_held", object_uid=object_uid, arm=arm),), + claims=( + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ), + ) + if node.get("role") == "cleanup": + return ResolvedActionContract( + requires=(StateAtom("arm_clear", arm=arm),), + effects=( + StateEffect("add", StateAtom("arm_home", arm=arm)), + StateEffect("add", StateAtom("arm_free", arm=arm)), + ), + claims=(ResourceClaim(f"arm:{arm}"),), + completion="terminal_barrier", + failure_policy="best_effort", + ) + return ResolvedActionContract( + requires=(StateAtom("arm_free", arm=arm),), + claims=(ResourceClaim(f"arm:{arm}"),), + ) + + +def _verify_arm_clearance( + *, + executor: Any, + step: Any, + arm: str, + outcome: Any, + attempted: torch.Tensor, +) -> torch.Tensor: + """Verify a released TCP is clear, plus the transfer side for handover.""" + policy = outcome.grounded.motion_policy + object_uid = policy.get("clearance_object_uid") + if not isinstance(object_uid, str) or not object_uid: + return attempted + transfer_arm = str(policy.get("transfer_arm", arm)) + if transfer_arm not in {"left_arm", "right_arm"}: + return torch.zeros_like(attempted) + entity = executor.env.sim.get_rigid_object(object_uid) + getter = getattr(executor.env, "get_current_xpos_agent", None) + if entity is None or not callable(getter): + return torch.zeros_like(attempted) + left, right = getter() + eef = torch.as_tensor( + left if transfer_arm == "left_arm" else right, + dtype=torch.float32, + device=executor.env.device, + ) + if eef.ndim == 2: + eef = eef.unsqueeze(0).repeat(int(executor.env.num_envs), 1, 1) + object_pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=executor.env.device, + ) + if object_pose.ndim == 2: + object_pose = object_pose.unsqueeze(0).repeat(int(executor.env.num_envs), 1, 1) + offset = eef[:, :3, 3] - object_pose[:, :3, 3] + distance = torch.linalg.vector_norm(offset, dim=1) + minimum_clearance = policy.get( + "minimum_clearance", + policy.get("minimum_transfer_clearance", 0.10), + ) + clear = distance >= float(minimum_clearance) + role_axis = policy.get("transfer_role_axis") + if role_axis is None: + return attempted & clear + role_axis = torch.as_tensor( + role_axis, + dtype=offset.dtype, + device=offset.device, + ) + if role_axis.ndim == 1: + role_axis = role_axis.unsqueeze(0).repeat(int(executor.env.num_envs), 1) + lateral = torch.sum(offset * role_axis, dim=1) + clear &= lateral >= float(policy.get("minimum_transfer_lateral_clearance", 0.06)) + return attempted & clear + + +def _actor_arms(actor: Mapping[str, Any]) -> tuple[str, ...]: + mode = str(actor.get("mode", "auto")) + if mode == "coordinated": + arms = actor.get("arms", ()) + if not isinstance(arms, (list, tuple)): + raise ValueError("Coordinated actor arms must be a sequence.") + result = tuple(str(item) for item in arms) + if len(result) < 2 or any(not item for item in result): + raise ValueError("Coordinated actor requires at least two named arms.") + return result + if mode in {"required", "preferred"}: + return (_required_string(actor.get("arm"), "actor.arm"),) + return ("auto",) + + +def _coordinated_arms(arms: tuple[str, ...], action: str) -> tuple[str, ...]: + if len(arms) < 2: + raise ValueError(f"{action} requires a coordinated actor.") + return arms + + +def _required_arm(arm: str | None, action: str) -> str: + if arm is None: + raise ValueError(f"{action} requires exactly one arm.") + return arm + + +def _required_string(value: Any, context: str) -> str: + if not isinstance(value, str) or not value: + raise ValueError(f"{context} must be a non-empty string.") + return value + + +def _callable_name(value: Callable[..., Any] | None) -> str | None: + if value is None: + return None + module = getattr(value, "__module__", "") + name = getattr( + value, "__qualname__", getattr(value, "__name__", type(value).__name__) + ) + return f"{module}.{name}" if module else str(name) diff --git a/embodichain/gen_sim/action_engine/capabilities/builtins.py b/embodichain/gen_sim/action_engine/capabilities/builtins.py new file mode 100644 index 000000000..839797e37 --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/builtins.py @@ -0,0 +1,1018 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Built-in semantic operators lowered to public atomic-action contracts.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.domain.motion import ( + motion_policy as build_motion_policy, +) +from embodichain.gen_sim.action_engine.domain.task_contracts import ( + PLACEMENT_RELATIONS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + normalize_placement_relation, +) +from embodichain.gen_sim.action_engine.orientation import ( + compile_orientation_constraint, +) + +from .registry import ( + ActionCapability, + ActionTemplate, + CapabilityRegistry, + OperatorCapability, + PhaseTemplate, +) + +__all__ = ["build_default_registry"] + +_SINGLE_ARM_PHASE_OPERATORS = frozenset( + {"build_stack", "hold_hover", "orient_object", "place_relative"} +) + + +def build_default_registry() -> CapabilityRegistry: + """Build a fresh registry containing all Action Engine v1 capabilities.""" + registry = CapabilityRegistry() + from .atomic import build_atomic_capability_registry + + for capability in build_atomic_capability_registry().catalog().values(): + registry.register_action( + ActionCapability( + str(capability["name"]), + frozenset(capability["binding_kinds"]), + frozenset(capability["controls"]), + ) + ) + + definitions = ( + OperatorCapability( + "arrange_line", + "Arrange two or more movable objects into one live-grounded line.", + _expand_arrange_line, + _build_arrange_line_phases, + expansion_topology="parallel_children", + ), + OperatorCapability( + "build_stack", + "Build one ordered vertical or nested stack.", + _expand_build_stack, + _build_single_arm_phases, + ), + OperatorCapability( + "place_relative", + "Place one object at a symbolic relation to another object.", + _expand_place_relative, + _build_single_arm_phases, + ), + OperatorCapability( + "orient_object", + "Reorient one object in place and release it in a stable pose.", + _expand_orient_object, + _build_orient_object_phases, + ), + OperatorCapability( + "coordinated_transport", + "Use both arms to pick and transport one shared object.", + _expand_coordinated_transport, + _build_coordinated_transport_phases, + ), + # These internal operators preserve runtime characterization coverage + # for public Atomic Actions. They are intentionally absent from the + # planner catalog during the five-skill first phase. + OperatorCapability( + "hold_hover", + "Internal terminal-hold compatibility operator.", + _expand_hold_hover, + _build_single_arm_phases, + lifecycle="terminal_hold", + planner_visible=False, + ), + OperatorCapability( + "press", + "Internal press compatibility operator.", + _expand_press, + _build_press_phases, + planner_visible=False, + ), + OperatorCapability( + "coordinated_place", + "Internal coordinated-placement compatibility operator.", + _expand_coordinated_place, + _build_coordinated_place_phases, + planner_visible=False, + ), + ) + for definition in definitions: + registry.register_operator(definition) + return registry + + +def _expand_arrange_line(step: Mapping[str, Any]) -> list[dict[str, Any]]: + objects = _collective_objects(step, "arrange_line", minimum=2) + goal = _goal( + step, + allowed={ + "anchor", + "axis", + "order_by", + "order_constraint", + "order_direction", + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "participation", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "arrange_line") + axis = str(goal.get("axis", "world_y")) + if axis not in {"world_x", "world_y", "table_long_axis"}: + raise ValueError("arrange_line goal.axis must be a symbolic table axis.") + anchor = str(goal.get("anchor", "table_center")) + if anchor != "table_center": + raise ValueError("arrange_line currently requires anchor='table_center'.") + order_constraint = str(goal.get("order_constraint", "free")) + if order_constraint not in {"free", "ordered"}: + raise ValueError( + "arrange_line goal.order_constraint must be 'free' or 'ordered'." + ) + order_by = str(goal.get("order_by", "explicit")) + if order_by not in {"explicit", "size", "color"}: + raise ValueError("arrange_line order_by must be explicit, size, or color.") + order_direction = str(goal.get("order_direction", "given")) + if order_direction not in {"given", "ascending", "descending"}: + raise ValueError( + "arrange_line order_direction must be given, ascending, or descending." + ) + participation = str(goal.get("participation", "auto")) + if participation not in {"auto", "both_arms"}: + raise ValueError("arrange_line participation must be auto or both_arms.") + + common_goal = { + "layout": "line", + "objects": objects, + "axis": axis, + "anchor": anchor, + "order_by": order_by, + "order_direction": order_direction, + "order_constraint": order_constraint, + "participation": participation, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + } + actor = _single_arm_actor(step) + expanded: list[dict[str, Any]] = [] + for slot_index, object_uid in enumerate(objects): + child_goal = { + **deepcopy(common_goal), + "nominal_slot_index": slot_index, + "slot_constraint": ( + "required" if order_constraint == "ordered" else "free_reassignable" + ), + } + expanded.append( + _execution_step( + step, + child_id=f"{step['id']}__{slot_index + 1:02d}", + object_uid=object_uid, + actor=( + { + **actor, + "allocation_group": f"{step['id']}_both_arms", + } + if participation == "both_arms" and slot_index < 2 + else actor + ), + goal=child_goal, + postcondition={ + "type": "line_member_placed", + "nominal_slot_index": slot_index, + "slot_constraint": child_goal["slot_constraint"], + "order_constraint": order_constraint, + }, + ) + ) + return expanded + + +def _expand_build_stack(step: Mapping[str, Any]) -> list[dict[str, Any]]: + objects = _collective_objects(step, "build_stack", minimum=1) + goal = _goal( + step, + allowed={ + "anchor", + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "stack_mode", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "build_stack") + stack_mode = str(goal.get("stack_mode", "on_top")) + if stack_mode not in {"on_top", "nested"}: + raise ValueError("build_stack goal.stack_mode must be 'on_top' or 'nested'.") + anchor = goal.get("anchor", "table_center") + if not isinstance(anchor, str) or not anchor: + raise ValueError("build_stack goal.anchor must be an object or table_center.") + + actor = _single_arm_actor(step) + expanded: list[dict[str, Any]] = [] + for layer_index, object_uid in enumerate(objects): + reference = objects[layer_index - 1] if layer_index else anchor + support_reference = "table" if reference == "table_center" else reference + child_goal: dict[str, Any] = { + "relation": ( + "inside" if stack_mode == "nested" and layer_index > 0 else "on" + ), + "reference_object": support_reference, + "reference_state": "live", + "layer_index": layer_index, + "stack_mode": stack_mode, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + } + expanded.append( + _execution_step( + step, + child_id=f"{step['id']}__{layer_index + 1:02d}", + object_uid=object_uid, + actor=actor, + goal=child_goal, + postcondition={ + "type": "stack_layer_supported", + "layer_index": layer_index, + "reference_object": support_reference, + }, + ) + ) + return expanded + + +def _expand_place_relative(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "place_relative") + goal = _goal( + step, + allowed={ + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "orientation_reference_object", + "payloads", + "reference_object", + "reference_state", + "relation", + "slot", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "place_relative") + reference = _required_string(goal, "reference_object", "place_relative") + relation = normalize_placement_relation(goal.get("relation", "on")) + normalized_goal = { + "reference_object": reference, + "reference_state": str(goal.get("reference_state", "live")), + "relation": relation, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + "slot": str(goal.get("slot", "auto")), + } + if normalized_goal["reference_state"] not in {"initial", "live"}: + raise ValueError("place_relative reference_state must be 'initial' or 'live'.") + if normalized_goal["slot"] not in {"auto", "left", "center", "right"}: + raise ValueError("place_relative slot must be left, center, right, or auto.") + if "orientation_reference_object" in goal: + normalized_goal["orientation_reference_object"] = goal[ + "orientation_reference_object" + ] + payloads = _normalize_payloads( + goal.get("payloads", []), + object_uid, + "place_relative", + ) + if payloads: + normalized_goal["payloads"] = payloads + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal=normalized_goal, + postcondition={ + "type": "semantic_goal", + "relation": relation, + "reference_object": reference, + }, + ) + ] + + +def _expand_hold_hover(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "hold_hover") + goal = _goal( + step, + allowed={ + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "reference_object", + "reference_state", + }, + ) + orientation_goal, orientation_axis = _orientation( + goal, + "hold_hover", + allow_change=False, + ) + reference = str(goal.get("reference_object", object_uid)) + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal={ + "relation": "held_above_initial", + "reference_object": reference, + "reference_state": str(goal.get("reference_state", "initial")), + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + }, + postcondition={"type": "object_held", "object": object_uid}, + ) + ] + + +def _expand_orient_object(step: Mapping[str, Any]) -> list[dict[str, Any]]: + """Normalize an in-place orientation request into one executable step. + + Keeping the target position symbolic is important: runtime observes the + object's live position immediately before grounding, so prior independent + operations and simulator settling cannot make this plan stale. + """ + object_uid = _single_object(step, "orient_object") + goal = _goal( + step, + allowed={ + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "position_anchor", + "support_object", + "upright_local_axis", + }, + ) + orientation_goal, orientation_axis = _orientation(goal, "orient_object") + if orientation_goal not in {"upright", "lay_flat", "axis_align"}: + raise ValueError( + "orient_object requires upright, lay_flat, or axis_align orientation." + ) + position_anchor = str(goal.get("position_anchor", "initial_xy")) + if position_anchor not in {"initial_xy", "live_xy"}: + raise ValueError( + "orient_object position_anchor must be 'initial_xy' or 'live_xy'." + ) + upright_local_axis = str(goal.get("upright_local_axis", "auto")) + if upright_local_axis not in {"auto", "long_axis", "x", "y", "z"}: + raise ValueError( + "orient_object upright_local_axis must be auto, long_axis, x, y, or z." + ) + support_object = str(goal.get("support_object", "table")) + if not support_object: + raise ValueError("orient_object support_object must be a non-empty string.") + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal={ + "relation": "none", + "reference_state": "live", + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + "position_anchor": position_anchor, + "support_object": support_object, + "upright_local_axis": upright_local_axis, + }, + postcondition={ + "type": "semantic_goal", + "relation": "none", + "orientation_goal": orientation_goal, + }, + ) + ] + + +def _expand_coordinated_transport( + step: Mapping[str, Any], +) -> list[dict[str, Any]]: + object_uid = _single_object(step, "coordinated_transport") + goal = _goal( + step, + allowed={ + "direction", + "orientation_axis", + "orientation_constraint", + "orientation_directed", + "orientation_goal", + "payloads", + "reference_object", + "relation", + "terminal_behavior", + }, + ) + orientation_goal, orientation_axis = _orientation( + goal, + "coordinated_transport", + ) + terminal_behavior = str(goal.get("terminal_behavior", "hold")) + if terminal_behavior not in TERMINAL_BEHAVIORS - {"none"}: + raise ValueError( + "coordinated_transport terminal_behavior must be 'hold' or 'place'." + ) + direction = str(goal.get("direction", "none")) + if direction not in TRANSPORT_DIRECTIONS: + raise ValueError( + f"coordinated_transport direction {direction!r} is unsupported." + ) + relation = goal.get("relation") + if relation is not None and str(relation) not in PLACEMENT_RELATIONS: + raise ValueError( + f"coordinated_transport relation {str(relation)!r} is unsupported." + ) + normalized_goal = { + "direction": direction, + "terminal_behavior": terminal_behavior, + "orientation_goal": orientation_goal, + "orientation_axis": orientation_axis, + **_orientation_extensions(goal), + } + normalized_payloads = _normalize_payloads( + goal.get("payloads", []), + object_uid, + "coordinated_transport", + ) + if normalized_payloads: + normalized_goal["payloads"] = normalized_payloads + for key in ("reference_object", "relation"): + if key in goal: + normalized_goal[key] = goal[key] + postcondition = ( + {"type": "semantic_goal", "relation": normalized_goal.get("relation", "at")} + if terminal_behavior == "place" + else {"type": "held_by_both_grippers", "object": object_uid} + ) + return [ + _execution_step( + step, + object_uid=object_uid, + actor={"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + goal=normalized_goal, + postcondition=postcondition, + ) + ] + + +def _expand_press(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "press") + goal = _goal( + step, + allowed={"interaction", "reference_object", "terminal_state"}, + ) + return [ + _execution_step( + step, + object_uid=object_uid, + actor=_single_arm_actor(step), + goal={ + "interaction": str(goal.get("interaction", "press")), + "terminal_state": str(goal.get("terminal_state", "activated")), + **( + {"reference_object": goal["reference_object"]} + if "reference_object" in goal + else {} + ), + }, + postcondition={ + "type": "pressed", + "object": object_uid, + "terminal_state": str(goal.get("terminal_state", "activated")), + }, + ) + ] + + +def _expand_coordinated_place(step: Mapping[str, Any]) -> list[dict[str, Any]]: + object_uid = _single_object(step, "coordinated_place") + goal = _goal( + step, + allowed={"relation", "release", "support_object"}, + ) + support_object = _required_string(goal, "support_object", "coordinated_place") + if support_object == object_uid: + raise ValueError("coordinated_place requires two distinct objects.") + relation = str(goal.get("relation", "on")) + if relation not in {"on", "inside"}: + raise ValueError("coordinated_place relation must be 'on' or 'inside'.") + release = goal.get("release", True) + if not isinstance(release, bool): + raise ValueError("coordinated_place goal.release must be boolean.") + return [ + _execution_step( + step, + object_uid=object_uid, + actor={"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + goal={ + "support_object": support_object, + "relation": relation, + "release": release, + }, + postcondition={ + "type": "coordinated_placed", + "object": object_uid, + "support_object": support_object, + "relation": relation, + }, + ) + ] + + +def _build_arrange_line_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + return ( + _pickup_phase(step), + _move_phase(step, "staging"), + _move_phase(step, "final"), + *_release_retreat_home(step), + ) + + +def _build_single_arm_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + if step["operator"] not in _SINGLE_ARM_PHASE_OPERATORS: + raise ValueError(f"Unexpected single-arm operator {step['operator']!r}.") + phases: tuple[PhaseTemplate, ...] = ( + _pickup_phase(step), + _move_phase(step), + ) + if step["operator"] == "hold_hover": + return phases + ( + PhaseTemplate( + name="keep_holding", + state_semantic=f"`{step['object']}` remains held", + actions=( + ActionTemplate( + "MoveJoints", + {"kind": "joint_state", "source": "gripper_closed"}, + build_motion_policy(), + control="hand", + ), + ), + ), + ) + return phases + _release_retreat_home(step) + + +def _build_orient_object_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + """Rotate at a clearance waypoint before descending to the support.""" + upright = build_motion_policy(("orientation", "upright")) + return ( + _pickup_phase(step, motion_policy=upright), + _move_phase( + step, + "staging", + motion_policy=upright, + ), + _move_phase( + step, + "final", + motion_policy=upright, + ), + *_release_retreat_home( + step, + release_policy=upright, + retreat_policy=upright, + ), + ) + + +def _build_coordinated_transport_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + phases: tuple[PhaseTemplate, ...] = ( + PhaseTemplate( + name="coordinated_transport", + state_semantic=f"`{step['object']}` reaches its coordinated goal", + actions=( + ActionTemplate( + "CoordinatedPickment", + { + "kind": "coordinated_goal", + "semantic_step": step["id"], + "object": step["object"], + "payloads": deepcopy(step["goal"].get("payloads", [])), + }, + build_motion_policy(), + control="coordinated", + ), + ), + ), + ) + if step["goal"]["terminal_behavior"] != "place": + return phases + return phases + ( + PhaseTemplate( + name="dual_release", + state_semantic="Both grippers release the transported object", + actions=tuple( + ActionTemplate( + "MoveJoints", + { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": release_role, + }, + build_motion_policy(), + control="hand", + actor={"mode": "required", "arm": arm}, + ) + for arm, release_role in ( + ("left_arm", "participant"), + ("right_arm", "commit"), + ) + ), + ), + ) + + +def _build_press_phases(step: Mapping[str, Any]) -> tuple[PhaseTemplate, ...]: + return ( + PhaseTemplate( + name="press", + state_semantic=f"`{step['object']}` has been pressed", + actions=( + ActionTemplate( + "Press", + { + "kind": "semantic_goal", + "semantic_step": step["id"], + "object": step["object"], + "interaction": "press", + }, + build_motion_policy(), + ), + ), + ), + ) + + +def _build_coordinated_place_phases( + step: Mapping[str, Any], +) -> tuple[PhaseTemplate, ...]: + return ( + PhaseTemplate( + name="dual_pick_up", + state_semantic=( + f"`{step['object']}` is held by the left arm and " + f"`{step['goal']['support_object']}` is held by the right arm" + ), + actions=( + ActionTemplate( + "PickUp", + { + "kind": "object", + "object": step["object"], + "affordance": "antipodal", + }, + build_motion_policy(), + actor={"mode": "required", "arm": "left_arm"}, + ), + ActionTemplate( + "PickUp", + { + "kind": "object", + "object": step["goal"]["support_object"], + "affordance": "antipodal", + }, + build_motion_policy(), + actor={"mode": "required", "arm": "right_arm"}, + ), + ), + ), + PhaseTemplate( + name="coordinated_place", + state_semantic=( + f"`{step['object']}` is coordinated with " + f"`{step['goal']['support_object']}`" + ), + actions=( + ActionTemplate( + "CoordinatedPlacement", + { + "kind": "coordinated_placement_goal", + "semantic_step": step["id"], + "placing_object": step["object"], + "support_object": step["goal"]["support_object"], + }, + build_motion_policy(), + control="coordinated", + ), + ), + ), + ) + + +def _pickup_phase( + step: Mapping[str, Any], + *, + motion_policy: Mapping[str, Any] | None = None, +) -> PhaseTemplate: + payloads = deepcopy(step["goal"].get("payloads", [])) + return PhaseTemplate( + name="pick_up", + state_semantic=f"Holding `{step['object']}`", + actions=( + ActionTemplate( + "PickUp", + { + "kind": "object", + "object": step["object"], + "affordance": "antipodal", + **({"payloads": payloads} if payloads else {}), + }, + motion_policy or build_motion_policy(), + ), + ), + ) + + +def _move_phase( + step: Mapping[str, Any], + phase: str | None = None, + *, + motion_policy: Mapping[str, Any] | None = None, +) -> PhaseTemplate: + target_binding = { + "kind": "semantic_goal", + "semantic_step": step["id"], + } + if phase is not None: + target_binding["phase"] = phase + payloads = deepcopy(step["goal"].get("payloads", [])) + if payloads: + target_binding["payloads"] = payloads + return PhaseTemplate( + name=f"move_to_{phase or 'semantic_goal'}", + state_semantic=f"`{step['object']}` is held at {phase or 'its semantic goal'}", + actions=( + ActionTemplate( + "MoveHeldObject", + target_binding, + motion_policy or build_motion_policy(), + ), + ), + ) + + +def _release_retreat_home( + step: Mapping[str, Any], + *, + release_policy: Mapping[str, Any] | None = None, + retreat_policy: Mapping[str, Any] | None = None, +) -> tuple[PhaseTemplate, ...]: + payloads = deepcopy(step["goal"].get("payloads", [])) + return ( + PhaseTemplate( + name="release", + state_semantic=f"`{step['object']}` is released at its semantic goal", + actions=( + ActionTemplate( + "Place", + { + "kind": "current_held_pose", + **({"payloads": payloads} if payloads else {}), + }, + release_policy or build_motion_policy(), + ), + ), + ), + PhaseTemplate( + name="retreat", + state_semantic=f"The end effector retreats from `{step['object']}`", + actions=( + ActionTemplate( + "MoveEndEffector", + { + "kind": "policy_pose", + "source": "release", + "operation": "retreat", + }, + retreat_policy or build_motion_policy(), + ), + ), + ), + PhaseTemplate( + name="home", + state_semantic="The selected arm returns to its initial state", + actions=( + ActionTemplate( + "MoveJoints", + {"kind": "joint_state", "source": "initial"}, + build_motion_policy(), + ), + ), + ), + ) + + +def _dual_arm_phase( + name: str, + state_semantic: str, + action_class: str, + target_binding: Mapping[str, Any], + motion_policy: Mapping[str, Any], + *, + control: str = "arm", +) -> PhaseTemplate: + return PhaseTemplate( + name=name, + state_semantic=state_semantic, + actions=tuple( + ActionTemplate( + action_class, + target_binding, + motion_policy, + control=control, + actor={"mode": "required", "arm": arm}, + ) + for arm in ("left_arm", "right_arm") + ), + ) + + +def _execution_step( + parent: Mapping[str, Any], + *, + object_uid: str, + actor: Mapping[str, Any], + goal: Mapping[str, Any], + postcondition: Mapping[str, Any], + child_id: str | None = None, +) -> dict[str, Any]: + return { + "id": child_id or parent["id"], + "parent_step_id": parent["id"], + "operator": parent["operator"], + "object": object_uid, + "actor": deepcopy(dict(actor)), + "goal": deepcopy(dict(goal)), + "depends_on": [], + "postcondition": deepcopy(dict(postcondition)), + "edge_ids": [], + } + + +def _single_object(step: Mapping[str, Any], operator: str) -> str: + if "object" not in step: + raise ValueError(f"{operator} requires one 'object', not 'objects'.") + return str(step["object"]) + + +def _collective_objects( + step: Mapping[str, Any], + operator: str, + *, + minimum: int, +) -> list[str]: + if "objects" not in step: + raise ValueError(f"{operator} requires an 'objects' list.") + objects = [str(value) for value in step["objects"]] + if len(objects) < minimum: + raise ValueError(f"{operator} requires at least {minimum} object(s).") + return objects + + +def _single_arm_actor(step: Mapping[str, Any]) -> dict[str, Any]: + actor = deepcopy(dict(step["actor"])) + if actor["mode"] == "coordinated": + raise ValueError(f"{step['operator']} requires one arm, not coordinated arms.") + if actor["mode"] == "required": + arm = str(actor["arm"]) + if arm in {"left", "right"}: + actor["arm"] = f"{arm}_arm" + return actor + + +def _goal(step: Mapping[str, Any], *, allowed: set[str]) -> dict[str, Any]: + goal = deepcopy(dict(step["goal"])) + unknown = sorted(set(goal) - allowed) + if unknown: + raise ValueError( + f"{step['operator']} goal contains unsupported fields: {unknown}." + ) + return goal + + +def _normalize_payloads( + value: Any, + carrier_uid: str, + operator: str, +) -> list[dict[str, str]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise ValueError(f"{operator} payloads must be a list.") + if len(value) > 4: + raise ValueError(f"{operator} supports at most four payloads.") + result = [] + for index, payload in enumerate(value): + item = {"object": payload} if isinstance(payload, str) else dict(payload) + uid = item.get("object") + slot = str(item.get("slot", "auto")) + if not isinstance(uid, str) or not uid: + raise ValueError(f"payloads[{index}] requires an object UID.") + if uid == carrier_uid: + raise ValueError(f"A {operator} carrier cannot be its own payload.") + if slot not in {"left", "right", "center", "auto"}: + raise ValueError(f"Unsupported payload slot {slot!r}.") + result.append({"object": uid, "slot": slot}) + payload_uids = [item["object"] for item in result] + if len(payload_uids) != len(set(payload_uids)): + raise ValueError(f"{operator} payload objects must be unique.") + return result + + +def _required_string( + value: Mapping[str, Any], + key: str, + operator: str, +) -> str: + result = value.get(key) + if not isinstance(result, str) or not result.strip(): + raise ValueError(f"{operator} goal.{key} must be a non-empty string.") + return result + + +def _orientation( + goal: Mapping[str, Any], + operator: str, + *, + allow_change: bool = True, +) -> tuple[str, str]: + default_goal = "none" if allow_change else "preserve" + orientation_goal = str(goal.get("orientation_goal", default_goal)) + orientation_axis = str(goal.get("orientation_axis", "none")) + allowed_goals = ( + {"none", "preserve", "upright", "lay_flat", "axis_align"} + if allow_change + else {"preserve"} + ) + if orientation_goal not in allowed_goals: + raise ValueError( + f"{operator} orientation_goal {orientation_goal!r} is unsupported." + ) + if orientation_axis not in {"none", "x", "y", "long_axis", "short_axis"}: + raise ValueError( + f"{operator} orientation_axis {orientation_axis!r} is unsupported." + ) + if orientation_goal == "axis_align" and orientation_axis == "none": + raise ValueError(f"{operator} axis_align requires an orientation_axis.") + compile_orientation_constraint(goal) + return orientation_goal, orientation_axis + + +def _orientation_extensions(goal: Mapping[str, Any]) -> dict[str, Any]: + """Copy optional composable fields after operator-level validation.""" + return { + key: deepcopy(goal[key]) + for key in ("orientation_constraint", "orientation_directed") + if key in goal + } diff --git a/embodichain/gen_sim/action_engine/capabilities/registry.py b/embodichain/gen_sim/action_engine/capabilities/registry.py new file mode 100644 index 000000000..83ade2d3a --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/registry.py @@ -0,0 +1,183 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Capability registry shared by planning metadata and compilation.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from embodichain.gen_sim.action_engine.domain.motion import validate_motion_policy + +__all__ = [ + "ActionCapability", + "ActionTemplate", + "CapabilityRegistry", + "OperatorCapability", + "PhaseTemplate", +] + + +@dataclass(frozen=True) +class ActionCapability: + """Describe one public AtomicAction class exposed to the compiler.""" + + class_name: str + target_binding_kinds: frozenset[str] + controls: frozenset[str] + + +@dataclass(frozen=True) +class ActionTemplate: + """Describe one symbolic atomic action before actor materialization.""" + + atomic_action_class: str + target_binding: Mapping[str, Any] + motion_policy: Mapping[str, Any] + control: str = "arm" + actor: Mapping[str, Any] | None = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "target_binding", + MappingProxyType(dict(self.target_binding)), + ) + object.__setattr__( + self, + "motion_policy", + MappingProxyType(validate_motion_policy(self.motion_policy)), + ) + if self.actor is not None: + object.__setattr__(self, "actor", MappingProxyType(dict(self.actor))) + + +@dataclass(frozen=True) +class PhaseTemplate: + """Group atomic actions that execute on one graph edge.""" + + name: str + state_semantic: str + actions: tuple[ActionTemplate, ...] + + +ExpandOperator = Callable[[Mapping[str, Any]], list[dict[str, Any]]] +BuildPhases = Callable[[Mapping[str, Any]], Sequence[PhaseTemplate]] + + +@dataclass(frozen=True) +class OperatorCapability: + """Bind a semantic operator to deterministic expansion and lowering.""" + + name: str + description: str + expand: ExpandOperator + build_phases: BuildPhases + expansion_topology: str = "serial" + lifecycle: str = "release" + planner_visible: bool = True + + def __post_init__(self) -> None: + if self.expansion_topology not in {"serial", "parallel_children"}: + raise ValueError( + "Operator expansion_topology must be 'serial' or " + "'parallel_children'." + ) + if self.lifecycle not in {"release", "terminal_hold"}: + raise ValueError("Operator lifecycle must be 'release' or 'terminal_hold'.") + + +class CapabilityRegistry: + """Store explicit operator and atomic-action capabilities. + + Registration is intentionally strict. Replacing a capability by accident + would silently change compilation semantics, so callers must construct a + new registry when they need a different definition. + """ + + def __init__(self) -> None: + self._operators: dict[str, OperatorCapability] = {} + self._actions: dict[str, ActionCapability] = {} + + def register_operator(self, capability: OperatorCapability) -> None: + """Register one semantic operator.""" + if capability.name in self._operators: + raise ValueError(f"Operator {capability.name!r} is already registered.") + self._operators[capability.name] = capability + + def register_action(self, capability: ActionCapability) -> None: + """Register one public AtomicAction contract.""" + if capability.class_name in self._actions: + raise ValueError( + f"Atomic action {capability.class_name!r} is already registered." + ) + self._actions[capability.class_name] = capability + + def operator(self, name: str) -> OperatorCapability: + """Return an operator or raise a capability-focused error.""" + try: + return self._operators[name] + except KeyError as exc: + raise ValueError( + f"Unknown semantic operator {name!r}; available operators are " + f"{sorted(self._operators)}." + ) from exc + + def action(self, class_name: str) -> ActionCapability: + """Return an atomic-action contract or raise a focused error.""" + try: + return self._actions[class_name] + except KeyError as exc: + raise ValueError( + f"Unknown atomic action {class_name!r}; available actions are " + f"{sorted(self._actions)}." + ) from exc + + def operator_names(self) -> tuple[str, ...]: + """Return only the semantic skills exposed to the LLM planner.""" + return tuple( + sorted( + name + for name, capability in self._operators.items() + if capability.planner_visible + ) + ) + + def operator_descriptions(self) -> dict[str, str]: + """Return JSON-safe operator descriptions.""" + return { + name: self._operators[name].description for name in self.operator_names() + } + + def validate_action_template(self, template: ActionTemplate) -> None: + """Validate one compiler-produced action against its registered API.""" + capability = self.action(template.atomic_action_class) + kind = template.target_binding.get("kind") + if kind not in capability.target_binding_kinds: + raise ValueError( + f"{template.atomic_action_class} does not accept target binding " + f"kind {kind!r}; expected one of " + f"{sorted(capability.target_binding_kinds)}." + ) + if template.control not in capability.controls: + raise ValueError( + f"{template.atomic_action_class} does not support control " + f"{template.control!r}; expected one of " + f"{sorted(capability.controls)}." + ) diff --git a/embodichain/gen_sim/action_engine/cli/__init__.py b/embodichain/gen_sim/action_engine/cli/__init__.py new file mode 100644 index 000000000..564654c85 --- /dev/null +++ b/embodichain/gen_sim/action_engine/cli/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Command-line entry points for Action Engine generation and execution.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py new file mode 100644 index 000000000..542dc1c70 --- /dev/null +++ b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py @@ -0,0 +1,240 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""CLI for generating Action Engine configs from a Prompt2Scene gym export.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from embodichain.gen_sim.action_engine.config import generation_defaults +from embodichain.gen_sim.action_engine.generation import ( + generate_action_engine_config, +) + +__all__ = ["build_parser", "cli"] + +_GENERATION_DEFAULTS = generation_defaults() +_TASK_DEFAULTS = _GENERATION_DEFAULTS["task"] +_SCENE_DEFAULTS = _GENERATION_DEFAULTS["scene"] + +_ROBOT_PROFILE_CHOICES = ( + "ur5", + "ur10", + "dual_ur5", + "dual_ur10", + "franka", + "dual_franka", +) + + +def build_parser() -> argparse.ArgumentParser: + """Build the standalone config-generation argument parser.""" + parser = argparse.ArgumentParser( + description=( + "Plan and compile an Action Engine task from an exported tabletop " + "gym project." + ) + ) + parser.add_argument( + "--gym_project", + "--gym-project", + required=True, + help=( + "Prompt2Scene task/export directory or gym_config.json/" + "scene_config.json path." + ), + ) + parser.add_argument( + "--output_dir", + "--output-dir", + required=True, + help="Directory receiving canonical JSON artifacts and the Seed PNG.", + ) + parser.add_argument( + "--task_name", + "--task-name", + required=True, + help="Stable task identifier stored in both programs.", + ) + parser.add_argument( + "--task_description", + "--task-description", + help="Natural-language goal passed to structured LLM interpretation.", + ) + parser.add_argument( + "--task_file", + "--task-file", + help="Optional UTF-8 file containing the natural-language goal.", + ) + parser.add_argument( + "--task-spec", + "--task_spec", + dest="task_spec", + help=( + "Optional existing Action Engine v2 TaskSpec JSON; bypasses text " + "LLM interpretation and uses its role_bindings hand-off." + ), + ) + parser.add_argument( + "--robot-profile", + "--robot_profile", + choices=_ROBOT_PROFILE_CHOICES, + default=str(_TASK_DEFAULTS["default_robot_profile"]), + help="Robot template used in fast_gym_config.json.", + ) + parser.add_argument( + "--llm_model", + "--llm-model", + default=None, + help="Optional planner model override.", + ) + parser.add_argument( + "--vlm_model", + "--vlm-model", + default=None, + help="Optional online visual/planner model override stored for A/B runs.", + ) + parser.add_argument( + "--planning-mode", + "--planning_mode", + choices=("offline", "ab"), + default="offline", + help="Generate one offline bundle or an offline/online A/B bundle.", + ) + parser.add_argument( + "--source_scene_z_rotation_degrees", + "--source-scene-z-rotation-degrees", + type=float, + default=None, + help=( + "World-frame scene rotation. Prompt2Scene exports default to -90 " + "degrees; other inputs default to zero." + ), + ) + parser.add_argument( + "--body-scale-policy", + choices=("preserve", "multiply", "absolute"), + default=str(_SCENE_DEFAULTS["body_scale_policy"]), + help="How the requested xyz scale combines with source body_scale.", + ) + parser.add_argument( + "--body-scale", + type=float, + nargs=3, + default=tuple(float(value) for value in _SCENE_DEFAULTS["body_scale"]), + metavar=("X", "Y", "Z"), + help="Positive xyz scale used by multiply or absolute policy.", + ) + parser.add_argument( + "--max_episodes", + "--max-episodes", + type=int, + default=int(_TASK_DEFAULTS["max_episodes"]), + help="Episode count written to fast_gym_config.json.", + ) + parser.add_argument( + "--max_episode_steps", + "--max-episode-steps", + type=int, + default=int(_TASK_DEFAULTS["max_episode_steps"]), + help="Per-episode step limit written to fast_gym_config.json.", + ) + parser.add_argument( + "--overwrite", + action="store_true", + help="Replace existing canonical artifacts in the output directory.", + ) + parser.add_argument( + "--randomize-scene", + action="store_true", + help="Randomize rigid-object poses and table height on every reset.", + ) + parser.add_argument( + "--randomize-table-material", + action="store_true", + help="Randomize the table material independently on every reset.", + ) + return parser + + +def cli() -> None: + """Generate and report the canonical Action Engine artifact bundle.""" + args = build_parser().parse_args() + task_description = _resolve_task_description(args) + paths = generate_action_engine_config( + args.gym_project, + args.output_dir, + task_name=args.task_name, + task_description=task_description, + task_spec=args.task_spec, + robot_profile=args.robot_profile, + llm_model=args.llm_model, + source_scene_z_rotation_degrees=args.source_scene_z_rotation_degrees, + body_scale_policy=args.body_scale_policy, + body_scale=args.body_scale, + overwrite=args.overwrite, + max_episodes=args.max_episodes, + max_episode_steps=args.max_episode_steps, + randomize_scene=args.randomize_scene, + randomize_table_material=args.randomize_table_material, + planning_mode=args.planning_mode, + vlm_model=args.vlm_model, + ) + + print(f"Generated gym config: {paths.gym_config}") + print(f"Generated agent config: {paths.agent_config}") + print(f"Generated TaskSpec: {paths.task_spec}") + print(f"Generated SceneRequirements: {paths.scene_requirements}") + print(f"Generated SeedGraph: {paths.seed_task_graph}") + print(f"Generated Seed graph PNG: {paths.seed_task_graph_png}") + print( + "Run with:\n" + "python -m embodichain.gen_sim.action_engine.cli.run_agent " + f"--task_name {args.task_name} " + f'--gym_config "{paths.gym_config}" ' + f'--agent_config "{paths.agent_config}" ' + "--regenerate" + ) + + +def _resolve_task_description(args: argparse.Namespace) -> str: + task_spec = getattr(args, "task_spec", None) + if task_spec: + if args.task_description or args.task_file: + raise ValueError( + "--task-spec cannot be combined with --task_description or " + "--task_file." + ) + return "" + if args.task_description and args.task_file: + raise ValueError("Use either --task_description or --task_file, not both.") + if args.task_file: + description = ( + Path(args.task_file).expanduser().read_text(encoding="utf-8").strip() + ) + else: + description = str(args.task_description or "").strip() + if not description: + raise ValueError( + "--task_description (or --task_file) must provide a non-empty goal." + ) + return description + + +if __name__ == "__main__": + cli() diff --git a/embodichain/gen_sim/action_engine/cli/run_agent.py b/embodichain/gen_sim/action_engine/cli/run_agent.py new file mode 100644 index 000000000..13cffe627 --- /dev/null +++ b/embodichain/gen_sim/action_engine/cli/run_agent.py @@ -0,0 +1,1620 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Run a generated Action Engine configuration.""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable, Mapping +from copy import deepcopy +from dataclasses import dataclass, replace +from datetime import datetime, timezone +import json +import multiprocessing as mp +import os +from pathlib import Path +import shutil +from types import SimpleNamespace +from typing import Any + +import gymnasium +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.config import generation_defaults +from embodichain.gen_sim.action_engine.environment import ( # noqa: F401 + ACTION_ENGINE_ENV_ID, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + load_agent_execution_program, + write_execution_report, +) +from embodichain.lab.gym.utils.gym_utils import ( + add_env_launcher_args_to_parser, + build_env_cfg_from_args, +) +from embodichain.utils import set_seed +from embodichain.utils.logger import log_info, log_warning +from embodichain.utils.utility import load_config + +__all__ = ["build_parser", "cli"] + +_DEFAULT_MAX_EPISODES = int(generation_defaults()["task"]["max_episodes"]) + + +def build_parser() -> argparse.ArgumentParser: + """Build the command-line parser used by generated demo commands.""" + parser = argparse.ArgumentParser(description="Execute an Action Engine task agent.") + add_env_launcher_args_to_parser(parser) + parser.add_argument("--task_name", required=True, help="Generated task name.") + parser.add_argument( + "--agent_config", + required=True, + help="Path to action_engine_config_v2 JSON.", + ) + parser.add_argument( + "--regenerate", + action="store_true", + help="Rebuild SeedGraph from TaskSpec in memory before execution.", + ) + parser.add_argument( + "--show-physical-collision", + action="store_true", + help="Show physical collision geometry after every reset.", + ) + parser.add_argument( + "--seed", + type=int, + default=None, + help="Base random seed; episode N uses seed + N.", + ) + parser.add_argument( + "--runtime-backend", + choices=("independent",), + default="independent", + help="Execution backend. Action Engine owns the production runtime.", + ) + parser.add_argument( + "--vlm-model", + default=None, + help="Optional runtime override for A/B visual facts and online planning.", + ) + parser.add_argument( + "--task-engine-report", + action="store_true", + help=argparse.SUPPRESS, + ) + return parser + + +def _validate_gym_id(config: dict[str, Any]) -> None: + if config.get("id") != ACTION_ENGINE_ENV_ID: + raise ValueError( + f"Gym config id must be {ACTION_ENGINE_ENV_ID!r}, " + f"got {config.get('id')!r}." + ) + + +def _validate_run_contract( + gym_config: dict[str, Any], + agent_config: dict[str, Any], + task_name: str, +) -> None: + """Validate the small cross-artifact contract before simulator startup.""" + configured_task = agent_config.get("task_name") + if configured_task != task_name: + raise ValueError( + f"--task_name {task_name!r} does not match agent_config task " + f"{configured_task!r}." + ) + extension = gym_config.get("env", {}).get("extensions", {}).get("action_engine", {}) + if extension.get("task_name") != task_name: + raise ValueError("Gym and agent configs describe different tasks.") + gym_hash = extension.get("seed_task_graph_hash") + agent_hash = agent_config.get("seed_task_graph_hash") + if not isinstance(agent_hash, str) or not agent_hash or gym_hash != agent_hash: + raise ValueError("Gym and agent configs have different program hashes.") + agent_mode = str(agent_config.get("planning_mode", "offline")) + gym_mode = str(extension.get("planning_mode", "offline")) + if agent_mode != gym_mode: + raise ValueError( + f"Gym and agent configs have different planning modes: " + f"gym={gym_mode!r}, agent={agent_mode!r}." + ) + + +def cli() -> int | None: + """Launch the environment and execute all configured episodes.""" + np.set_printoptions(precision=5, suppress=True) + torch.set_printoptions(precision=5, sci_mode=False) + args = build_parser().parse_args() + if args.seed is not None: + set_seed(args.seed) + env_cfg, gym_config, _ = build_env_cfg_from_args(args) + if args.seed is not None: + env_cfg.seed = args.seed + _validate_gym_id(gym_config) + agent_config = load_config(args.agent_config) + if not isinstance(agent_config, dict): + raise ValueError("agent_config must contain a JSON object.") + _validate_run_contract(gym_config, agent_config, args.task_name) + planning_mode = str(agent_config.get("planning_mode", "offline")) + if planning_mode == "ab": + _run_ab( + args, + env_cfg=env_cfg, + gym_config=gym_config, + agent_config=agent_config, + ) + return 0 if args.task_engine_report else None + if planning_mode != "offline": + raise ValueError(f"Unsupported Action Engine planning_mode {planning_mode!r}.") + execution_program = load_agent_execution_program( + agent_config, + agent_config_path=args.agent_config, + regenerate=bool(args.regenerate), + ) + grounded_plan = _load_grounded_task_plan(args.agent_config) + action_reporter = None + if grounded_plan is not None: + from embodichain.gen_sim.action_engine.agent import ActionAgent + + action_reporter = ActionAgent() + + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + episodes = int(gym_config.get("max_episodes", _DEFAULT_MAX_EPISODES)) + runtime_arguments = { + "agent_config": str(Path(args.agent_config).expanduser().resolve()), + "base_seed": args.seed, + "gym_config": str(Path(args.gym_config).expanduser().resolve()), + "max_episodes": episodes, + "planning_mode": planning_mode, + "regenerate": bool(args.regenerate), + "runtime_backend": str(args.runtime_backend), + "task_name": str(args.task_name), + } + any_failed = False + task_engine_reports: list[ExecutionReport] = [] + episode_index = 0 + episode_seed = None + seed_graph = getattr(execution_program, "seed_graph", None) + env = None + try: + env = gymnasium.make( + id=gym_config["id"], + cfg=env_cfg, + agent_config=agent_config, + agent_config_path=args.agent_config, + task_name=args.task_name, + runtime_backend=args.runtime_backend, + ) + for episode_index in range(episodes): + episode_seed = None if args.seed is None else int(args.seed) + episode_index + env.reset(seed=episode_seed) + if args.show_physical_collision: + _show_physical_collision(env) + execute = env.get_wrapper_attr("create_demo_action_list") + result = execute( + regenerate=bool(args.regenerate), + runtime_run_id=run_id, + episode_index=episode_index, + ) + if not getattr(result, "already_executed", False): + raise RuntimeError( + "Action Engine env returned an offline action sequence." + ) + success = torch.as_tensor( + getattr(result, "runtime_success"), + dtype=torch.bool, + ) + any_failed = any_failed or not bool(success.all()) + log_info( + "Action Engine episode " + f"{episode_index}: {int(success.sum())}/{success.numel()} " + "environments succeeded.", + color="green", + ) + record_dir = getattr(result, "runtime_graph_output_dir", None) + if record_dir: + log_info(f"Runtime records: {record_dir}", color="green") + if action_reporter is not None and isinstance(seed_graph, Mapping): + report = action_reporter.report_execution_result( + result, + action_graph=seed_graph, + grounded_plan=grounded_plan, + run_id=run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + task_engine_reports.append(report) + _publish_task_engine_report( + args.agent_config, + report, + enabled=bool(args.task_engine_report), + ) + log_info( + "Execution report: " + f"status={report.status}, actions={report.action_count}", + color="green" if report.status == "succeeded" else "yellow", + ) + # EmbodiedEnv publishes the just-finished rollout during reset. Flush + # the final episode as well; otherwise only episodes followed by a next + # iteration reach the configured dataset recorder. + env.reset(options={"final": True}) + except KeyboardInterrupt: + log_warning("Action Engine run interrupted by user.") + return 130 if args.task_engine_report else None + except Exception as exc: + if action_reporter is not None and isinstance(seed_graph, Mapping): + report = action_reporter.abortion_report( + seed_graph, + exc, + grounded_plan=grounded_plan, + environment_count=_runtime_environment_count(env), + run_id=run_id, + episode_index=episode_index, + episode_seed=episode_seed, + runtime_arguments=runtime_arguments, + ) + write_execution_report(Path(args.agent_config).resolve().parent, report) + if args.task_engine_report: + log_warning(f"Action Engine execution aborted: {type(exc).__name__}: {exc}") + return 3 + raise + finally: + close = getattr(env, "close", None) if env is not None else None + if callable(close): + close() + if not args.task_engine_report: + return None + return _task_engine_exit_code(any_failed, task_engine_reports) + + +def _publish_task_engine_report( + agent_config_path: str | Path, + report: ExecutionReport, + *, + enabled: bool, +) -> Path | None: + """Mirror one normal execution report into its Task Engine bundle.""" + if not enabled: + return None + return write_execution_report(Path(agent_config_path).resolve().parent, report) + + +def _task_engine_exit_code( + any_failed: bool, + reports: list[ExecutionReport], +) -> int: + """Return a report-authoritative exit code for Task Engine execution.""" + return int( + bool(any_failed) or any(report.status != "succeeded" for report in reports) + ) + + +def _load_grounded_task_plan(agent_config_path: str | Path) -> dict[str, Any] | None: + """Load the optional Task Engine hand-off beside an agent config.""" + path = ( + Path(agent_config_path).expanduser().resolve().parent + / "grounded_task_plan.json" + ) + if not path.is_file(): + return None + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Unable to read GroundedTaskPlan at {path}: {exc}") from exc + if not isinstance(value, Mapping): + raise ValueError("grounded_task_plan.json must contain a JSON object.") + from embodichain.gen_sim.task_engine.orchestration.contracts import ( + validate_grounded_task_plan, + ) + + return validate_grounded_task_plan(value) + + +def _runtime_environment_count(env: Any) -> int: + value = getattr(getattr(env, "unwrapped", env), "num_envs", 1) + try: + return max(1, int(value)) + except (TypeError, ValueError): + return 1 + + +class _BranchExecutor: + def __init__( + self, + graph: dict[str, Any], + env: gymnasium.Env, + *, + record_root: Path, + ) -> None: + self.graph = graph + self.env = env + self.record_root = record_root + + def preflight(self) -> bool: + """Compile and capability-check the branch without sending motion.""" + route = getattr(self.env.unwrapped, "action_engine_ab_route", None) + if route in {"offline", "online"} and self.graph.get("planner_route") != route: + raise ValueError( + f"A/B branch route {route!r} cannot execute graph route " + f"{self.graph.get('planner_route')!r}." + ) + try: + preflight = self.env.get_wrapper_attr("preflight_seed_graph") + except AttributeError: + preflight = None + if callable(preflight): + value = preflight(self.graph) + return value is not False + # Older generated environments expose only execute_seed_graph. The + # loader is still a useful structural/capability preflight and does + # not step the simulator. + from embodichain.gen_sim.action_engine.runtime import load_execution_program + + source = self.env.unwrapped.agent_config.get("source") + if source is None: + source = {} + if not isinstance(source, dict): + raise ValueError("agent_config.source must be a mapping when provided.") + uid_map = source.get("uid_map", {}) + if not isinstance(uid_map, dict): + raise ValueError( + "agent_config.source.uid_map must be a mapping when provided." + ) + known_objects = {str(uid) for uid in uid_map.values()} + load_execution_program(self.graph, known_objects=known_objects or None) + return True + + def run(self, *, run_id: str, episode_index: int) -> Any: + execute = self.env.get_wrapper_attr("execute_seed_graph") + return execute( + self.graph, + runtime_run_id=run_id, + episode_index=episode_index, + record_root=self.record_root.as_posix(), + ) + + +@dataclass(frozen=True) +class _ABWorkerConfig: + """Serializable startup contract for one process-isolated A/B branch.""" + + route: str + gym_config: dict[str, Any] + env_options: dict[str, Any] + gym_id: str + agent_config: dict[str, Any] + agent_config_path: str + task_name: str + runtime_backend: str + seed: int + camera_uids: tuple[str, ...] + staging_dir: str + + +class _ABBranchWorker: + """Small RPC proxy for one simulator process. + + DexSim entities resolve through a process-global default world. Keeping + each branch in a separate process is therefore a correctness requirement, + not merely a way to parallelize A/B execution. + """ + + _STARTUP_TIMEOUT_SECONDS = 300.0 + _COMMAND_TIMEOUT_SECONDS = 1800.0 + _SHUTDOWN_TIMEOUT_SECONDS = 30.0 + + def __init__(self, config: _ABWorkerConfig) -> None: + self.action_engine_ab_route = config.route + self._config = config + self._closed = False + self._context = mp.get_context("spawn") + self._connection, child_connection = self._context.Pipe(duplex=True) + self._process = self._context.Process( + target=_ab_worker_main, + args=(child_connection, config), + name=f"action-engine-ab-{config.route}", + ) + try: + self._process.start() + except BaseException: + child_connection.close() + self._connection.close() + raise + child_connection.close() + try: + startup = self._receive( + "startup", timeout_seconds=self._STARTUP_TIMEOUT_SECONDS + ) + except Exception: + self.close() + raise + if not isinstance(startup, dict): + self.close() + raise RuntimeError( + f"A/B {config.route} worker returned an invalid startup payload." + ) + snapshot = startup.get("snapshot") + if not isinstance(snapshot, dict): + self.close() + raise RuntimeError( + f"A/B {config.route} worker did not return its reset snapshot." + ) + self.startup_snapshot = snapshot + self.startup_observation = startup.get("observation") + + def snapshot(self) -> dict[str, Any]: + value = self._request("snapshot") + if not isinstance(value, dict): + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned an invalid snapshot." + ) + return value + + def preflight(self, graph: dict[str, Any]) -> bool: + value = self._request("preflight", graph=graph) + return value is not False + + def run( + self, + graph: dict[str, Any], + *, + run_id: str, + episode_index: int, + record_root: Path, + ) -> Any: + value = self._request( + "run", + graph=graph, + run_id=run_id, + episode_index=int(episode_index), + record_root=record_root.as_posix(), + ) + if not isinstance(value, dict): + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned an invalid result." + ) + return _execution_result_from_wire(value) + + def finalize(self, branch_dir: Path, *, episode_index: int) -> list[str]: + value = self._request( + "finalize", + branch_dir=branch_dir.as_posix(), + episode_index=int(episode_index), + ) + if not isinstance(value, list) or not all( + isinstance(path, str) and path for path in value + ): + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned invalid video paths." + ) + return value + + def close(self) -> None: + """Ask the worker to clean up, then force-stop only if it is stuck.""" + if self._closed: + return + self._closed = True + try: + if self._process.is_alive(): + try: + self._connection.send({"op": "shutdown"}) + self._receive( + "shutdown", timeout_seconds=self._SHUTDOWN_TIMEOUT_SECONDS + ) + except Exception: + # The process is still joined/terminated below. Cleanup + # errors cannot justify leaking a simulator child. + pass + finally: + try: + self._connection.close() + finally: + self._process.join(timeout=self._SHUTDOWN_TIMEOUT_SECONDS) + if self._process.is_alive(): + self._process.terminate() + self._process.join(timeout=self._SHUTDOWN_TIMEOUT_SECONDS) + + def _request(self, operation: str, **payload: Any) -> Any: + if self._closed: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker is already closed." + ) + try: + self._connection.send({"op": operation, **payload}) + except (BrokenPipeError, EOFError, OSError) as exc: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker could not receive " + f"{operation!r}." + ) from exc + return self._receive(operation, timeout_seconds=self._COMMAND_TIMEOUT_SECONDS) + + def _receive(self, operation: str, *, timeout_seconds: float) -> Any: + try: + ready = self._connection.poll(timeout_seconds) + except (EOFError, OSError) as exc: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker closed during {operation}." + ) from exc + if not ready: + exit_code = self._process.exitcode + if exit_code is not None: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker exited with code " + f"{exit_code} during {operation}." + ) + raise TimeoutError( + f"A/B {self.action_engine_ab_route} worker timed out during " + f"{operation} after {timeout_seconds:.0f}s." + ) + try: + response = self._connection.recv() + except (EOFError, OSError) as exc: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker closed during {operation}." + ) from exc + if not isinstance(response, dict) or "ok" not in response: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker returned a malformed " + f"response during {operation}." + ) + if response["ok"] is True: + return response.get("value") + message = response.get("error") + if not isinstance(message, str) or not message: + message = "unknown worker error" + raise RuntimeError( + f"A/B {self.action_engine_ab_route} worker failed during {operation}: " + f"{message}" + ) + + +class _SerializedABBranch: + """Run one branch in fresh isolated workers when two worlds do not fit. + + Each worker still owns a separate DexSim process and is reset from the + exact same seed. The proxy only serializes their GPU residency: it probes + a reset for planning, starts a fresh worker for preflight, then starts one + more fresh worker for execution. Every startup digest must match the + planning reset before an RPC is allowed to progress. + """ + + def __init__( + self, + config: _ABWorkerConfig, + *, + startup_snapshot: Mapping[str, Any], + startup_observation: Any, + expected_initial_state_digest: str, + worker_factory: Callable[[_ABWorkerConfig], Any] | None = None, + ) -> None: + self.action_engine_ab_route = config.route + self.startup_snapshot = deepcopy(dict(startup_snapshot)) + self.startup_observation = startup_observation + self._config = config + self._expected_initial_state_digest = expected_initial_state_digest + self._worker_factory = worker_factory or _ABBranchWorker + self._active_worker: Any | None = None + self._closed = False + + def snapshot(self) -> dict[str, Any]: + """Return the verified reset snapshot without rehydrating a GPU world.""" + return deepcopy(self.startup_snapshot) + + def preflight(self, graph: dict[str, Any]) -> bool: + worker = self._start_worker("preflight") + try: + return worker.preflight(graph) + finally: + worker.close() + + def run( + self, + graph: dict[str, Any], + *, + run_id: str, + episode_index: int, + record_root: Path, + ) -> Any: + if self._active_worker is not None: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} execution worker is already active." + ) + worker = self._start_worker("execute") + self._active_worker = worker + return worker.run( + graph, + run_id=run_id, + episode_index=episode_index, + record_root=record_root, + ) + + def finalize(self, branch_dir: Path, *, episode_index: int) -> list[str]: + worker = self._active_worker + if worker is None: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} has no execution worker to finalize." + ) + try: + return worker.finalize(branch_dir, episode_index=episode_index) + finally: + self._active_worker = None + worker.close() + + def close(self) -> None: + self._closed = True + worker = self._active_worker + self._active_worker = None + if worker is not None: + worker.close() + + def _start_worker(self, phase: str) -> Any: + if self._closed: + raise RuntimeError( + f"A/B {self.action_engine_ab_route} serialized branch is closed." + ) + worker = self._worker_factory(_ab_phase_worker_config(self._config, phase)) + try: + from embodichain.gen_sim.action_engine.evaluation import state_digest + + snapshot = worker.startup_snapshot + actual_digest = state_digest(snapshot) + if actual_digest != self._expected_initial_state_digest: + raise RuntimeError( + "Strict A/B serialized reset mismatch before " + f"{phase}: route={self.action_engine_ab_route}, " + f"expected={self._expected_initial_state_digest}, " + f"actual={actual_digest}." + ) + return worker + except BaseException: + worker.close() + raise + + +def _ab_phase_worker_config(config: _ABWorkerConfig, phase: str) -> _ABWorkerConfig: + """Give serial lifecycle phases distinct recorder and dataset roots.""" + if not phase: + return config + staging_dir = Path(config.staging_dir) + return replace( + config, + staging_dir=(staging_dir.parent / phase / staging_dir.name).as_posix(), + ) + + +def _prepare_ab_branches( + configs: Mapping[str, _ABWorkerConfig], + *, + worker_factory: Callable[[_ABWorkerConfig], Any] = _ABBranchWorker, + prefer_serial: bool | None = None, + gpu_id: int | None = None, +) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """Start concurrent worlds, with a digest-checked serialized fallback. + + A renderer can consume several GiB per DexSim process. On smaller GPUs, + starting the second isolated branch may fail before any action is sent. In + that case keeping one world resident is not a semantic requirement, while + the reset digest is; use fresh one-at-a-time workers instead. + """ + if prefer_serial is None: + prefer_serial = _prefer_serial_ab_startup(gpu_id=gpu_id) + if prefer_serial: + log_warning( + "A/B GPU capacity is below the concurrent-world budget; " + "using serialized isolated workers with reset-digest checks." + ) + return _prepare_serial_ab_branches(configs, worker_factory=worker_factory) + + workers: dict[str, Any] = {} + try: + for route in ("offline", "online"): + workers[route] = worker_factory(configs[route]) + except Exception as error: + for worker in workers.values(): + worker.close() + if not _is_gpu_memory_error(error): + raise + log_warning( + "A/B concurrent simulator startup exhausted GPU memory; " + "using serialized isolated workers with reset-digest checks." + ) + return _prepare_serial_ab_branches(configs, worker_factory=worker_factory) + return ( + workers, + {route: worker.startup_snapshot for route, worker in workers.items()}, + ) + + +def _prepare_serial_ab_branches( + configs: Mapping[str, _ABWorkerConfig], + *, + worker_factory: Callable[[_ABWorkerConfig], Any], +) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]: + """Probe one branch at a time and return lazy serialized branch proxies.""" + + snapshots: dict[str, dict[str, Any]] = {} + observations: dict[str, Any] = {} + for route in ("offline", "online"): + worker = worker_factory(_ab_phase_worker_config(configs[route], "probe")) + try: + snapshots[route] = worker.startup_snapshot + observations[route] = worker.startup_observation + finally: + worker.close() + from embodichain.gen_sim.action_engine.evaluation import state_digest + + expected_digest = state_digest(snapshots["offline"]) + return ( + { + route: _SerializedABBranch( + configs[route], + startup_snapshot=snapshots[route], + startup_observation=observations[route], + expected_initial_state_digest=expected_digest, + worker_factory=worker_factory, + ) + for route in ("offline", "online") + }, + snapshots, + ) + + +def _is_gpu_memory_error(error: BaseException) -> bool: + """Recognize the process-startup failures where serialization is safe.""" + message = str(error).lower() + memory_markers = ( + "out of memory", + "out_of_memory", + "out_of_device_memory", + "outofmemory", + "resource exhausted", + ) + return any(marker in message for marker in memory_markers) and ( + "cuda" in message + or "gpu" in message + or "vulkan" in message + or "device" in message + ) + + +def _prefer_serial_ab_startup(*, gpu_id: int | None = None) -> bool: + """Avoid a known OOM trial on GPUs too small for two renderer worlds.""" + if not torch.cuda.is_available(): + return False + try: + device = torch.device(f"cuda:{int(gpu_id)}" if gpu_id is not None else "cuda") + free, _ = torch.cuda.mem_get_info(device=device) + except (RuntimeError, ValueError): + return False + # One hybrid DexSim world with the four VLM cameras can occupy roughly + # 11--13 GiB on the supported RTX setup. Reserve 24 GiB for two worlds; + # larger cards still attempt concurrent startup and retain the OOM fallback + # for unusually heavy scenes. + return int(free) < 24 * 1024**3 + + +class _RemoteBranchExecutor: + """Executor adapter which keeps simulator calls inside the branch worker.""" + + def __init__( + self, + graph: dict[str, Any], + worker: Any, + *, + record_root: Path, + ) -> None: + self.graph = graph + self.worker = worker + self.record_root = record_root + + def preflight(self) -> bool: + return self.worker.preflight(self.graph) + + def run(self, *, run_id: str, episode_index: int) -> Any: + return self.worker.run( + self.graph, + run_id=run_id, + episode_index=episode_index, + record_root=self.record_root, + ) + + +def _ab_worker_main(connection: Any, config: _ABWorkerConfig) -> None: + """Create and drive exactly one real environment in a child process.""" + # SimulationManager otherwise exits the whole worker with os._exit(0) + # during environment cleanup, bypassing the artifact/RPC shutdown contract. + os.environ["EMBODICHAIN_SIM_EXIT_PROCESS"] = "0" + env: gymnasium.Env | None = None + try: + from embodichain.lab.gym.utils.gym_utils import ( + config_to_cfg, + get_manager_modules, + ) + + # ``config_to_cfg`` creates local component-config classes which are + # intentionally not picklable. Send the merged JSON contract over IPC + # and reconstruct it inside each worker instead of pickling ``env_cfg``. + branch_cfg = config_to_cfg( + deepcopy(config.gym_config), + manager_modules=get_manager_modules(), + ) + _apply_ab_env_options(branch_cfg, config.env_options) + branch_cfg.seed = int(config.seed) + _configure_ab_branch_cfg( + branch_cfg, + staging_dir=Path(config.staging_dir), + dataset_dir=Path(config.staging_dir).parent / ".dataset", + ) + set_seed(int(config.seed)) + env = gymnasium.make( + id=config.gym_id, + cfg=branch_cfg, + agent_config=deepcopy(config.agent_config), + agent_config_path=config.agent_config_path, + task_name=config.task_name, + runtime_backend=config.runtime_backend, + ) + setattr(env.unwrapped, "action_engine_ab_route", config.route) + env.reset(seed=int(config.seed)) + startup: dict[str, Any] = { + "snapshot": _snapshot_environment(env, list(config.camera_uids)), + } + if config.route == "online": + from embodichain.gen_sim.action_engine.planning import ( + collect_scene_observation, + ) + + startup["observation"] = collect_scene_observation( + env.unwrapped, + camera_uids=config.camera_uids, + env_id=0, + ) + # The recorder normally receives its first frame from an interval + # event during ``env.step``. Capture one reset-time, no-motion frame so + # an execution branch that fails before its first action still has a + # valid video artifact after the mandatory final reset. + _capture_ab_initial_frame(env) + _worker_send(connection, ok=True, value=startup) + while True: + try: + request = connection.recv() + except EOFError: + break + if not isinstance(request, dict): + raise ValueError("A/B worker request must be a mapping.") + operation = request.get("op") + try: + if operation == "snapshot": + value = _snapshot_environment(env, list(config.camera_uids)) + elif operation == "preflight": + graph = _worker_graph(request.get("graph")) + value = _BranchExecutor( + graph, + env, + record_root=Path(config.staging_dir).parent / "runtime", + ).preflight() + elif operation == "run": + graph = _worker_graph(request.get("graph")) + run_id = request.get("run_id") + record_root = request.get("record_root") + if not isinstance(run_id, str) or not run_id: + raise ValueError( + "A/B worker run_id must be a non-empty string." + ) + if not isinstance(record_root, str) or not record_root: + raise ValueError( + "A/B worker record_root must be a non-empty path string." + ) + result = _BranchExecutor( + graph, + env, + record_root=Path(record_root).expanduser().resolve(), + ).run( + run_id=run_id, + episode_index=int(request.get("episode_index", 0)), + ) + value = _execution_result_to_wire(result) + elif operation == "finalize": + branch_dir = request.get("branch_dir") + if not isinstance(branch_dir, str) or not branch_dir: + raise ValueError( + "A/B worker branch_dir must be a non-empty path string." + ) + value = _finalize_ab_branch_video( + env, + staging_dir=Path(config.staging_dir), + branch_dir=Path(branch_dir).expanduser().resolve(), + ) + elif operation == "shutdown": + _worker_send(connection, ok=True, value=True) + break + else: + raise ValueError(f"Unknown A/B worker operation {operation!r}.") + except BaseException as exc: + _worker_send(connection, ok=False, error=_worker_error(exc)) + except BaseException as exc: + _worker_send(connection, ok=False, error=_worker_error(exc)) + finally: + if env is not None: + try: + env.close() + except BaseException: + pass + try: + from embodichain.lab.sim import SimulationManager + + SimulationManager.flush_cleanup_queue() + except BaseException: + pass + try: + connection.close() + except OSError: + pass + + +def _worker_graph(value: Any) -> dict[str, Any]: + if not isinstance(value, dict): + raise ValueError("A/B worker SeedGraph must be a JSON object.") + return value + + +def _worker_send( + connection: Any, + *, + ok: bool, + value: Any = None, + error: str | None = None, +) -> None: + try: + payload: dict[str, Any] = {"ok": bool(ok)} + if ok: + payload["value"] = value + else: + payload["error"] = error or "unknown worker error" + connection.send(payload) + except (BrokenPipeError, EOFError, OSError): + pass + + +def _worker_error(error: BaseException) -> str: + return f"{type(error).__name__}: {error}" + + +def _capture_ab_initial_frame(env: gymnasium.Env) -> None: + """Append one audience-camera frame without advancing simulation state.""" + base = env.unwrapped + manager = getattr(base, "event_manager", None) + mode_cfgs = getattr(manager, "_mode_functor_cfgs", {}) + candidates: list[tuple[Any, dict[str, Any]]] = [] + for configured in mode_cfgs.values(): + for functor_cfg in configured: + functor = _config_member(functor_cfg, "func") + class_name = getattr(type(functor), "__name__", "") + if not callable(functor) or class_name not in { + "record_camera_data", + "record_camera_data_async", + }: + continue + params = _config_member(functor_cfg, "params") or {} + if not isinstance(params, Mapping): + raise ValueError("A/B record_camera params must be a mapping.") + params = dict(params) + if params.get("name") == "record_cam_audience_view": + candidates.insert(0, (functor, params)) + else: + candidates.append((functor, params)) + if candidates: + # Prefer the explicitly generated audience recorder. A single + # unnamed legacy recorder remains a compatible fallback; selecting + # among multiple non-audience recorders would silently produce the + # wrong camera view, so fail instead. + if len(candidates) > 1 and candidates[0][1].get("name") != ( + "record_cam_audience_view" + ): + raise RuntimeError( + "A/B environment has multiple camera recorders but none is " + "named 'record_cam_audience_view'." + ) + functor, params = candidates[0] + functor(base, None, **params) + return + raise RuntimeError( + "A/B environment must define a record_camera_data audience recorder." + ) + + +def _execution_result_to_wire(result: Any) -> dict[str, Any]: + """Strip simulator-owned state from an execution result before IPC.""" + + actions = [_wire_tensor(action) for action in list(getattr(result, "actions", ()))] + success = _wire_tensor(getattr(result, "success", False)) + return { + "actions": actions, + "success": success, + "record_dir": getattr(result, "record_dir", None), + "already_executed": bool(getattr(result, "already_executed", True)), + "retry_count": int(getattr(result, "retry_count", 0)), + "recovery_count": int(getattr(result, "recovery_count", 0)), + "revision_count": int(getattr(result, "revision_count", 0)), + "failure_events": list(getattr(result, "failure_events", ())), + "runtime_revisions": list(getattr(result, "runtime_revisions", ())), + } + + +def _wire_tensor(value: Any) -> Any: + if isinstance(value, torch.Tensor): + return value.detach().cpu() + return value + + +def _execution_result_from_wire(value: dict[str, Any]) -> SimpleNamespace: + required = { + "actions", + "success", + "retry_count", + "recovery_count", + "revision_count", + "failure_events", + "runtime_revisions", + } + missing = sorted(required - set(value)) + if missing: + raise RuntimeError(f"A/B worker result is missing fields: {missing}.") + return SimpleNamespace(**value) + + +def _run_ab( + args: argparse.Namespace, + *, + env_cfg: Any, + gym_config: dict[str, Any], + agent_config: dict[str, Any], +) -> None: + """Plan and execute strict offline/online branches for every episode.""" + from embodichain.gen_sim.action_engine.evaluation import run_strict_ab, state_digest + from embodichain.gen_sim.action_engine.planning import ( + plan_candidates_parallel, + plan_online_seed_graph, + ) + from embodichain.gen_sim.action_engine.generation import VLM_CAMERA_UIDS + + config_path = Path(args.agent_config).expanduser().resolve() + task_path = _resolve_artifact_path( + agent_config, + config_path, + "task_spec", + "task_spec_path", + ) + task_spec = _read_json(task_path, "TaskSpec") + reference_program = load_agent_execution_program( + agent_config, + agent_config_path=config_path, + regenerate=bool(getattr(args, "regenerate", False)), + require_executable=False, + ) + if reference_program.seed_graph is None: + raise ValueError("A/B execution requires an immutable offline SeedGraph.") + reference_graph = reference_program.seed_graph + source = agent_config.get("source") + if not isinstance(source, dict): + source = {} + uid_map = source.get("uid_map", {}) + if not isinstance(uid_map, dict): + raise ValueError("agent_config.source.uid_map must be a mapping when provided.") + known_objects = {str(uid) for uid in uid_map.values() if str(uid)} + online_config = agent_config.get("online_planning", {}) + if online_config is None: + online_config = {} + if not isinstance(online_config, dict): + raise ValueError("agent_config.online_planning must be a mapping.") + camera_uids = online_config.get("camera_uids") or agent_config.get( + "vlm_camera_uids", [] + ) + if camera_uids != list(VLM_CAMERA_UIDS): + raise ValueError( + "A/B execution requires the canonical VLM cameras " + f"{list(VLM_CAMERA_UIDS)}." + ) + vlm_model = ( + getattr(args, "vlm_model", None) + or online_config.get("vlm_model") + or agent_config.get("vlm_model") + ) + robot_profile = str(agent_config.get("robot_profile", "dual_ur10")) + base_seed = 0 if args.seed is None else int(args.seed) + if args.seed is None: + set_seed(base_seed) + run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + output_root = config_path.parent / "ab_runs" / run_id + episodes = int(gym_config.get("max_episodes", _DEFAULT_MAX_EPISODES)) + summaries = [] + env_options = _ab_env_options(env_cfg) + + for episode_index in range(episodes): + episode_seed = base_seed + episode_index + episode_root = output_root / f"episode_{episode_index:04d}" + branch_envs: dict[str, Any] = {} + ownership_transferred = False + try: + # The online VLM observes the same reset that its branch executes. + # Each branch owns one simulator process because DexSim entities + # resolve through a process-global default world. + worker_gym_config = _ab_runtime_gym_config(gym_config, env_cfg) + worker_configs = { + route: _ABWorkerConfig( + route=route, + gym_config=worker_gym_config, + env_options=deepcopy(env_options), + gym_id=str(gym_config["id"]), + agent_config=agent_config, + agent_config_path=config_path.as_posix(), + task_name=args.task_name, + runtime_backend=getattr(args, "runtime_backend", "independent"), + seed=episode_seed, + camera_uids=tuple(str(uid) for uid in camera_uids), + staging_dir=(episode_root / ".work" / route / "video").as_posix(), + ) + for route in ("offline", "online") + } + branch_envs, snapshots = _prepare_ab_branches( + worker_configs, + gpu_id=getattr(getattr(env_cfg, "sim_cfg", None), "gpu_id", None), + ) + planning_digest = state_digest(snapshots["offline"]) + if planning_digest != state_digest(snapshots["online"]): + raise RuntimeError( + "Strict A/B initial state mismatch before planning: " + f"offline={planning_digest}, " + f"online={state_digest(snapshots['online'])}." + ) + observation = branch_envs["online"].startup_observation + if observation is None: + raise RuntimeError( + "Online A/B worker did not return scene observation." + ) + visual_facts: dict[str, Any] = {} + + def offline_planner(*, task_spec: dict[str, Any]) -> dict[str, Any]: + # Generation has already materialized the fixed recipe from the + # shared TaskSpec. Reuse that immutable artifact verbatim so A/B + # also supports legacy v2 bundles whose graph metadata predates + # ``role_bindings``. + del task_spec + return deepcopy(reference_graph) + + def online_planner(*, task_spec: dict[str, Any]) -> dict[str, Any]: + graph, facts = plan_online_seed_graph( + task_spec, + observation, + vlm_model=vlm_model, + robot_profile=robot_profile, + ) + visual_facts.update(facts) + return graph + + candidates = plan_candidates_parallel( + task_spec, + offline_planner=offline_planner, + online_planner=online_planner, + known_objects=known_objects or None, + robot_profile=robot_profile, + ) + if candidates.offline != reference_graph: + from embodichain.gen_sim.action_engine.domain import seed_graph_hash + + if seed_graph_hash(candidates.offline) != seed_graph_hash( + reference_graph + ): + raise RuntimeError( + "A/B offline recipe no longer matches the generated " + "reference graph." + ) + + # Visual evidence is an online-planning artifact, not an execution + # artifact. Persist it as soon as both candidates have passed their + # static checks so it remains auditable even if a later preflight, + # runtime action, or video flush fails. + _write_json(episode_root / "online" / "visual_facts.json", visual_facts) + + # Rendering and external planning must be side-effect free. Take + # a second full snapshot immediately before executor preflight so + # an accidental simulation advance cannot be hidden behind the + # reset-time digest used for visual planning. + snapshots = { + route: worker.snapshot() for route, worker in branch_envs.items() + } + execution_digests = { + route: state_digest(snapshot) for route, snapshot in snapshots.items() + } + if ( + execution_digests["offline"] != planning_digest + or execution_digests["online"] != planning_digest + ): + raise RuntimeError( + "Strict A/B initial state changed during visual planning: " + f"offline={execution_digests['offline']}, " + f"online={execution_digests['online']}, " + f"expected={planning_digest}." + ) + + def executor_factory(graph: dict[str, Any], worker: Any) -> Any: + route = worker.action_engine_ab_route + if route not in branch_envs: + raise ValueError(f"Unknown A/B worker route {route!r}.") + return _RemoteBranchExecutor( + graph, + worker, + record_root=episode_root / ".work" / route / "runtime", + ) + + def branch_finalizer(**kwargs: Any) -> list[str]: + worker = kwargs["env"] + branch_dir = Path(kwargs["branch_dir"]) + return worker.finalize( + branch_dir, + episode_index=int(kwargs.get("episode_index", episode_index)), + ) + + result = run_strict_ab( + task_spec, + candidates.offline, + candidates.online, + executor_factory=executor_factory, + snapshot_reader=lambda env: _snapshot_environment(env, camera_uids), + output_dir=episode_root, + seed=episode_seed, + shared_config={ + "robot_profile": robot_profile, + "camera_uids": camera_uids, + "vlm_model": vlm_model, + "strict_state_digest": True, + }, + planning_metrics=candidates.planning_metrics, + known_objects=known_objects or None, + expected_initial_state_digest=planning_digest, + branch_finalizer=branch_finalizer, + episode_index=episode_index, + strict_state_digest=True, + prepared_environments=branch_envs, + prepared_snapshots=snapshots, + require_branch_videos=True, + ) + # run_strict_ab owns and closes prepared workers on both its normal + # and exceptional execution paths. Do not claim ownership until + # it has entered/returned from that cleanup boundary; this also + # closes workers when graph validation fails before its try/finally. + ownership_transferred = True + finally: + if not ownership_transferred: + for worker in branch_envs.values(): + worker.close() + summaries.append( + { + "episode_index": episode_index, + "seed": episode_seed, + "comparison": result.comparison_path.as_posix(), + "initial_state_digest": result.initial_state_digest, + } + ) + log_info( + "Action Engine A/B episode " + f"{episode_index}: offline=" + f"{result.comparison['branches']['offline']['success_rate']:.3f}, " + f"online={result.comparison['branches']['online']['success_rate']:.3f}.", + color="green", + ) + + summary_path = output_root / "run_summary.json" + _write_json( + summary_path, + { + "schema_version": "action_engine_ab_run_v1", + "task_id": args.task_name, + "run_id": run_id, + "episodes": summaries, + }, + ) + log_info(f"A/B comparison artifacts: {output_root}", color="green") + + +def _configure_ab_branch_cfg( + env_cfg: Any, + *, + staging_dir: Path, + dataset_dir: Path, +) -> None: + """Give one worker exclusive recorder paths before env construction.""" + staging_dir.mkdir(parents=True, exist_ok=True) + dataset_dir.mkdir(parents=True, exist_ok=True) + events = _config_member(env_cfg, "events") + recorder = _config_member(events, "record_camera") + if recorder is None: + raise ValueError("A/B environment config must define record_camera.") + _set_config_param(recorder, "save_path", staging_dir.as_posix()) + + # Dataset output is not part of the A/B contract, but leaving the + # generated path shared would still let two workers overwrite each other. + dataset = _config_member(env_cfg, "dataset") + if dataset is not None: + for name in ("lerobot", "record", "dataset"): + term = _config_member(dataset, name) + if term is not None: + _set_config_param(term, "save_path", dataset_dir.as_posix()) + + +def _ab_runtime_gym_config( + gym_config: Mapping[str, Any], env_cfg: Any +) -> dict[str, Any]: + """Carry launcher-resolved simulation settings into spawned workers.""" + result = deepcopy(dict(gym_config)) + sim_cfg = getattr(env_cfg, "sim_cfg", None) + if sim_cfg is None: + return result + result.update( + { + "device": str(getattr(sim_cfg, "sim_device", "cpu")), + "gpu_id": int(getattr(sim_cfg, "gpu_id", 0)), + "headless": bool(getattr(sim_cfg, "headless", False)), + "arena_space": float(getattr(sim_cfg, "arena_space", 5.0)), + "num_envs": int(getattr(sim_cfg, "num_envs", result.get("num_envs", 1))), + } + ) + render_cfg = getattr(sim_cfg, "render_cfg", None) + renderer = getattr(render_cfg, "renderer", None) + if renderer is not None: + result["renderer"] = str(renderer) + return result + + +def _ab_env_options(env_cfg: Any) -> dict[str, Any]: + """Extract the non-JSON flags applied after gym config parsing.""" + profiler = getattr(env_cfg, "profiler", None) + return { + "filter_visual_rand": bool(getattr(env_cfg, "filter_visual_rand", False)), + "filter_dataset_saving": bool(getattr(env_cfg, "filter_dataset_saving", False)), + "record_trajectory": bool(getattr(env_cfg, "record_trajectory", False)), + "trajectory_save_dir": getattr(env_cfg, "trajectory_save_dir", None), + "profile": bool(getattr(profiler, "enable_time", False)), + "profile_output": getattr(profiler, "output_path", None), + } + + +def _apply_ab_env_options(env_cfg: Any, options: Mapping[str, Any]) -> None: + """Apply launcher flags after reconstructing a worker's config.""" + env_cfg.filter_visual_rand = bool(options.get("filter_visual_rand", False)) + env_cfg.filter_dataset_saving = bool(options.get("filter_dataset_saving", False)) + env_cfg.record_trajectory = bool(options.get("record_trajectory", False)) + trajectory_dir = options.get("trajectory_save_dir") + if trajectory_dir: + env_cfg.trajectory_save_dir = str(trajectory_dir) + if bool(options.get("profile", False)): + from embodichain.lab.gym.utils.profiler import EnvProfilerCfg + + env_cfg.profiler = EnvProfilerCfg( + enable_time=True, + output_path=options.get("profile_output"), + ) + + +def _config_member(value: Any, name: str) -> Any: + if isinstance(value, dict): + return value.get(name) + return getattr(value, name, None) if value is not None else None + + +def _set_config_param(term: Any, name: str, value: Any) -> None: + params = _config_member(term, "params") + if params is None: + params = {} + if isinstance(term, dict): + term["params"] = params + else: + setattr(term, "params", params) + if not isinstance(params, dict): + raise ValueError( + f"A/B recorder params must be a mapping, got {type(params)!r}." + ) + params[name] = value + + +def _finalize_ab_branch_video( + env: gymnasium.Env, + *, + staging_dir: Path, + branch_dir: Path, +) -> list[str]: + """Flush the final episode and publish exactly this worker's video.""" + before = { + path: (path.stat().st_mtime_ns, path.stat().st_size) + for path in staging_dir.glob("episode_*_record_cam_audience_view.mp4") + if path.is_file() + } + env.reset(options={"final": True}) + video = _publish_branch_video(staging_dir, branch_dir, before=before) + return [video.as_posix()] + + +def _snapshot_environment( + env: gymnasium.Env, + camera_uids: list[str], +) -> dict[str, Any]: + base = env.unwrapped + sim = base.sim + object_poses = {} + for uid in sim.get_rigid_object_uid_list(): + entity = sim.get_rigid_object(uid) + if entity is not None: + object_poses[str(uid)] = _snapshot_tensor( + entity.get_local_pose(to_matrix=True) + ) + articulation_state = {} + for uid in getattr(sim, "get_articulation_uid_list", lambda: [])(): + entity = sim.get_articulation(uid) + if entity is None: + continue + articulation_state[str(uid)] = { + "pose": _snapshot_tensor(entity.get_local_pose(to_matrix=True)), + "qpos": _snapshot_tensor(entity.get_qpos()), + } + camera_calibration = {} + available_sensor_uids = getattr(sim, "get_sensor_uid_list", lambda: [])() + snapshot_camera_uids = sorted( + {str(uid) for uid in [*camera_uids, *available_sensor_uids] if str(uid)} + ) + for uid in snapshot_camera_uids: + sensor = sim.get_sensor(uid) + if sensor is None: + raise ValueError(f"A/B snapshot cannot find camera {uid!r}.") + camera_calibration[uid] = { + "intrinsics": _snapshot_sensor_value(sensor, "get_intrinsics"), + "extrinsics": _snapshot_sensor_value( + sensor, "get_arena_pose", to_matrix=True + ), + } + return { + "robot_qpos": _snapshot_tensor(base.robot.get_qpos()), + "object_poses": object_poses, + "articulation_state": articulation_state, + "camera_calibration": camera_calibration, + } + + +def _snapshot_tensor(value: Any) -> torch.Tensor: + """Normalize simulator values for deterministic digesting.""" + tensor = torch.as_tensor(value) + return tensor.detach().cpu().contiguous() + + +def _snapshot_sensor_value( + sensor: Any, + method_name: str, + **kwargs: Any, +) -> torch.Tensor: + method = getattr(sensor, method_name, None) + if not callable(method): + raise ValueError(f"A/B snapshot sensor lacks {method_name}().") + try: + value = method(**kwargs) + except TypeError: + value = method() + return _snapshot_tensor(value) + + +def _publish_branch_video( + staging_dir: Path, + branch_dir: Path, + *, + before: dict[Path, tuple[int, int]] | None = None, +) -> Path: + candidates = sorted( + ( + path + for path in staging_dir.glob("episode_*_record_cam_audience_view.mp4") + if path.is_file() + and ( + before is None + or path not in before + or (path.stat().st_mtime_ns, path.stat().st_size) != before[path] + ) + ), + key=lambda path: path.stat().st_mtime_ns, + ) + if not candidates: + raise RuntimeError(f"No completed A/B audience video found in {staging_dir}.") + source = candidates[-1] + destination = branch_dir / "video.mp4" + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + if destination.stat().st_size == 0: + raise RuntimeError(f"A/B audience video is empty: {destination}.") + return destination + + +def _resolve_artifact_path( + config: dict[str, Any], + config_path: Path, + *keys: str, +) -> Path: + for key in keys: + value = config.get(key) + if value is None: + continue + if not isinstance(value, str) or not value: + raise ValueError(f"agent_config.{key} must be a non-empty path string.") + path = Path(value).expanduser() + return ( + path.resolve() + if path.is_absolute() + else (config_path.parent / path).resolve() + ) + joined = " or ".join(f"agent_config.{key}" for key in keys) + raise ValueError(f"A/B execution requires {joined}.") + + +def _read_json(path: Path, label: str) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Unable to read {label} at {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"{label} must contain a JSON object.") + return value + + +def _write_json(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def _show_physical_collision(env: gymnasium.Env) -> None: + """Enable physical-shape visualization for all supported scene assets.""" + sim = env.get_wrapper_attr("sim") + uids: list[str] = [] + for getter_name in ( + "get_rigid_object_uid_list", + "get_rigid_object_group_uid_list", + "get_articulation_uid_list", + ): + getter = getattr(sim, getter_name, None) + if callable(getter): + uids.extend(getter()) + visible = 0 + for uid in uids: + asset = sim.get_asset(uid) + if asset is None or not hasattr(asset, "set_physical_visible"): + continue + try: + asset.set_physical_visible( + visible=True, + rgba=[1.0, 0.15, 0.1, 0.35], + ) + visible += 1 + except Exception as exc: + log_warning(f"Unable to show collision geometry for {uid!r}: {exc}") + log_info(f"Physical collision geometry visible for {visible} assets.") + + +if __name__ == "__main__": + raise SystemExit(cli()) diff --git a/embodichain/gen_sim/action_engine/compiler/__init__.py b/embodichain/gen_sim/action_engine/compiler/__init__.py new file mode 100644 index 000000000..fbd36a858 --- /dev/null +++ b/embodichain/gen_sim/action_engine/compiler/__init__.py @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stable deterministic compiler API.""" + +from __future__ import annotations + +from .core import compile_task_agent +from .v2 import ( + compile_task_agent_v2, + execution_program_to_seed_graph, + seed_graph_to_execution_program, +) + +__all__ = [ + "compile_task_agent", + "compile_task_agent_v2", + "execution_program_to_seed_graph", + "seed_graph_to_execution_program", +] diff --git a/embodichain/gen_sim/action_engine/compiler/core.py b/embodichain/gen_sim/action_engine/compiler/core.py new file mode 100644 index 000000000..aa1a36617 --- /dev/null +++ b/embodichain/gen_sim/action_engine/compiler/core.py @@ -0,0 +1,594 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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) + for content in goal.get("contents", []): + value = content.get("object") if isinstance(content, Mapping) else content + if isinstance(value, str): + result.add(value) + return result + + +def _ancestor_sets( + semantic_steps: Sequence[Mapping[str, Any]], +) -> dict[str, set[str]]: + direct = {step["id"]: set(step["depends_on"]) for step in semantic_steps} + ancestors: dict[str, set[str]] = {} + for step in semantic_steps: + pending = list(direct[step["id"]]) + result: set[str] = set() + while pending: + dependency = pending.pop() + if dependency in result: + continue + result.add(dependency) + pending.extend(direct[dependency]) + ancestors[step["id"]] = result + return ancestors + + +def _starts_with_pickup(phases: Sequence[PhaseTemplate]) -> bool: + return bool( + phases + and phases[0].actions + and phases[0].actions[0].atomic_action_class == "PickUp" + ) + + +def _actors_request_distinct_arms( + first: Mapping[str, Any], + second: Mapping[str, Any], +) -> bool: + """Return whether actors explicitly request a distinct-arm assignment.""" + first_group = first.get("allocation_group") + same_group = first_group is not None and first_group == second.get( + "allocation_group" + ) + required_opposite = ( + first["mode"] == "required" + and second["mode"] == "required" + and _canonical_arm(first["arm"]) != _canonical_arm(second["arm"]) + ) + if same_group and not required_opposite: + both_required = first["mode"] == second["mode"] == "required" + if both_required: + raise ValueError( + f"Allocation group {first_group!r} requires distinct arms, " + "but both steps require the same arm." + ) + return same_group or required_opposite + + +def _canonical_arm(value: Any) -> str: + arm = str(value) + return f"{arm}_arm" if arm in {"left", "right"} else arm + + +def _stable_topological_steps( + semantic_steps: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + original = [deepcopy(dict(step)) for step in semantic_steps] + order = {step["id"]: index for index, step in enumerate(original)} + by_id = {step["id"]: step for step in original} + indegree = {step["id"]: len(step["depends_on"]) for step in original} + dependents: dict[str, list[str]] = {step["id"]: [] for step in original} + for step in original: + for dependency in step["depends_on"]: + dependents[dependency].append(step["id"]) + + ready = deque( + sorted( + (step_id for step_id, degree in indegree.items() if degree == 0), + key=order.__getitem__, + ) + ) + result: list[dict[str, Any]] = [] + while ready: + step_id = ready.popleft() + result.append(by_id[step_id]) + newly_ready: list[str] = [] + for dependent in dependents[step_id]: + indegree[dependent] -= 1 + if indegree[dependent] == 0: + newly_ready.append(dependent) + ready.extend(sorted(newly_ready, key=order.__getitem__)) + return result + + +def _slug(value: Any) -> str: + slug = _UNSAFE_ID_RE.sub("_", str(value).lower()).strip("_") + return slug[:64].rstrip("_") or "step" diff --git a/embodichain/gen_sim/action_engine/compiler/v2.py b/embodichain/gen_sim/action_engine/compiler/v2.py new file mode 100644 index 000000000..13942afe7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/compiler/v2.py @@ -0,0 +1,424 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Bridge mature v1 task recipes to the direct AtomicAction SeedGraph v3.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +import re +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, + capability_precondition, +) +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + MOTION_POLICY_VERSION, + validate_execution_program, + validate_seed_graph, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA +from embodichain.gen_sim.action_engine.planning.linker import ( + link_seed_graph, + validate_persisted_contracts, +) + +__all__ = [ + "compile_task_agent_v2", + "execution_program_to_seed_graph", + "seed_graph_to_execution_program", +] + +_UNSAFE_ID_RE = re.compile(r"[^0-9a-z]+") +_OPERATOR_TASK_TYPES = { + "arrange_line": "E1", + "build_stack": "E1", + "coordinated_place": "E5", + "coordinated_transport": "E5", + "hold_hover": "E1", + "orient_object": "E2", + "place_in_line": "E1", + "place_relative": "E1", + "press": "E9", +} + + +def compile_task_agent_v2( + program: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Compile a mature semantic recipe directly to the persisted v3 graph.""" + from .core import compile_task_agent + + legacy = compile_task_agent(program, known_objects=known_objects) + return execution_program_to_seed_graph( + legacy, + known_objects=known_objects, + registry=registry, + ) + + +def execution_program_to_seed_graph( + program: Mapping[str, Any], + *, + planner_route: str = "offline", + known_objects: Collection[str] | None = None, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Convert a mature v1 result without changing its AtomicAction topology.""" + legacy = validate_execution_program(program) + capabilities = registry or build_atomic_capability_registry() + steps = {str(step["id"]): step for step in legacy["semantic_steps"]} + node_ids_by_edge: dict[str, list[str]] = {} + nodes: list[dict[str, Any]] = [] + + for edge in legacy["edges"]: + edge_id = str(edge["id"]) + step = steps[str(edge["semantic_step_id"])] + task_type = _task_type(str(step["operator"])) + dependencies = [ + node_id + for dependency in edge.get("depends_on", []) + for node_id in node_ids_by_edge[str(dependency)] + ] + edge_nodes: list[str] = [] + actions = list(edge["actions"]) + for action_index, action in enumerate(actions): + action_name = str(action["atomic_action_class"]) + descriptor_view = { + "atomic_action": action_name, + "control": action.get("control", "arm"), + "target_binding": action["target_binding"], + } + capabilities.validate_binding(descriptor_view) + capability = capabilities.get(action_name) + node_id = _node_id(edge_id, action_name, action_index, len(actions)) + postcondition = ( + deepcopy(step["postcondition"]) + if edge_id == step["edge_ids"][-1] + else {} + ) + node = { + "id": node_id, + "atomic_action": action_name, + "object_uid": str(step["object"]), + "actor": _v2_actor(action["actor"]), + "control": str(action.get("control", "arm")), + "target_binding": deepcopy(dict(action["target_binding"])), + "depends_on": list(dict.fromkeys(dependencies)), + "task_instance_id": str(step["id"]), + "task_type": task_type, + "role": _node_role(action_name, action["target_binding"]), + "precondition": capability_precondition( + capability, + object_uid=str(step["object"]), + actor=_v2_actor(action["actor"]), + target_binding=action["target_binding"], + ), + "postcondition": postcondition, + "motion_policy": deepcopy(dict(action["motion_policy"])), + } + if len(actions) > 1: + node["sync_group"] = edge_id + nodes.append(node) + edge_nodes.append(node_id) + node_ids_by_edge[edge_id] = edge_nodes + + groups = [] + for step in legacy["semantic_steps"]: + group_node_ids = [ + node_id + for edge_id in step["edge_ids"] + for node_id in node_ids_by_edge[str(edge_id)] + ] + groups.append( + { + "id": str(step["id"]), + "task_type": _task_type(str(step["operator"])), + "role": "primary", + "operator": str(step["operator"]), + "object_uid": str(step["object"]), + "actor": _v2_actor(step["actor"]), + "goal": deepcopy(dict(step.get("goal", {}))), + "depends_on": list(step.get("depends_on", [])), + "parent_task_instance_id": str(step.get("parent_step_id", step["id"])), + "node_ids": group_node_ids, + "success": deepcopy(dict(step["postcondition"])), + } + ) + + level = _level(groups) + graph = { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": str(legacy["task"]), + "instruction": str(legacy["goal_description"]), + "level": level, + "reasoning_type": "none", + "planner_route": planner_route, + "nodes": nodes, + "task_groups": groups, + "success": { + "op": "all", + "terms": [deepcopy(group["success"]) for group in groups], + }, + "capability_catalog_hash": capabilities.catalog_hash(), + "metadata": { + "source_schema": EXECUTION_PROGRAM_SCHEMA, + "legacy_allocation_groups": deepcopy(legacy.get("allocation_groups", [])), + "planning_latency_seconds": 0.0, + "vlm_call_count": 0, + }, + } + return link_seed_graph( + graph, + registry=capabilities, + task_order=[str(step["id"]) for step in legacy["semantic_steps"]], + known_objects=known_objects, + ) + + +def seed_graph_to_execution_program( + graph: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, + registry: AtomicCapabilityRegistry | None = None, + require_executable: bool = True, +) -> dict[str, Any]: + """Materialize the v3 DAG as the existing in-memory runtime view.""" + capabilities = registry or build_atomic_capability_registry() + seed = validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=require_executable, + ) + if seed["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError( + "SeedGraph capability_catalog_hash does not match the runtime catalog." + ) + validate_persisted_contracts(seed, capabilities) + for node in seed["nodes"]: + capabilities.validate_binding(node) + + node_by_id = {str(node["id"]): node for node in seed["nodes"]} + unit_by_node, units = _execution_units(seed["nodes"]) + ordered_units = _topological_units(units) + edge_id_by_unit = {unit_id: f"edge_{_slug(unit_id)}" for unit_id in ordered_units} + target_by_unit = { + unit_id: f"state_{index + 1:04d}_{_slug(unit_id)}" + for index, unit_id in enumerate(ordered_units) + } + start = "state_start" + edges = [] + graph_nodes = [{"id": start, "semantic": "Initial live simulator state"}] + for unit_id in ordered_units: + unit = units[unit_id] + dependencies = sorted(unit["depends_on"]) + source = start if not dependencies else target_by_unit[dependencies[0]] + target = target_by_unit[unit_id] + graph_nodes.append( + { + "id": target, + "semantic": f"Completed AtomicAction unit {unit_id}", + } + ) + unit_nodes = [node_by_id[node_id] for node_id in unit["node_ids"]] + edges.append( + { + "id": edge_id_by_unit[unit_id], + "source": source, + "target": target, + "semantic_step_id": str(unit_nodes[0]["task_instance_id"]), + "actions": [ + { + "atomic_action_class": node["atomic_action"], + "actor": deepcopy(node["actor"]), + "control": node["control"], + "target_binding": deepcopy(node["target_binding"]), + "motion_policy": node["motion_policy"], + "seed_node_id": node["id"], + "failure_policy": node["contract"]["failure_policy"], + } + for node in unit_nodes + ], + "depends_on": [edge_id_by_unit[item] for item in dependencies], + "resources": sorted( + { + str(claim["resource"]) + for node in unit_nodes + for claim in node["contract"]["claims"] + } + ), + } + ) + + group_by_id = {str(group["id"]): group for group in seed["task_groups"]} + semantic_steps = [] + for group_id in _topological_groups(seed["task_groups"]): + group = group_by_id[group_id] + group_units = [ + unit_id + for unit_id in ordered_units + if any( + node_by_id[node_id]["task_instance_id"] == group_id + for node_id in units[unit_id]["node_ids"] + ) + ] + semantic_steps.append( + { + "id": group_id, + "parent_step_id": str(group.get("parent_task_instance_id", group_id)), + "operator": str(group["operator"]), + "object": str(group["object_uid"]), + "actor": deepcopy(group["actor"]), + "goal": deepcopy(group["goal"]), + "depends_on": list(group["depends_on"]), + "postcondition": deepcopy(group["success"]), + "edge_ids": [edge_id_by_unit[unit_id] for unit_id in group_units], + } + ) + + metadata = seed.get("metadata", {}) + allocation_groups = ( + deepcopy(metadata.get("legacy_allocation_groups", [])) + if isinstance(metadata, Mapping) + else [] + ) + program = { + "schema_version": EXECUTION_PROGRAM_SCHEMA, + "task": seed["task_id"], + "goal_description": seed["instruction"], + "start": start, + "goal": target_by_unit[ordered_units[-1]], + "nodes": graph_nodes, + "edges": edges, + "semantic_steps": semantic_steps, + "allocation_groups": allocation_groups, + "motion_policy_version": MOTION_POLICY_VERSION, + } + return validate_execution_program(program) + + +def _execution_units( + nodes: Sequence[Mapping[str, Any]], +) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: + unit_by_node = { + str(node["id"]): str(node.get("sync_group", node["id"])) for node in nodes + } + units: dict[str, dict[str, Any]] = {} + for node in nodes: + node_id = str(node["id"]) + unit_id = unit_by_node[node_id] + unit = units.setdefault(unit_id, {"node_ids": [], "depends_on": set()}) + unit["node_ids"].append(node_id) + for dependency in node["depends_on"]: + dependency_unit = unit_by_node[str(dependency)] + if dependency_unit == unit_id: + raise ValueError( + f"Synchronized unit {unit_id!r} has an internal dependency." + ) + unit["depends_on"].add(dependency_unit) + for unit_id, unit in units.items(): + groups = { + str( + next(node for node in nodes if node["id"] == node_id)[ + "task_instance_id" + ] + ) + for node_id in unit["node_ids"] + } + if len(groups) != 1: + raise ValueError(f"Synchronized unit {unit_id!r} crosses task groups.") + return unit_by_node, units + + +def _topological_units(units: Mapping[str, Mapping[str, Any]]) -> list[str]: + return _topological_ids( + {unit_id: list(unit["depends_on"]) for unit_id, unit in units.items()} + ) + + +def _topological_groups(groups: Sequence[Mapping[str, Any]]) -> list[str]: + return _topological_ids( + {str(group["id"]): list(group["depends_on"]) for group in groups} + ) + + +def _topological_ids(dependencies: Mapping[str, Sequence[str]]) -> list[str]: + outgoing = {item_id: [] for item_id in dependencies} + indegree = {item_id: 0 for item_id in dependencies} + for item_id, parents in dependencies.items(): + for parent in parents: + outgoing[parent].append(item_id) + indegree[item_id] += 1 + ready = deque( + sorted(item_id for item_id, degree in indegree.items() if degree == 0) + ) + ordered = [] + while ready: + item_id = ready.popleft() + ordered.append(item_id) + for child in sorted(outgoing[item_id]): + indegree[child] -= 1 + if indegree[child] == 0: + ready.append(child) + if len(ordered) != len(dependencies): + raise ValueError("Graph contains a dependency cycle.") + return ordered + + +def _v2_actor(value: Mapping[str, Any]) -> dict[str, Any]: + actor = deepcopy(dict(value)) + actor.pop("allocation_group", None) + if actor.get("mode") == "required" and actor.get("arm") in {"left", "right"}: + actor["arm"] = f"{actor['arm']}_arm" + return actor + + +def _task_type(operator: str) -> str: + return _OPERATOR_TASK_TYPES.get(operator, "E1") + + +def _level(groups: Sequence[Mapping[str, Any]]) -> str: + types = {str(group["task_type"]) for group in groups} + if len(groups) == 1: + return "L1" + return "L2" if len(types) == 1 else "L3" + + +def _node_role(action_name: str, binding: Mapping[str, Any]) -> str: + if ( + action_name == "MoveJoints" and binding.get("source") == "initial" + ) or binding.get("kind") == "policy_pose": + return "cleanup" + return "primary" + + +def _node_id(edge_id: str, action: str, index: int, count: int) -> str: + base = f"{_slug(edge_id)}_{_slug(action)}" + return base if count == 1 else f"{base}_{index + 1}" + + +def _slug(value: str) -> str: + return _UNSAFE_ID_RE.sub("_", value.lower()).strip("_") or "node" diff --git a/embodichain/gen_sim/action_engine/config/__init__.py b/embodichain/gen_sim/action_engine/config/__init__.py new file mode 100644 index 000000000..d9759d702 --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/__init__.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Validated package policy for Action Engine generation and runtime.""" + +from __future__ import annotations + +from .runtime_policy import ( + ACTION_ENGINE_DEFAULTS_SCHEMA, + RUNTIME_POLICY_SCHEMA, + ArmSelectionPolicyCfg, + RuntimePolicyCfg, + default_runtime_policy, + generation_defaults, + resolve_agent_runtime_policy, + runtime_policy_hash, +) + +__all__ = [ + "ACTION_ENGINE_DEFAULTS_SCHEMA", + "RUNTIME_POLICY_SCHEMA", + "ArmSelectionPolicyCfg", + "RuntimePolicyCfg", + "default_runtime_policy", + "generation_defaults", + "resolve_agent_runtime_policy", + "runtime_policy_hash", +] diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml new file mode 100644 index 000000000..117c76bac --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -0,0 +1,344 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +schema_version: action_engine_defaults_v1 + +# Generation policy is materialized into fast_gym_config.json. It is not part +# of the coordinate-free Execution Program or the runtime-policy hash. +generation: + task: + default_robot_profile: ur10 + max_episodes: 1 + max_episode_steps: 2000 + environment: + viewer_camera_uid: cam_high + ignore_terminations_during_agent: true + recording: + enabled: true + resolution: [640, 360] + interval_step: 1 + arm_aim_yaw_offset: + left: 0.0 + right: 0.0 + scene: + prompt2scene_z_rotation_degrees: -90.0 + default_tabletop_z: 0.7 + body_scale_policy: preserve + body_scale: [1.0, 1.0, 1.0] + object_length_sample_points: 5000 + physics: + background: + mass: 10.0 + static_friction: 0.95 + dynamic_friction: 0.9 + restitution: 0.01 + max_convex_hull_num: 1 + rigid_object: + mass: 0.1 + static_friction: 0.95 + dynamic_friction: 0.9 + linear_damping: 0.9 + angular_damping: 0.9 + contact_offset: 0.003 + rest_offset: 0.001 + restitution: 0.05 + max_depenetration_velocity: 0.8 + max_linear_velocity: 5.0 + max_angular_velocity: 5.0 + min_position_iters: 32 + min_velocity_iters: 8 + max_convex_hull_num: 16 + acd_method: vhacd + randomization: + rigid_object_position_range: [[-0.04, -0.04, 0.0], [0.04, 0.04, 0.0]] + rigid_object_rotation_range: [[0.0, 0.0, -30.0], [0.0, 0.0, 30.0]] + table_height_delta_range: [[-0.05], [0.05]] + table_material: + random_texture_prob: 0.0 + base_color_range: [[0.55, 0.55, 0.55], [0.95, 0.95, 0.95]] + metallic_range: [0.0, 0.15] + roughness_range: [0.45, 0.95] + dataset: + control_frequency: 25 + save_failed_episodes: true + use_videos: true + +# Runtime policy is resolved per robot profile, snapshotted in agent_config, +# hash-verified at startup, and recorded with every execution. +runtime: + common: + execution: + max_transitions: 1000 + semantic_step_settle_steps: 10 + max_retries_per_action: 2 + max_graph_revisions: 8 + max_recovery_actions: 12 + support_stability_samples: 3 + support_stability_interval_steps: 5 + support_linear_velocity_tolerance: 0.02 + support_angular_velocity_tolerance: 0.20 + + planner: + backend: curobo + single_arm_strategy: motion_gen + coordinated_strategy: ik_interp + fallback_strategy: ik_interp + allow_fallback: true + dynamic_collision: false + static_obstacle_uids: [] + dynamic_obstacle_uids: [] + curobo: + log_level: error + obstacle_representation: cuboid + multi_env: false + use_cuda_graph: true + preserve_plan_samples: false + max_attempts: 5 + collision_activation_distance: 0.01 + + # Crossing is measured along the live right-to-left arm-base axis so the + # same-side constraint follows translated and rotated robot workspaces. + arm_selection: + crossing_deadband_ratio: 0.08 + allow_cross_side_fallback: false + pickup_crossing_weight: 1.0 + placement_crossing_weight: 1.5 + motion_cost_scale: 3.141592653589793 + fallback_workspace_half_width: 0.5 + orient_object_preferred_arm_deadband: 0.02 + + grounding: + semantic_defaults: + surface_clearance: 0.003 + transport_clearance: 0.10 + staging_lift_height: 0.12 + relation_distance: 0.16 + hover_height: 0.10 + press_depth: 0.004 + retreat_height: 0.10 + maximum_eef_height: 0.80 + arrangement: + slot_margin: 0.08 + minimum_spacing: 0.07 + layout_clearance: 0.025 + row_search_step: 0.025 + row_search_radius: 0.25 + placement: + clearance: 0.012 + candidate_count: 5 + candidate_offset_fraction: 0.50 + support_margin: 0.002 + recovery_attempts: 2 + coordinated_grasp: + inset_fraction: 0.15 + minimum_inset: 0.01 + handover: + retreat_height: 0.10 + retreat_distance: 0.10 + maximum_eef_height: 1.50 + minimum_transfer_clearance: 0.10 + minimum_transfer_lateral_clearance: 0.06 + joint_state: + hand_close_sample_interval: 10 + hand_open_sample_interval: 15 + + grasp: + antipodal_n_sample: 10000 + antipodal_max_angle: 0.2617993877991494 + max_open_length: 0.15 + min_open_length: 0.01 + finger_length: 0.13 + point_sample_dense: 0.012 + max_deviation_angle: 0.3490658503988659 + n_deviated_approach_directions: 4 + viser_port: 11801 + max_decomposition_hulls: 16 + force_grasp_reannotate: false + + motion_defaults: + PickUp: + pre_grasp_distance: 0.08 + lift_height: 0.30 + sample_interval: 45 + MoveHeldObject: + sample_interval: 45 + relation_distance: 0.18 + robot_relative_distance: 0.10 + relation_clearance: 0.02 + exchange_clearance: 0.06 + exchange_candidate_offset: 0.16 + exchange_obstacle_clearance: 0.04 + exchange_gripper_horizontal_envelope: 0.035 + exchange_wrist_horizontal_envelope: 0.055 + exchange_gripper_vertical_envelope: 0.025 + exchange_wrist_vertical_envelope: 0.04 + exchange_minimum_reach: 0.10 + exchange_maximum_reach: 1.00 + exchange_candidate_count: 4 + hover_height: 0.10 + line_spacing: 0.14 + transport_clearance: 0.10 + staging_lift_height: 0.30 + surface_clearance: 0.005 + postcondition_tolerance: 0.08 + line_axis_tolerance: 0.06 + line_perpendicular_tolerance: 0.06 + preserve_orientation_tolerance: 0.2617993877991494 + Place: + sample_interval: 15 + lift_height: 0.0 + post_hold_steps: 0 + cartesian_waypoint_count: 2 + MoveEndEffector: + sample_interval: 20 + retreat_height: 0.30 + minimum_retreat_height: 0.05 + maximum_eef_height: 1.10 + postcondition_tolerance: 0.05 + MoveJoints: + sample_interval: 30 + postcondition_tolerance: 0.05 + Press: + sample_interval: 80 + press_depth: 0.004 + postcondition_tolerance: 0.03 + CoordinatedPickment: + sample_interval: 120 + object_motion_keyframes: 6 + pre_grasp_distance: 0.10 + lift_height: 0.08 + middle_empty_ratio: 0.4 + is_filter_ground_collision: false + postcondition_tolerance: 0.06 + HandOver: + sample_interval: 140 + pre_grasp_distance: 0.08 + lift_height: 0.08 + receiver_hold_joint_tolerance: 0.002 + receive_pick_object_part: bottom + exchange_clearance: 0.06 + exchange_candidate_offset: 0.16 + exchange_obstacle_clearance: 0.04 + exchange_gripper_horizontal_envelope: 0.035 + exchange_wrist_horizontal_envelope: 0.055 + exchange_gripper_vertical_envelope: 0.025 + exchange_wrist_vertical_envelope: 0.04 + exchange_minimum_reach: 0.10 + exchange_maximum_reach: 1.00 + exchange_candidate_count: 4 + held_position_tolerance: 0.03 + hand_interp_steps: 10 + hold_steps: 4 + retreat_steps: 28 + postcondition_tolerance: 0.06 + CoordinatedPlacement: + sample_interval: 100 + hand_interp_steps: 10 + hold_steps: 4 + retreat_steps: 16 + postcondition_tolerance: 0.06 + + motion_modifiers: + orientation: + upright: + PickUp: + rotate_upright: 0.7853981633974483 + upright_yaw_samples: 8 + MoveHeldObject: + staging_lift_height: 0.25 + surface_clearance: 0.05 + upright_yaw_samples: 8 + upright_xy_tolerance: 0.05 + upright_max_tilt: 0.2617993877991494 + Place: + sample_interval: 64 + post_hold_steps: 12 + hand_interp_steps: 12 + MoveEndEffector: + sample_interval: 30 + retreat_height: 0.30 + retreat_distance: 0.10 + maximum_eef_height: 1.50 + handover_role: + transfer: + PickUp: + sample_interval: 80 + hand_interp_steps: 5 + pick_object_part: top + + predicate_fallbacks: + held_position_tolerance: 0.06 + held_gripper_tolerance: 0.01 + position_tolerance: 0.05 + xy_tolerance: 0.05 + container_xy_radius: 0.20 + container_min_z_offset: -0.05 + container_max_z_offset: 0.35 + support_xy_radius: 0.08 + support_com_margin: 0.002 + support_max_vertical_gap: 0.03 + support_max_penetration: 0.01 + support_min_overlap_ratio: 0.25 + not_fallen_max_tilt: 0.7853981633974483 + upright_max_tilt: 0.2617993877991494 + axis_tolerance: 0.03 + collinearity_tolerance: 0.03 + ordering_tolerance: 0.02 + minimum_lift_height: 0.08 + arm_initial_qpos_tolerance: 0.05 + gripper_state_tolerance: 0.001 + gripper_clear_min_distance: 0.08 + line_axis_tolerance: 0.06 + line_perpendicular_tolerance: 0.06 + preserve_orientation_tolerance: 0.2617993877991494 + payload_minimum_upright_cosine: 0.94 + payload_position_tolerance: 0.08 + payload_support_margin: 0.015 + + profiles: + dual_franka: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 0.85 + MoveEndEffector: + retreat_height: 0.10 + HandOver: + exchange_maximum_reach: 0.85 + dual_ur3: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 0.55 + HandOver: + exchange_maximum_reach: 0.55 + dual_ur5: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 0.85 + HandOver: + exchange_maximum_reach: 0.85 + motion_modifiers: + orientation: + upright: + PickUp: + lift_height: 0.12 + MoveHeldObject: + staging_lift_height: 0.12 + dual_ur10: + motion_defaults: + MoveHeldObject: + exchange_maximum_reach: 1.25 + HandOver: + exchange_maximum_reach: 1.25 diff --git a/embodichain/gen_sim/action_engine/config/runtime_policy.py b/embodichain/gen_sim/action_engine/config/runtime_policy.py new file mode 100644 index 000000000..56198211b --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/runtime_policy.py @@ -0,0 +1,775 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Load, resolve, snapshot, and hash package-owned Action Engine defaults.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +import hashlib +import json +import math +from pathlib import Path +from typing import Any, Final + +from embodichain.gen_sim.action_engine.domain.motion import MOTION_MODIFIER_MODES +from embodichain.utils import configclass +from embodichain.utils.utility import load_config + +__all__ = [ + "ACTION_ENGINE_DEFAULTS_SCHEMA", + "RUNTIME_POLICY_SCHEMA", + "ArmSelectionPolicyCfg", + "RuntimePolicyCfg", + "default_runtime_policy", + "generation_defaults", + "resolve_agent_runtime_policy", + "runtime_policy_hash", +] + +ACTION_ENGINE_DEFAULTS_SCHEMA: Final = "action_engine_defaults_v1" +RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v6" +_PREVIOUS_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v5" +_PRE_GRASP_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v4" +_PRE_PLANNER_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v3" +_LEGACY_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v1" +_DEFAULTS_PATH = Path(__file__).with_name("defaults.yaml") +_ARM_SELECTION_KEYS = ( + "crossing_deadband_ratio", + "pickup_crossing_weight", + "placement_crossing_weight", + "motion_cost_scale", + "fallback_workspace_half_width", + "orient_object_preferred_arm_deadband", +) +_ARM_SELECTION_OPTIONAL_KEYS = {"allow_cross_side_fallback"} +_GROUNDING_KEYS = { + "semantic_defaults": { + "surface_clearance", + "transport_clearance", + "staging_lift_height", + "relation_distance", + "hover_height", + "press_depth", + "retreat_height", + "maximum_eef_height", + }, + "arrangement": { + "slot_margin", + "minimum_spacing", + "layout_clearance", + "row_search_step", + "row_search_radius", + }, + "placement": { + "clearance", + "candidate_count", + "candidate_offset_fraction", + "support_margin", + "recovery_attempts", + }, + "coordinated_grasp": {"inset_fraction", "minimum_inset"}, + "handover": { + "retreat_height", + "retreat_distance", + "maximum_eef_height", + "minimum_transfer_clearance", + "minimum_transfer_lateral_clearance", + }, + "joint_state": { + "hand_close_sample_interval", + "hand_open_sample_interval", + }, +} +_GRASP_KEYS = { + "antipodal_n_sample", + "antipodal_max_angle", + "max_open_length", + "min_open_length", + "finger_length", + "point_sample_dense", + "max_deviation_angle", + "n_deviated_approach_directions", + "viser_port", + "max_decomposition_hulls", + "force_grasp_reannotate", +} +_PLANNER_KEYS = { + "backend", + "single_arm_strategy", + "coordinated_strategy", + "fallback_strategy", + "allow_fallback", + "dynamic_collision", + "static_obstacle_uids", + "dynamic_obstacle_uids", + "curobo", +} +_CUROBO_KEYS = { + "log_level", + "obstacle_representation", + "multi_env", + "use_cuda_graph", + "preserve_plan_samples", + "max_attempts", + "collision_activation_distance", +} +_MOTION_DEFAULT_ACTIONS = { + "CoordinatedPickment", + "CoordinatedPlacement", + "HandOver", + "MoveEndEffector", + "MoveHeldObject", + "MoveJoints", + "PickUp", + "Place", + "Press", +} +_PREDICATE_KEYS = { + "held_position_tolerance", + "held_gripper_tolerance", + "position_tolerance", + "xy_tolerance", + "container_xy_radius", + "container_min_z_offset", + "container_max_z_offset", + "support_xy_radius", + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + "not_fallen_max_tilt", + "upright_max_tilt", + "axis_tolerance", + "collinearity_tolerance", + "ordering_tolerance", + "minimum_lift_height", + "arm_initial_qpos_tolerance", + "gripper_state_tolerance", + "gripper_clear_min_distance", + "line_axis_tolerance", + "line_perpendicular_tolerance", + "preserve_orientation_tolerance", + "payload_minimum_upright_cosine", + "payload_position_tolerance", + "payload_support_margin", +} +_DEPRECATED_PREDICATE_KEYS = { + "support_min_z_offset", + "support_max_z_offset", +} + + +@configclass +class ArmSelectionPolicyCfg: + """Arm-allocation constraints and costs resolved for one robot profile.""" + + crossing_deadband_ratio: float = 0.08 + allow_cross_side_fallback: bool = False + pickup_crossing_weight: float = 1.0 + placement_crossing_weight: float = 1.5 + motion_cost_scale: float = math.pi + fallback_workspace_half_width: float = 0.5 + orient_object_preferred_arm_deadband: float = 0.02 + + def __post_init__(self) -> None: + if not isinstance(self.allow_cross_side_fallback, bool): + raise TypeError("allow_cross_side_fallback must be a bool.") + for name in _ARM_SELECTION_KEYS: + value = float(getattr(self, name)) + if not math.isfinite(value): + raise ValueError(f"{name} must be finite.") + if not 0.0 <= float(self.crossing_deadband_ratio) < 1.0: + raise ValueError("crossing_deadband_ratio must be in [0, 1).") + for name in _ARM_SELECTION_KEYS[1:3]: + if float(getattr(self, name)) < 0.0: + raise ValueError(f"{name} must be non-negative.") + for name in _ARM_SELECTION_KEYS[3:]: + if float(getattr(self, name)) <= 0.0: + raise ValueError(f"{name} must be positive.") + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> ArmSelectionPolicyCfg: + """Build a strict policy from a JSON/YAML mapping.""" + keys = frozenset(value) + if keys not in { + frozenset(_ARM_SELECTION_KEYS), + frozenset((*_ARM_SELECTION_KEYS, *_ARM_SELECTION_OPTIONAL_KEYS)), + }: + raise ValueError("arm_selection fields do not match the policy schema.") + fields: dict[str, Any] = {key: float(value[key]) for key in _ARM_SELECTION_KEYS} + if "allow_cross_side_fallback" in value: + fields["allow_cross_side_fallback"] = value["allow_cross_side_fallback"] + return cls(**fields) + + def as_mapping(self) -> dict[str, float | bool]: + """Return a stable JSON-compatible representation.""" + return { + "crossing_deadband_ratio": float(self.crossing_deadband_ratio), + "allow_cross_side_fallback": bool(self.allow_cross_side_fallback), + "pickup_crossing_weight": float(self.pickup_crossing_weight), + "placement_crossing_weight": float(self.placement_crossing_weight), + "motion_cost_scale": float(self.motion_cost_scale), + "fallback_workspace_half_width": float(self.fallback_workspace_half_width), + "orient_object_preferred_arm_deadband": float( + self.orient_object_preferred_arm_deadband + ), + } + + +@configclass +class RuntimePolicyCfg: + """Effective runtime policy persisted in generated agent artifacts.""" + + schema_version: str = RUNTIME_POLICY_SCHEMA + arm_selection: ArmSelectionPolicyCfg = ArmSelectionPolicyCfg() + execution: dict[str, Any] = {} + planner: dict[str, Any] = {} + grounding: dict[str, Any] = {} + grasp: dict[str, Any] = {} + motion_defaults: dict[str, dict[str, Any]] = {} + motion_modifiers: dict[str, dict[str, dict[str, dict[str, Any]]]] = {} + predicate_fallbacks: dict[str, Any] = {} + + def __post_init__(self) -> None: + if self.schema_version != RUNTIME_POLICY_SCHEMA: + raise ValueError( + f"Unsupported runtime policy schema {self.schema_version!r}." + ) + if not isinstance(self.arm_selection, ArmSelectionPolicyCfg): + raise TypeError("arm_selection must be an ArmSelectionPolicyCfg.") + for name in ( + "execution", + "planner", + "grounding", + "grasp", + "motion_defaults", + "motion_modifiers", + "predicate_fallbacks", + ): + if not isinstance(getattr(self, name), dict): + raise TypeError(f"{name} must be a mapping.") + _validate_finite_numbers(getattr(self, name), name) + if int(self.execution.get("max_transitions", 0)) <= 0: + raise ValueError("execution.max_transitions must be positive.") + if int(self.execution.get("semantic_step_settle_steps", -1)) < 0: + raise ValueError( + "execution.semantic_step_settle_steps must be non-negative." + ) + for name in ( + "max_retries_per_action", + "max_graph_revisions", + "max_recovery_actions", + "support_stability_interval_steps", + ): + if int(self.execution.get(name, -1)) < 0: + raise ValueError(f"execution.{name} must be non-negative.") + _require_keys( + self.execution, + { + "max_transitions", + "semantic_step_settle_steps", + "max_retries_per_action", + "max_graph_revisions", + "max_recovery_actions", + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + }, + "execution", + ) + if int(self.execution["support_stability_samples"]) <= 0: + raise ValueError("execution.support_stability_samples must be positive.") + for name in ( + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + if float(self.execution[name]) < 0.0: + raise ValueError(f"execution.{name} must be non-negative.") + _validate_planner(self.planner) + _require_keys(self.grounding, set(_GROUNDING_KEYS), "grounding") + for name, keys in _GROUNDING_KEYS.items(): + section = self.grounding.get(name) + if not isinstance(section, Mapping): + raise ValueError(f"grounding.{name} must be a mapping.") + _require_keys(section, keys, f"grounding.{name}") + placement = self.grounding["placement"] + if not 1 <= int(placement["candidate_count"]) <= 9: + raise ValueError("grounding.placement.candidate_count must be in [1, 9].") + if int(placement["recovery_attempts"]) < 0: + raise ValueError( + "grounding.placement.recovery_attempts must be non-negative." + ) + if not 0.0 <= float(placement["candidate_offset_fraction"]) <= 1.0: + raise ValueError( + "grounding.placement.candidate_offset_fraction must be in [0, 1]." + ) + if float(placement["support_margin"]) < 0.0: + raise ValueError("grounding.placement.support_margin must be non-negative.") + _require_keys(self.grasp, _GRASP_KEYS, "grasp") + _require_keys( + self.motion_defaults, + _MOTION_DEFAULT_ACTIONS, + "motion_defaults", + ) + if not all( + isinstance(policy, Mapping) and policy + for policy in self.motion_defaults.values() + ): + raise ValueError("Every motion default must be a non-empty mapping.") + _validate_motion_modifiers(self.motion_modifiers) + _require_keys( + self.predicate_fallbacks, + _PREDICATE_KEYS, + "predicate_fallbacks", + ) + if float(self.grasp.get("min_open_length", -1.0)) < 0.0: + raise ValueError("grasp.min_open_length must be non-negative.") + if float(self.grasp.get("max_open_length", 0.0)) <= float( + self.grasp.get("min_open_length", 0.0) + ): + raise ValueError("grasp.max_open_length must exceed min_open_length.") + direction_count = self.grasp.get("n_deviated_approach_directions") + if ( + isinstance(direction_count, bool) + or not isinstance(direction_count, int) + or not 1 <= direction_count <= 16 + ): + raise ValueError("grasp.n_deviated_approach_directions must be in [1, 16].") + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> RuntimePolicyCfg: + """Parse one fully resolved policy snapshot.""" + fields = { + "schema_version", + "execution", + "planner", + "arm_selection", + "grounding", + "grasp", + "motion_defaults", + "motion_modifiers", + "predicate_fallbacks", + } + if set(value) != fields: + raise ValueError("Runtime policy fields do not match the policy schema.") + if value.get("schema_version") != RUNTIME_POLICY_SCHEMA: + raise ValueError("Runtime policy has an unexpected schema_version.") + arm_selection = value.get("arm_selection") + if not isinstance(arm_selection, Mapping): + raise ValueError("Runtime policy requires an arm_selection mapping.") + sections = { + name: value.get(name) + for name in fields + if name not in {"schema_version", "arm_selection"} + } + if not all(isinstance(section, Mapping) for section in sections.values()): + raise ValueError("Runtime policy sections must be mappings.") + resolved_sections = { + name: deepcopy(dict(section)) for name, section in sections.items() + } + predicate_fallbacks = resolved_sections["predicate_fallbacks"] + for key in _DEPRECATED_PREDICATE_KEYS: + predicate_fallbacks.pop(key, None) + return cls( + schema_version=RUNTIME_POLICY_SCHEMA, + arm_selection=ArmSelectionPolicyCfg.from_mapping(arm_selection), + **resolved_sections, + ) + + def as_mapping(self) -> dict[str, Any]: + """Return the canonical artifact snapshot.""" + return { + "schema_version": self.schema_version, + "execution": deepcopy(self.execution), + "planner": deepcopy(self.planner), + "arm_selection": self.arm_selection.as_mapping(), + "grounding": deepcopy(self.grounding), + "grasp": deepcopy(self.grasp), + "motion_defaults": deepcopy(self.motion_defaults), + "motion_modifiers": deepcopy(self.motion_modifiers), + "predicate_fallbacks": deepcopy(self.predicate_fallbacks), + } + + +def default_runtime_policy(robot_profile: str) -> RuntimePolicyCfg: + """Resolve a package policy for one canonical robot profile.""" + document = _load_defaults() + runtime = document.get("runtime") + if not isinstance(runtime, Mapping) or set(runtime) != {"common", "profiles"}: + raise ValueError("Runtime defaults require common and profiles mappings.") + common, profiles = runtime["common"], runtime["profiles"] + if not isinstance(common, Mapping) or not isinstance(profiles, Mapping): + raise ValueError("Runtime common and profiles must be mappings.") + override = profiles.get(str(robot_profile)) + if not isinstance(override, Mapping): + raise ValueError(f"Unknown runtime robot profile {robot_profile!r}.") + resolved = _deep_merge(common, override) + return RuntimePolicyCfg.from_mapping( + { + "schema_version": RUNTIME_POLICY_SCHEMA, + **resolved, + } + ) + + +def generation_defaults() -> dict[str, Any]: + """Return a detached generation-policy mapping.""" + value = _load_defaults().get("generation") + if not isinstance(value, Mapping): + raise ValueError("Action Engine defaults require a generation mapping.") + required = { + "task", + "environment", + "scene", + "physics", + "randomization", + "dataset", + } + if set(value) != required: + raise ValueError("Generation defaults do not match the expected sections.") + return deepcopy(dict(value)) + + +def _load_defaults() -> dict[str, Any]: + document = load_config(_DEFAULTS_PATH) + if not isinstance(document, dict) or set(document) != { + "schema_version", + "generation", + "runtime", + }: + raise ValueError("Action Engine defaults do not match the package schema.") + if document.get("schema_version") != ACTION_ENGINE_DEFAULTS_SCHEMA: + raise ValueError("Action Engine defaults have an unexpected schema_version.") + return document + + +def _deep_merge( + base: Mapping[str, Any], + override: Mapping[str, Any], +) -> dict[str, Any]: + result = deepcopy(dict(base)) + for key, value in override.items(): + current = result.get(key) + result[key] = ( + _deep_merge(current, value) + if isinstance(current, Mapping) and isinstance(value, Mapping) + else deepcopy(value) + ) + return result + + +def _validate_finite_numbers(value: Any, path: str) -> None: + if isinstance(value, Mapping): + for key, item in value.items(): + _validate_finite_numbers(item, f"{path}.{key}") + elif isinstance(value, (list, tuple)): + for index, item in enumerate(value): + _validate_finite_numbers(item, f"{path}[{index}]") + elif isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{path} must be finite.") + + +def _require_keys( + value: Mapping[str, Any], + expected: set[str], + path: str, +) -> None: + if set(value) != expected: + raise ValueError(f"{path} fields do not match the defaults schema.") + + +def _validate_string_sequence(value: Any, path: str) -> None: + if not isinstance(value, (list, tuple)): + raise ValueError(f"{path} must be a list of object UIDs.") + normalized = [str(item) for item in value] + if any(not item.strip() for item in normalized): + raise ValueError(f"{path} entries must be non-empty strings.") + if any(not isinstance(item, str) for item in value): + raise ValueError(f"{path} entries must be strings.") + if len(set(normalized)) != len(normalized): + raise ValueError(f"{path} must not contain duplicate object UIDs.") + + +def _validate_planner(value: Mapping[str, Any]) -> None: + _require_keys(value, _PLANNER_KEYS, "planner") + backend = value.get("backend") + if backend not in {"curobo", "toppra"}: + raise ValueError("planner.backend must be 'curobo' or 'toppra'.") + for name in ("single_arm_strategy", "coordinated_strategy"): + if value.get(name) not in {"motion_gen", "ik_interp"}: + raise ValueError(f"planner.{name} must be 'motion_gen' or 'ik_interp'.") + if value.get("fallback_strategy") != "ik_interp": + raise ValueError("planner.fallback_strategy must be 'ik_interp'.") + if value.get("coordinated_strategy") == "motion_gen" and backend == "curobo": + raise ValueError( + "planner.coordinated_strategy must be 'ik_interp' with cuRobo." + ) + for name in ("allow_fallback", "dynamic_collision"): + if not isinstance(value.get(name), bool): + raise ValueError(f"planner.{name} must be a boolean.") + if value.get("dynamic_collision") and backend != "curobo": + raise ValueError("planner.dynamic_collision requires the cuRobo backend.") + _validate_string_sequence( + value.get("static_obstacle_uids"), + "planner.static_obstacle_uids", + ) + _validate_string_sequence( + value.get("dynamic_obstacle_uids"), + "planner.dynamic_obstacle_uids", + ) + overlap = set(value["static_obstacle_uids"]) & set(value["dynamic_obstacle_uids"]) + if overlap: + raise ValueError( + "Planner obstacle UIDs cannot be both static and dynamic: " + f"{sorted(overlap)}." + ) + + curobo = value.get("curobo") + if not isinstance(curobo, Mapping): + raise ValueError("planner.curobo must be a mapping.") + _require_keys(curobo, _CUROBO_KEYS, "planner.curobo") + if curobo.get("log_level") not in { + "debug", + "info", + "warning", + "warn", + "error", + }: + raise ValueError("planner.curobo.log_level is unsupported.") + if curobo.get("obstacle_representation") not in {"sphere", "cuboid", "mesh"}: + raise ValueError( + "planner.curobo.obstacle_representation must be sphere, cuboid, or mesh." + ) + for name in ("multi_env", "use_cuda_graph", "preserve_plan_samples"): + if not isinstance(curobo.get(name), bool): + raise ValueError(f"planner.curobo.{name} must be a boolean.") + max_attempts = curobo.get("max_attempts") + if ( + isinstance(max_attempts, bool) + or not isinstance(max_attempts, int) + or max_attempts <= 0 + ): + raise ValueError("planner.curobo.max_attempts must be positive.") + activation_distance = curobo.get("collision_activation_distance") + if ( + isinstance(activation_distance, bool) + or not isinstance(activation_distance, (int, float)) + or float(activation_distance) < 0.0 + ): + raise ValueError( + "planner.curobo.collision_activation_distance must be non-negative." + ) + + +def _validate_motion_modifiers(value: Mapping[str, Any]) -> None: + _require_keys(value, set(MOTION_MODIFIER_MODES), "motion_modifiers") + for modifier_type, modes in MOTION_MODIFIER_MODES.items(): + configured_modes = value.get(modifier_type) + if not isinstance(configured_modes, Mapping): + raise ValueError(f"motion_modifiers.{modifier_type} must be a mapping.") + _require_keys( + configured_modes, + set(modes), + f"motion_modifiers.{modifier_type}", + ) + for mode, patches in configured_modes.items(): + path = f"motion_modifiers.{modifier_type}.{mode}" + if not isinstance(patches, Mapping) or not patches: + raise ValueError(f"{path} must contain action-specific patches.") + unknown_actions = set(patches) - _MOTION_DEFAULT_ACTIONS + if unknown_actions: + raise ValueError( + f"{path} references unknown actions: {sorted(unknown_actions)}." + ) + if not all( + isinstance(patch, Mapping) and patch for patch in patches.values() + ): + raise ValueError(f"Every {path} action patch must be non-empty.") + + +def runtime_policy_hash(policy: RuntimePolicyCfg | Mapping[str, Any]) -> str: + """Hash the canonical effective policy independently of the Seed graph.""" + resolved = ( + policy + if isinstance(policy, RuntimePolicyCfg) + else RuntimePolicyCfg.from_mapping(policy) + ) + return _mapping_hash(resolved.as_mapping()) + + +def resolve_agent_runtime_policy(agent_config: Mapping[str, Any]) -> RuntimePolicyCfg: + """Resolve a generated snapshot or fall back for a legacy v1 artifact.""" + snapshot = agent_config.get("runtime_policy") + expected_hash = agent_config.get("runtime_policy_hash") + if snapshot is None: + if expected_hash is not None: + raise ValueError("runtime_policy_hash requires a runtime_policy snapshot.") + return default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + if not isinstance(snapshot, Mapping): + raise ValueError("agent_config.runtime_policy must be a mapping.") + if not isinstance(expected_hash, str) or not expected_hash: + raise ValueError( + "agent_config.runtime_policy requires a non-empty runtime_policy_hash." + ) + if _mapping_hash(snapshot) != expected_hash: + raise ValueError( + "agent_config runtime policy hash does not match its snapshot." + ) + if snapshot.get("schema_version") == _LEGACY_RUNTIME_POLICY_SCHEMA: + if set(snapshot) != {"schema_version", "arm_selection"} or not isinstance( + snapshot.get("arm_selection"), Mapping + ): + raise ValueError("Legacy runtime policy snapshot is malformed.") + policy = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + merged = policy.arm_selection.as_mapping() + merged.update(snapshot["arm_selection"]) + policy.arm_selection = ArmSelectionPolicyCfg.from_mapping(merged) + return policy + if snapshot.get("schema_version") == _PREVIOUS_RUNTIME_POLICY_SCHEMA: + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated_execution = deepcopy(dict(migrated.get("execution", {}))) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + migrated_execution[key] = defaults.execution[key] + migrated["execution"] = migrated_execution + migrated_grounding = deepcopy(dict(migrated.get("grounding", {}))) + migrated_placement = deepcopy(dict(migrated_grounding.get("placement", {}))) + for key, value in defaults.grounding["placement"].items(): + migrated_placement.setdefault(key, value) + migrated_grounding["placement"] = migrated_placement + migrated["grounding"] = migrated_grounding + migrated_predicates = deepcopy(dict(migrated.get("predicate_fallbacks", {}))) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + migrated_predicates[key] = defaults.predicate_fallbacks[key] + migrated["predicate_fallbacks"] = migrated_predicates + return RuntimePolicyCfg.from_mapping(migrated) + if snapshot.get("schema_version") == _PRE_GRASP_RUNTIME_POLICY_SCHEMA: + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated_grasp = deepcopy(dict(migrated.get("grasp", {}))) + migrated_grasp["n_deviated_approach_directions"] = defaults.grasp[ + "n_deviated_approach_directions" + ] + migrated["grasp"] = migrated_grasp + migrated_execution = deepcopy(dict(migrated.get("execution", {}))) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + migrated_execution[key] = defaults.execution[key] + migrated["execution"] = migrated_execution + migrated_grounding = deepcopy(dict(migrated.get("grounding", {}))) + migrated_placement = deepcopy(dict(migrated_grounding.get("placement", {}))) + for key, value in defaults.grounding["placement"].items(): + migrated_placement.setdefault(key, value) + migrated_grounding["placement"] = migrated_placement + migrated["grounding"] = migrated_grounding + migrated_predicates = deepcopy(dict(migrated.get("predicate_fallbacks", {}))) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + migrated_predicates[key] = defaults.predicate_fallbacks[key] + migrated["predicate_fallbacks"] = migrated_predicates + return RuntimePolicyCfg.from_mapping(migrated) + if snapshot.get("schema_version") == _PRE_PLANNER_RUNTIME_POLICY_SCHEMA: + expected_fields = { + "schema_version", + "execution", + "arm_selection", + "grounding", + "grasp", + "motion_defaults", + "motion_modifiers", + "predicate_fallbacks", + } + if set(snapshot) != expected_fields: + raise ValueError("Previous runtime policy snapshot is malformed.") + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated["planner"] = deepcopy(defaults.planner) + migrated_execution = deepcopy(dict(migrated.get("execution", {}))) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + migrated_execution[key] = defaults.execution[key] + migrated["execution"] = migrated_execution + migrated_grounding = deepcopy(dict(migrated.get("grounding", {}))) + migrated_placement = deepcopy(dict(migrated_grounding.get("placement", {}))) + for key, value in defaults.grounding["placement"].items(): + migrated_placement.setdefault(key, value) + migrated_grounding["placement"] = migrated_placement + migrated["grounding"] = migrated_grounding + migrated_grasp = deepcopy(dict(migrated["grasp"])) + migrated_grasp["n_deviated_approach_directions"] = defaults.grasp[ + "n_deviated_approach_directions" + ] + migrated["grasp"] = migrated_grasp + migrated_predicates = deepcopy(dict(migrated["predicate_fallbacks"])) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + migrated_predicates[key] = defaults.predicate_fallbacks[key] + migrated["predicate_fallbacks"] = migrated_predicates + return RuntimePolicyCfg.from_mapping(migrated) + policy = RuntimePolicyCfg.from_mapping(snapshot) + return policy + + +def _mapping_hash(value: Mapping[str, Any]) -> str: + payload = json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() diff --git a/embodichain/gen_sim/action_engine/domain/__init__.py b/embodichain/gen_sim/action_engine/domain/__init__.py new file mode 100644 index 000000000..5f2c303ad --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/__init__.py @@ -0,0 +1,93 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stable public contracts for Action Engine programs.""" + +from __future__ import annotations + +from .motion import ( + MOTION_MODIFIER_MODES, + MOTION_POLICY_VERSION, + motion_policy, + validate_motion_policy, +) +from .programs import ( + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, + execution_program_hash, + validate_execution_program, + validate_task_agent, +) +from .task_contracts import ( + PLACEMENT_RELATIONS, + RELATIONS, + TASK_CONTRACTS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + TaskContract, + normalize_placement_relation, + task_contract, + task_success_type, +) +from .v2 import ( + REASONING_TYPES, + TASK_LEVELS, + TASK_TYPES, + public_task_spec, + seed_graph_hash, + validate_public_task_spec, + validate_scene_requirements, + validate_seed_graph, + validate_task_spec, +) +from .visual_contracts import ( + OCCLUSION_RELATION, + VISUAL_RELATION_PARTICIPANTS, + requested_visual_task_predicates, +) + +__all__ = [ + "EXECUTION_PROGRAM_SCHEMA", + "MOTION_POLICY_VERSION", + "MOTION_MODIFIER_MODES", + "OCCLUSION_RELATION", + "REASONING_TYPES", + "RELATIONS", + "PLACEMENT_RELATIONS", + "TASK_CONTRACTS", + "TASK_LEVELS", + "TASK_TYPES", + "TASK_AGENT_SCHEMA", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "VISUAL_RELATION_PARTICIPANTS", + "TaskContract", + "execution_program_hash", + "motion_policy", + "normalize_placement_relation", + "public_task_spec", + "requested_visual_task_predicates", + "seed_graph_hash", + "task_contract", + "task_success_type", + "validate_public_task_spec", + "validate_scene_requirements", + "validate_seed_graph", + "validate_task_spec", + "validate_execution_program", + "validate_motion_policy", + "validate_task_agent", +] diff --git a/embodichain/gen_sim/action_engine/domain/motion.py b/embodichain/gen_sim/action_engine/domain/motion.py new file mode 100644 index 000000000..ae9ced35d --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/motion.py @@ -0,0 +1,111 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Typed, composable motion-policy references persisted in symbolic graphs.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from typing import Any, Final + +__all__ = [ + "MOTION_MODIFIER_MODES", + "MOTION_POLICY_VERSION", + "motion_policy", + "validate_motion_policy", +] + +MOTION_POLICY_VERSION: Final = "action_engine_motion_policy_v3" +MOTION_MODIFIER_MODES: Final = { + "orientation": frozenset({"upright"}), + "handover_role": frozenset({"transfer"}), +} + +_POLICY_KEYS = frozenset({"modifiers"}) +_MODIFIER_KEYS = frozenset({"type", "mode"}) + + +def motion_policy(*modifiers: tuple[str, str]) -> dict[str, Any]: + """Build one canonical policy reference from typed modifier pairs.""" + return validate_motion_policy( + { + "modifiers": [ + {"type": modifier_type, "mode": mode} + for modifier_type, mode in modifiers + ] + } + ) + + +def validate_motion_policy( + value: Any, + context: str = "motion_policy", +) -> dict[str, Any]: + """Validate and detach one symbolic motion-policy reference.""" + if not isinstance(value, Mapping): + raise ValueError( + f"{context} must be a mapping with typed modifiers; named string " + "policies are no longer supported. Regenerate the graph." + ) + if set(value) != _POLICY_KEYS: + raise ValueError(f"{context} fields must be {sorted(_POLICY_KEYS)}.") + raw_modifiers = value.get("modifiers") + if not isinstance(raw_modifiers, (list, tuple)): + raise ValueError(f"{context}.modifiers must be a sequence.") + + modifiers: list[dict[str, str]] = [] + seen: set[tuple[str, str]] = set() + seen_types: set[str] = set() + for index, raw_modifier in enumerate(raw_modifiers): + modifier_context = f"{context}.modifiers[{index}]" + if not isinstance(raw_modifier, Mapping): + raise ValueError(f"{modifier_context} must be a mapping.") + if set(raw_modifier) != _MODIFIER_KEYS: + raise ValueError( + f"{modifier_context} fields must be {sorted(_MODIFIER_KEYS)}." + ) + modifier_type = raw_modifier.get("type") + mode = raw_modifier.get("mode") + if not isinstance(modifier_type, str) or not modifier_type: + raise ValueError(f"{modifier_context}.type must be a non-empty string.") + if modifier_type not in MOTION_MODIFIER_MODES: + raise ValueError( + f"{modifier_context}.type {modifier_type!r} is unsupported; " + f"expected one of {sorted(MOTION_MODIFIER_MODES)}." + ) + if ( + not isinstance(mode, str) + or mode not in MOTION_MODIFIER_MODES[modifier_type] + ): + raise ValueError( + f"{modifier_context}.mode {mode!r} is unsupported for " + f"{modifier_type!r}; expected one of " + f"{sorted(MOTION_MODIFIER_MODES[modifier_type])}." + ) + key = (modifier_type, mode) + if key in seen: + raise ValueError(f"{modifier_context} duplicates modifier {key!r}.") + if modifier_type in seen_types: + raise ValueError( + f"{context} may select only one mode for modifier type " + f"{modifier_type!r}." + ) + seen.add(key) + seen_types.add(modifier_type) + modifiers.append({"type": modifier_type, "mode": mode}) + + return {"modifiers": deepcopy(modifiers)} diff --git a/embodichain/gen_sim/action_engine/domain/programs.py b/embodichain/gen_sim/action_engine/domain/programs.py new file mode 100644 index 000000000..a1243f1ea --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/programs.py @@ -0,0 +1,848 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Coordinate-free task and execution program contracts. + +The task agent is the only structure an LLM is allowed to influence. The +execution program is produced deterministically and contains the complete +symbolic action DAG consumed by runtime. Neither representation may contain +poses, trajectories, joint values, or other environment-specific geometry. +""" + +from __future__ import annotations + +import hashlib +import json +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, +) +from .motion import MOTION_POLICY_VERSION, validate_motion_policy + +__all__ = [ + "EXECUTION_PROGRAM_SCHEMA", + "MOTION_POLICY_VERSION", + "TASK_AGENT_SCHEMA", + "execution_program_hash", + "validate_execution_program", + "validate_task_agent", +] + +_ACTOR_MODES = frozenset({"auto", "required", "coordinated"}) +_CONTROL_MODES = frozenset({"arm", "hand", "coordinated"}) +_TASK_KEYS = frozenset( + {"schema_version", "task", "goal", "semantic_steps", "allocation_groups"} +) +_TASK_STEP_KEYS = frozenset( + {"id", "operator", "object", "objects", "actor", "goal", "depends_on"} +) +_EXECUTION_KEYS = frozenset( + { + "schema_version", + "task", + "goal_description", + "start", + "goal", + "nodes", + "edges", + "semantic_steps", + "allocation_groups", + "motion_policy_version", + } +) +_EXECUTION_STEP_KEYS = frozenset( + { + "id", + "parent_step_id", + "operator", + "object", + "actor", + "goal", + "depends_on", + "postcondition", + "edge_ids", + } +) +_EDGE_KEYS = frozenset( + { + "id", + "source", + "target", + "semantic_step_id", + "actions", + "depends_on", + "resources", + } +) +_ACTION_KEYS = frozenset( + { + "atomic_action_class", + "actor", + "control", + "target_binding", + "motion_policy", + "seed_node_id", + "failure_policy", + } +) +_FAILURE_POLICIES = frozenset({"task_required", "safety_required", "best_effort"}) +_TASK_ALLOCATION_GROUP_KEYS = frozenset({"id", "semantic_step_ids", "arm_constraint"}) +_BINDING_REQUIREMENTS = { + "articulation_goal": frozenset({"object"}), + "coordinated_goal": frozenset({"object"}), + "coordinated_placement_goal": frozenset({"placing_object", "support_object"}), + "current_held_pose": frozenset(), + "handover_goal": frozenset({"object"}), + "handover_staging": frozenset({"object"}), + "joint_state": frozenset({"source"}), + "object": frozenset({"object"}), + "policy_pose": frozenset(), + "pour_goal": frozenset({"object", "reference_object"}), + "semantic_goal": frozenset({"semantic_step"}), + "visual_constraint": frozenset({"camera_uid", "normalized_keypoint"}), +} +_POSTCONDITION_TYPES = frozenset( + { + "both_arms_at_initial_qpos", + "both_grippers_open", + "coordinated_placed", + "grippers_clear_of_object", + "held_by_both_grippers", + "line_member_placed", + "object_axis_near", + "object_axis_offset_near", + "object_held", + "object_held_by_both_grippers", + "object_held_by_gripper", + "object_in_container", + "object_lifted", + "object_not_fallen", + "object_on_object", + "object_supported_by", + "object_position_near", + "object_upright", + "object_xy_near", + "objects_collinear", + "objects_ordered", + "pressed", + "poured", + "articulation_joint_near", + "handover_complete", + "visual_relation", + "stable_unobstructed", + "sum_equals", + "semantic_goal", + "stack_layer_supported", + } +) +_OBJECT_REFERENCE_KEYS = frozenset( + { + "anchor", + "container", + "object", + "object_uid", + "orientation_reference_object", + "placing_object", + "reference", + "reference_object", + "support", + "support_object", + } +) + +# These fields indicate that planning-time or runtime geometry leaked into a +# symbolic program. Integers such as slot and layer indices remain valid. +_GROUNDED_FIELD_NAMES = frozenset( + { + "absolute_position", + "coordinates", + "joint_positions", + "object_target_pose", + "position", + "positions", + "pose", + "qpos", + "release_position", + "staging_position", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def validate_task_agent( + program: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, +) -> dict[str, Any]: + """Validate and return a detached canonical TaskAgent mapping. + + Defaults are added only for structural fields that have one unambiguous + meaning: ``actor={"mode": "auto"}``, an empty goal, and no dependencies. + Operator-specific semantics are validated by the capability registry. + + Args: + program: Candidate route-free task agent. + known_objects: Optional runtime scene UIDs. When supplied, every object + reference is validated before compilation. + + Returns: + A deep-copied, canonical mapping safe for compilation. + + Raises: + ValueError: If the program violates the TaskAgent contract. + """ + value = _mapping_copy(program, "TaskAgent") + _reject_unknown_keys(value, _TASK_KEYS, "TaskAgent") + _require_schema(value, TASK_AGENT_SCHEMA, "TaskAgent") + _require_nonempty_string(value.get("task"), "TaskAgent.task") + _require_nonempty_string(value.get("goal"), "TaskAgent.goal") + + raw_steps = _sequence(value.get("semantic_steps"), "TaskAgent.semantic_steps") + if not raw_steps: + raise ValueError("TaskAgent.semantic_steps must not be empty.") + + steps: list[dict[str, Any]] = [] + for index, raw_step in enumerate(raw_steps): + context = f"TaskAgent.semantic_steps[{index}]" + step = _mapping_copy(raw_step, context) + _reject_unknown_keys(step, _TASK_STEP_KEYS, context) + _require_nonempty_string(step.get("id"), f"{context}.id") + _require_nonempty_string(step.get("operator"), f"{context}.operator") + + has_object = "object" in step + has_objects = "objects" in step + if has_object == has_objects: + raise ValueError( + f"{context} must contain exactly one of 'object' or 'objects'." + ) + if has_object: + _require_nonempty_string(step["object"], f"{context}.object") + else: + objects = _string_list(step["objects"], f"{context}.objects") + if not objects: + raise ValueError(f"{context}.objects must not be empty.") + _require_unique(objects, f"{context}.objects") + step["objects"] = objects + + step["actor"] = _validate_actor( + step.get("actor", {"mode": "auto"}), + f"{context}.actor", + ) + step["goal"] = _mapping_copy(step.get("goal", {}), f"{context}.goal") + step["depends_on"] = _string_list( + step.get("depends_on", []), + f"{context}.depends_on", + ) + _require_unique(step["depends_on"], f"{context}.depends_on") + steps.append(step) + + step_ids = [step["id"] for step in steps] + _require_unique(step_ids, "TaskAgent semantic step IDs") + dependencies = {step["id"]: step["depends_on"] for step in steps} + _validate_dependency_dag(dependencies, "TaskAgent semantic steps") + value["semantic_steps"] = steps + value["allocation_groups"] = _validate_task_allocation_groups( + value.get("allocation_groups", []), + set(step_ids), + ) + if known_objects is not None: + _validate_known_objects(value, known_objects) + _reject_grounded_values(value) + return value + + +def validate_execution_program(program: Mapping[str, Any]) -> dict[str, Any]: + """Validate and return a detached canonical ExecutionProgram mapping. + + Validation covers both dependency DAGs: semantic-step dependencies and + executable edge dependencies. It also proves node reachability, edge + ownership, resource declarations, and the symbolic action envelope. + + Args: + program: Candidate deterministic execution program. + + Returns: + A deep-copied mapping safe for runtime consumption or hashing. + + Raises: + ValueError: If the program violates the ExecutionProgram contract. + """ + value = _mapping_copy(program, "ExecutionProgram") + _reject_unknown_keys(value, _EXECUTION_KEYS, "ExecutionProgram") + _require_schema(value, EXECUTION_PROGRAM_SCHEMA, "ExecutionProgram") + _require_nonempty_string(value.get("task"), "ExecutionProgram.task") + _require_nonempty_string( + value.get("goal_description"), + "ExecutionProgram.goal_description", + ) + _require_nonempty_string(value.get("start"), "ExecutionProgram.start") + _require_nonempty_string(value.get("goal"), "ExecutionProgram.goal") + if value.get("motion_policy_version") != MOTION_POLICY_VERSION: + raise ValueError( + "ExecutionProgram.motion_policy_version must be " + f"{MOTION_POLICY_VERSION!r}." + ) + + nodes = _validate_nodes(value.get("nodes")) + node_ids = {node["id"] for node in nodes} + if value["start"] not in node_ids or value["goal"] not in node_ids: + raise ValueError("ExecutionProgram start and goal must reference nodes.") + + edges = _validate_edges(value.get("edges"), node_ids) + edge_by_id = {edge["id"]: edge for edge in edges} + _validate_dependency_dag( + {edge_id: edge["depends_on"] for edge_id, edge in edge_by_id.items()}, + "ExecutionProgram edges", + ) + _validate_node_reachability( + start=value["start"], + goal=value["goal"], + node_ids=node_ids, + edges=edges, + ) + + semantic_steps = _validate_execution_steps( + value.get("semantic_steps"), + edge_by_id, + ) + step_ids = {step["id"] for step in semantic_steps} + _validate_allocation_groups(value.get("allocation_groups"), step_ids) + + value["nodes"] = nodes + value["edges"] = edges + value["semantic_steps"] = semantic_steps + value["allocation_groups"] = deepcopy(list(value.get("allocation_groups", []))) + _reject_grounded_values(value) + return value + + +def execution_program_hash(program: Mapping[str, Any]) -> str: + """Return the stable SHA-256 hash of a validated ExecutionProgram.""" + canonical = validate_execution_program(program) + try: + payload = json.dumps( + canonical, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise ValueError( + "ExecutionProgram must contain JSON-serializable values." + ) from exc + return hashlib.sha256(payload).hexdigest() + + +def _validate_nodes(value: Any) -> list[dict[str, Any]]: + raw_nodes = _sequence(value, "ExecutionProgram.nodes") + if len(raw_nodes) < 2: + raise ValueError("ExecutionProgram.nodes must contain start and goal nodes.") + nodes: list[dict[str, Any]] = [] + for index, raw_node in enumerate(raw_nodes): + context = f"ExecutionProgram.nodes[{index}]" + node = _mapping_copy(raw_node, context) + _reject_unknown_keys(node, frozenset({"id", "semantic"}), context) + _require_nonempty_string(node.get("id"), f"{context}.id") + _require_nonempty_string(node.get("semantic"), f"{context}.semantic") + nodes.append(node) + _require_unique([node["id"] for node in nodes], "ExecutionProgram node IDs") + return nodes + + +def _validate_edges(value: Any, node_ids: set[str]) -> list[dict[str, Any]]: + raw_edges = _sequence(value, "ExecutionProgram.edges") + if not raw_edges: + raise ValueError("ExecutionProgram.edges must not be empty.") + edges: list[dict[str, Any]] = [] + for index, raw_edge in enumerate(raw_edges): + context = f"ExecutionProgram.edges[{index}]" + edge = _mapping_copy(raw_edge, context) + _reject_unknown_keys(edge, _EDGE_KEYS, context) + for key in ("id", "source", "target", "semantic_step_id"): + _require_nonempty_string(edge.get(key), f"{context}.{key}") + if edge["source"] not in node_ids or edge["target"] not in node_ids: + raise ValueError(f"{context} references an unknown graph node.") + + edge["depends_on"] = _string_list( + edge.get("depends_on", []), + f"{context}.depends_on", + ) + edge["resources"] = _string_list( + edge.get("resources", []), + f"{context}.resources", + ) + _require_unique(edge["depends_on"], f"{context}.depends_on") + _require_unique(edge["resources"], f"{context}.resources") + edge["actions"] = _validate_actions(edge.get("actions"), context) + edges.append(edge) + + edge_ids = [edge["id"] for edge in edges] + _require_unique(edge_ids, "ExecutionProgram edge IDs") + known_edges = set(edge_ids) + for edge in edges: + unknown = set(edge["depends_on"]) - known_edges + if unknown: + raise ValueError( + f"Edge {edge['id']!r} depends on unknown edges: {sorted(unknown)}." + ) + return edges + + +def _validate_actions(value: Any, edge_context: str) -> list[dict[str, Any]]: + raw_actions = _sequence(value, f"{edge_context}.actions") + if not raw_actions: + raise ValueError(f"{edge_context}.actions must not be empty.") + actions: list[dict[str, Any]] = [] + for index, raw_action in enumerate(raw_actions): + context = f"{edge_context}.actions[{index}]" + action = _mapping_copy(raw_action, context) + _reject_unknown_keys(action, _ACTION_KEYS, context) + _require_nonempty_string( + action.get("atomic_action_class"), + f"{context}.atomic_action_class", + ) + action["actor"] = _validate_actor(action.get("actor"), f"{context}.actor") + control = _require_nonempty_string(action.get("control"), f"{context}.control") + if control not in _CONTROL_MODES: + raise ValueError( + f"{context}.control must be one of {sorted(_CONTROL_MODES)}." + ) + binding = _mapping_copy( + action.get("target_binding"), + f"{context}.target_binding", + ) + _require_nonempty_string( + binding.get("kind"), + f"{context}.target_binding.kind", + ) + kind = binding["kind"] + required = _BINDING_REQUIREMENTS.get(kind) + if required is None: + raise ValueError(f"{context}.target_binding.kind {kind!r} is unsupported.") + missing = sorted( + key for key in required if not _is_present_binding_value(binding.get(key)) + ) + if missing: + raise ValueError( + f"{context}.target_binding is missing required fields: {missing}." + ) + action["target_binding"] = binding + action["motion_policy"] = validate_motion_policy( + action.get("motion_policy"), + f"{context}.motion_policy", + ) + if "seed_node_id" in action: + _require_nonempty_string( + action.get("seed_node_id"), + f"{context}.seed_node_id", + ) + failure_policy = action.get("failure_policy", "task_required") + if failure_policy not in _FAILURE_POLICIES: + raise ValueError( + f"{context}.failure_policy must be one of " + f"{sorted(_FAILURE_POLICIES)}." + ) + action["failure_policy"] = str(failure_policy) + actions.append(action) + return actions + + +def _validate_execution_steps( + value: Any, + edge_by_id: Mapping[str, Mapping[str, Any]], +) -> list[dict[str, Any]]: + raw_steps = _sequence(value, "ExecutionProgram.semantic_steps") + if not raw_steps: + raise ValueError("ExecutionProgram.semantic_steps must not be empty.") + steps: list[dict[str, Any]] = [] + for index, raw_step in enumerate(raw_steps): + context = f"ExecutionProgram.semantic_steps[{index}]" + step = _mapping_copy(raw_step, context) + _reject_unknown_keys(step, _EXECUTION_STEP_KEYS, context) + for key in ("id", "parent_step_id", "operator", "object"): + _require_nonempty_string(step.get(key), f"{context}.{key}") + step["actor"] = _validate_actor(step.get("actor"), f"{context}.actor") + step["goal"] = _mapping_copy(step.get("goal"), f"{context}.goal") + step["postcondition"] = _mapping_copy( + step.get("postcondition"), + f"{context}.postcondition", + ) + _require_nonempty_string( + step["postcondition"].get("type"), + f"{context}.postcondition.type", + ) + if step["postcondition"]["type"] not in _POSTCONDITION_TYPES: + raise ValueError( + f"{context}.postcondition.type " + f"{step['postcondition']['type']!r} is unsupported." + ) + step["depends_on"] = _string_list( + step.get("depends_on", []), + f"{context}.depends_on", + ) + step["edge_ids"] = _string_list( + step.get("edge_ids"), + f"{context}.edge_ids", + ) + if not step["edge_ids"]: + raise ValueError(f"{context}.edge_ids must not be empty.") + _require_unique(step["depends_on"], f"{context}.depends_on") + _require_unique(step["edge_ids"], f"{context}.edge_ids") + steps.append(step) + + step_ids = [step["id"] for step in steps] + _require_unique(step_ids, "ExecutionProgram semantic step IDs") + _validate_dependency_dag( + {step["id"]: step["depends_on"] for step in steps}, + "ExecutionProgram semantic steps", + ) + + covered_edges: list[str] = [] + for step in steps: + for edge_id in step["edge_ids"]: + edge = edge_by_id.get(edge_id) + if edge is None: + raise ValueError( + f"Semantic step {step['id']!r} owns unknown edge {edge_id!r}." + ) + if edge["semantic_step_id"] != step["id"]: + raise ValueError( + f"Edge {edge_id!r} is assigned to {edge['semantic_step_id']!r}, " + f"not {step['id']!r}." + ) + covered_edges.append(edge_id) + _require_unique(covered_edges, "ExecutionProgram semantic edge ownership") + if set(covered_edges) != set(edge_by_id): + missing = sorted(set(edge_by_id) - set(covered_edges)) + raise ValueError(f"ExecutionProgram has unowned edges: {missing}.") + return steps + + +def _validate_allocation_groups(value: Any, step_ids: set[str]) -> None: + groups = _sequence(value, "ExecutionProgram.allocation_groups") + group_ids: list[str] = [] + for index, raw_group in enumerate(groups): + context = f"ExecutionProgram.allocation_groups[{index}]" + group = _mapping_copy(raw_group, context) + allowed = frozenset( + { + "id", + "semantic_step_ids", + "arm_constraint", + "execution_policy", + "parallel_action_classes", + "workspace_policy", + } + ) + _reject_unknown_keys(group, allowed, context) + group_ids.append(_require_nonempty_string(group.get("id"), f"{context}.id")) + members = _string_list( + group.get("semantic_step_ids"), + f"{context}.semantic_step_ids", + ) + if len(members) < 2: + raise ValueError(f"{context} must contain at least two semantic steps.") + _require_unique(members, f"{context}.semantic_step_ids") + unknown = set(members) - step_ids + if unknown: + raise ValueError(f"{context} references unknown steps: {sorted(unknown)}.") + for key in ("arm_constraint", "execution_policy", "workspace_policy"): + _require_nonempty_string(group.get(key), f"{context}.{key}") + if group["arm_constraint"] != "distinct_arms": + raise ValueError(f"{context}.arm_constraint must be 'distinct_arms'.") + if group["execution_policy"] != "parallel_if_feasible": + raise ValueError( + f"{context}.execution_policy must be 'parallel_if_feasible'." + ) + if group["workspace_policy"] != "shared_target_serial": + raise ValueError( + f"{context}.workspace_policy must be 'shared_target_serial'." + ) + action_classes = _string_list( + group.get("parallel_action_classes"), + f"{context}.parallel_action_classes", + ) + if not action_classes: + raise ValueError(f"{context}.parallel_action_classes must not be empty.") + _require_unique(group_ids, "ExecutionProgram allocation group IDs") + + +def _validate_task_allocation_groups( + value: Any, + step_ids: set[str], +) -> list[dict[str, Any]]: + groups = _sequence(value, "TaskAgent.allocation_groups") + result: list[dict[str, Any]] = [] + ids: list[str] = [] + members_seen: set[str] = set() + for index, raw_group in enumerate(groups): + context = f"TaskAgent.allocation_groups[{index}]" + group = _mapping_copy(raw_group, context) + _reject_unknown_keys(group, _TASK_ALLOCATION_GROUP_KEYS, context) + group_id = _require_nonempty_string(group.get("id"), f"{context}.id") + members = _string_list( + group.get("semantic_step_ids"), + f"{context}.semantic_step_ids", + ) + if len(members) < 2: + raise ValueError(f"{context} must contain at least two semantic steps.") + _require_unique(members, f"{context}.semantic_step_ids") + unknown = set(members) - step_ids + if unknown: + raise ValueError(f"{context} references unknown steps: {sorted(unknown)}.") + overlap = set(members) & members_seen + if overlap: + raise ValueError( + f"TaskAgent allocation groups overlap at steps: {sorted(overlap)}." + ) + constraint = _require_nonempty_string( + group.get("arm_constraint"), + f"{context}.arm_constraint", + ) + if constraint != "distinct_arms": + raise ValueError(f"{context}.arm_constraint must be 'distinct_arms'.") + ids.append(group_id) + members_seen.update(members) + result.append( + { + "id": group_id, + "semantic_step_ids": members, + "arm_constraint": constraint, + } + ) + _require_unique(ids, "TaskAgent allocation group IDs") + return result + + +def _validate_known_objects( + program: Mapping[str, Any], + known_objects: Collection[str], +) -> None: + known = {str(uid) for uid in known_objects} + if not known: + raise ValueError("known_objects must not be empty when supplied.") + allowed_sentinels = {"self", "table", "table_center"} + references: list[tuple[str, str]] = [] + for index, step in enumerate(program["semantic_steps"]): + if "object" in step: + references.append((f"semantic_steps[{index}].object", step["object"])) + for item_index, uid in enumerate(step.get("objects", [])): + references.append((f"semantic_steps[{index}].objects[{item_index}]", uid)) + _collect_object_references( + step["goal"], + f"semantic_steps[{index}].goal", + references, + ) + unknown = [ + f"{path}={uid!r}" + for path, uid in references + if uid not in known and uid not in allowed_sentinels + ] + if unknown: + raise ValueError( + "TaskAgent references objects not present in the scene: " + + ", ".join(unknown) + + "." + ) + + +def _collect_object_references( + value: Any, + path: str, + output: list[tuple[str, str]], +) -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + child_path = f"{path}.{key}" + if key in _OBJECT_REFERENCE_KEYS and isinstance(child, str): + output.append((child_path, child)) + elif ( + key in {"objects", "object_uids", "payloads"} + and isinstance(child, Sequence) + and not isinstance(child, (str, bytes, bytearray)) + ): + for index, item in enumerate(child): + uid = item.get("object") if isinstance(item, Mapping) else item + if isinstance(uid, str): + output.append((f"{child_path}[{index}]", uid)) + else: + _collect_object_references(child, child_path, output) + + +def _is_present_binding_value(value: Any) -> bool: + return value is not None and value != "" and value != [] + + +def _validate_actor(value: Any, context: str) -> dict[str, Any]: + actor = _mapping_copy(value, context) + mode = _require_nonempty_string(actor.get("mode"), f"{context}.mode") + if mode not in _ACTOR_MODES: + raise ValueError(f"{context}.mode must be one of {sorted(_ACTOR_MODES)}.") + if mode == "auto": + _reject_unknown_keys(actor, frozenset({"mode", "allocation_group"}), context) + elif mode == "required": + _reject_unknown_keys( + actor, + frozenset({"mode", "arm", "allocation_group"}), + context, + ) + _require_nonempty_string(actor.get("arm"), f"{context}.arm") + else: + _reject_unknown_keys(actor, frozenset({"mode", "arms"}), context) + arms = _string_list(actor.get("arms"), f"{context}.arms") + if len(arms) < 2: + raise ValueError(f"{context}.arms must contain at least two arms.") + _require_unique(arms, f"{context}.arms") + actor["arms"] = arms + if "allocation_group" in actor: + _require_nonempty_string( + actor["allocation_group"], + f"{context}.allocation_group", + ) + return actor + + +def _validate_dependency_dag( + dependencies: Mapping[str, Sequence[str]], + context: str, +) -> None: + known = set(dependencies) + outgoing: dict[str, list[str]] = {item_id: [] for item_id in known} + indegree = {item_id: 0 for item_id in known} + for item_id, required_ids in dependencies.items(): + unknown = set(required_ids) - known + if unknown: + raise ValueError(f"{context} reference unknown IDs: {sorted(unknown)}.") + if item_id in required_ids: + raise ValueError(f"{context} contain a self-dependency at {item_id!r}.") + for required_id in required_ids: + outgoing[required_id].append(item_id) + indegree[item_id] += 1 + + ready = deque( + sorted(item_id for item_id, degree in indegree.items() if degree == 0) + ) + visited = 0 + while ready: + item_id = ready.popleft() + visited += 1 + for dependent_id in sorted(outgoing[item_id]): + indegree[dependent_id] -= 1 + if indegree[dependent_id] == 0: + ready.append(dependent_id) + if visited != len(known): + cyclic = sorted(item_id for item_id, degree in indegree.items() if degree > 0) + raise ValueError(f"{context} contain a dependency cycle: {cyclic}.") + + +def _validate_node_reachability( + *, + start: str, + goal: str, + node_ids: set[str], + edges: Sequence[Mapping[str, Any]], +) -> None: + outgoing: dict[str, list[str]] = {node_id: [] for node_id in node_ids} + for edge in edges: + outgoing[edge["source"]].append(edge["target"]) + reachable = {start} + ready = deque([start]) + while ready: + node_id = ready.popleft() + for target_id in outgoing[node_id]: + if target_id not in reachable: + reachable.add(target_id) + ready.append(target_id) + if goal not in reachable: + raise ValueError("ExecutionProgram goal is unreachable from start.") + unreachable = sorted(node_ids - reachable) + if unreachable: + raise ValueError(f"ExecutionProgram contains unreachable nodes: {unreachable}.") + + +def _reject_grounded_values(value: Any, path: str = "program") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + normalized = str(key).strip().lower() + if normalized in _GROUNDED_FIELD_NAMES: + raise ValueError( + f"{path}.{key} is grounded runtime data and is not allowed." + ) + _reject_grounded_values(child, f"{path}.{key}") + return + if isinstance(value, list): + for index, child in enumerate(value): + _reject_grounded_values(child, f"{path}[{index}]") + return + if isinstance(value, float): + raise ValueError( + f"{path} contains a floating-point runtime value; use a named " + "motion policy or symbolic relation instead." + ) + + +def _mapping_copy(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise ValueError(f"{context} must be a list.") + return list(value) + + +def _string_list(value: Any, context: str) -> list[str]: + items = _sequence(value, context) + result: list[str] = [] + for index, item in enumerate(items): + result.append(_require_nonempty_string(item, f"{context}[{index}]")) + return result + + +def _require_schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +def _require_nonempty_string(value: Any, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{context} must be a non-empty string.") + return value + + +def _require_unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must not contain duplicates.") + + +def _reject_unknown_keys( + value: Mapping[str, Any], + allowed: frozenset[str], + context: str, +) -> None: + unknown = sorted(set(value) - allowed) + if unknown: + raise ValueError(f"{context} contains unknown fields: {unknown}.") diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py new file mode 100644 index 000000000..6d7a66cad --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -0,0 +1,125 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Action Engine recipes layered over the Task Engine semantic ontology.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +from embodichain.gen_sim.task_engine.ontology import ( + RELATIONS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + TaskContract as SemanticTaskContract, + task_success_type, +) +from embodichain.gen_sim.task_engine.ontology import ( + TASK_CONTRACTS as SEMANTIC_TASK_CONTRACTS, +) + +__all__ = [ + "PLACEMENT_RELATIONS", + "RELATIONS", + "TASK_CONTRACTS", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "TaskContract", + "task_contract", + "task_success_type", + "normalize_placement_relation", +] + +PLACEMENT_RELATIONS = RELATIONS - {"none"} +_SUPPORTED_PLACEMENT_ALIASES = frozenset({"above", "on_top", "on_top_of"}) + + +def normalize_placement_relation(value: Any) -> str: + """Lower task-language relations to physically executable release goals. + + A released object cannot remain freely hovering. Task-language ``above`` + therefore lowers to the supported ``on`` relation for placement operators; + non-placement operators such as pouring retain their distinct ``above`` + semantics. + """ + relation = str(value) + if ( + relation not in PLACEMENT_RELATIONS + and relation not in _SUPPORTED_PLACEMENT_ALIASES + ): + raise ValueError(f"Unsupported placement relation {relation!r}.") + return "on" if relation in _SUPPORTED_PLACEMENT_ALIASES else relation + + +_CORE_ACTIONS: Mapping[str, tuple[str, ...]] = MappingProxyType( + { + "E1": ("PickUp", "MoveHeldObject", "Place"), + "E2": ("AxisAlign",), + "E3": ("Pour",), + "E4": ("PickUp", "MoveHeldObject", "HandOver"), + "E5": ("CoordinatedPickment",), + "E6": ("PullArticulatedPart",), + "E7": ("PushArticulatedPart",), + "E8": ("TurnKnob",), + "E9": ("Press",), + } +) + + +@dataclass(frozen=True, slots=True) +class TaskContract: + """Action-facing view of one Task Engine semantic contract.""" + + task_type: str + semantics: str + core_actions: tuple[str, ...] + applicable_intent_fields: frozenset[str] + source_structure: str + required_affordances: frozenset[str] + success_type: str + scene_affordances: frozenset[str] + + +def _action_contract(value: SemanticTaskContract) -> TaskContract: + return TaskContract( + task_type=value.task_type, + semantics=value.semantics, + core_actions=_CORE_ACTIONS[value.task_type], + applicable_intent_fields=value.applicable_intent_fields, + source_structure=value.source_structure, + required_affordances=value.required_affordances, + success_type=value.success_type, + scene_affordances=value.scene_affordances, + ) + + +TASK_CONTRACTS: Mapping[str, TaskContract] = MappingProxyType( + { + task_type: _action_contract(contract) + for task_type, contract in SEMANTIC_TASK_CONTRACTS.items() + } +) + + +def task_contract(task_type: str) -> TaskContract: + """Return the Action Engine view of one canonical task contract.""" + try: + return TASK_CONTRACTS[str(task_type)] + except KeyError as exc: + raise ValueError(f"Unsupported task type {task_type!r}.") from exc diff --git a/embodichain/gen_sim/action_engine/domain/v2.py b/embodichain/gen_sim/action_engine/domain/v2.py new file mode 100644 index 000000000..bae0cbbbe --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/v2.py @@ -0,0 +1,1297 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict coordinate-free contracts for Action Engine SeedGraph v3.""" + +from __future__ import annotations + +from collections import deque +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +import hashlib +import json +import math +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from .motion import validate_motion_policy + +__all__ = [ + "REASONING_TYPES", + "TASK_LEVELS", + "TASK_TYPES", + "public_task_spec", + "seed_graph_hash", + "validate_public_task_spec", + "validate_scene_requirements", + "validate_seed_graph", + "validate_task_spec", +] + +TASK_LEVELS = frozenset({"L1", "L2", "L3", "L4"}) +TASK_TYPES = frozenset({f"E{index}" for index in range(1, 10)}) +REASONING_TYPES = frozenset( + { + "none", + "memory", + "visual_semantics", + "pattern", + "logic", + "common_sense", + "constraint", + } +) + +_TASK_SPEC_KEYS = frozenset( + { + "schema_version", + "task_id", + "level", + "instruction", + "reasoning_type", + "task_instances", + "success", + "oracle", + "metadata", + } +) +_TASK_INSTANCE_KEYS = frozenset({"id", "task_type", "params", "depends_on", "role"}) +_SCENE_REQUIREMENTS_KEYS = frozenset( + { + "schema_version", + "task_id", + "objects", + "cameras", + "spatial_constraints", + "distractor_count", + "metadata", + } +) +_OBJECT_REQUIREMENT_KEYS = frozenset( + { + "role_id", + "category", + "count", + "affordances", + "initial_state", + "attributes", + } +) +_SEED_GRAPH_KEYS = frozenset( + { + "schema_version", + "task_id", + "instruction", + "level", + "reasoning_type", + "planner_route", + "nodes", + "task_groups", + "success", + "capability_catalog_hash", + "metadata", + } +) +_ACTION_NODE_KEYS = frozenset( + { + "id", + "atomic_action", + "object_uid", + "actor", + "control", + "target_binding", + "depends_on", + "contract", + "task_instance_id", + "task_type", + "role", + "precondition", + "postcondition", + "motion_policy", + "sync_group", + } +) +_TASK_GROUP_KEYS = frozenset( + { + "id", + "task_type", + "role", + "operator", + "object_uid", + "actor", + "goal", + "depends_on", + "parent_task_instance_id", + "node_ids", + "success", + "contract", + } +) +_ACTOR_MODES = frozenset({"auto", "required", "preferred", "coordinated"}) +_NODE_ROLES = frozenset({"primary", "recovery", "cleanup"}) +_GROUP_ROLES = frozenset({"primary", "recovery"}) +_PLANNER_ROUTES = frozenset({"offline", "online", "selected", "fused"}) +_ACTION_CONTRACT_KEYS = frozenset( + { + "version", + "requires", + "effects", + "claims", + "completion", + "failure_policy", + } +) +_TASK_GROUP_CONTRACT_KEYS = frozenset( + { + "entry_requires", + "exit_effects", + "claims", + "entry_node_ids", + "terminal_node_ids", + "completion", + } +) +_STATE_ATOM_KEYS = frozenset({"predicate", "object_uid", "arm"}) +_STATE_PREDICATES = frozenset( + { + "arm_free", + "object_free", + "object_held", + "object_coordinated_held", + "handover_complete", + "arm_clear", + "arm_home", + } +) +_EFFECT_KEYS = frozenset({"op", "atom"}) +_EFFECT_OPERATIONS = frozenset({"add", "delete"}) +_CLAIM_KEYS = frozenset({"resource", "access", "lifetime"}) +_CLAIM_ACCESS = frozenset({"shared_read", "exclusive"}) +_CLAIM_LIFETIMES = frozenset({"action", "until_release"}) +_ACTION_COMPLETION = frozenset({"ordinary", "cleanup", "terminal_barrier"}) +_FAILURE_POLICIES = frozenset({"task_required", "safety_required", "best_effort"}) +_GROUP_COMPLETION = frozenset({"ordinary", "terminal_barrier"}) +_OBJECT_REFERENCE_KEYS = frozenset( + { + "anchor", + "container", + "object", + "object_uid", + "placing_object", + "reference", + "reference_object", + "support", + "support_object", + } +) +_GROUNDED_FIELD_NAMES = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "object_target_pose", + "position", + "positions", + "pose", + "qpos", + "release_position", + "staging_position", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def validate_task_spec(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate one task-first, scene-independent task specification.""" + result = _mapping(value, "TaskSpec") + _keys(result, _TASK_SPEC_KEYS, "TaskSpec") + _schema(result, TASK_SPEC_SCHEMA, "TaskSpec") + _string(result.get("task_id"), "TaskSpec.task_id") + level = _enum(result.get("level"), TASK_LEVELS, "TaskSpec.level") + _string(result.get("instruction"), "TaskSpec.instruction") + reasoning = _enum( + result.get("reasoning_type", "none"), + REASONING_TYPES, + "TaskSpec.reasoning_type", + ) + if level == "L4" and reasoning == "none": + raise ValueError("TaskSpec L4 tasks require a non-'none' reasoning_type.") + if level != "L4" and reasoning != "none": + raise ValueError("Only TaskSpec L4 tasks may declare reasoning_type.") + + instances: list[dict[str, Any]] = [] + for index, item in enumerate( + _sequence(result.get("task_instances"), "TaskSpec.task_instances") + ): + context = f"TaskSpec.task_instances[{index}]" + instance = _mapping(item, context) + _keys(instance, _TASK_INSTANCE_KEYS, context) + _string(instance.get("id"), f"{context}.id") + _enum(instance.get("task_type"), TASK_TYPES, f"{context}.task_type") + instance["params"] = _mapping(instance.get("params", {}), f"{context}.params") + instance["depends_on"] = _strings( + instance.get("depends_on", []), f"{context}.depends_on" + ) + instance["role"] = _enum( + instance.get("role", "primary"), _GROUP_ROLES, f"{context}.role" + ) + instances.append(instance) + if not instances: + raise ValueError("TaskSpec.task_instances must not be empty.") + _unique([item["id"] for item in instances], "TaskSpec task instance IDs") + _dag( + {item["id"]: item["depends_on"] for item in instances}, + "TaskSpec task instances", + ) + _validate_level_shape(level, instances) + result["task_instances"] = instances + result["success"] = _mapping(result.get("success"), "TaskSpec.success") + result["oracle"] = _mapping(result.get("oracle", {}), "TaskSpec.oracle") + result["metadata"] = _mapping(result.get("metadata", {}), "TaskSpec.metadata") + _reject_grounded(result) + _finite(result) + return result + + +def public_task_spec(value: Mapping[str, Any]) -> dict[str, Any]: + """Return the validated TaskSpec view safe to expose to an online agent.""" + if "task_instances" not in value and value.get("level") == "L4": + result = dict(value) + _strip_public_private_metadata(result) + return validate_public_task_spec(result) + result = validate_task_spec(value) + result.pop("oracle", None) + if result["level"] == "L4": + # L4 task instances are the hidden reference plan, not public intent. + result.pop("task_instances", None) + _strip_public_private_metadata(result) + return validate_public_task_spec(result) + + +def _strip_public_private_metadata(value: dict[str, Any]) -> None: + """Remove role/UID bindings that would turn the public view into an oracle.""" + metadata = value.get("metadata") + if not isinstance(metadata, Mapping): + return + private_keys = { + "role_bindings", + "uid_map", + "source_uid_map", + "reference_seed_graph", + "oracle", + } + + def strip(child: Any) -> Any: + if isinstance(child, Mapping): + return { + key: strip(nested) + for key, nested in child.items() + if str(key).lower() not in private_keys + } + if isinstance(child, list): + return [strip(item) for item in child] + if isinstance(child, tuple): + return [strip(item) for item in child] + return deepcopy(child) + + value["metadata"] = strip(metadata) + + +def validate_public_task_spec(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the oracle-free TaskSpec projection consumed online.""" + result = _mapping(value, "PublicTaskSpec") + allowed = _TASK_SPEC_KEYS - {"oracle"} + _keys(result, allowed, "PublicTaskSpec") + if "oracle" in result: + raise ValueError("PublicTaskSpec must not contain oracle data.") + _schema(result, TASK_SPEC_SCHEMA, "PublicTaskSpec") + _string(result.get("task_id"), "PublicTaskSpec.task_id") + level = _enum(result.get("level"), TASK_LEVELS, "PublicTaskSpec.level") + _string(result.get("instruction"), "PublicTaskSpec.instruction") + reasoning = _enum( + result.get("reasoning_type", "none"), + REASONING_TYPES, + "PublicTaskSpec.reasoning_type", + ) + if (level == "L4") != (reasoning != "none"): + raise ValueError( + "PublicTaskSpec reasoning_type must be non-'none' exactly for L4." + ) + if level == "L4": + if "task_instances" in result: + raise ValueError("Public L4 TaskSpec must hide reference task instances.") + else: + # Reuse the complete structural validator for explicit L1-L3 tasks. + normalized = validate_task_spec({**result, "oracle": {}}) + normalized.pop("oracle", None) + return normalized + result["success"] = _mapping(result.get("success"), "PublicTaskSpec.success") + result["metadata"] = _mapping(result.get("metadata", {}), "PublicTaskSpec.metadata") + _reject_grounded(result) + _finite(result) + return result + + +def validate_scene_requirements(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the structured hand-off contract consumed by a Scene Engine.""" + result = _mapping(value, "SceneRequirements") + _keys(result, _SCENE_REQUIREMENTS_KEYS, "SceneRequirements") + _schema(result, SCENE_REQUIREMENTS_SCHEMA, "SceneRequirements") + _string(result.get("task_id"), "SceneRequirements.task_id") + objects: list[dict[str, Any]] = [] + for index, item in enumerate( + _sequence(result.get("objects"), "SceneRequirements.objects") + ): + context = f"SceneRequirements.objects[{index}]" + requirement = _mapping(item, context) + _keys(requirement, _OBJECT_REQUIREMENT_KEYS, context) + _string(requirement.get("role_id"), f"{context}.role_id") + _string(requirement.get("category"), f"{context}.category") + count = requirement.get("count", 1) + if not isinstance(count, int) or isinstance(count, bool) or count < 1: + raise ValueError(f"{context}.count must be a positive integer.") + requirement["count"] = count + requirement["affordances"] = _strings( + requirement.get("affordances", []), f"{context}.affordances" + ) + requirement["initial_state"] = _mapping( + requirement.get("initial_state", {}), f"{context}.initial_state" + ) + requirement["attributes"] = _mapping( + requirement.get("attributes", {}), f"{context}.attributes" + ) + objects.append(requirement) + if not objects: + raise ValueError("SceneRequirements.objects must not be empty.") + _unique([item["role_id"] for item in objects], "SceneRequirements role IDs") + result["objects"] = objects + result["cameras"] = [ + _mapping(item, f"SceneRequirements.cameras[{index}]") + for index, item in enumerate( + _sequence(result.get("cameras", []), "SceneRequirements.cameras") + ) + ] + result["spatial_constraints"] = [ + _mapping(item, f"SceneRequirements.spatial_constraints[{index}]") + for index, item in enumerate( + _sequence( + result.get("spatial_constraints", []), + "SceneRequirements.spatial_constraints", + ) + ) + ] + distractors = result.get("distractor_count", 0) + if ( + not isinstance(distractors, int) + or isinstance(distractors, bool) + or distractors < 0 + ): + raise ValueError("SceneRequirements.distractor_count must be non-negative.") + result["distractor_count"] = distractors + result["metadata"] = _mapping( + result.get("metadata", {}), "SceneRequirements.metadata" + ) + _finite(result) + return result + + +def validate_seed_graph( + value: Mapping[str, Any], + *, + known_objects: Collection[str] | None = None, + known_actions: Collection[str] | None = None, + executable_actions: Collection[str] | None = None, + require_executable: bool = False, +) -> dict[str, Any]: + """Validate a direct, coordinate-free AtomicAction DAG.""" + result = _mapping(value, "SeedGraph") + _keys(result, _SEED_GRAPH_KEYS, "SeedGraph") + _schema(result, SEED_GRAPH_SCHEMA, "SeedGraph") + _string(result.get("task_id"), "SeedGraph.task_id") + _string(result.get("instruction"), "SeedGraph.instruction") + level = _enum(result.get("level"), TASK_LEVELS, "SeedGraph.level") + reasoning = _enum( + result.get("reasoning_type", "none"), + REASONING_TYPES, + "SeedGraph.reasoning_type", + ) + if (level == "L4") != (reasoning != "none"): + raise ValueError("SeedGraph reasoning_type must be non-'none' exactly for L4.") + result["planner_route"] = _enum( + result.get("planner_route"), _PLANNER_ROUTES, "SeedGraph.planner_route" + ) + _string(result.get("capability_catalog_hash"), "SeedGraph.capability_catalog_hash") + + nodes: list[dict[str, Any]] = [] + for index, item in enumerate(_sequence(result.get("nodes"), "SeedGraph.nodes")): + context = f"SeedGraph.nodes[{index}]" + node = _mapping(item, context) + _keys(node, _ACTION_NODE_KEYS, context) + _string(node.get("id"), f"{context}.id") + action = _string(node.get("atomic_action"), f"{context}.atomic_action") + if known_actions is not None and action not in set(known_actions): + raise ValueError(f"{context} references unknown AtomicAction {action!r}.") + if ( + require_executable + and executable_actions is not None + and action not in set(executable_actions) + ): + raise ValueError( + f"AtomicAction {action!r} is planning-only and cannot be executed." + ) + object_uid = _string(node.get("object_uid"), f"{context}.object_uid") + if known_objects is not None and object_uid not in set(known_objects): + raise ValueError(f"{context} references unknown object {object_uid!r}.") + node["actor"] = _actor(node.get("actor", {"mode": "auto"}), f"{context}.actor") + node["control"] = _string(node.get("control", "arm"), f"{context}.control") + binding = _mapping(node.get("target_binding"), f"{context}.target_binding") + _string(binding.get("kind"), f"{context}.target_binding.kind") + node["target_binding"] = binding + node["depends_on"] = _strings( + node.get("depends_on", []), f"{context}.depends_on" + ) + node["contract"] = _action_contract(node.get("contract"), f"{context}.contract") + node["task_instance_id"] = _string( + node.get("task_instance_id"), f"{context}.task_instance_id" + ) + node["task_type"] = _enum( + node.get("task_type"), TASK_TYPES, f"{context}.task_type" + ) + node["role"] = _enum( + node.get("role", "primary"), _NODE_ROLES, f"{context}.role" + ) + node["precondition"] = _mapping( + node.get("precondition", {}), f"{context}.precondition" + ) + node["postcondition"] = _mapping( + node.get("postcondition", {}), f"{context}.postcondition" + ) + node["motion_policy"] = validate_motion_policy( + node.get("motion_policy"), f"{context}.motion_policy" + ) + if "sync_group" in node: + node["sync_group"] = _string(node["sync_group"], f"{context}.sync_group") + nodes.append(node) + if not nodes: + raise ValueError("SeedGraph.nodes must not be empty.") + node_ids = [node["id"] for node in nodes] + _unique(node_ids, "SeedGraph node IDs") + _dag({node["id"]: node["depends_on"] for node in nodes}, "SeedGraph nodes") + + groups: list[dict[str, Any]] = [] + for index, item in enumerate( + _sequence(result.get("task_groups"), "SeedGraph.task_groups") + ): + context = f"SeedGraph.task_groups[{index}]" + group = _mapping(item, context) + _keys(group, _TASK_GROUP_KEYS, context) + group["id"] = _string(group.get("id"), f"{context}.id") + group["task_type"] = _enum( + group.get("task_type"), TASK_TYPES, f"{context}.task_type" + ) + group["role"] = _enum( + group.get("role", "primary"), _GROUP_ROLES, f"{context}.role" + ) + group["operator"] = _string(group.get("operator"), f"{context}.operator") + group["object_uid"] = _string(group.get("object_uid"), f"{context}.object_uid") + group["actor"] = _actor( + group.get("actor", {"mode": "auto"}), f"{context}.actor" + ) + group["goal"] = _mapping(group.get("goal", {}), f"{context}.goal") + group["depends_on"] = _strings( + group.get("depends_on", []), f"{context}.depends_on" + ) + if "parent_task_instance_id" in group: + group["parent_task_instance_id"] = _string( + group["parent_task_instance_id"], + f"{context}.parent_task_instance_id", + ) + group["node_ids"] = _strings(group.get("node_ids"), f"{context}.node_ids") + if not group["node_ids"]: + raise ValueError(f"{context}.node_ids must not be empty.") + group["success"] = _mapping(group.get("success"), f"{context}.success") + group["contract"] = _task_group_contract( + group.get("contract"), f"{context}.contract" + ) + groups.append(group) + if not groups: + raise ValueError("SeedGraph.task_groups must not be empty.") + _validate_groups(nodes, groups) + _validate_group_contract_topology(nodes, groups) + _validate_cleanup_barriers(nodes, groups) + _validate_task_group_semantics(nodes, groups) + _validate_ownership_transitions(nodes, groups) + _validate_resource_conflicts(nodes, groups, result.get("metadata", {})) + _dag( + {group["id"]: group["depends_on"] for group in groups}, + "SeedGraph task groups", + ) + _validate_group_dependency_alignment(nodes, groups) + result["nodes"] = nodes + result["task_groups"] = groups + result["success"] = _mapping(result.get("success"), "SeedGraph.success") + result["metadata"] = _mapping(result.get("metadata", {}), "SeedGraph.metadata") + _reject_grounded(result) + _finite(result) + if known_objects is not None: + _known_object_references(result, set(known_objects)) + return result + + +def seed_graph_hash(value: Mapping[str, Any]) -> str: + """Return a stable SHA-256 hash for one validated SeedGraph.""" + canonical = validate_seed_graph(value) + payload = json.dumps( + canonical, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _validate_level_shape(level: str, instances: Sequence[Mapping[str, Any]]) -> None: + primary = [item for item in instances if item["role"] == "primary"] + types = {str(item["task_type"]) for item in primary} + if level == "L1" and len(primary) != 1: + raise ValueError("L1 requires exactly one primary task instance.") + if level == "L2" and (len(primary) < 2 or len(types) != 1): + raise ValueError("L2 requires at least two primary instances of one E type.") + if level == "L3" and (len(primary) < 2 or len(types) < 2): + raise ValueError( + "L3 requires at least two primary instances of different E types." + ) + + +def _validate_groups( + nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]] +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + _unique([str(group["id"]) for group in groups], "SeedGraph task group IDs") + memberships: dict[str, str] = {} + for group in groups: + group_id = str(group["id"]) + for node_id in group["node_ids"]: + if node_id not in node_by_id: + raise ValueError( + f"SeedGraph task group {group_id!r} references unknown node {node_id!r}." + ) + if node_id in memberships: + raise ValueError( + f"SeedGraph node {node_id!r} belongs to multiple task groups." + ) + node = node_by_id[node_id] + if node["task_instance_id"] != group_id: + raise ValueError( + f"SeedGraph node {node_id!r} task_instance_id does not match {group_id!r}." + ) + if node["task_type"] != group["task_type"]: + raise ValueError( + f"SeedGraph node {node_id!r} task_type does not match its group." + ) + memberships[node_id] = group_id + missing = sorted(set(node_by_id) - set(memberships)) + if missing: + raise ValueError( + f"SeedGraph nodes are missing task group membership: {missing}." + ) + + +def _validate_task_group_semantics( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + required_actions = { + "E1": set(), + "E2": set(), + "E3": {"Pour"}, + "E4": {"HandOver"}, + "E5": set(), + "E6": {"PullArticulatedPart"}, + "E7": {"PushArticulatedPart"}, + "E8": {"TurnKnob"}, + "E9": {"Press"}, + } + for group in groups: + task_type = str(group["task_type"]) + group_nodes = [node_by_id[node_id] for node_id in group["node_ids"]] + actions = {str(node["atomic_action"]) for node in group_nodes} + missing = required_actions[task_type] - actions + if task_type == "E1": + if not actions.intersection({"MoveHeldObject", "Place"}): + missing = {"MoveHeldObject|Place"} + elif "PickUp" not in actions: + first = group_nodes[0] + precondition = first.get("precondition", {}) + if precondition.get("type") != "object_held": + missing = {"PickUp|object_held precondition"} + if task_type == "E2" and "AxisAlign" not in actions: + missing = {"MoveHeldObject", "Place"} - actions + if "PickUp" not in actions: + first = group_nodes[0] + precondition = first.get("precondition", {}) + if precondition.get("type") != "object_held": + missing.add("PickUp|object_held precondition") + # Recovery may explicitly preserve a verified downstream hold. Ordinary + # E2 groups always complete their supported world state with Place. + if ( + task_type == "E2" + and "Place" not in actions + and group.get("goal", {}).get("terminal_behavior") == "hold" + and "MoveHeldObject" in actions + and group.get("role") == "recovery" + ): + missing.discard("Place") + if task_type == "E5" and not actions.intersection( + {"CoordinatedPickment", "CoordinatedPlacement"} + ): + missing = {"CoordinatedPickment|CoordinatedPlacement"} + if missing: + raise ValueError( + f"SeedGraph TaskGroup {group['id']!r} is missing {task_type} " + f"core actions: {sorted(missing)}." + ) + + +def _validate_ownership_transitions( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], +) -> None: + """Check release/reacquire and explicit single-arm hold transitions. + + An ordinary E2 -> E4 transition persists the supported upright state, ends + the predecessor resource lease, and lets E4 acquire a fresh transfer grasp. + E4 -> E1 keeps receiver ownership because the exchanged object is not yet + supported. Recovery groups may preserve an explicitly requested hold. + """ + node_by_id = {str(node["id"]): node for node in nodes} + group_by_id = {str(group["id"]): group for group in groups} + nodes_by_group = { + group_id: [node_by_id[node_id] for node_id in group["node_ids"]] + for group_id, group in group_by_id.items() + } + + def direct_predecessor( + group: Mapping[str, Any], task_type: str + ) -> Mapping[str, Any] | None: + for dependency in group.get("depends_on", []): + candidate = group_by_id.get(str(dependency)) + if candidate is not None and candidate.get("task_type") == task_type: + return candidate + return None + + def held_arm(node: Mapping[str, Any]) -> str | None: + precondition = node.get("precondition", {}) + if ( + isinstance(precondition, Mapping) + and precondition.get("type") == "object_held" + ): + arm = str(precondition.get("arm", "")) + if arm in {"left_arm", "right_arm"}: + return arm + actor = node.get("actor", {}) + if isinstance(actor, Mapping) and actor.get("mode") == "required": + arm = str(actor.get("arm", "")) + if arm in {"left_arm", "right_arm"}: + return arm + return None + + for group_id, group in group_by_id.items(): + task_type = str(group.get("task_type")) + group_nodes = nodes_by_group[group_id] + actions = [str(node.get("atomic_action")) for node in group_nodes] + object_uid = str(group.get("object_uid")) + + if task_type == "E2": + handover = next( + ( + candidate + for candidate in groups + if candidate.get("task_type") == "E4" + and group_id + in {str(item) for item in candidate.get("depends_on", [])} + and str(candidate.get("object_uid")) == object_uid + ), + None, + ) + if handover is None: + continue + if ( + group.get("goal", {}).get("terminal_behavior") == "hold" + and group.get("role") != "recovery" + ): + raise ValueError( + f"SeedGraph E2 group {group_id!r} may not preserve a holder " + "across an ordinary E2->E4 TaskGroup boundary." + ) + if ( + group.get("role") != "recovery" + and "Place" not in actions + and "AxisAlign" not in actions + ): + raise ValueError( + f"SeedGraph E2 group {group_id!r} must release its supported " + "object before E4 reacquires it." + ) + + if task_type == "E4": + predecessor = direct_predecessor(group, "E2") + if ( + predecessor is not None + and str(predecessor.get("object_uid")) == object_uid + ): + predecessor_nodes = nodes_by_group[str(predecessor["id"])] + preserves_hold = ( + predecessor.get("role") == "recovery" + and predecessor.get("goal", {}).get("terminal_behavior") == "hold" + ) + if preserves_hold: + if "PickUp" in actions or not group_nodes: + raise ValueError( + f"SeedGraph E4 group {group_id!r} must consume the " + "recovery-held object without PickUp." + ) + first = group_nodes[0] + holder_arm = next( + ( + held_arm(node) + for node in reversed(predecessor_nodes) + if held_arm(node) is not None + ), + None, + ) + if ( + str(first.get("atomic_action")) != "MoveHeldObject" + or held_arm(first) is None + or held_arm(first) != holder_arm + ): + raise ValueError( + f"SeedGraph E2->E4 recovery holder mismatch for object " + f"{object_uid!r}." + ) + else: + predecessor_actions = { + str(node.get("atomic_action")) for node in predecessor_nodes + } + if not predecessor_actions.intersection({"Place", "AxisAlign"}): + raise ValueError( + f"SeedGraph E2 predecessor {predecessor['id']!r} must " + "release its object before E4." + ) + if ( + not group_nodes + or str(group_nodes[0].get("atomic_action")) != "PickUp" + ): + raise ValueError( + f"SeedGraph E4 group {group_id!r} must reacquire the " + "supported E2 object with PickUp." + ) + + if task_type == "E1": + predecessor = direct_predecessor(group, "E4") + if ( + predecessor is not None + and str(predecessor.get("object_uid")) == object_uid + ): + if "PickUp" in actions or not group_nodes: + raise ValueError( + f"SeedGraph E1 group {group_id!r} must preserve the E4 receiver hold " + "without PickUp." + ) + first = group_nodes[0] + if ( + str(first.get("atomic_action")) != "MoveHeldObject" + or held_arm(first) is None + ): + raise ValueError( + f"SeedGraph E1 group {group_id!r} must start with MoveHeldObject " + "from the receiver hold." + ) + + +def _validate_group_dependency_alignment( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], +) -> None: + group_by_node = { + str(node_id): str(group["id"]) + for group in groups + for node_id in group["node_ids"] + } + group_dependencies = { + str(group["id"]): set(str(parent) for parent in group["depends_on"]) + for group in groups + } + + def group_reaches(child: str, parent: str) -> bool: + pending = list(group_dependencies[child]) + visited = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(group_dependencies[current]) + return False + + for node in nodes: + child_group = group_by_node[str(node["id"])] + for dependency in node["depends_on"]: + parent_group = group_by_node[str(dependency)] + if parent_group != child_group and not group_reaches( + child_group, parent_group + ): + raise ValueError( + f"SeedGraph node {node['id']!r} depends on TaskGroup " + f"{parent_group!r}, but TaskGroup {child_group!r} does not." + ) + + +def _validate_resource_conflicts( + nodes: Sequence[Mapping[str, Any]], + groups: Sequence[Mapping[str, Any]], + metadata: Any, +) -> None: + dependencies = { + str(node["id"]): set(str(item) for item in node["depends_on"]) for node in nodes + } + + def reaches(child: str, parent: str) -> bool: + pending = list(dependencies[child]) + visited = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(dependencies[current]) + return False + + group_by_node = { + str(node_id): str(group["id"]) + for group in groups + for node_id in group["node_ids"] + } + distinct_arm_pairs = _distinct_arm_pairs(metadata) + for index, first in enumerate(nodes): + for second in nodes[index + 1 :]: + first_id = str(first["id"]) + second_id = str(second["id"]) + if reaches(first_id, second_id) or reaches(second_id, first_id): + continue + if ( + first.get("sync_group") == second.get("sync_group") + and first.get("sync_group") is not None + ): + continue + first_claims = { + str(claim["resource"]): str(claim["access"]) + for claim in first["contract"]["claims"] + } + second_claims = { + str(claim["resource"]): str(claim["access"]) + for claim in second["contract"]["claims"] + } + conflicts = sorted( + resource + for resource in set(first_claims) & set(second_claims) + if "exclusive" in {first_claims[resource], second_claims[resource]} + ) + if ( + frozenset({group_by_node[first_id], group_by_node[second_id]}) + in distinct_arm_pairs + ): + conflicts = [item for item in conflicts if item != "arm:auto"] + if conflicts: + raise ValueError( + f"SeedGraph concurrent nodes {first_id!r} and {second_id!r} " + f"have resource conflicts: {conflicts}." + ) + + +def _distinct_arm_pairs(value: Any) -> set[frozenset[str]]: + if not isinstance(value, Mapping): + return set() + groups = value.get("legacy_allocation_groups", value.get("allocation_groups", ())) + if not isinstance(groups, Sequence) or isinstance(groups, (str, bytes, bytearray)): + return set() + result: set[frozenset[str]] = set() + for group in groups: + if ( + not isinstance(group, Mapping) + or group.get("arm_constraint") != "distinct_arms" + ): + continue + members = group.get("semantic_step_ids", group.get("task_instance_ids", ())) + if not isinstance(members, Sequence) or isinstance( + members, (str, bytes, bytearray) + ): + continue + member_ids = [str(item) for item in members] + for index, first in enumerate(member_ids): + for second in member_ids[index + 1 :]: + result.add(frozenset({first, second})) + return result + + +def _action_contract(value: Any, context: str) -> dict[str, Any]: + contract = _mapping(value, context) + _keys(contract, _ACTION_CONTRACT_KEYS, context) + if set(contract) != _ACTION_CONTRACT_KEYS: + missing = sorted(_ACTION_CONTRACT_KEYS - set(contract)) + raise ValueError(f"{context} is missing required fields: {missing}.") + if contract["version"] != "action_contract_v2": + raise ValueError(f"{context}.version must be 'action_contract_v2'.") + contract["requires"] = [ + _state_atom(item, f"{context}.requires[{index}]") + for index, item in enumerate( + _sequence(contract["requires"], f"{context}.requires") + ) + ] + contract["effects"] = [ + _state_effect(item, f"{context}.effects[{index}]") + for index, item in enumerate( + _sequence(contract["effects"], f"{context}.effects") + ) + ] + contract["claims"] = [ + _resource_claim(item, f"{context}.claims[{index}]") + for index, item in enumerate(_sequence(contract["claims"], f"{context}.claims")) + ] + _unique( + [str(item["resource"]) for item in contract["claims"]], + f"{context}.claims resources", + ) + contract["completion"] = _enum( + contract["completion"], _ACTION_COMPLETION, f"{context}.completion" + ) + contract["failure_policy"] = _enum( + contract["failure_policy"], + _FAILURE_POLICIES, + f"{context}.failure_policy", + ) + return contract + + +def _task_group_contract(value: Any, context: str) -> dict[str, Any]: + contract = _mapping(value, context) + _keys(contract, _TASK_GROUP_CONTRACT_KEYS, context) + if set(contract) != _TASK_GROUP_CONTRACT_KEYS: + missing = sorted(_TASK_GROUP_CONTRACT_KEYS - set(contract)) + raise ValueError(f"{context} is missing required fields: {missing}.") + contract["entry_requires"] = [ + _state_atom(item, f"{context}.entry_requires[{index}]") + for index, item in enumerate( + _sequence(contract["entry_requires"], f"{context}.entry_requires") + ) + ] + contract["exit_effects"] = [ + _state_effect(item, f"{context}.exit_effects[{index}]") + for index, item in enumerate( + _sequence(contract["exit_effects"], f"{context}.exit_effects") + ) + ] + contract["claims"] = [ + _resource_claim(item, f"{context}.claims[{index}]") + for index, item in enumerate(_sequence(contract["claims"], f"{context}.claims")) + ] + _unique( + [str(item["resource"]) for item in contract["claims"]], + f"{context}.claims resources", + ) + contract["entry_node_ids"] = _strings( + contract["entry_node_ids"], f"{context}.entry_node_ids" + ) + contract["terminal_node_ids"] = _strings( + contract["terminal_node_ids"], f"{context}.terminal_node_ids" + ) + if not contract["entry_node_ids"] or not contract["terminal_node_ids"]: + raise ValueError(f"{context} requires entry and terminal node IDs.") + contract["completion"] = _enum( + contract["completion"], _GROUP_COMPLETION, f"{context}.completion" + ) + return contract + + +def _state_atom(value: Any, context: str) -> dict[str, str]: + atom = _mapping(value, context) + _keys(atom, _STATE_ATOM_KEYS, context) + predicate = _enum(atom.get("predicate"), _STATE_PREDICATES, f"{context}.predicate") + required = { + "arm_free": {"arm"}, + "object_free": {"object_uid"}, + "object_held": {"object_uid", "arm"}, + "object_coordinated_held": {"object_uid"}, + "handover_complete": {"object_uid"}, + "arm_clear": {"arm"}, + "arm_home": {"arm"}, + }[predicate] + present = set(atom) - {"predicate"} + if present != required: + raise ValueError( + f"{context} predicate {predicate!r} requires exactly {sorted(required)}." + ) + for field in required: + atom[field] = _string(atom.get(field), f"{context}.{field}") + return atom + + +def _state_effect(value: Any, context: str) -> dict[str, Any]: + effect = _mapping(value, context) + _keys(effect, _EFFECT_KEYS, context) + if set(effect) != _EFFECT_KEYS: + raise ValueError(f"{context} requires op and atom.") + effect["op"] = _enum(effect["op"], _EFFECT_OPERATIONS, f"{context}.op") + effect["atom"] = _state_atom(effect["atom"], f"{context}.atom") + return effect + + +def _resource_claim(value: Any, context: str) -> dict[str, str]: + claim = _mapping(value, context) + _keys(claim, _CLAIM_KEYS, context) + if set(claim) != _CLAIM_KEYS: + raise ValueError(f"{context} requires resource, access, and lifetime.") + claim["resource"] = _string(claim["resource"], f"{context}.resource") + claim["access"] = _enum(claim["access"], _CLAIM_ACCESS, f"{context}.access") + claim["lifetime"] = _enum( + claim["lifetime"], _CLAIM_LIFETIMES, f"{context}.lifetime" + ) + return claim + + +def _validate_group_contract_topology( + nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]] +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + children: dict[str, set[str]] = {node_id: set() for node_id in node_by_id} + for node in nodes: + for dependency in node["depends_on"]: + children[str(dependency)].add(str(node["id"])) + for group in groups: + group_id = str(group["id"]) + node_ids = {str(item) for item in group["node_ids"]} + expected_entries = { + node_id + for node_id in node_ids + if not any( + str(parent) in node_ids for parent in node_by_id[node_id]["depends_on"] + ) + } + expected_terminals = { + node_id for node_id in node_ids if not (children[node_id] & node_ids) + } + contract = group["contract"] + if set(contract["entry_node_ids"]) != expected_entries: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} contract entry_node_ids do not " + "match its internal topology." + ) + if set(contract["terminal_node_ids"]) != expected_terminals: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} contract terminal_node_ids do not " + "match its internal topology." + ) + if contract["completion"] == "terminal_barrier" and not all( + node_by_id[node_id]["contract"]["completion"] == "terminal_barrier" + for node_id in expected_terminals + ): + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} terminal barrier must end in " + "terminal_barrier AtomicActions." + ) + + +def _validate_cleanup_barriers( + nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]] +) -> None: + node_by_id = {str(node["id"]): node for node in nodes} + for group in groups: + group_id = str(group["id"]) + group_nodes = [node_by_id[str(node_id)] for node_id in group["node_ids"]] + cleanup = [ + node for node in group_nodes if node["contract"]["completion"] == "cleanup" + ] + has_handover = any(node["atomic_action"] == "HandOver" for node in group_nodes) + recovery_successor = any( + candidate.get("role") == "recovery" + and candidate.get("parent_task_instance_id") == group_id + and group_id in candidate.get("depends_on", ()) + for candidate in groups + ) + if has_handover and not cleanup and recovery_successor: + continue + if has_handover and not cleanup: + raise ValueError( + f"SeedGraph HandOver TaskGroup {group_id!r} is missing retreat cleanup." + ) + if not cleanup and not has_handover: + continue + terminal_ids = group["contract"]["terminal_node_ids"] + if group["contract"]["completion"] != "terminal_barrier" or not all( + node_by_id[node_id]["contract"]["completion"] == "terminal_barrier" + for node_id in terminal_ids + ): + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} cleanup must end at a home " + "terminal barrier." + ) + + +def _actor(value: Any, context: str) -> dict[str, Any]: + actor = _mapping(value, context) + mode = _enum(actor.get("mode"), _ACTOR_MODES, f"{context}.mode") + allowed = {"mode"} + if mode in {"required", "preferred"}: + allowed.add("arm") + _string(actor.get("arm"), f"{context}.arm") + elif mode == "coordinated": + allowed.add("arms") + arms = _strings(actor.get("arms"), f"{context}.arms") + if len(arms) < 2: + raise ValueError(f"{context}.arms must contain at least two arms.") + actor["arms"] = arms + _keys(actor, frozenset(allowed), context) + return actor + + +def _known_object_references( + value: Any, known: set[str], path: str = "SeedGraph" +) -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + child_path = f"{path}.{key}" + if key in _OBJECT_REFERENCE_KEYS and isinstance(child, str): + if child not in known and child not in {"table_center", "world"}: + raise ValueError( + f"{child_path} references unknown object {child!r}." + ) + _known_object_references(child, known, child_path) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _known_object_references(child, known, f"{path}[{index}]") + + +def _reject_grounded(value: Any, path: str = "document") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _GROUNDED_FIELD_NAMES: + raise ValueError(f"{path}.{key} contains grounded motion data.") + _reject_grounded(child, f"{path}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _reject_grounded(child, f"{path}[{index}]") + + +def _finite(value: Any, path: str = "document") -> None: + if isinstance(value, Mapping): + for key, child in value.items(): + _finite(child, f"{path}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _finite(child, f"{path}[{index}]") + elif isinstance(value, float) and not math.isfinite(value): + raise ValueError(f"{path} must be finite.") + + +def _dag(dependencies: Mapping[str, Sequence[str]], context: str) -> None: + known = set(dependencies) + outgoing = {item_id: [] for item_id in known} + indegree = {item_id: 0 for item_id in known} + for item_id, required in dependencies.items(): + unknown = set(required) - known + if unknown: + raise ValueError(f"{context} reference unknown IDs: {sorted(unknown)}.") + if item_id in required: + raise ValueError(f"{context} contain a self-dependency at {item_id!r}.") + for parent in required: + outgoing[parent].append(item_id) + indegree[item_id] += 1 + ready = deque( + sorted(item_id for item_id, degree in indegree.items() if degree == 0) + ) + visited = 0 + while ready: + item_id = ready.popleft() + visited += 1 + for child in sorted(outgoing[item_id]): + indegree[child] -= 1 + if indegree[child] == 0: + ready.append(child) + if visited != len(known): + cyclic = sorted(item_id for item_id, degree in indegree.items() if degree) + raise ValueError(f"{context} contain a dependency cycle: {cyclic}.") + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise TypeError(f"{context} must be a list.") + return list(value) + + +def _strings(value: Any, context: str) -> list[str]: + result = [ + _string(item, f"{context}[{index}]") + for index, item in enumerate(_sequence(value, context)) + ] + _unique(result, context) + return result + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{context} must be a non-empty string.") + return value.strip() + + +def _enum(value: Any, allowed: Collection[str], context: str) -> str: + result = _string(value, context) + if result not in allowed: + raise ValueError(f"{context} must be one of {sorted(allowed)}.") + return result + + +def _unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must be unique.") + + +def _keys(value: Mapping[str, Any], allowed: frozenset[str], context: str) -> None: + unknown = sorted(set(value) - allowed) + if unknown: + raise ValueError(f"{context} contains unsupported fields: {unknown}.") + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") diff --git a/embodichain/gen_sim/action_engine/domain/visual_contracts.py b/embodichain/gen_sim/action_engine/domain/visual_contracts.py new file mode 100644 index 000000000..78ca4a0f1 --- /dev/null +++ b/embodichain/gen_sim/action_engine/domain/visual_contracts.py @@ -0,0 +1,60 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Canonical visual-fact contracts shared by planning and evaluation.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Any + +__all__ = [ + "OCCLUSION_RELATION", + "VISUAL_RELATION_PARTICIPANTS", + "requested_visual_task_predicates", +] + + +OCCLUSION_RELATION = "occludes" + +# Participant order is semantic. For ``occludes`` it is +# ``[occluder_uid, occluded_uid]``. +VISUAL_RELATION_PARTICIPANTS: Mapping[str, tuple[str, ...]] = MappingProxyType( + {OCCLUSION_RELATION: ("occluder", "occluded")} +) + + +def requested_visual_task_predicates(task_spec: Mapping[str, Any]) -> frozenset[str]: + """Return task-level visual predicates explicitly requested by a TaskSpec.""" + result: set[str] = set() + + def collect(value: Any) -> None: + if isinstance(value, Mapping): + if value.get("type") == "visual_relation": + relation = value.get("relation") + if isinstance(relation, str) and relation: + result.add(relation) + for child in value.values(): + collect(child) + elif isinstance(value, Sequence) and not isinstance( + value, (str, bytes, bytearray) + ): + for child in value: + collect(child) + + collect(task_spec.get("success", {})) + return frozenset(result) diff --git a/embodichain/gen_sim/action_engine/environment/__init__.py b/embodichain/gen_sim/action_engine/environment/__init__.py new file mode 100644 index 000000000..269a419c1 --- /dev/null +++ b/embodichain/gen_sim/action_engine/environment/__init__.py @@ -0,0 +1,23 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tracked Action Engine environment package.""" + +from __future__ import annotations + +from .agent_env import ACTION_ENGINE_ENV_ID, ActionEngineEnv + +__all__ = ["ACTION_ENGINE_ENV_ID", "ActionEngineEnv"] diff --git a/embodichain/gen_sim/action_engine/environment/agent_env.py b/embodichain/gen_sim/action_engine/environment/agent_env.py new file mode 100644 index 000000000..ed927cfce --- /dev/null +++ b/embodichain/gen_sim/action_engine/environment/agent_env.py @@ -0,0 +1,535 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Gym environment that executes Action Engine programs against live state.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + generation_defaults, + resolve_agent_runtime_policy, +) +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import validate_seed_graph +from embodichain.gen_sim.action_engine.protocol import ACTION_ENGINE_ENV_ID +from embodichain.gen_sim.action_engine.runtime import ( + ProgramExecutor, + evaluate_predicate, + load_agent_execution_program, + load_execution_program, +) +from embodichain.gen_sim.action_engine.runtime.solver_compat import ( + install_action_engine_solver_compat, + repair_action_engine_ur5_solver_cfg, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) +from embodichain.lab.gym.envs import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env + +__all__ = ["ACTION_ENGINE_ENV_ID", "ActionEngineEnv"] + +_MAX_EPISODE_STEPS = int(generation_defaults()["task"]["max_episode_steps"]) + + +@register_env(ACTION_ENGINE_ENV_ID, max_episode_steps=_MAX_EPISODE_STEPS) +class ActionEngineEnv(EmbodiedEnv): + """EmbodiedEnv adapter for in-memory compiled execution programs.""" + + def __init__(self, cfg: EmbodiedEnvCfg | None = None, **kwargs: Any) -> None: + agent_config = kwargs.pop("agent_config", None) + task_name = kwargs.pop("task_name", None) + agent_config_path = kwargs.pop("agent_config_path", None) + runtime_backend = kwargs.pop("runtime_backend", "independent") + runtime_policy = kwargs.pop("runtime_policy", None) + if not isinstance(agent_config, Mapping): + raise ValueError("ActionEngineEnv requires an agent_config mapping.") + if not isinstance(task_name, str) or not task_name: + raise ValueError("ActionEngineEnv requires a non-empty task_name.") + if not isinstance(agent_config_path, str) or not agent_config_path: + raise ValueError("ActionEngineEnv requires agent_config_path.") + self.agent_config = dict(agent_config) + self.agent_config_path = agent_config_path + self.task_name = task_name + if runtime_policy is None: + runtime_policy = resolve_agent_runtime_policy(self.agent_config) + if not isinstance(runtime_policy, RuntimePolicyCfg): + raise TypeError("ActionEngineEnv runtime_policy must be RuntimePolicyCfg.") + self.runtime_policy = runtime_policy + if runtime_backend != "independent": + raise ValueError( + "ActionEngineEnv only supports its independent runtime, got " + f"{runtime_backend!r}." + ) + self.runtime_backend = str(runtime_backend) + self.last_execution: Any | None = None + self._runtime_state_ready = False + repair_action_engine_ur5_solver_cfg(getattr(cfg, "robot", None)) + super().__init__(cfg, **kwargs) + install_action_engine_solver_compat(self.robot) + if bool(getattr(self, "ignore_terminations_during_agent", False)): + # Atomic trajectories execute online through env.step(). Prevent a + # transient task signal from resetting an environment mid-program. + self.cfg.ignore_terminations = True + self._capture_runtime_state() + + def reset( + self, + seed: int | None = None, + options: dict[str, Any] | None = None, + ) -> tuple[Any, dict[str, Any]]: + self._runtime_state_ready = False + observation, info = super().reset(seed=seed, options=options) + self.last_execution = None + self._capture_runtime_state() + return observation, info + + def _capture_runtime_state(self) -> None: + """Capture reset-relative robot and object state used by symbolic bindings.""" + self.init_qpos = self.robot.get_qpos().clone() + self._agent_arm_slots = self._resolve_arm_slots() + for side in ("left", "right"): + self._initialize_arm(side, self._agent_arm_slots.get(side)) + + default_open = getattr(self, "gripper_open_state", (0.04, 0.04)) + default_close = getattr(self, "gripper_close_state", (0.0, 0.0)) + self.open_state = torch.as_tensor( + getattr(self, "agent_open_state", default_open), + dtype=self.init_qpos.dtype, + device=self.init_qpos.device, + ).flatten() + self.close_state = torch.as_tensor( + getattr(self, "agent_close_state", default_close), + dtype=self.init_qpos.dtype, + device=self.init_qpos.device, + ).flatten() + self.left_arm_current_gripper_state = self._hand_qpos("left") + self.right_arm_current_gripper_state = self._hand_qpos("right") + self.update_obj_info() + self.agent_initial_object_poses = { + uid: item["pose"].clone() for uid, item in self.obj_info.items() + } + self.agent_initial_object_heights = { + uid: item["height"].clone() for uid, item in self.obj_info.items() + } + self._runtime_state_ready = True + + def _resolve_arm_slots(self) -> dict[str, dict[str, str | None] | None]: + configured = getattr(self, "agent_arm_slots", None) + if isinstance(configured, Mapping): + result: dict[str, dict[str, str | None] | None] = { + "left": None, + "right": None, + } + for side in result: + value = configured.get(side) + if isinstance(value, str): + result[side] = {"arm": value, "eef": None} + elif isinstance(value, Mapping): + result[side] = { + "arm": value.get("arm", value.get("arm_control_part")), + "eef": value.get( + "eef", + value.get("hand", value.get("eef_control_part")), + ), + } + return result + parts = getattr(self.robot, "control_parts", {}) or {} + if "left_arm" in parts or "right_arm" in parts: + return { + "left": {"arm": "left_arm", "eef": "left_eef"}, + "right": {"arm": "right_arm", "eef": "right_eef"}, + } + if "arm" in parts: + side = str(getattr(self, "agent_single_arm_slot", "right")) + result = {"left": None, "right": None} + result[side] = {"arm": "arm", "eef": "hand"} + return result + raise ValueError("Robot exposes no arm control part for Action Engine.") + + def _initialize_arm( + self, + side: str, + slot: dict[str, str | None] | None, + ) -> None: + arm = None if slot is None else slot.get("arm") + eef = None if slot is None else slot.get("eef") + arm_ids = self._control_part_ids(arm) + eef_ids = self._control_part_ids(eef) + setattr(self, f"{side}_arm_joints", arm_ids) + setattr(self, f"{side}_eef_joints", eef_ids) + arm_qpos = self.init_qpos[:, arm_ids] + setattr(self, f"{side}_arm_init_qpos", arm_qpos.clone()) + setattr(self, f"{side}_arm_current_qpos", arm_qpos.clone()) + if arm is None or not arm_ids: + setattr(self, f"{side}_arm_init_xpos", None) + setattr(self, f"{side}_arm_current_xpos", None) + return + xpos = self.robot.compute_fk(arm_qpos, name=arm, to_matrix=True) + setattr(self, f"{side}_arm_init_xpos", xpos.clone()) + setattr(self, f"{side}_arm_current_xpos", xpos.clone()) + + def _control_part_ids(self, name: str | None) -> list[int]: + if name is None: + return [] + parts = getattr(self.robot, "control_parts", {}) or {} + if name not in parts: + return [] + return list(self.robot.get_joint_ids(name=name)) + + def _hand_qpos(self, side: str) -> torch.Tensor: + ids = list(getattr(self, f"{side}_eef_joints", ())) + return self.init_qpos[:, ids].clone() + + def get_agent_arm_control_part(self, is_left: bool) -> str: + value = self._agent_arm_slots["left" if is_left else "right"] + arm = None if value is None else value.get("arm") + if not isinstance(arm, str) or not arm: + raise ValueError(f"{'left' if is_left else 'right'} arm is not configured.") + return arm + + def get_agent_eef_control_part(self, is_left: bool) -> str | None: + value = self._agent_arm_slots["left" if is_left else "right"] + eef = None if value is None else value.get("eef") + return str(eef) if eef else None + + def get_current_qpos_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return tuple( + qpos[:, list(getattr(self, f"{side}_arm_joints", ()))].clone() + for side in ("left", "right") + ) + + def is_object_pressed( + self, + uid: str, + terminal_state: str = "activated", + ) -> torch.Tensor: + """Verify a calibrated button state from its live articulation qpos.""" + articulation = self.sim.get_articulation(uid) + if articulation is None: + return torch.zeros(int(self.num_envs), dtype=torch.bool, device=self.device) + settings = self.agent_config.get("articulation_settings", {}) + per_articulation = ( + settings.get(uid, {}) if isinstance(settings, Mapping) else {} + ) + matches = [ + (name, values) + for name, values in per_articulation.items() + if name in articulation.joint_names + and isinstance(values, Sequence) + and not isinstance(values, (str, bytes, bytearray)) + and len(values) >= 2 + ] + if len(matches) != 1: + return torch.zeros(int(self.num_envs), dtype=torch.bool, device=self.device) + joint_name, raw_values = matches[0] + values = torch.as_tensor( + raw_values, dtype=torch.float32, device=self.device + ).flatten() + if not torch.isfinite(values).all(): + return torch.zeros(int(self.num_envs), dtype=torch.bool, device=self.device) + if terminal_state == "activated": + target = values[-1] + elif terminal_state == "inactive": + target = values[0] + else: + return torch.zeros(int(self.num_envs), dtype=torch.bool, device=self.device) + unique = torch.unique(values, sorted=True) + if unique.numel() < 2: + return torch.zeros(int(self.num_envs), dtype=torch.bool, device=self.device) + minimum_spacing = torch.diff(unique).abs().min() + fallback = float(self.runtime_policy.predicate_fallbacks["axis_tolerance"]) + tolerance = min(fallback, float(minimum_spacing) * 0.25) + joint_id = articulation.joint_names.index(joint_name) + qpos = articulation.get_qpos()[:, joint_id] + return torch.isfinite(qpos) & (torch.abs(qpos - target) <= tolerance) + + def set_current_qpos_agent( + self, + arm_qpos: torch.Tensor, + is_left: bool, + ) -> None: + side = "left" if is_left else "right" + setattr(self, f"{side}_arm_current_qpos", arm_qpos) + + def get_current_xpos_agent( + self, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + qpos = self.robot.get_qpos() + result = [] + for side in ("left", "right"): + slot = self._agent_arm_slots.get(side) + arm = None if slot is None else slot.get("arm") + arm_ids = list(getattr(self, f"{side}_arm_joints", ())) + if not arm or not arm_ids: + result.append(None) + continue + result.append( + self.robot.compute_fk( + qpos[:, arm_ids], + name=arm, + to_matrix=True, + ) + ) + return result[0], result[1] + + def set_current_xpos_agent( + self, + arm_xpos: torch.Tensor, + is_left: bool, + ) -> None: + side = "left" if is_left else "right" + setattr(self, f"{side}_arm_current_xpos", arm_xpos) + + def get_current_gripper_state_agent( + self, + ) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return tuple( + qpos[:, list(getattr(self, f"{side}_eef_joints", ()))].clone() + for side in ("left", "right") + ) + + def set_current_gripper_state_agent( + self, + arm_gripper_state: torch.Tensor, + is_left: bool, + ) -> None: + side = "left" if is_left else "right" + setattr(self, f"{side}_arm_current_gripper_state", arm_gripper_state) + + def get_arm_fk(self, qpos: torch.Tensor, is_left: bool) -> torch.Tensor: + return self.robot.compute_fk( + name=self.get_agent_arm_control_part(is_left), + qpos=torch.as_tensor(qpos, device=self.robot.device), + to_matrix=True, + ) + + def sync_agent_state_from_qpos(self, qpos: torch.Tensor) -> None: + """Keep arm-selection seeds synchronized with the command sent to sim.""" + qpos = torch.as_tensor( + qpos, + dtype=self.init_qpos.dtype, + device=self.init_qpos.device, + ) + for side in ("left", "right"): + arm_ids = list(getattr(self, f"{side}_arm_joints", ())) + hand_ids = list(getattr(self, f"{side}_eef_joints", ())) + arm_qpos = qpos[:, arm_ids] + setattr(self, f"{side}_arm_current_qpos", arm_qpos.clone()) + slot = self._agent_arm_slots.get(side) + arm = None if slot is None else slot.get("arm") + if arm and arm_ids: + xpos = self.robot.compute_fk(arm_qpos, name=arm, to_matrix=True) + setattr(self, f"{side}_arm_current_xpos", xpos) + setattr( + self, + f"{side}_arm_current_gripper_state", + qpos[:, hand_ids].clone(), + ) + + def get_arm_ik( + self, + target_xpos: torch.Tensor, + is_left: bool, + qpos_seed: torch.Tensor | None = None, + env_ids: list[int] | None = None, + ) -> tuple[bool, torch.Tensor]: + success, qpos = self.robot.compute_ik( + name=self.get_agent_arm_control_part(is_left), + pose=target_xpos, + joint_seed=qpos_seed, + env_ids=env_ids, + ) + success_value = ( + bool(torch.as_tensor(success).all().item()) + if isinstance(success, torch.Tensor) + else bool(success) + ) + return success_value, qpos + + def update_obj_info(self) -> None: + info = getattr(self, "obj_info", {}) + for uid in self.sim.get_rigid_object_uid_list(): + entity = self.sim.get_rigid_object(uid) + if entity is None: + continue + pose = entity.get_local_pose(to_matrix=True) + info[uid] = {"pose": pose, "height": pose[:, 2, 3]} + self.obj_info = info + + def create_demo_action_list( + self, + regenerate: bool = False, + **kwargs: Any, + ) -> Any: + """Compile in memory when requested, then execute the program online.""" + program = load_agent_execution_program( + self.agent_config, + agent_config_path=self.agent_config_path, + regenerate=regenerate, + ) + executor = ProgramExecutor( + program, + self, + max_transitions=( + int(self.action_engine_max_transitions) + if hasattr(self, "action_engine_max_transitions") + else None + ), + settle_steps=( + int(self.action_engine_settle_steps) + if hasattr(self, "action_engine_settle_steps") + else None + ), + record_runtime=bool(getattr(self, "action_engine_record_runtime", True)), + record_root=getattr(self, "action_engine_record_root", None), + runtime_policy=self.runtime_policy, + ) + self.last_execution = executor.run( + run_id=kwargs.get("runtime_run_id"), + episode_index=int(kwargs.get("episode_index", 0)), + ) + return self.last_execution + + def execute_seed_graph( + self, + seed_graph: Mapping[str, Any], + *, + runtime_run_id: str, + episode_index: int, + record_root: str | None = None, + ) -> Any: + """Execute one already validated branch graph without rewriting config.""" + program = self.preflight_seed_graph(seed_graph) + route = getattr(self, "action_engine_ab_route", None) + graph_route = seed_graph.get("planner_route") + if route in {"offline", "online"} and graph_route != route: + raise ValueError( + f"A/B branch route {route!r} cannot execute graph route " + f"{graph_route!r}." + ) + executor = ProgramExecutor( + program, + self, + max_transitions=( + int(self.action_engine_max_transitions) + if hasattr(self, "action_engine_max_transitions") + else None + ), + settle_steps=( + int(self.action_engine_settle_steps) + if hasattr(self, "action_engine_settle_steps") + else None + ), + record_runtime=bool(getattr(self, "action_engine_record_runtime", True)), + record_root=record_root, + runtime_policy=self.runtime_policy, + ) + self.last_execution = executor.run( + run_id=runtime_run_id, + episode_index=episode_index, + ) + return self.last_execution + + def preflight_seed_graph(self, seed_graph: Mapping[str, Any]) -> Any: + """Validate/compile one branch graph without stepping the simulator. + + This hook is intentionally separate from :meth:`execute_seed_graph` so + strict A/B can preflight both branches before either executor sends a + command to the robot. + """ + source = self.agent_config.get("source", {}) + if not isinstance(source, Mapping): + source = {} + uid_map = source.get("uid_map", {}) + if not isinstance(uid_map, Mapping): + uid_map = {} + known_objects = {str(uid) for uid in uid_map.values() if str(uid)} + registry = build_atomic_capability_registry() + graph = validate_seed_graph( + seed_graph, + known_objects=known_objects or None, + known_actions=registry.names(), + executable_actions=registry.executable_names(), + require_executable=True, + ) + for node in graph["nodes"]: + registry.validate_binding(node) + resolve_motion_policy( + str( + self.agent_config.get( + "robot_profile", + getattr(self, "agent_robot_profile", "dual_ur10"), + ) + ), + node["motion_policy"], + ) + return load_execution_program( + graph, + known_objects=known_objects or None, + registry=registry, + ) + + def _normalize_demo_action_list(self, action_list: Any) -> Any: + """Preserve metadata on action streams that already ran online. + + ``EmbodiedEnv`` normally rebuilds returned sequences after validating + their action width. Rebuilding an ``ExecutionResult`` would discard its + success masks and runtime-record location, and its commands have + already been sent to the simulator, so no replay normalization is + needed. + """ + if getattr(action_list, "already_executed", False): + return action_list + return super()._normalize_demo_action_list(action_list) + + def is_task_success(self, **_: Any) -> torch.Tensor: + configured = getattr(self, "agent_success", None) + if isinstance(configured, Mapping): + return evaluate_predicate(self, configured) + if self.last_execution is not None: + return torch.as_tensor( + getattr( + self.last_execution, + "runtime_success", + getattr(self.last_execution, "success", False), + ), + dtype=torch.bool, + device=self.device, + ) + return torch.zeros( + int(self.num_envs), + dtype=torch.bool, + device=self.device, + ) + + def compute_task_state( + self, + **_: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + success = self.is_task_success() + return success, torch.zeros_like(success), {} diff --git a/embodichain/gen_sim/action_engine/evaluation/__init__.py b/embodichain/gen_sim/action_engine/evaluation/__init__.py new file mode 100644 index 000000000..2d2c15190 --- /dev/null +++ b/embodichain/gen_sim/action_engine/evaluation/__init__.py @@ -0,0 +1,27 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict offline/online comparison utilities.""" + +from __future__ import annotations + +from .ab import ABExecutionResult, run_strict_ab, state_digest + +__all__ = [ + "ABExecutionResult", + "run_strict_ab", + "state_digest", +] diff --git a/embodichain/gen_sim/action_engine/evaluation/ab.py b/embodichain/gen_sim/action_engine/evaluation/ab.py new file mode 100644 index 000000000..1c95a61ea --- /dev/null +++ b/embodichain/gen_sim/action_engine/evaluation/ab.py @@ -0,0 +1,831 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Execute offline and online SeedGraphs from strictly identical resets.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass +import hashlib +import json +from pathlib import Path +from time import perf_counter +from typing import Any + +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.domain import ( + seed_graph_hash, + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.protocol import ( + COMPARISON_FILENAME, + EXECUTION_PROGRAM_FILENAME, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +__all__ = ["ABExecutionResult", "run_strict_ab", "state_digest"] + +EnvFactory = Callable[..., Any] +ExecutorFactory = Callable[[Mapping[str, Any], Any], Any] +SnapshotReader = Callable[[Any], Mapping[str, Any]] +SuccessEvaluator = Callable[..., Any] +BranchFinalizer = Callable[..., list[str]] + +_FULL_SNAPSHOT_KEYS = frozenset( + {"robot_qpos", "object_poses", "articulation_state", "camera_calibration"} +) +_PRIVATE_OR_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +@dataclass(frozen=True) +class ABExecutionResult: + """Paths and summaries from one strict A/B run.""" + + comparison_path: Path + offline_dir: Path + online_dir: Path + initial_state_digest: str + comparison: dict[str, Any] + + +def state_digest(snapshot: Mapping[str, Any]) -> str: + """Hash nested tensors/arrays/mappings without lossy JSON conversion.""" + digest = hashlib.sha256() + _update_digest(digest, snapshot) + return digest.hexdigest() + + +def run_strict_ab( + task_spec: Mapping[str, Any], + offline_graph: Mapping[str, Any], + online_graph: Mapping[str, Any], + *, + env_factory: EnvFactory | None = None, + executor_factory: ExecutorFactory, + snapshot_reader: SnapshotReader, + output_dir: str | Path, + seed: int, + shared_config: Mapping[str, Any] | None = None, + planning_metrics: Mapping[str, Mapping[str, Any]] | None = None, + success_evaluator: SuccessEvaluator | None = None, + known_objects: set[str] | None = None, + expected_initial_state_digest: str | None = None, + branch_finalizer: BranchFinalizer | None = None, + episode_index: int = 0, + strict_state_digest: bool | None = None, + prepared_environments: Mapping[str, Any] | None = None, + prepared_snapshots: Mapping[str, Mapping[str, Any]] | None = None, + require_branch_videos: bool = False, +) -> ABExecutionResult: + """Run both planners in isolated environments after exact state checks. + + Callers that need visual observations before planning may supply two already + reset environments and their snapshots. The environments remain owned by + this function once supplied and are closed on every exit path. + + Set ``require_branch_videos`` for production A/B runs. In that mode each + finalizer must publish one non-empty ``video.mp4`` in its branch directory. + """ + supplied_environments = ( + tuple(prepared_environments.values()) + if prepared_environments is not None + else () + ) + try: + task, config, offline, online = _validate_ab_inputs( + task_spec, + offline_graph, + online_graph, + shared_config=shared_config, + success_evaluator=success_evaluator, + known_objects=known_objects, + ) + except BaseException: + # Prepared environments are already live before graph validation. The + # caller transfers ownership at function entry, including invalid-input + # paths that return before the normal environment scope below. + _close_environments(supplied_environments) + raise + + metrics = dict(planning_metrics or {}) + environments: dict[str, Any] = {} + try: + routes = ("offline", "online") + if prepared_environments is not None: + # Keep every supplied object in ``environments`` until after shape + # validation so the finally block closes extras on this error path. + environments = dict(prepared_environments) + if set(environments) != set(routes): + raise ValueError( + "prepared_environments must contain exactly offline and online." + ) + environments = {route: prepared_environments[route] for route in routes} + else: + if not callable(env_factory): + raise TypeError( + "env_factory is required when prepared_environments is not supplied." + ) + for route in routes: + environments[route] = env_factory( + route=route, + seed=int(seed), + config=config, + ) + if id(environments["offline"]) == id(environments["online"]): + raise RuntimeError( + "Strict A/B requires two isolated environment instances; " + "env_factory returned the same object twice." + ) + for route, env in environments.items(): + marker = getattr(env, "action_engine_ab_route", None) + if marker is not None and str(marker) != route: + raise RuntimeError( + f"A/B environment route marker {marker!r} does not match {route!r}." + ) + snapshots: dict[str, Mapping[str, Any]] = {} + if prepared_snapshots is not None and prepared_environments is None: + raise ValueError( + "prepared_snapshots requires prepared_environments so the state " + "being compared is unambiguous." + ) + if prepared_snapshots is not None: + if set(prepared_snapshots) != set(routes): + raise ValueError( + "prepared_snapshots must contain exactly offline and online." + ) + snapshots = {route: prepared_snapshots[route] for route in routes} + digests = {} + for route, env in environments.items(): + if prepared_snapshots is None: + env.reset(seed=int(seed)) + snapshots[route] = snapshot_reader(env) + _validate_snapshot(snapshots[route], route=route, require_full=False) + if strict_state_digest is None: + strict_state_digest = bool(config.get("strict_state_digest", False)) + # Automatically enforce the expanded contract whenever a caller + # supplies any of the new state components, while retaining the + # two-field v1 test helper compatibility. + strict_state_digest = strict_state_digest or any( + set(snapshot) & (_FULL_SNAPSHOT_KEYS - {"robot_qpos", "object_poses"}) + for snapshot in snapshots.values() + ) + if strict_state_digest: + for route, snapshot in snapshots.items(): + _validate_snapshot(snapshot, route=route, require_full=True) + for route, snapshot in snapshots.items(): + digests[route] = state_digest(snapshots[route]) + if digests["offline"] != digests["online"]: + raise RuntimeError( + "Strict A/B initial state mismatch: " + f"offline={digests['offline']}, online={digests['online']}." + ) + if ( + expected_initial_state_digest is not None + and digests["offline"] != expected_initial_state_digest + ): + raise RuntimeError( + "Strict A/B execution state does not match the online-planning " + f"snapshot: planning={expected_initial_state_digest}, " + f"execution={digests['offline']}." + ) + + root = Path(output_dir).expanduser().resolve() + branch_dirs = {route: root / route for route in environments} + for branch_dir in branch_dirs.values(): + branch_dir.mkdir(parents=True, exist_ok=True) + # Construct and preflight both executors before invoking either run. + # A route-specific executor may perform capability/robot checks that + # cannot be expressed in the serializable SeedGraph validator. + executors: dict[str, Any] = {} + for route, graph in (("offline", offline), ("online", online)): + _write_json(branch_dirs[route] / EXECUTION_PROGRAM_FILENAME, graph) + executors[route] = executor_factory(graph, environments[route]) + preflight_errors: dict[str, Exception] = {} + for route, executor in executors.items(): + preflight = getattr(executor, "preflight", None) + if not callable(preflight): + preflight = getattr(executor, "validate", None) + if not callable(preflight): + continue + try: + outcome = _call_preflight( + preflight, + route=route, + graph=(offline if route == "offline" else online), + env=environments[route], + ) + if outcome is not None: + try: + preflight_ok = bool(outcome) + except (TypeError, ValueError, RuntimeError) as exc: + raise RuntimeError( + f"{route} executor preflight returned a non-scalar result." + ) from exc + if not preflight_ok: + raise RuntimeError( + f"{route} executor preflight returned false." + ) + except Exception as exc: + preflight_errors[route] = exc + if preflight_errors: + detail = "; ".join( + f"{route}: {type(error).__name__}: {error}" + for route, error in sorted(preflight_errors.items()) + ) + raise RuntimeError( + "Strict A/B preflight failed; no branch was allowed to move. " + detail + ) + results = {} + finalization_errors: dict[str, Exception] = {} + for route, graph in (("offline", offline), ("online", online)): + result = None + started = perf_counter() + try: + executor = executors[route] + result = executor.run( + run_id=f"ab-{seed}-{route}", + episode_index=episode_index, + ) + elapsed = perf_counter() - started + success_override = ( + success_evaluator( + task_spec=task, + graph=graph, + env=environments[route], + result=result, + route=route, + ) + if success_evaluator is not None + else None + ) + results[route] = _result_summary( + result, + elapsed, + graph, + metrics.get(route, {}), + success_override=success_override, + ) + except Exception as exc: + elapsed = perf_counter() - started + results[route] = _error_result_summary( + exc, + elapsed, + graph, + metrics.get(route, {}), + ) + try: + raw_video_paths = ( + branch_finalizer( + route=route, + env=environments[route], + result=result, + branch_dir=branch_dirs[route], + episode_index=episode_index, + ) + if branch_finalizer is not None + else list(getattr(result, "video_paths", ())) + ) + video_paths = [str(path) for path in raw_video_paths] + if require_branch_videos: + _validate_branch_video_paths( + video_paths, + route=route, + branch_dir=branch_dirs[route], + ) + except Exception as exc: + video_paths = [] + results[route]["video_error"] = f"{type(exc).__name__}: {exc}" + finalization_errors[route] = exc + results[route]["video_paths"] = video_paths + results[route]["initial_state_digest"] = digests[route] + results[route]["seed_graph_hash"] = seed_graph_hash(graph) + _write_json( + branch_dirs[route] / "runtime_revisions.json", + { + "schema_version": "action_engine_runtime_revisions_v1", + "task_id": task["task_id"], + "route": route, + "revisions": list(getattr(result, "runtime_revisions", ())), + }, + ) + _write_json(branch_dirs[route] / "result.json", results[route]) + + comparison = { + "schema_version": "action_engine_ab_comparison_v1", + "task_id": task["task_id"], + "seed": int(seed), + "shared_config": config, + "initial_state_digest": digests["offline"], + "initial_state_digests": dict(digests), + "strict_state_digest": bool(strict_state_digest), + "graph_hashes": { + "offline": seed_graph_hash(offline), + "online": seed_graph_hash(online), + }, + "branches": { + "offline": { + **results["offline"], + }, + "online": { + **results["online"], + }, + }, + "graph_difference": _graph_difference(offline, online), + "video_finalization_errors": { + route: f"{type(error).__name__}: {error}" + for route, error in sorted(finalization_errors.items()) + }, + } + comparison_path = root / COMPARISON_FILENAME + _write_json(comparison_path, comparison) + if finalization_errors: + detail = "; ".join( + f"{route}: {type(error).__name__}: {error}" + for route, error in sorted(finalization_errors.items()) + ) + first_error = next(iter(finalization_errors.values())) + raise RuntimeError( + "Strict A/B branch video finalization failed; comparison report " + "was written with artifact errors. " + detail + ) from first_error + return ABExecutionResult( + comparison_path, + branch_dirs["offline"], + branch_dirs["online"], + digests["offline"], + comparison, + ) + finally: + _close_environments(environments.values()) + + +def _validate_ab_inputs( + task_spec: Mapping[str, Any], + offline_graph: Mapping[str, Any], + online_graph: Mapping[str, Any], + *, + shared_config: Mapping[str, Any] | None, + success_evaluator: SuccessEvaluator | None, + known_objects: set[str] | None, +) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any], dict[str, Any]]: + """Validate every serializable input before either branch may move.""" + task = validate_task_spec(task_spec) + if task["level"] == "L4" and not callable(success_evaluator): + raise ValueError( + "Strict L4 A/B requires a path-independent private-oracle " + "success_evaluator." + ) + config = dict(shared_config or {}) + capabilities = build_atomic_capability_registry() + offline = validate_seed_graph( + offline_graph, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=True, + ) + online = validate_seed_graph( + online_graph, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=True, + ) + robot_profile = str(config.get("robot_profile", "dual_ur10")) + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + for graph in (offline, online): + if graph["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError("A/B SeedGraph capability catalog does not match runtime.") + validate_persisted_contracts(graph, capabilities) + for node in graph["nodes"]: + capabilities.validate_binding(node) + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + _reject_private_or_live_fields(graph, "A/B SeedGraph") + if offline["task_id"] != task["task_id"] or online["task_id"] != task["task_id"]: + raise ValueError("A/B graphs and TaskSpec must have the same task_id.") + for route, graph in (("offline", offline), ("online", online)): + if ( + graph["level"] != task["level"] + or graph["reasoning_type"] != task["reasoning_type"] + ): + raise ValueError( + f"A/B {route} SeedGraph level/reasoning does not match TaskSpec." + ) + _validate_task_group_coverage(task, graph, route=route) + if offline["planner_route"] != "offline" or online["planner_route"] != "online": + raise ValueError( + "Strict A/B requires explicit offline and online graph routes." + ) + return task, config, offline, online + + +def _validate_branch_video_paths( + video_paths: list[str], *, route: str, branch_dir: Path +) -> None: + """Require the normalized video artifact used by strict production A/B.""" + expected = (branch_dir / "video.mp4").resolve() + if len(video_paths) != 1: + raise RuntimeError( + f"Strict A/B {route} branch must publish exactly one video.mp4." + ) + published = Path(video_paths[0]).expanduser().resolve() + if published != expected: + raise RuntimeError( + f"Strict A/B {route} video must be published as {expected.as_posix()}." + ) + if not expected.is_file() or expected.stat().st_size <= 0: + raise RuntimeError(f"Strict A/B {route} video.mp4 is missing or empty.") + + +def _close_environments(environments: Any) -> None: + """Best-effort close every distinct supplied environment exactly once.""" + seen: set[int] = set() + for env in environments: + if id(env) in seen: + continue + seen.add(id(env)) + close = getattr(env, "close", None) + if not callable(close): + continue + try: + close() + except Exception: + # Preserve the validation/execution failure that triggered cleanup, + # but continue closing the other independent branch. + continue + + +def _result_summary( + result: Any, + elapsed: float, + graph: Mapping[str, Any], + planning_metrics: Mapping[str, Any], + *, + success_override: Any | None, +) -> dict[str, Any]: + success = torch.as_tensor( + ( + getattr(result, "success", False) + if success_override is None + else success_override + ), + dtype=torch.bool, + ) + actions = list(getattr(result, "actions", ())) + retries = int(getattr(result, "retry_count", 0)) + recoveries = int(getattr(result, "recovery_count", 0)) + revisions = int(getattr(result, "revision_count", 0)) + metadata = graph.get("metadata", {}) + if not isinstance(metadata, Mapping): + metadata = {} + return { + "route": str(graph.get("planner_route", "")), + "seed_graph_hash": seed_graph_hash(graph), + "planning_seconds": float( + planning_metrics.get( + "planning_seconds", + metadata.get("planning_latency_seconds", 0.0), + ) + ), + "execution_seconds": float(elapsed), + "vlm_call_count": int( + planning_metrics.get( + "vlm_call_count", + metadata.get("vlm_call_count", 0), + ) + ), + "success": success.tolist(), + "success_source": ( + "runtime_postconditions" if success_override is None else "private_oracle" + ), + "success_rate": float(success.float().mean()) if success.numel() else 0.0, + "action_command_count": len(actions), + "path_length": _path_length(actions), + "retry_count": retries, + "recovery_count": recoveries, + "revision_count": revisions, + "failure_events": list(getattr(result, "failure_events", ())), + "ik_failure_count": sum( + item.get("failure_type") in {"plan_failed", "search_exhausted"} + for item in getattr(result, "failure_events", ()) + ), + "record_dir": getattr(result, "record_dir", None), + "video_paths": list(getattr(result, "video_paths", ())), + } + + +def _error_result_summary( + error: Exception, + elapsed: float, + graph: Mapping[str, Any], + planning_metrics: Mapping[str, Any], +) -> dict[str, Any]: + metadata = graph.get("metadata", {}) + if not isinstance(metadata, Mapping): + metadata = {} + return { + "route": str(graph.get("planner_route", "")), + "seed_graph_hash": seed_graph_hash(graph), + "planning_seconds": float( + planning_metrics.get( + "planning_seconds", + metadata.get("planning_latency_seconds", 0.0), + ) + ), + "execution_seconds": float(elapsed), + "vlm_call_count": int( + planning_metrics.get("vlm_call_count", metadata.get("vlm_call_count", 0)) + ), + "success": [False], + "success_source": "runtime_exception", + "success_rate": 0.0, + "action_command_count": 0, + "path_length": 0.0, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failure_events": [], + "ik_failure_count": 0, + "record_dir": None, + "video_paths": [], + "error": f"{type(error).__name__}: {error}", + } + + +def _path_length(actions: list[Any]) -> float: + if len(actions) < 2: + return 0.0 + tensors = [torch.as_tensor(action, dtype=torch.float32) for action in actions] + return float( + sum( + torch.linalg.vector_norm(current - previous, dim=-1).sum() + for previous, current in zip(tensors, tensors[1:]) + ) + ) + + +def _graph_difference( + offline: Mapping[str, Any], online: Mapping[str, Any] +) -> dict[str, Any]: + offline_nodes = {str(node["id"]): node for node in offline["nodes"]} + online_nodes = {str(node["id"]): node for node in online["nodes"]} + offline_actions = [node["atomic_action"] for node in offline["nodes"]] + online_actions = [node["atomic_action"] for node in online["nodes"]] + common_ids = sorted(set(offline_nodes) & set(online_nodes)) + node_changes = [] + for node_id in common_ids: + left = offline_nodes[node_id] + right = online_nodes[node_id] + changed_fields = sorted( + key for key in set(left) | set(right) if left.get(key) != right.get(key) + ) + if changed_fields: + node_changes.append({"id": node_id, "changed_fields": changed_fields}) + offline_groups = {str(group["id"]): group for group in offline["task_groups"]} + online_groups = {str(group["id"]): group for group in online["task_groups"]} + common_group_ids = sorted(set(offline_groups) & set(online_groups)) + group_changes = [] + for group_id in common_group_ids: + left = offline_groups[group_id] + right = online_groups[group_id] + changed_fields = sorted( + key for key in set(left) | set(right) if left.get(key) != right.get(key) + ) + if changed_fields: + group_changes.append({"id": group_id, "changed_fields": changed_fields}) + atomic_action_difference = { + "offline": offline_actions, + "online": online_actions, + "same_sequence": offline_actions == online_actions, + "added_node_ids": sorted(set(online_nodes) - set(offline_nodes)), + "removed_node_ids": sorted(set(offline_nodes) - set(online_nodes)), + "changed_nodes": node_changes, + } + task_group_difference = { + "offline_ids": sorted(offline_groups), + "online_ids": sorted(online_groups), + "added_ids": sorted(set(online_groups) - set(offline_groups)), + "removed_ids": sorted(set(offline_groups) - set(online_groups)), + "same_ids": set(offline_groups) == set(online_groups), + "changed_groups": group_changes, + } + return { + "offline_node_count": len(offline_actions), + "online_node_count": len(online_actions), + "same_action_sequence": offline_actions == online_actions, + "offline_actions": offline_actions, + "online_actions": online_actions, + "added_node_ids": sorted(set(online_nodes) - set(offline_nodes)), + "removed_node_ids": sorted(set(offline_nodes) - set(online_nodes)), + "changed_nodes": node_changes, + "offline_task_group_ids": sorted(offline_groups), + "online_task_group_ids": sorted(online_groups), + "added_task_group_ids": sorted(set(online_groups) - set(offline_groups)), + "removed_task_group_ids": sorted(set(offline_groups) - set(online_groups)), + "same_task_group_ids": set(offline_groups) == set(online_groups), + "changed_task_groups": group_changes, + "atomic_action_difference": atomic_action_difference, + "task_group_difference": task_group_difference, + } + + +def _validate_task_group_coverage( + task: Mapping[str, Any], graph: Mapping[str, Any], *, route: str +) -> None: + """Check explicit TaskSpec instances are neither dropped nor duplicated.""" + if task.get("level") == "L4": + return + expected = { + str(item["id"]) + for item in task.get("task_instances", ()) + if isinstance(item, Mapping) + } + actual = {str(group["id"]) for group in graph.get("task_groups", ())} + if expected != actual: + raise ValueError( + f"A/B {route} TaskGroup coverage mismatch; " + f"missing={sorted(expected - actual)}, " + f"unexpected={sorted(actual - expected)}." + ) + + +def _validate_snapshot( + snapshot: Mapping[str, Any], *, route: str, require_full: bool = False +) -> None: + if not isinstance(snapshot, Mapping): + raise TypeError(f"A/B {route} snapshot must be a mapping.") + required = {"robot_qpos", "object_poses"} + if require_full: + required = set(_FULL_SNAPSHOT_KEYS) + missing = required - set(snapshot) + if missing: + raise ValueError( + f"A/B {route} snapshot is missing required state {sorted(missing)}." + ) + qpos = torch.as_tensor(snapshot["robot_qpos"]) + if qpos.numel() == 0 or not bool(torch.isfinite(qpos).all()): + raise ValueError(f"A/B {route} robot_qpos must be finite and non-empty.") + object_poses = snapshot["object_poses"] + if not isinstance(object_poses, Mapping) or not object_poses: + raise ValueError(f"A/B {route} object_poses must be a non-empty mapping.") + for uid, pose in object_poses.items(): + tensor = torch.as_tensor(pose) + if not isinstance(uid, str) or not uid or tensor.numel() == 0: + raise ValueError(f"A/B {route} object_poses contains an invalid entry.") + if not bool(torch.isfinite(tensor).all()): + raise ValueError(f"A/B {route} pose for {uid!r} must be finite.") + if "articulation_state" in snapshot: + articulation_state = snapshot["articulation_state"] + if not isinstance(articulation_state, Mapping): + raise ValueError(f"A/B {route} articulation_state must be a mapping.") + for uid, state in articulation_state.items(): + if not isinstance(uid, str) or not uid: + raise ValueError( + f"A/B {route} articulation_state contains invalid UID." + ) + if not isinstance(state, Mapping): + raise ValueError( + f"A/B {route} articulation state for {uid!r} must be a mapping." + ) + if not state: + raise ValueError( + f"A/B {route} articulation state for {uid!r} is empty." + ) + for name, value in state.items(): + tensor = torch.as_tensor(value) + if tensor.numel() == 0 or not bool(torch.isfinite(tensor).all()): + raise ValueError( + f"A/B {route} articulation {uid!r}.{name} must be finite." + ) + if "camera_calibration" in snapshot: + calibrations = snapshot["camera_calibration"] + if not isinstance(calibrations, Mapping): + raise ValueError(f"A/B {route} camera_calibration must be a mapping.") + if require_full and not calibrations: + raise ValueError(f"A/B {route} camera_calibration must not be empty.") + for uid, calibration in calibrations.items(): + if not isinstance(uid, str) or not uid: + raise ValueError( + f"A/B {route} camera_calibration contains invalid UID." + ) + if not isinstance(calibration, Mapping): + raise ValueError( + f"A/B {route} calibration for {uid!r} must be a mapping." + ) + for name in ("intrinsics", "extrinsics"): + if name not in calibration: + raise ValueError( + f"A/B {route} calibration for {uid!r} is missing {name}." + ) + tensor = torch.as_tensor(calibration[name]) + if tensor.numel() == 0 or not bool(torch.isfinite(tensor).all()): + raise ValueError( + f"A/B {route} calibration {uid!r}.{name} must be finite." + ) + + +def _update_digest(digest: Any, value: Any) -> None: + if isinstance(value, Mapping): + digest.update(b"mapping{") + for key in sorted(value, key=str): + _update_digest(digest, str(key)) + _update_digest(digest, value[key]) + digest.update(b"}") + return + if isinstance(value, (list, tuple)): + digest.update(b"sequence[") + for item in value: + _update_digest(digest, item) + digest.update(b"]") + return + if isinstance(value, torch.Tensor): + value = value.detach().cpu().contiguous().numpy() + if isinstance(value, np.ndarray): + digest.update(str(value.dtype).encode("ascii")) + digest.update(str(tuple(value.shape)).encode("ascii")) + digest.update(value.tobytes(order="C")) + return + digest.update(type(value).__name__.encode("ascii")) + digest.update(repr(value).encode("utf-8")) + + +def _reject_private_or_live_fields(value: Any, context: str) -> None: + """Reject oracle/grounded fields before either branch can execute.""" + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _PRIVATE_OR_LIVE_KEYS: + raise ValueError(f"{context} contains private/live field {key!r}.") + _reject_private_or_live_fields(child, f"{context}.{key}") + elif isinstance(value, (list, tuple)): + for index, child in enumerate(value): + _reject_private_or_live_fields(child, f"{context}[{index}]") + + +def _call_preflight( + callback: Callable[..., Any], *, route: str, graph: Mapping[str, Any], env: Any +) -> Any: + """Call executor preflight hooks across the small supported API variants.""" + try: + return callback() + except TypeError as first_error: + # Third-party branch executors often expose contextual keyword-only + # arguments. Retry only for an argument-binding TypeError; if the + # callback itself raised TypeError, preserve that original failure. + try: + return callback(route=route, graph=graph, env=env) + except TypeError: + raise first_error + + +def _write_json(path: Path, value: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.py b/embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.py new file mode 100644 index 000000000..b4b256806 --- /dev/null +++ b/embodichain/gen_sim/action_engine/evaluation/e1_e2_scene_action.py @@ -0,0 +1,448 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Measure deterministic E1/E2 scene feasibility and graph compilation. + +This CPU benchmark exercises the contract path before simulator motion. It +checks that Scene Engine v1 output adapts successfully, required capabilities +are executable, and E1/E2 compile to action graphs containing pickup, +held-object motion, and placement. + +Run this module with ``--iterations 100`` for the default benchmark. +""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from time import perf_counter +import tracemalloc +from types import SimpleNamespace +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + TASK_CONTRACTS, + task_success_type, + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph +from embodichain.gen_sim.task_engine.scene import ( + FeasibilityBroker, + SceneEngineV1Adapter, +) + +__all__ = ["BenchmarkResult", "run_benchmark"] + + +@dataclass(frozen=True) +class BenchmarkResult: + """One scenario's contract latency, memory, and correctness metrics.""" + + scenario: str + iterations: int + elapsed_seconds: float + peak_bytes: int + success_count: int + feasibility_status: str + action_count: int + unknown_checks: int + runtime_probe_checks: int + + @property + def success_rate(self) -> float: + """Return successful iterations divided by all iterations.""" + return self.success_count / self.iterations + + @property + def mean_milliseconds(self) -> float: + """Return mean contract latency in milliseconds.""" + return self.elapsed_seconds * 1000.0 / self.iterations + + +def run_benchmark( + *, + iterations: int = 100, + output_dir: str | Path = "outputs/benchmarks", +) -> tuple[tuple[BenchmarkResult, ...], Path]: + """Run E1/E2 contract regressions and write one Markdown report.""" + if ( + isinstance(iterations, bool) + or not isinstance(iterations, int) + or iterations < 1 + ): + raise ValueError("iterations must be a positive integer.") + results = tuple( + _benchmark_scenario(task_type, iterations=iterations) + for task_type in ("E1", "E2") + ) + report = _write_report(results, output_dir=Path(output_dir)) + return results, report + + +def _benchmark_scenario(task_type: str, *, iterations: int) -> BenchmarkResult: + task, requirements = _generated_task(task_type) + bindings = { + item["role_id"]: f"{task_type.lower()}_{item['role_id']}" + for item in requirements["objects"] + } + manifest = _static_manifest(task_type, requirements, bindings) + candidate, reference_bindings = _candidate( + task_type, + task, + requirements, + bindings, + ) + registry = build_atomic_capability_registry() + broker = FeasibilityBroker() + task_actions = { + name: contract.core_actions for name, contract in TASK_CONTRACTS.items() + } + success_count = 0 + last_report: dict[str, Any] = {} + last_graph: dict[str, Any] = {} + + tracemalloc.start() + start = perf_counter() + try: + for _ in range(iterations): + last_report = broker.assess( + candidate, + reference_bindings, + manifest, + capability_catalog=registry.catalog(), + task_actions=task_actions, + ) + last_graph = instantiate_seed_graph(task, bindings, registry=registry) + actions = { + str(node.get("atomic_action")) + for node in last_graph["nodes"] + if node.get("atomic_action") + } + if ( + last_report["status"] != "contradicted" + and {"PickUp", "MoveHeldObject", "Place"} <= actions + ): + success_count += 1 + finally: + elapsed = perf_counter() - start + _, peak_bytes = tracemalloc.get_traced_memory() + tracemalloc.stop() + + action_count = sum( + bool(node.get("atomic_action")) for node in last_graph.get("nodes", ()) + ) + return BenchmarkResult( + scenario=task_type, + iterations=iterations, + elapsed_seconds=elapsed, + peak_bytes=peak_bytes, + success_count=success_count, + feasibility_status=str(last_report.get("status", "unknown")), + action_count=action_count, + unknown_checks=int(last_report.get("summary", {}).get("unknown", 0)), + runtime_probe_checks=int( + last_report.get("summary", {}).get("runtime_probe", 0) + ), + ) + + +def _generated_task(task_type: str) -> tuple[dict[str, Any], dict[str, Any]]: + if task_type not in {"E1", "E2"}: + raise ValueError("This benchmark supports only E1 and E2 fixtures.") + params: dict[str, Any] = {"object_role": "object"} + initial_state = {} + if task_type == "E1": + params.update({"target_role": "target", "relation": "inside"}) + else: + params.update( + { + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + ) + initial_state = {"orientation": "fallen"} + task_id = f"benchmark-{task_type.lower()}" + task = validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": task_id, + "level": "L1", + "instruction": "benchmark-instruction", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": task_type, + "params": params, + "depends_on": [], + "role": "primary", + } + ], + "success": { + "type": task_success_type(task_type, params), + "task_instance_id": "task_01", + }, + "oracle": {}, + "metadata": {"benchmark_fixture": True}, + } + ) + objects = [ + { + "role_id": "object", + "category": "can", + "count": 1, + "affordances": sorted(TASK_CONTRACTS[task_type].scene_affordances), + "initial_state": initial_state, + "attributes": {}, + } + ] + if task_type == "E1": + objects.append( + { + "role_id": "target", + "category": "container", + "count": 1, + "affordances": ["container", "support_surface"], + "initial_state": {}, + "attributes": {}, + } + ) + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": task_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [ + {"type": "reachable", "roles": "all_interaction_objects"} + ], + "distractor_count": 0, + "metadata": {"benchmark_fixture": True}, + } + ) + return task, requirements + + +def _static_manifest( + task_type: str, + requirements: dict[str, Any], + bindings: dict[str, str], +) -> dict[str, Any]: + planner_objects = [] + runtime_objects = [] + for index, requirement in enumerate(requirements["objects"]): + role_id = str(requirement["role_id"]) + uid = bindings[role_id] + role = "rigid_object" + planner_objects.append( + { + "uid": uid, + "source_uid": f"{uid}_0", + "role": role, + "name": uid, + "description": f"Synthetic {task_type} benchmark object.", + "category": str(requirement["category"]), + "color": requirement.get("attributes", {}).get("color"), + "shape": {"shape_type": "Mesh", "fpath": f"/{uid}.glb"}, + "init_pos": [0.15 * index, 0.0, 0.7], + "init_rot": [90.0, 0.0, 0.0] if task_type == "E2" else [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "affordances": list(requirement["affordances"]), + "initial_state": dict(requirement["initial_state"]), + "attributes": dict(requirement["attributes"]), + } + ) + runtime_objects.append( + { + "uid": uid, + "shape": {"shape_type": "Mesh", "fpath": f"/{uid}.glb"}, + "attrs": {"mass": 0.1}, + "body_type": "dynamic", + } + ) + prepared = SimpleNamespace( + source_config_path=Path("/synthetic/scene_config.json"), + planner_objects=tuple(planner_objects), + background=(), + rigid_objects=tuple(runtime_objects), + articulations=(), + asset_hashes={ + uid: uid.encode().hex().ljust(64, "0")[:64] for uid in bindings.values() + }, + ) + return SceneEngineV1Adapter().adapt_prepared_scene( + prepared, + source_format="benchmark", + robot_profile="dual_franka", + ) + + +def _candidate( + task_type: str, + task: dict[str, Any], + requirements: dict[str, Any], + bindings: dict[str, str], +) -> tuple[dict[str, Any], dict[str, list[str]]]: + references = [] + reference_bindings = {} + for index, requirement in enumerate(requirements["objects"]): + role_id = str(requirement["role_id"]) + role = "object" if index == 0 else "target" + reference_id = f"task_01.{role}" + references.append( + { + "reference_id": reference_id, + "role": role, + "source_structure": "rigid_object", + "affordances": list(requirement["affordances"]), + "initial_state": dict(requirement["initial_state"]), + "attributes": dict(requirement["attributes"]), + } + ) + reference_bindings[reference_id] = [bindings[role_id]] + return ( + { + "candidate_id": "candidate_01", + "draft": { + "task_id": task["task_id"], + "steps": [{"id": "task_01", "task_type": task_type}], + }, + "scene_request": {"references": references}, + }, + reference_bindings, + ) + + +def _write_report( + results: tuple[BenchmarkResult, ...], + *, + output_dir: Path, +) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + path = output_dir / f"e1_e2_scene_action_{timestamp}.md" + performance_rows = [ + { + "Scenario": item.scenario, + "Iterations": item.iterations, + "Total ms": f"{item.elapsed_seconds * 1000.0:.3f}", + "Mean ms": f"{item.mean_milliseconds:.3f}", + "Peak KiB": f"{item.peak_bytes / 1024.0:.1f}", + } + for item in results + ] + metric_rows = [ + { + "Scenario": item.scenario, + "Success rate": f"{item.success_rate:.3f}", + "Feasibility": item.feasibility_status, + "Actions": item.action_count, + "Unknown checks": item.unknown_checks, + "Runtime probes": item.runtime_probe_checks, + } + for item in results + ] + leaderboard_rows = [ + { + "Rank": rank, + "Scenario": item.scenario, + "Success rate": f"{item.success_rate:.3f}", + "Mean ms": f"{item.mean_milliseconds:.3f}", + } + for rank, item in enumerate( + sorted( + results, + key=lambda value: (-value.success_rate, value.mean_milliseconds), + ), + start=1, + ) + ] + lines = [ + "# E1/E2 Scene-Action Contract Benchmark", + "", + f"Generated at: {datetime.now().isoformat(timespec='seconds')}", + "", + "## Time & Memory", + "", + *_table(performance_rows), + "", + "## Success & Other Metrics", + "", + *_table(metric_rows), + "", + "## Leaderboard", + "", + *_table(leaderboard_rows), + "", + "## Notes", + "", + "- This benchmark covers deterministic contracts and graph compilation, not GPU motion execution.", + ] + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return path + + +def _table(rows: list[dict[str, object]]) -> list[str]: + headers = list(rows[0]) + return [ + "| " + " | ".join(headers) + " |", + "| " + " | ".join("---" for _ in headers) + " |", + *[ + "| " + " | ".join(str(row[header]) for header in headers) + " |" + for row in rows + ], + ] + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Benchmark E1/E2 scene-action contract stability." + ) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--output-dir", type=Path, default=Path("outputs/benchmarks")) + return parser + + +def main() -> int: + """Run from the command line and print the generated report path.""" + args = _build_parser().parse_args() + results, report = run_benchmark( + iterations=args.iterations, + output_dir=args.output_dir, + ) + for result in results: + print( + f"{result.scenario}: success={result.success_rate:.3f}, " + f"mean={result.mean_milliseconds:.3f} ms, " + f"peak={result.peak_bytes / 1024.0:.1f} KiB" + ) + print(f"Report: {report}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/action_engine/generation/__init__.py b/embodichain/gen_sim/action_engine/generation/__init__.py new file mode 100644 index 000000000..4446010d6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/__init__.py @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Independent config generation for Action Engine.""" + +from __future__ import annotations + +from .config_builder import VLM_CAMERA_UIDS, canonical_robot_profile +from .assets import normalize_scene_assets +from .generator import generate_action_engine_config +from .models import GeneratedConfigPaths, PreparedScene + +__all__ = [ + "GeneratedConfigPaths", + "PreparedScene", + "VLM_CAMERA_UIDS", + "canonical_robot_profile", + "generate_action_engine_config", + "normalize_scene_assets", +] diff --git a/embodichain/gen_sim/action_engine/generation/artifacts.py b/embodichain/gen_sim/action_engine/generation/artifacts.py new file mode 100644 index 000000000..0af961051 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/artifacts.py @@ -0,0 +1,141 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Publish canonical generation artifacts without intermediate copies.""" + +from __future__ import annotations + +from collections.abc import Mapping +import json +import os +from pathlib import Path +import tempfile +from typing import Any + +from embodichain.gen_sim.action_engine.protocol import ( + AGENT_CONFIG_FILENAME, + EXECUTION_PROGRAM_FILENAME, + FAST_GYM_CONFIG_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + SEED_TASK_GRAPH_PNG_FILENAME, + TASK_SPEC_FILENAME, +) + +from .models import GeneratedConfigPaths + +__all__ = ["artifact_paths", "write_generation_artifacts"] + + +def artifact_paths( + output_dir: str | Path, + *, + planning_mode: str = "offline", +) -> GeneratedConfigPaths: + """Return canonical resolved paths for one output directory.""" + directory = Path(output_dir).expanduser().resolve() + _validate_planning_mode(planning_mode) + graph_directory = directory if planning_mode == "offline" else directory / "offline" + return GeneratedConfigPaths( + gym_config=directory / FAST_GYM_CONFIG_FILENAME, + agent_config=directory / AGENT_CONFIG_FILENAME, + task_spec=directory / TASK_SPEC_FILENAME, + scene_requirements=directory / SCENE_REQUIREMENTS_FILENAME, + seed_task_graph=graph_directory / EXECUTION_PROGRAM_FILENAME, + seed_task_graph_png=graph_directory / SEED_TASK_GRAPH_PNG_FILENAME, + planning_mode=planning_mode, + ) + + +def write_generation_artifacts( + output_dir: str | Path, + *, + gym_config: Mapping[str, Any], + agent_config: Mapping[str, Any], + task_spec: Mapping[str, Any], + scene_requirements: Mapping[str, Any], + seed_task_graph: Mapping[str, Any], + seed_task_graph_png: bytes, + overwrite: bool, + planning_mode: str = "offline", +) -> GeneratedConfigPaths: + """Serialize validated artifacts and replace their destinations atomically.""" + paths = artifact_paths(output_dir, planning_mode=planning_mode) + if not isinstance(seed_task_graph_png, (bytes, bytearray)): + raise TypeError("seed_task_graph_png must be bytes.") + payloads = { + paths.gym_config: _serialize_json(gym_config), + paths.agent_config: _serialize_json(agent_config), + paths.task_spec: _serialize_json(task_spec), + paths.scene_requirements: _serialize_json(scene_requirements), + paths.seed_task_graph: _serialize_json(seed_task_graph), + paths.seed_task_graph_png: bytes(seed_task_graph_png), + } + existing = sorted(path for path in payloads if path.exists()) + if existing and not overwrite: + names = ", ".join(path.name for path in existing) + raise FileExistsError( + f"Generated artifacts already exist in {paths.gym_config.parent}: " + f"{names}. Pass --overwrite to replace them." + ) + + paths.gym_config.parent.mkdir(parents=True, exist_ok=True) + temporary: dict[Path, Path] = {} + try: + for destination, payload in payloads.items(): + destination.parent.mkdir(parents=True, exist_ok=True) + temporary[destination] = _write_temporary(destination.parent, payload) + for destination, temporary_path in temporary.items(): + os.replace(temporary_path, destination) + finally: + for temporary_path in temporary.values(): + temporary_path.unlink(missing_ok=True) + return paths + + +def _serialize_json(value: Mapping[str, Any]) -> str: + try: + return ( + json.dumps( + dict(value), + ensure_ascii=False, + indent=2, + sort_keys=False, + allow_nan=False, + ) + + "\n" + ) + except (TypeError, ValueError) as exc: + raise ValueError("Generated artifact is not strict JSON data.") from exc + + +def _write_temporary(directory: Path, payload: str | bytes) -> Path: + data = payload if isinstance(payload, bytes) else payload.encode("utf-8") + with tempfile.NamedTemporaryFile( + mode="wb", + dir=directory, + prefix=".action_engine_", + suffix=".tmp", + delete=False, + ) as stream: + stream.write(data) + stream.flush() + os.fsync(stream.fileno()) + return Path(stream.name) + + +def _validate_planning_mode(value: Any) -> None: + if value not in {"offline", "ab"}: + raise ValueError("planning_mode must be 'offline' or 'ab'.") diff --git a/embodichain/gen_sim/action_engine/generation/assets.py b/embodichain/gen_sim/action_engine/generation/assets.py new file mode 100644 index 000000000..d1b4da830 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/assets.py @@ -0,0 +1,158 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Normalize GLB node transforms and body scale into reusable runtime assets.""" + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import replace +import hashlib +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from .models import PreparedScene + +__all__ = ["normalize_scene_assets"] + +_POLICY = "action_engine_glb_geometry_v2" + + +def normalize_scene_assets( + scene: PreparedScene, + output_dir: str | Path, +) -> PreparedScene: + """Return a scene whose valid GLB meshes have flattened runtime geometry. + + Source files are never modified. Cache names derive from source bytes, + object scale, and the normalization policy, so repeated generation reuses + identical assets. + """ + sections = { + "background": [deepcopy(value) for value in scene.background], + "rigid_object": [deepcopy(value) for value in scene.rigid_objects], + "articulation": [deepcopy(value) for value in scene.articulations], + } + cache_dir = Path(output_dir).expanduser().resolve() / "mesh_assets" / "normalized" + reports: list[dict[str, Any]] = [] + hashes = dict(scene.asset_hashes) + normalized_by_uid: dict[str, dict[str, Any]] = {} + for section in ("background", "rigid_object"): + for config in sections[section]: + report = _normalize_object(config, cache_dir) + if report is not None: + reports.append(report) + hashes[str(config["uid"])] = str(report["runtime_sha256"]) + normalized_by_uid[str(config["uid"])] = config + + planner = [deepcopy(value) for value in scene.planner_objects] + for item in planner: + runtime = normalized_by_uid.get(str(item["runtime_uid"])) + if runtime is None: + continue + item["shape"] = deepcopy(runtime.get("shape", {})) + item["body_scale"] = list(runtime.get("body_scale", [1.0, 1.0, 1.0])) + return replace( + scene, + planner_objects=tuple(planner), + background=tuple(sections["background"]), + rigid_objects=tuple(sections["rigid_object"]), + articulations=tuple(sections["articulation"]), + asset_hashes=hashes, + asset_provenance=tuple(reports), + ) + + +def _normalize_object( + config: dict[str, Any], + cache_dir: Path, +) -> dict[str, Any] | None: + shape = config.get("shape") + if not isinstance(shape, dict) or not shape.get("fpath"): + return None + source = Path(str(shape["fpath"])).expanduser().resolve() + if source.suffix.lower() not in {".glb", ".gltf"}: + return None + source_hash = _file_hash(source) + scale = [float(value) for value in config.get("body_scale", [1.0, 1.0, 1.0])] + key = hashlib.sha256( + json.dumps( + {"source": source_hash, "scale": scale, "policy": _POLICY}, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + destination = cache_dir / f"{source.stem[:32]}_{key[:16]}.glb" + status = "reused" if destination.is_file() else "generated" + if status == "generated": + try: + _bake_glb(source, destination, scale) + except Exception as exc: + return { + "uid": str(config.get("uid", "")), + "source_path": source.as_posix(), + "source_sha256": source_hash, + "runtime_path": source.as_posix(), + "runtime_sha256": source_hash, + "body_scale": scale, + "status": "preserved_invalid_source", + "error": f"{type(exc).__name__}: {exc}", + "policy_version": _POLICY, + } + shape["fpath"] = destination.as_posix() + config["body_scale"] = [1.0, 1.0, 1.0] + return { + "uid": str(config.get("uid", "")), + "source_path": source.as_posix(), + "source_sha256": source_hash, + "runtime_path": destination.as_posix(), + "runtime_sha256": _file_hash(destination), + "body_scale": scale, + "status": status, + "policy_version": _POLICY, + } + + +def _bake_glb(source: Path, destination: Path, sim_scale: list[float]) -> None: + import trimesh + + source_scene = trimesh.load(source.as_posix(), force="scene") + baked = trimesh.Scene() + scale = np.diag([sim_scale[0], sim_scale[2], sim_scale[1], 1.0]) + for node_name in source_scene.graph.nodes_geometry: + node_transform, geometry_name = source_scene.graph.get(node_name) + mesh = source_scene.geometry[geometry_name].copy() + mesh.apply_transform(scale @ node_transform) + baked.add_geometry( + mesh, + node_name=str(node_name), + geom_name=f"geometry_{len(baked.geometry)}", + ) + if not baked.geometry: + raise ValueError(f"GLB contains no mesh geometry: {source}") + destination.parent.mkdir(parents=True, exist_ok=True) + baked.export(destination.as_posix(), file_type="glb") + + +def _file_hash(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/embodichain/gen_sim/action_engine/generation/config_builder.py b/embodichain/gen_sim/action_engine/generation/config_builder.py new file mode 100644 index 000000000..f1193494a --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/config_builder.py @@ -0,0 +1,852 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Build the simulator and Action Engine artifact manifests.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from functools import lru_cache +import json +import math +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.action_engine.config import ( + ACTION_ENGINE_DEFAULTS_SCHEMA, + RuntimePolicyCfg, + default_runtime_policy, + generation_defaults, + runtime_policy_hash, +) +from embodichain.gen_sim.action_engine.protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + ACTION_ENGINE_ENV_ID, + EXECUTION_PROGRAM_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + TASK_SPEC_FILENAME, +) + +from .models import PreparedScene + +__all__ = [ + "build_agent_config", + "build_fast_gym_config", + "canonical_robot_profile", + "VLM_CAMERA_UIDS", + "validate_fast_gym_config", +] + +_TEMPLATE_DIR = Path(__file__).resolve().parent / "templates" +_GENERATION_DEFAULTS = generation_defaults() +_DEFAULT_TABLETOP_Z = float(_GENERATION_DEFAULTS["scene"]["default_tabletop_z"]) + +_ARM_SLOTS = { + "left": {"arm": "left_arm", "eef": "left_eef"}, + "right": {"arm": "right_arm", "eef": "right_eef"}, +} + +# These IDs are part of the A/B runtime contract. Keep the order stable so +# visual-fact payloads and comparison reports are reproducible across runs. +VLM_CAMERA_UIDS = ( + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", +) + + +def canonical_robot_profile(profile: str) -> str: + """Normalize the supported CLI aliases to one runtime profile ID.""" + normalized = str(profile).strip().lower().replace("-", "_") + profiles = _robot_profiles() + if normalized in profiles: + return normalized + for profile_id, value in profiles.items(): + if normalized in value["aliases"]: + return profile_id + raise ValueError( + f"Unsupported robot profile {profile!r}; expected one of: " + f"{', '.join(sorted(profiles))}" + ) + + +def build_agent_config( + *, + task_name: str, + robot_profile: str, + execution_program_hash: str, + source_config_path: Path, + uid_map: dict[str, str], + static_obstacle_uids: Sequence[str] | None = None, + dynamic_obstacle_uids: Sequence[str] | None = None, + table_top_z: float | None = None, + articulation_settings: Mapping[str, Mapping[str, Sequence[float]]] | None = None, + planning_mode: str = "offline", + seed_task_graph_path: str | Path | None = EXECUTION_PROGRAM_FILENAME, + vlm_model: str | None = None, + vlm_camera_uids: Sequence[str] | None = None, +) -> dict[str, Any]: + """Build the small manifest consumed by ``run_agent``.""" + profile = canonical_robot_profile(robot_profile) + runtime_policy = default_runtime_policy(profile) + if ( + static_obstacle_uids is not None + or dynamic_obstacle_uids is not None + or table_top_z is not None + ): + policy = runtime_policy.as_mapping() + planner = policy["planner"] + if static_obstacle_uids is not None: + planner["static_obstacle_uids"] = [str(uid) for uid in static_obstacle_uids] + if dynamic_obstacle_uids is not None: + planner["dynamic_obstacle_uids"] = [ + str(uid) for uid in dynamic_obstacle_uids + ] + planner["dynamic_collision"] = bool(dynamic_obstacle_uids) + if table_top_z is not None: + tabletop = float(table_top_z) + if not math.isfinite(tabletop): + raise ValueError("table_top_z must be finite when provided.") + height_offset = tabletop - _DEFAULT_TABLETOP_Z + height_policies = ( + policy["grounding"]["semantic_defaults"], + policy["grounding"]["handover"], + policy["motion_defaults"]["MoveEndEffector"], + policy["motion_modifiers"]["orientation"]["upright"]["MoveEndEffector"], + ) + for height_policy in height_policies: + height_policy["maximum_eef_height"] = round( + float(height_policy["maximum_eef_height"]) + height_offset, + 6, + ) + runtime_policy = RuntimePolicyCfg.from_mapping(policy) + _validate_planning_mode(planning_mode) + graph_path = _validate_seed_graph_path(seed_task_graph_path) + if planning_mode == "ab" and graph_path == EXECUTION_PROGRAM_FILENAME: + graph_path = f"offline/{EXECUTION_PROGRAM_FILENAME}" + result = { + "schema_version": ACTION_ENGINE_CONFIG_SCHEMA, + "task_name": task_name, + "robot_profile": profile, + "planning_mode": planning_mode, + "task_spec": TASK_SPEC_FILENAME, + "scene_requirements": SCENE_REQUIREMENTS_FILENAME, + "seed_task_graph": graph_path, + "seed_task_graph_hash": execution_program_hash, + "runtime_policy": runtime_policy.as_mapping(), + "runtime_policy_hash": runtime_policy_hash(runtime_policy), + "source": { + "gym_config": source_config_path.as_posix(), + "uid_map": dict(sorted(uid_map.items())), + }, + "articulation_settings": _normalize_articulation_settings( + articulation_settings or {} + ), + } + if planning_mode == "ab": + camera_uids = _normalize_vlm_camera_uids(vlm_camera_uids) + configured_model = _optional_model(vlm_model) + # Retain concise top-level aliases for early A/B bundles while keeping + # the nested section as the canonical runtime namespace. + result["offline_seed_task_graph"] = graph_path + result["vlm_model"] = configured_model + result["vlm_camera_uids"] = list(camera_uids) + result["online_planning"] = { + # Model names are deliberately persisted only when explicitly + # supplied by the generator. Runtime resolution can then apply + # the documented ACTION_ENGINE_VLM_MODEL/OPENAI_MODEL fallback. + "vlm_model": configured_model, + "camera_uids": camera_uids, + } + return result + + +def _normalize_articulation_settings( + value: Mapping[str, Mapping[str, Sequence[float]]], +) -> dict[str, dict[str, list[float]]]: + """Own finite per-joint ordinal setting calibrations for runtime grounding.""" + result: dict[str, dict[str, list[float]]] = {} + for uid, joints in value.items(): + if not isinstance(uid, str) or not uid or not isinstance(joints, Mapping): + raise ValueError("articulation_settings must map UIDs to joint mappings.") + normalized_joints = {} + for joint_name, settings in joints.items(): + if ( + not isinstance(joint_name, str) + or not joint_name + or not isinstance(settings, Sequence) + or isinstance(settings, (str, bytes, bytearray)) + or not settings + ): + raise ValueError( + "articulation_settings joints require non-empty setting lists." + ) + normalized = [float(item) for item in settings] + if any(not math.isfinite(item) for item in normalized): + raise ValueError("articulation setting values must be finite.") + normalized_joints[joint_name] = normalized + result[uid] = normalized_joints + return dict(sorted(result.items())) + + +def build_fast_gym_config( + scene: PreparedScene, + *, + task_name: str, + task_description: str, + robot_profile: str, + execution_program_hash: str, + max_episodes: int, + max_episode_steps: int, + randomize_scene: bool = False, + randomize_table_material: bool = False, + planning_mode: str = "offline", + seed_task_graph_path: str | Path | None = EXECUTION_PROGRAM_FILENAME, +) -> dict[str, Any]: + """Build a runnable EmbodiChain gym config from a prepared source scene.""" + if max_episodes < 1: + raise ValueError("max_episodes must be at least 1.") + if max_episode_steps < 1: + raise ValueError("max_episode_steps must be at least 1.") + _validate_planning_mode(planning_mode) + graph_path = _validate_seed_graph_path(seed_task_graph_path) + if planning_mode == "ab" and graph_path == EXECUTION_PROGRAM_FILENAME: + graph_path = f"offline/{EXECUTION_PROGRAM_FILENAME}" + profile = canonical_robot_profile(robot_profile) + + profile_config = _profile(profile) + robot = _make_robot(profile, profile_config, scene.table_top_z) + observations = _make_observations(robot) + # These two template fields describe serialization order to generation, not + # RobotCfg. Remove them after deriving observation IDs to avoid parser noise. + robot.pop("observation_joint_parts", None) + robot.pop("qpos_control_part_order", None) + sensors = _load_template("default_sensors.json") + if not isinstance(sensors, list) or not sensors: + raise ValueError("Default sensor template must define at least one camera.") + environment_policy = _GENERATION_DEFAULTS["environment"] + viewer_camera_uid = str(environment_policy["viewer_camera_uid"]) + sensors[0]["uid"] = viewer_camera_uid + if planning_mode == "ab": + vlm_sensors = _load_template("vlm_sensors.json") + if not isinstance(vlm_sensors, list) or len(vlm_sensors) != len( + VLM_CAMERA_UIDS + ): + raise ValueError("A/B planning requires exactly four VLM cameras.") + _validate_vlm_sensors(vlm_sensors) + _anchor_vlm_sensors(vlm_sensors, scene) + sensors.extend(vlm_sensors) + light = _load_template("default_lights.json") + + rigid_uids = [str(config["uid"]) for config in scene.rigid_objects] + background_uids = [str(config["uid"]) for config in scene.background] + engine_extension = { + "schema_version": "action_engine_runtime_v2", + "defaults_schema_version": ACTION_ENGINE_DEFAULTS_SCHEMA, + "task_name": task_name, + "robot_profile": profile, + "planning_mode": planning_mode, + "task_spec": TASK_SPEC_FILENAME, + "scene_requirements": SCENE_REQUIREMENTS_FILENAME, + "seed_task_graph": graph_path, + "seed_task_graph_hash": execution_program_hash, + "source_gym_config": scene.source_config_path.as_posix(), + "source_scene_z_rotation_degrees": scene.z_rotation_degrees, + "source_scene_xy_translation": list(scene.source_scene_xy_translation), + "body_scale_policy": scene.body_scale_policy, + "body_scale": list(scene.body_scale), + "asset_hashes": dict(sorted(scene.asset_hashes.items())), + "asset_provenance": [deepcopy(value) for value in scene.asset_provenance], + "uid_map": dict(sorted(scene.uid_map.items())), + } + extensions = { + "action_engine": engine_extension, + "agent_robot_profile": profile, + "agent_arm_slots": deepcopy(_ARM_SLOTS), + "agent_static_obstacle_uids": background_uids, + "agent_dynamic_obstacle_uids": rigid_uids, + "gripper_open_state": list(profile_config["gripper_open_state"]), + "gripper_close_state": list(profile_config["gripper_close_state"]), + "arm_aim_yaw_offset": deepcopy(environment_policy["arm_aim_yaw_offset"]), + "ignore_terminations_during_agent": bool( + environment_policy["ignore_terminations_during_agent"] + ), + "viewer_camera_uid": viewer_camera_uid, + } + + config: dict[str, Any] = { + "id": ACTION_ENGINE_ENV_ID, + "max_episodes": int(max_episodes), + "max_episode_steps": int(max_episode_steps), + "env": { + "extensions": extensions, + "events": _make_events( + sensors[0], + rigid_uids, + planning_mode=planning_mode, + randomize_scene=randomize_scene, + randomize_table_material=randomize_table_material, + ), + "observations": observations, + "dataset": _make_dataset( + task_name=task_name, + task_description=task_description, + source_config_path=scene.source_config_path, + robot_type=str(robot["uid"]), + ), + }, + "robot": robot, + "sensor": sensors, + "light": light, + "background": [deepcopy(obj_config) for obj_config in scene.background], + "rigid_object": [deepcopy(obj_config) for obj_config in scene.rigid_objects], + } + if scene.articulations: + config["articulation"] = [ + { + key: deepcopy(value) + for key, value in articulation.items() + if key not in {"attributes", "role"} + } + for articulation in scene.articulations + ] + validate_fast_gym_config(config) + return config + + +def validate_fast_gym_config(config: dict[str, Any]) -> None: + """Check the cross-file and simulator-facing invariants generation owns.""" + if config.get("id") != ACTION_ENGINE_ENV_ID: + raise ValueError(f"Gym config id must be {ACTION_ENGINE_ENV_ID!r}.") + if not isinstance(config.get("robot"), dict) or not config["robot"].get("uid"): + raise ValueError("Gym config requires a concrete robot template.") + if not config.get("sensor"): + raise ValueError("Gym config requires at least one sensor.") + if not all(isinstance(sensor, dict) for sensor in config["sensor"]): + raise ValueError("Generated sensors must be object mappings.") + sensor_uids = [str(sensor.get("uid", "")) for sensor in config["sensor"]] + if not all(sensor_uids) or len(sensor_uids) != len(set(sensor_uids)): + raise ValueError("Generated sensor UIDs must be non-empty and unique.") + if not config.get("background"): + raise ValueError("Gym config requires at least one background object.") + + objects = [ + *config.get("background", []), + *config.get("rigid_object", []), + *config.get("articulation", []), + ] + uids = [str(obj.get("uid", "")) for obj in objects] + if not all(uids) or len(uids) != len(set(uids)): + raise ValueError("Generated scene object UIDs must be non-empty and unique.") + if "table" not in uids: + raise ValueError("Generated tabletop scene must expose runtime UID 'table'.") + + for obj in objects: + shape = obj.get("shape") + fpath = shape.get("fpath") if isinstance(shape, dict) else obj.get("fpath") + if fpath is None: + continue + path = Path(str(fpath)) + if not path.is_absolute() or not path.is_file(): + raise ValueError( + f"Generated asset path for {obj.get('uid')!r} is not an " + f"existing absolute file: {path}" + ) + + action_engine = config.get("env", {}).get("extensions", {}).get("action_engine", {}) + if action_engine.get("defaults_schema_version") != ACTION_ENGINE_DEFAULTS_SCHEMA: + raise ValueError("Gym config has an unexpected defaults schema version.") + if action_engine.get("task_spec") != TASK_SPEC_FILENAME: + raise ValueError("Gym config points to an unexpected TaskSpec artifact.") + if action_engine.get("scene_requirements") != SCENE_REQUIREMENTS_FILENAME: + raise ValueError("Gym config points to unexpected SceneRequirements.") + graph_path = action_engine.get("seed_task_graph") + if ( + not isinstance(graph_path, str) + or Path(graph_path).name != EXECUTION_PROGRAM_FILENAME + ): + raise ValueError("Gym config points to an unexpected SeedGraph artifact.") + + planning_mode = action_engine.get("planning_mode", "offline") + _validate_planning_mode(planning_mode) + if planning_mode == "ab": + sensors = config["sensor"] + vlm_sensors = [ + sensor + for sensor in sensors + if isinstance(sensor, dict) + and str(sensor.get("uid", "")).startswith("vlm_") + ] + _validate_vlm_sensors(vlm_sensors) + + registered = { + entry.get("entity_cfg", {}).get("uid") + for entry in ( + config.get("env", {}) + .get("events", {}) + .get("register_info_to_env", {}) + .get("params", {}) + .get("registry", []) + ) + } + rigid_uids = {obj["uid"] for obj in config.get("rigid_object", [])} + if registered != rigid_uids: + raise ValueError("Every rigid object must have one live-pose registry entry.") + + +def _make_robot( + profile_id: str, + profile: dict[str, Any], + table_top_z: float | None, +) -> dict[str, Any]: + robot = _load_template(str(profile["template"])) + tabletop_z = _DEFAULT_TABLETOP_Z if table_top_z is None else float(table_top_z) + robot["init_pos"][2] = round( + tabletop_z + + float(profile["tabletop_clearance"]) + - float(profile["arm_component_z"]), + 6, + ) + family = str(profile["robot_family"]) + if family.startswith("ur"): + display = family.upper() + urdf_dir = display + robot["uid"] = f"Dual{display}" + robot["urdf_cfg"]["fname"] = f"dual_{family}_robotiq_arg2f_140_basket" + for component in robot["urdf_cfg"]["components"]: + if str(component.get("component_type", "")).endswith("_arm"): + component["urdf_path"] = f"UniversalRobots/{urdf_dir}/{urdf_dir}.urdf" + component["transform"][0][3] = float(profile["arm_base_x"]) + component["transform"][2][3] = float(profile["arm_component_z"]) + for arm in ("left_arm", "right_arm"): + robot["solver_cfg"][arm]["ur_type"] = family + robot["drive_pros"]["max_effort"][arm] = float(profile["max_effort"]) + robot["qpos_control_part_order"] = [ + "left_arm", + "right_arm", + "left_eef", + "right_eef", + ] + robot["observation_joint_parts"] = ["left_eef", "right_eef"] + if profile_id != canonical_robot_profile(profile_id): + raise ValueError(f"Invalid canonical robot profile {profile_id!r}.") + return robot + + +@lru_cache(maxsize=1) +def _robot_profiles() -> dict[str, dict[str, Any]]: + value = _read_template("robot_profiles.json") + if not isinstance(value, dict) or not value: + raise ValueError("robot_profiles.json must contain a non-empty object.") + return value + + +def _profile(profile_id: str) -> dict[str, Any]: + profile = deepcopy(_robot_profiles()[profile_id]) + required = { + "aliases", + "template", + "robot_family", + "tabletop_clearance", + "arm_component_z", + "gripper_open_state", + "gripper_close_state", + } + missing = sorted(required - set(profile)) + if missing: + raise ValueError(f"Robot profile {profile_id!r} is missing fields: {missing}.") + return profile + + +def _make_events( + camera: dict[str, Any], + rigid_uids: list[str], + *, + planning_mode: str, + randomize_scene: bool = False, + randomize_table_material: bool = False, +) -> dict[str, Any]: + extrinsics = camera["extrinsics"] + eye = list(extrinsics["eye"]) + target = list(extrinsics["target"]) + # The recording view mirrors the interactive viewer around its target. + audience_eye = [ + 2.0 * float(target[0]) - float(eye[0]), + 2.0 * float(target[1]) - float(eye[1]), + float(eye[2]), + ] + recording_enabled, recording_resolution, recording_interval = _recording_policy( + planning_mode + ) + source_width = int(camera["width"]) + source_height = int(camera["height"]) + if source_width <= 0 or source_height <= 0: + raise ValueError("Recording source camera resolution must be positive.") + intrinsics = camera.get("intrinsics") + if ( + not isinstance(intrinsics, Sequence) + or isinstance(intrinsics, (str, bytes, bytearray)) + or len(intrinsics) != 4 + ): + raise ValueError("Recording source camera intrinsics must be a 4-vector.") + scale_x = recording_resolution[0] / source_width + scale_y = recording_resolution[1] / source_height + recording_intrinsics = [ + float(intrinsics[0]) * scale_x, + float(intrinsics[1]) * scale_y, + float(intrinsics[2]) * scale_x, + float(intrinsics[3]) * scale_y, + ] + events = { + "record_camera": { + "func": "record_camera_data", + "mode": "interval", + "interval_step": recording_interval, + "params": { + "name": "record_cam_audience_view", + "resolution": list(recording_resolution), + "intrinsics": recording_intrinsics, + "eye": audience_eye, + "target": target, + "up": [ + -float(extrinsics["up"][0]), + -float(extrinsics["up"][1]), + float(extrinsics["up"][2]), + ], + }, + }, + "validation_cameras": { + "func": "validation_cameras", + "mode": "trigger", + "params": {}, + }, + "prepare_extra_attr": { + "func": "prepare_extra_attr", + "mode": "reset", + "params": { + "attrs": [ + { + "name": "object_lengths", + "mode": "callable", + "entity_uids": list(rigid_uids), + "func_name": "compute_object_length", + "func_kwargs": { + "is_svd_frame": True, + "sample_points": int( + _GENERATION_DEFAULTS["scene"][ + "object_length_sample_points" + ] + ), + }, + } + ] + }, + }, + "register_info_to_env": { + "func": "register_info_to_env", + "mode": "reset", + "params": { + "registry": [ + { + "entity_cfg": {"uid": uid}, + "pose_register_params": { + "compute_relative": False, + "compute_pose_object_to_arena": True, + "to_matrix": True, + }, + } + for uid in sorted(rigid_uids) + ], + "registration": "affordance_datas", + "sim_update": True, + }, + }, + } + if not recording_enabled: + events.pop("record_camera") + if randomize_table_material: + material = _GENERATION_DEFAULTS["randomization"]["table_material"] + events["randomize_table_material"] = { + "func": "randomize_visual_material", + "mode": "reset", + "params": { + "entity_cfg": {"uid": "table"}, + "random_texture_prob": float(material["random_texture_prob"]), + "base_color_range": deepcopy(material["base_color_range"]), + "metallic_range": list(material["metallic_range"]), + "roughness_range": list(material["roughness_range"]), + }, + } + if randomize_scene: + randomization = _GENERATION_DEFAULTS["randomization"] + for uid in sorted(rigid_uids): + events[f"randomize_{uid}_pose"] = { + "func": "randomize_rigid_object_pose", + "mode": "reset", + "params": { + "entity_cfg": {"uid": uid}, + "position_range": deepcopy( + randomization["rigid_object_position_range"] + ), + "rotation_range": deepcopy( + randomization["rigid_object_rotation_range"] + ), + "relative_position": True, + "relative_rotation": True, + }, + } + events["randomize_table_height"] = { + "func": "randomize_anchor_height", + "mode": "reset", + "params": { + "anchor_uid": "table", + "height_delta_range": deepcopy( + randomization["table_height_delta_range"] + ), + }, + } + return events + + +def _recording_policy(planning_mode: str) -> tuple[bool, tuple[int, int], int]: + """Resolve the bounded GenSim audience-recording policy.""" + value = _GENERATION_DEFAULTS["environment"].get("recording") + required = {"enabled", "resolution", "interval_step"} + if not isinstance(value, dict) or set(value) != required: + raise ValueError( + "generation.environment.recording must define enabled, resolution, " + "and interval_step." + ) + enabled = value["enabled"] + if not isinstance(enabled, bool): + raise ValueError("generation.environment.recording.enabled must be a boolean.") + resolution = value["resolution"] + if ( + not isinstance(resolution, Sequence) + or isinstance(resolution, (str, bytes, bytearray)) + or len(resolution) != 2 + or any( + isinstance(item, bool) or not isinstance(item, int) for item in resolution + ) + or any(int(item) <= 0 for item in resolution) + ): + raise ValueError( + "generation.environment.recording.resolution must contain two " + "positive integers." + ) + interval_step = value["interval_step"] + if ( + isinstance(interval_step, bool) + or not isinstance(interval_step, int) + or interval_step <= 0 + ): + raise ValueError( + "generation.environment.recording.interval_step must be positive." + ) + return ( + bool(enabled or planning_mode == "ab"), + (int(resolution[0]), int(resolution[1])), + int(interval_step), + ) + + +def _make_observations(robot: dict[str, Any]) -> dict[str, Any]: + control_parts = robot["control_parts"] + qpos_order = robot["qpos_control_part_order"] + observed_parts = set(robot["observation_joint_parts"]) + offset = 0 + joint_ids: list[int] = [] + for part in qpos_order: + count = len(control_parts[part]) + if part in observed_parts: + joint_ids.extend(range(offset, offset + count)) + offset += count + return { + "norm_robot_eef_joint": { + "func": "normalize_robot_joint_data", + "mode": "modify", + "name": "robot/qpos", + "params": {"joint_ids": joint_ids}, + } + } + + +def _make_dataset( + *, + task_name: str, + task_description: str, + source_config_path: Path, + robot_type: str, +) -> dict[str, Any]: + dataset_policy = _GENERATION_DEFAULTS["dataset"] + return { + "lerobot": { + "func": "LeRobotRecorder", + "mode": "save", + "save_failed_episodes": bool(dataset_policy["save_failed_episodes"]), + "params": { + "robot_meta": { + "robot_type": robot_type, + "control_freq": int(dataset_policy["control_frequency"]), + }, + "instruction": {"lang": task_description}, + "extra": { + "scene_type": source_config_path.parent.name, + "task_name": task_name, + # LeRobotRecorder uses this legacy field as a directory label. + "task_description": task_name, + "data_type": "sim", + }, + "use_videos": bool(dataset_policy["use_videos"]), + }, + } + } + + +def _load_template(name: str) -> Any: + return deepcopy(_read_template(name)) + + +@lru_cache(maxsize=None) +def _read_template(name: str) -> Any: + path = _TEMPLATE_DIR / name + if not path.is_file(): + raise FileNotFoundError(f"Action Engine template not found: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def _validate_planning_mode(value: Any) -> str: + """Validate and return the two supported generation/runtime modes.""" + if value not in {"offline", "ab"}: + raise ValueError("planning_mode must be 'offline' or 'ab'.") + return str(value) + + +def _validate_seed_graph_path(value: str | Path | None) -> str: + """Validate a relative or absolute path while preserving caller spelling.""" + if value is None: + return EXECUTION_PROGRAM_FILENAME + if not isinstance(value, (str, Path)): + raise ValueError("seed_task_graph_path must be a non-empty path string.") + path = str(value).strip() + if not path: + raise ValueError("seed_task_graph_path must be a non-empty path string.") + if Path(path).name != EXECUTION_PROGRAM_FILENAME: + raise ValueError("seed_task_graph_path must point to seed_task_graph.json.") + return path + + +def _optional_model(value: Any) -> str | None: + """Normalize optional model names without serializing blank strings.""" + if value is None: + return None + if not isinstance(value, str): + raise TypeError("Model name must be a string or None.") + normalized = value.strip() + return normalized or None + + +def _normalize_vlm_camera_uids(value: Sequence[str] | None) -> list[str]: + """Return the canonical four-camera list used by A/B execution.""" + if value is None: + return list(VLM_CAMERA_UIDS) + if isinstance(value, (str, bytes, bytearray)) or not isinstance(value, Sequence): + raise TypeError("vlm_camera_uids must be a list of strings.") + if not all(isinstance(item, str) for item in value): + raise TypeError("vlm_camera_uids must be a list of strings.") + normalized = [item.strip() for item in value] + if normalized != list(VLM_CAMERA_UIDS): + raise ValueError( + "A/B planning requires VLM cameras in canonical order: " + f"{list(VLM_CAMERA_UIDS)}." + ) + return normalized + + +def _validate_vlm_sensors(value: list[dict[str, Any]]) -> None: + """Validate camera template fields needed by visual fact extraction.""" + if len(value) != len(VLM_CAMERA_UIDS): + raise ValueError("A/B planning requires exactly four VLM cameras.") + if not all(isinstance(sensor, dict) for sensor in value): + raise ValueError("VLM sensors must be object mappings.") + uids = [str(sensor.get("uid", "")) for sensor in value] + if uids != list(VLM_CAMERA_UIDS): + raise ValueError("VLM camera UIDs must be exactly " f"{list(VLM_CAMERA_UIDS)}.") + for sensor in value: + if sensor.get("sensor_type", "Camera") != "Camera": + raise ValueError(f"VLM sensor {sensor.get('uid')!r} must be a Camera.") + if int(sensor.get("width", 0)) != 640 or int(sensor.get("height", 0)) != 480: + raise ValueError("VLM cameras must use 640x480 resolution.") + if not bool(sensor.get("enable_color")) or not bool(sensor.get("enable_depth")): + raise ValueError("VLM cameras must enable RGB and depth.") + extrinsics = sensor.get("extrinsics") + if not isinstance(extrinsics, dict) or not all( + key in extrinsics for key in ("eye", "target", "up") + ): + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} requires eye/target/up extrinsics." + ) + for name in ("eye", "target", "up"): + vector = extrinsics[name] + if ( + not isinstance(vector, Sequence) + or isinstance(vector, (str, bytes, bytearray)) + or len(vector) != 3 + ): + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} {name} must be a 3-vector." + ) + try: + values = [float(item) for item in vector] + except (TypeError, ValueError) as exc: + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} {name} must be numeric." + ) from exc + if not all(math.isfinite(item) for item in values): + raise ValueError( + f"VLM sensor {sensor.get('uid')!r} {name} must be finite." + ) + + +def _anchor_vlm_sensors(sensors: list[dict[str, Any]], scene: PreparedScene) -> None: + """Aim the fixed high views at the normalized tabletop center.""" + table = next( + ( + item + for item in scene.background + if isinstance(item, dict) and str(item.get("uid")) == "table" + ), + None, + ) + init_pos = table.get("init_pos", [0.0, 0.0, 0.0]) if table else [0.0, 0.0, 0.0] + if not isinstance(init_pos, Sequence) or len(init_pos) != 3: + init_pos = [0.0, 0.0, 0.0] + center = [ + float(init_pos[0]), + float(init_pos[1]), + float(scene.table_top_z if scene.table_top_z is not None else 0.75), + ] + for sensor in sensors: + extrinsics = sensor["extrinsics"] + eye = [float(value) for value in extrinsics["eye"]] + target = [float(value) for value in extrinsics["target"]] + offset = [target[index] - 0.0 for index in range(3)] + extrinsics["target"] = list(center) + extrinsics["eye"] = [ + center[index] + eye[index] - offset[index] for index in range(3) + ] diff --git a/embodichain/gen_sim/action_engine/generation/generator.py b/embodichain/gen_sim/action_engine/generation/generator.py new file mode 100644 index 000000000..ac9da7be2 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/generator.py @@ -0,0 +1,877 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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, + source_scene_xy_translation: Sequence[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, + source_scene_xy_translation=source_scene_xy_translation, + 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, + articulation_settings={ + str(config["uid"]): deepcopy( + config.get("attributes", {}).get("joint_settings", {}) + ) + for config in scene.planner_objects + if config.get("role") == "articulation" + and isinstance(config.get("attributes"), Mapping) + and config.get("attributes", {}).get("joint_settings") + }, + planning_mode=planning_mode, + seed_task_graph_path=graph_relative_path, + vlm_model=vlm_model, + vlm_camera_uids=vlm_camera_uids, + ) + gym_config = build_fast_gym_config( + scene, + task_name=task_name, + task_description=task_description, + robot_profile=robot_profile, + execution_program_hash=program_hash, + max_episodes=max_episodes, + max_episode_steps=max_episode_steps, + randomize_scene=randomize_scene, + randomize_table_material=randomize_table_material, + planning_mode=planning_mode, + seed_task_graph_path=graph_relative_path, + ) + if planning_mode == "ab": + output_root = Path(output_dir).expanduser().resolve() + gym_config["env"]["events"]["record_camera"]["params"]["save_path"] = ( + output_root / ".ab_video_staging" + ).as_posix() + gym_config["env"]["dataset"]["lerobot"]["params"]["save_path"] = ( + output_root / ".ab_datasets" + ).as_posix() + _validate_agent_config(agent_config) + return write_generation_artifacts( + output_dir, + gym_config=gym_config, + agent_config=agent_config, + task_spec=task_spec, + scene_requirements=scene_requirements, + seed_task_graph=execution_program, + seed_task_graph_png=seed_task_graph_png, + overwrite=overwrite, + planning_mode=planning_mode, + ) + + +def _read_task_spec( + source: Mapping[str, Any] | str | Path, +) -> tuple[dict[str, Any], Path | None]: + """Read one existing v2 TaskSpec without invoking a text planner.""" + if isinstance(source, Mapping): + return deepcopy(dict(source)), None + path = Path(source).expanduser().resolve() + return _read_json_mapping(path, label="TaskSpec"), path + + +def _read_json_mapping(path: Path, *, label: str) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except OSError as exc: + raise ValueError(f"Unable to read {label} at {path}: {exc}") from exc + except json.JSONDecodeError as exc: + raise ValueError(f"{label} at {path} is not valid JSON: {exc}") from exc + if not isinstance(value, Mapping): + raise ValueError(f"{label} JSON must contain an object.") + return deepcopy(dict(value)) + + +def _read_sibling_scene_requirements( + task_spec_path: Path | None, + task_name: str, +) -> dict[str, Any] | None: + """Load the canonical sidecar when a task-first batch supplied one.""" + if task_spec_path is None: + return None + candidate = task_spec_path.parent / SCENE_REQUIREMENTS_FILENAME + if not candidate.is_file(): + return None + requirements = _read_json_mapping(candidate, label="SceneRequirements") + if requirements.get("task_id") != task_name: + raise ValueError( + "Sibling SceneRequirements task_id does not match the requested " + "task_name." + ) + return requirements + + +def _require_matching_task_spec(task_spec: Mapping[str, Any], task_name: str) -> None: + if task_spec.get("task_id") != task_name: + raise ValueError( + f"TaskSpec task_id {task_spec.get('task_id')!r} does not match " + f"requested task_name {task_name!r}." + ) + + +def _task_spec_role_bindings( + task_spec: Mapping[str, Any], + known_objects: Sequence[str], + *, + scene_requirements: Mapping[str, Any] | None = None, + scene_objects: Sequence[Mapping[str, Any]] | None = None, + robot_profile: str = "dual_ur10", +) -> dict[str, str]: + """Resolve v2 roles from explicit hand-off data or a strict sidecar match. + + Task-first artifacts may contain abstract role IDs rather than scene UIDs. + When their sibling SceneRequirements is available, match + every still-unbound role against the source scene's static category, + attributes, state, and affordance metadata. This is a deterministic + Scene-Engine hand-off, not a text-model fallback: missing or ambiguous + evidence remains an error. + """ + metadata = task_spec.get("metadata", {}) + if not isinstance(metadata, Mapping): + raise ValueError("TaskSpec.metadata must be a mapping.") + metadata_bindings = metadata.get("role_bindings", {}) + if not isinstance(metadata_bindings, Mapping): + raise ValueError("TaskSpec.metadata.role_bindings must be a mapping.") + candidates: list[tuple[str, Mapping[str, Any]]] = [] + if metadata_bindings: + candidates.append(("TaskSpec.metadata", metadata_bindings)) + + # Older grounded v2 TaskSpecs kept this private hand-off in ``oracle`` + # rather than metadata. Accept that representation while publishing the + # normalized binding in metadata for runtime regeneration. + oracle = task_spec.get("oracle", {}) + if isinstance(oracle, Mapping) and oracle.get("role_bindings"): + oracle_bindings = oracle["role_bindings"] + if not isinstance(oracle_bindings, Mapping): + raise ValueError("TaskSpec.oracle.role_bindings must be a mapping.") + candidates.append(("TaskSpec.oracle", oracle_bindings)) + if isinstance(oracle, Mapping): + reference = oracle.get("reference_seed_graph") + if isinstance(reference, Mapping): + graph_metadata = reference.get("metadata", {}) + if isinstance(graph_metadata, Mapping) and graph_metadata.get( + "role_bindings" + ): + graph_bindings = graph_metadata["role_bindings"] + if not isinstance(graph_bindings, Mapping): + raise ValueError( + "SeedGraph.metadata.role_bindings must be a mapping." + ) + candidates.append(("SeedGraph.metadata", graph_bindings)) + + if scene_requirements is not None: + requirement_metadata = scene_requirements.get("metadata", {}) + if isinstance(requirement_metadata, Mapping) and requirement_metadata.get( + "role_bindings" + ): + requirement_bindings = requirement_metadata["role_bindings"] + if not isinstance(requirement_bindings, Mapping): + raise ValueError( + "SceneRequirements.metadata.role_bindings must be a mapping." + ) + candidates.append(("SceneRequirements.metadata", requirement_bindings)) + + supplied: dict[str, Any] = {} + supplied_sources: dict[str, str] = {} + for source, candidate in candidates: + for raw_role, uid in candidate.items(): + if not isinstance(raw_role, str) or not raw_role.strip(): + raise ValueError(f"{source}.role_bindings must use non-empty role IDs.") + role = raw_role.strip() + if role in supplied and supplied[role] != uid: + raise ValueError( + "Conflicting role_bindings were supplied for " + f"{role!r} by {supplied_sources[role]} and {source}." + ) + supplied[role] = uid + supplied_sources[role] = source + + known = {str(uid) for uid in known_objects} + required = _task_spec_role_references(task_spec.get("task_instances", [])) + required.discard("table") + if not required: + raise ValueError("TaskSpec must reference at least one non-table object role.") + + bindings: dict[str, str] = {} + missing: list[str] = [] + for role in sorted(required): + raw_uid = supplied.get(role, role if role in known else None) + if raw_uid is None: + missing.append(role) + continue + if not isinstance(raw_uid, str) or not raw_uid.strip(): + raise ValueError( + "TaskSpec.metadata.role_bindings must map role IDs to non-empty " + "runtime UIDs." + ) + uid = raw_uid.strip() + if uid not in known: + raise ValueError(f"TaskSpec role {role!r} binds unknown scene UID {uid!r}.") + bindings[role] = uid + if missing and scene_requirements is not None and scene_objects is not None: + bindings.update( + _infer_role_bindings_from_scene_requirements( + missing, + known_objects=known, + scene_objects=scene_objects, + scene_requirements=scene_requirements, + existing_bindings=bindings, + robot_profile=robot_profile, + ) + ) + missing = [role for role in missing if role not in bindings] + if missing: + raise ValueError( + "TaskSpec requires explicit role_bindings or an unambiguous sibling " + f"SceneRequirements match for roles {missing}; a task-first spec must " + "be grounded by a Scene Engine before it can be compiled for this gym " + "project." + ) + if len(bindings.values()) != len(set(bindings.values())): + raise ValueError("TaskSpec role bindings must resolve to unique scene UIDs.") + if scene_requirements is not None and scene_objects is not None: + _validate_bound_role_requirements( + bindings, + scene_requirements=scene_requirements, + scene_objects=scene_objects, + robot_profile=robot_profile, + ) + return bindings + + +def _infer_role_bindings_from_scene_requirements( + roles: Sequence[str], + *, + known_objects: set[str], + scene_objects: Sequence[Mapping[str, Any]], + scene_requirements: Mapping[str, Any], + existing_bindings: Mapping[str, str], + robot_profile: str, +) -> dict[str, str]: + """Bind abstract task roles only when static evidence is unique.""" + from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + requirements = _requirements_by_role(scene_requirements) + inventory = SceneInventory(scene_objects, robot_profile=robot_profile) + entities = [entity for entity in inventory.entities if entity.uid in known_objects] + used_uids = set(existing_bindings.values()) + inferred: dict[str, str] = {} + for role in sorted(roles): + requirement = requirements.get(role) + if requirement is None: + raise ValueError( + "Sibling SceneRequirements is missing TaskSpec role " f"{role!r}." + ) + count = requirement.get("count", 1) + if count != 1: + raise ValueError( + f"TaskSpec role {role!r} has count={count}; direct SeedGraph " + "binding requires exactly one concrete scene UID." + ) + matches = [ + entity + for entity in entities + if entity.uid not in used_uids + and _entity_matches_requirement( + entity, + requirement, + require_complete_static_evidence=True, + ) + ] + if len(matches) != 1: + raise ValueError( + "TaskSpec role " + f"{role!r} requires one unambiguous scene match, found " + f"{[entity.uid for entity in matches]}." + ) + uid = matches[0].uid + inferred[role] = uid + used_uids.add(uid) + return inferred + + +def _validate_bound_role_requirements( + bindings: Mapping[str, str], + *, + scene_requirements: Mapping[str, Any], + scene_objects: Sequence[Mapping[str, Any]], + robot_profile: str, +) -> None: + """Ensure an explicit binding does not contradict its static sidecar.""" + from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + requirements = _requirements_by_role(scene_requirements) + entities = SceneInventory(scene_objects, robot_profile=robot_profile).by_uid + for role, uid in bindings.items(): + requirement = requirements.get(role) + if requirement is None: + raise ValueError( + "Sibling SceneRequirements is missing TaskSpec role " f"{role!r}." + ) + entity = entities.get(uid) + if entity is None: + raise ValueError( + f"TaskSpec role {role!r} binds unavailable scene UID {uid!r}." + ) + if not _entity_matches_requirement( + entity, + requirement, + require_complete_static_evidence=False, + ): + raise ValueError( + f"TaskSpec role {role!r} binding {uid!r} conflicts with its " + "SceneRequirements category, attributes, state, or affordances." + ) + + +def _requirements_by_role( + scene_requirements: Mapping[str, Any], +) -> dict[str, Mapping[str, Any]]: + objects = scene_requirements.get("objects", []) + if not isinstance(objects, Sequence) or isinstance(objects, (str, bytes)): + raise ValueError("SceneRequirements.objects must be a list.") + result: dict[str, Mapping[str, Any]] = {} + for requirement in objects: + if not isinstance(requirement, Mapping): + raise ValueError("SceneRequirements.objects must contain mappings.") + role = requirement.get("role_id") + if not isinstance(role, str) or not role: + raise ValueError("SceneRequirements role_id must be a non-empty string.") + result[role] = requirement + return result + + +def _entity_matches_requirement( + entity: Any, + requirement: Mapping[str, Any], + *, + require_complete_static_evidence: bool, +) -> bool: + """Match explicit metadata; UID inference requires complete evidence.""" + category = requirement.get("category") + expected_category = category.strip().casefold() if isinstance(category, str) else "" + actual_category = str(entity.category).strip().casefold() + if expected_category: + if not actual_category: + if require_complete_static_evidence: + return False + elif expected_category != actual_category: + return False + required_affordances = requirement.get("affordances", []) + if not isinstance(required_affordances, Sequence) or isinstance( + required_affordances, (str, bytes) + ): + return False + expected_affordances = { + str(value).strip().casefold() for value in required_affordances + } + if ( + expected_affordances + and (require_complete_static_evidence or entity.affordances) + and not expected_affordances.issubset(entity.affordances) + ): + return False + expected_attributes = requirement.get("attributes", {}) + if not isinstance(expected_attributes, Mapping): + return False + for name, expected in expected_attributes.items(): + if not _static_attribute_matches( + entity, + str(name), + expected, + require_complete_static_evidence=require_complete_static_evidence, + ): + return False + expected_state = requirement.get("initial_state", {}) + if not isinstance(expected_state, Mapping): + return False + missing = object() + for name, expected in expected_state.items(): + actual = entity.initial_state.get(str(name), missing) + if actual is missing: + if require_complete_static_evidence: + return False + elif actual != expected: + return False + return True + + +def _static_attribute_matches( + entity: Any, + name: str, + expected: Any, + *, + require_complete_static_evidence: bool, +) -> bool: + """Compare one requirement against explicit exported metadata only.""" + marker = object() + actual = entity.color if name == "color" else entity.attributes.get(name, marker) + if actual is marker or actual is None or actual == "": + return not require_complete_static_evidence + if name == "color" and isinstance(actual, str) and isinstance(expected, str): + return actual.strip().casefold() == expected.strip().casefold() + return actual == expected + + +def _task_spec_role_references(value: Any, key: str = "") -> set[str]: + if isinstance(value, Mapping): + return { + role + for child_key, child in value.items() + for role in _task_spec_role_references(child, str(child_key)) + } + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return { + role for child in value for role in _task_spec_role_references(child, key) + } + if isinstance(value, str) and (key.endswith("_role") or key.endswith("_roles")): + return {value} + return set() + + +def _with_role_bindings( + task_spec: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> dict[str, Any]: + result = deepcopy(dict(task_spec)) + metadata = result.setdefault("metadata", {}) + if not isinstance(metadata, dict): + raise ValueError("TaskSpec.metadata must be a mapping.") + metadata["role_bindings"] = dict(sorted(role_bindings.items())) + return result + + +def _validate_requirement_roles( + requirements: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> None: + requirement_roles = { + str(item["role_id"]) + for item in requirements["objects"] + if isinstance(item, Mapping) + } + missing = sorted(set(role_bindings) - requirement_roles) + if missing: + raise ValueError( + "SceneRequirements is missing TaskSpec role bindings for " f"{missing}." + ) + + +def _scene_requirements_from_bindings( + task_id: str, + planner_objects: Sequence[Mapping[str, Any]], + role_bindings: Mapping[str, str], +) -> dict[str, Any]: + """Derive a minimal concrete SceneRequirements view for grounded roles.""" + source = _scene_requirements_from_scene(task_id, planner_objects) + by_uid = {str(item["role_id"]): item for item in source["objects"]} + objects = [] + for role, uid in sorted(role_bindings.items()): + requirement = by_uid.get(uid) + if requirement is None: + raise ValueError( + f"TaskSpec role {role!r} binds UID {uid!r}, which has no " + "source-scene requirement." + ) + resolved = deepcopy(requirement) + resolved["role_id"] = role + objects.append(resolved) + return { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": task_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [{"type": "preserve_source_scene"}], + "distractor_count": 0, + "metadata": {"source": "task_spec_role_bindings"}, + } + + +def _validated_mapping( + value: Any, + *, + validator: Any, + label: str, +) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError( + f"{label} producer returned {type(value).__name__}, not a mapping." + ) + candidate = deepcopy(dict(value)) + validated = validator(candidate) + if validated is None: + # Validators may either return a normalized mapping or validate in place. + validated = candidate + if not isinstance(validated, Mapping): + raise TypeError(f"{label} validator must return a mapping or None.") + return deepcopy(dict(validated)) + + +def _validate_agent_config(config: Mapping[str, Any]) -> None: + if config.get("schema_version") != ACTION_ENGINE_CONFIG_SCHEMA: + raise ValueError("Agent config has an unexpected schema_version.") + if config.get("task_spec") != TASK_SPEC_FILENAME: + raise ValueError("Agent config must point to the canonical TaskSpec.") + if config.get("scene_requirements") != SCENE_REQUIREMENTS_FILENAME: + raise ValueError("Agent config must point to canonical SceneRequirements.") + graph_path = config.get("seed_task_graph") + if ( + not isinstance(graph_path, str) + or Path(graph_path).name != EXECUTION_PROGRAM_FILENAME + ): + raise ValueError("Agent config must point to the canonical SeedGraph.") + planning_mode = config.get("planning_mode", "offline") + if planning_mode not in {"offline", "ab"}: + raise ValueError("Agent config planning_mode must be 'offline' or 'ab'.") + if planning_mode == "ab": + online = config.get("online_planning") + if not isinstance(online, Mapping): + raise ValueError("A/B agent config requires online_planning settings.") + camera_uids = online.get("camera_uids") + if camera_uids != list(VLM_CAMERA_UIDS): + raise ValueError( + "A/B agent config must list the canonical four VLM cameras." + ) + model = online.get("vlm_model") + if model is not None and (not isinstance(model, str) or not model.strip()): + raise ValueError("online_planning.vlm_model must be a string or null.") + if config.get("offline_seed_task_graph") != graph_path: + raise ValueError( + "A/B agent config offline_seed_task_graph must match seed_task_graph." + ) + if config.get("vlm_camera_uids") != camera_uids: + raise ValueError( + "A/B agent config vlm_camera_uids must match online_planning." + ) + if config.get("vlm_model") != model: + raise ValueError("A/B agent config vlm_model must match online_planning.") + resolve_agent_runtime_policy(config) + + +def _raise_if_outputs_exist( + output_dir: str | Path, + *, + overwrite: bool, + planning_mode: str = "offline", +) -> None: + if overwrite: + return + paths = artifact_paths(output_dir, planning_mode=planning_mode) + existing = [ + path + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.seed_task_graph, + paths.seed_task_graph_png, + ) + if path.exists() + ] + if existing: + names = ", ".join(path.name for path in existing) + raise FileExistsError( + f"Generated artifacts already exist in {paths.gym_config.parent}: " + f"{names}. Pass --overwrite to replace them." + ) + + +def _scene_requirements_from_scene( + task_id: str, + planner_objects: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + objects = [] + for item in planner_objects: + uid = str(item.get("runtime_uid", item.get("uid", ""))).strip() + if not uid: + raise ValueError("Planner scene object is missing a runtime UID.") + role = str(item.get("role", "object")).strip().lower() + raw_category = item.get("category", item.get("object_category", "")) + category = str(raw_category).strip().lower() or role or "object" + raw_attributes = item.get("attributes", {}) + attributes = ( + deepcopy(dict(raw_attributes)) + if isinstance(raw_attributes, Mapping) + else {} + ) + color = item.get("color") + if color not in (None, ""): + attributes.setdefault("color", color) + objects.append( + { + "role_id": uid, + "category": category, + "count": 1, + "affordances": [], + "initial_state": {}, + "attributes": attributes, + } + ) + return { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": task_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [{"type": "preserve_source_scene"}], + "distractor_count": 0, + "metadata": {"source": "existing_gym_project"}, + } + + +def _add_ab_camera_requirements( + requirements: Mapping[str, Any], +) -> dict[str, Any]: + """Declare fixed multi-view inputs in the shared A/B hand-off.""" + from embodichain.gen_sim.action_engine.domain import validate_scene_requirements + + result = deepcopy(dict(requirements)) + cameras = result.get("cameras", []) + if not isinstance(cameras, list): + raise ValueError("SceneRequirements.cameras must be a list.") + existing_uids = { + str(item.get("uid")) + for item in cameras + if isinstance(item, Mapping) and item.get("uid") + } + for uid in VLM_CAMERA_UIDS: + if uid in existing_uids: + continue + cameras.append( + { + "uid": uid, + "role": "vlm_view", + "modalities": ["rgb", "depth"], + "coverage": "all_interaction_objects", + "resolution": [640, 480], + } + ) + result["cameras"] = cameras + metadata = result.setdefault("metadata", {}) + if not isinstance(metadata, dict): + metadata = {} + result["metadata"] = metadata + metadata["planning_mode"] = "ab" + metadata["vlm_camera_uids"] = list(VLM_CAMERA_UIDS) + return validate_scene_requirements(result) diff --git a/embodichain/gen_sim/action_engine/generation/models.py b/embodichain/gen_sim/action_engine/generation/models.py new file mode 100644 index 000000000..f12ae6440 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/models.py @@ -0,0 +1,83 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Small value objects used by Action Engine config generation.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +__all__ = ["GeneratedConfigPaths", "PreparedScene"] + + +@dataclass(frozen=True) +class GeneratedConfigPaths: + """Paths written by one successful generation transaction.""" + + gym_config: Path + agent_config: Path + task_spec: Path + scene_requirements: Path + seed_task_graph: Path + seed_task_graph_png: Path + planning_mode: str = "offline" + + @property + def execution_program(self) -> Path: + """Retain the Python API alias for callers migrating to SeedGraph v3.""" + return self.seed_task_graph + + @property + def offline_seed_task_graph(self) -> Path: + """Explicit A/B alias for the immutable offline SeedGraph artifact.""" + return self.seed_task_graph + + @property + def seed_task_graph_path(self) -> Path: + """Path-style alias used by runtime config loaders.""" + return self.seed_task_graph + + @property + def offline_seed_task_graph_path(self) -> Path: + """Verbose alias for callers that distinguish A/B graph branches.""" + return self.seed_task_graph + + @property + def offline_seed_task_graph_png(self) -> Path: + """Explicit A/B alias for the review rendering of the offline graph.""" + return self.seed_task_graph_png + + +@dataclass(frozen=True) +class PreparedScene: + """A source scene normalized for both planning and simulator loading.""" + + source_config_path: Path + scene_dir: Path + planner_objects: tuple[dict[str, Any], ...] + background: tuple[dict[str, Any], ...] + rigid_objects: tuple[dict[str, Any], ...] + articulations: tuple[dict[str, Any], ...] + uid_map: dict[str, str] + table_top_z: float | None + z_rotation_degrees: float + body_scale_policy: str + body_scale: tuple[float, float, float] + asset_hashes: dict[str, str] + source_scene_xy_translation: tuple[float, float] = (0.0, 0.0) + asset_provenance: tuple[dict[str, Any], ...] = () diff --git a/embodichain/gen_sim/action_engine/generation/source_scene.py b/embodichain/gen_sim/action_engine/generation/source_scene.py new file mode 100644 index 000000000..fd2f8845b --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/source_scene.py @@ -0,0 +1,644 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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.data import get_data_path +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, + source_scene_xy_translation: Sequence[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) + source_robot = source.get("robot") + source_has_robot = isinstance(source_robot, Mapping) and bool(source_robot) + source_table = next( + ( + item + for role, item in source_entries + if role == "background" and str(item.get("uid", "")) == table_source_uid + ), + None, + ) + if source_scene_xy_translation is not None: + if len(source_scene_xy_translation) != 2 or any( + not math.isfinite(float(value)) for value in source_scene_xy_translation + ): + raise ValueError( + "source_scene_xy_translation must contain two finite values." + ) + resolved_xy_translation = tuple( + float(value) for value in source_scene_xy_translation + ) + elif source_has_robot and source_table is not None: + table_anchor = _vector3(source_table.get("init_pos", (0.0, 0.0, 0.0))) + resolved_xy_translation = (-table_anchor[0], -table_anchor[1]) + else: + resolved_xy_translation = (0.0, 0.0) + 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) + normalized["init_pos"][0] += resolved_xy_translation[0] + normalized["init_pos"][1] += resolved_xy_translation[1] + _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, + source_scene_xy_translation=resolved_xy_translation, + ) + + +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: + source = _read_json_object(path) + if not any( + isinstance(source.get(section), Sequence) for section in _SCENE_SECTIONS + ): + expected = ", ".join(_CONFIG_FILENAMES) + raise ValueError( + f"Expected one of {expected} or an explicit legacy scene JSON, " + f"got: {path}" + ) + return ResolvedSceneSource( + path=path, + source_format=_LEGACY_GYM_FORMAT, + is_prompt2scene=False, + ) + 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() and not raw.is_absolute(): + resolved = Path(get_data_path(fpath)).expanduser().resolve() + if not resolved.is_file(): + raise FileNotFoundError(f"Scene asset does not exist: {resolved}") + return resolved + + +def _normalize_pose_fields(config: dict[str, Any]) -> None: + config["init_pos"] = _vector3(config.get("init_pos", [0.0, 0.0, 0.0])) + config["init_rot"] = _vector3(config.get("init_rot", [0.0, 0.0, 0.0])) + if "body_scale" in config: + scale = _vector3(config["body_scale"]) + if any(value <= 0.0 for value in scale): + raise ValueError( + f"Object {config.get('uid')!r} has non-positive body_scale." + ) + config["body_scale"] = scale + + +def _apply_body_scale_policy( + config: dict[str, Any], + *, + policy: str, + requested: Sequence[float], +) -> None: + source = _vector3(config.get("body_scale", [1.0, 1.0, 1.0])) + if policy == "preserve": + result = source + elif policy == "multiply": + result = [left * right for left, right in zip(source, requested)] + else: + result = list(requested) + config["body_scale"] = [_clean_float(value) for value in result] + + +def _file_hash(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _apply_world_z_rotation(config: dict[str, Any], degrees: float) -> None: + if math.isclose(degrees, 0.0, abs_tol=1e-12): + return + theta = math.radians(degrees) + cos_theta, sin_theta = math.cos(theta), math.sin(theta) + x, y, z = _vector3(config["init_pos"]) + config["init_pos"] = [ + _clean_float(x * cos_theta - y * sin_theta), + _clean_float(x * sin_theta + y * cos_theta), + _clean_float(z), + ] + + # EmbodiChain and Prompt2Scene both interpret these values as intrinsic XYZ. + from scipy.spatial.transform import Rotation + + original = Rotation.from_euler("XYZ", config["init_rot"], degrees=True) + world_z = Rotation.from_rotvec([0.0, 0.0, theta]) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="Gimbal lock detected") + rotated = (world_z * original).as_euler("XYZ", degrees=True) + config["init_rot"] = [_clean_float(value) for value in rotated] + if "init_local_pose" in config: + # Keeping two pose representations risks the stale local matrix + # overriding the rotated Euler pose in ObjectBaseCfg.from_dict. + del config["init_local_pose"] + + +def _planner_object( + config: Mapping[str, Any], + *, + source_uid: str, + role: str, +) -> dict[str, Any]: + description = str(config.get("description", "")).strip() + shape = deepcopy(dict(config.get("shape", {}))) + raw_attributes = config.get("attributes", {}) + if not isinstance(raw_attributes, Mapping): + raw_attributes = {} + raw_initial_state = config.get("initial_state", config.get("state", {})) + if not isinstance(raw_initial_state, Mapping): + raw_initial_state = {} + raw_affordances = config.get("affordances", config.get("capabilities", [])) + affordances = ( + [str(value) for value in raw_affordances] + if isinstance(raw_affordances, Sequence) + and not isinstance(raw_affordances, (str, bytes)) + else [] + ) + return { + "uid": str(config["uid"]), + "runtime_uid": str(config["uid"]), + "source_uid": source_uid, + "role": role, + "name": str(config.get("name", "")).strip(), + "description": description, + "shape": shape, + "init_pos": list(config["init_pos"]), + "init_rot": list(config["init_rot"]), + "body_scale": list(config.get("body_scale", [1.0, 1.0, 1.0])), + "category": config.get("category", config.get("object_category", "")), + "color": config.get("color", raw_attributes.get("color")), + "attributes": deepcopy(dict(raw_attributes)), + "initial_state": deepcopy(dict(raw_initial_state)), + "affordances": affordances, + } + + +def _runtime_object(config: Mapping[str, Any], *, role: str) -> dict[str, Any]: + if role == "articulation": + # Articulation schemas vary by asset; preserve their source fields after + # path and pose normalization instead of guessing a reduced schema. + result = deepcopy(dict(config)) + result.pop("description", None) + return result + + result = { + key: deepcopy(config[key]) + for key in ( + "uid", + "shape", + "init_pos", + "init_rot", + "body_scale", + ) + if key in config + } + result.setdefault("body_scale", [1.0, 1.0, 1.0]) + source_attrs = dict(config.get("attrs", {})) + if role == "background": + result["attrs"] = {**source_attrs, **_BACKGROUND_ATTRS} + result["body_type"] = "kinematic" + result["max_convex_hull_num"] = int(_BACKGROUND_POLICY["max_convex_hull_num"]) + else: + result["attrs"] = {**source_attrs, **_RIGID_ATTRS} + result["body_type"] = "dynamic" + hull_limit = int(_RIGID_POLICY["max_convex_hull_num"]) + max_hulls = max( + 1, + min(int(config.get("max_convex_hull_num", hull_limit)), hull_limit), + ) + result["max_convex_hull_num"] = max_hulls + result["acd_method"] = str(_RIGID_POLICY["acd_method"]) + shape = result.get("shape") + if isinstance(shape, dict): + shape["acd_method"] = str(_RIGID_POLICY["acd_method"]) + shape["max_convex_hull_num"] = max_hulls + return result + + +def _estimate_mesh_top_z(config: Mapping[str, Any]) -> float | None: + shape = config.get("shape", {}) + if not isinstance(shape, Mapping) or not shape.get("fpath"): + return None + try: + import numpy as np + import trimesh + from scipy.spatial.transform import Rotation + + loaded = trimesh.load(str(shape["fpath"]), force="scene") + geometry = ( + loaded.to_geometry() + if hasattr(loaded, "to_geometry") + else loaded.dump(concatenate=True) + ) + vertices = np.asarray(geometry.vertices, dtype=np.float64) + if vertices.size == 0: + return None + # DexSim converts glTF Y-up vertices to its Z-up basis at load time. + sim_vertices = np.column_stack( + (vertices[:, 0], -vertices[:, 2], vertices[:, 1]) + ) + sim_vertices *= np.asarray( + config.get("body_scale", [1.0, 1.0, 1.0]), dtype=np.float64 + ) + rotated = Rotation.from_euler( + "XYZ", config.get("init_rot", [0.0, 0.0, 0.0]), degrees=True + ).apply(sim_vertices) + rotated += np.asarray(config.get("init_pos", [0.0, 0.0, 0.0]), dtype=np.float64) + return float(rotated[:, 2].max()) + except Exception: + # Mesh bounds improve robot placement but are not needed to preserve the + # exported scene. The robot builder has a conservative tabletop fallback. + return None + + +def _vector3(value: Any) -> list[float]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"Expected a finite xyz vector, got: {value!r}") + values = [float(item) for item in value] + if len(values) != 3 or not all(math.isfinite(item) for item in values): + raise ValueError(f"Expected a finite xyz vector, got: {value!r}") + return values + + +def _clean_float(value: float) -> float: + rounded = round(float(value), 12) + return 0.0 if abs(rounded) < 1e-12 else rounded diff --git a/embodichain/gen_sim/action_engine/generation/templates/default_lights.json b/embodichain/gen_sim/action_engine/generation/templates/default_lights.json new file mode 100644 index 000000000..5ea73ee5b --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/default_lights.json @@ -0,0 +1,3 @@ +{ + "direct": [] +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/default_sensors.json b/embodichain/gen_sim/action_engine/generation/templates/default_sensors.json new file mode 100644 index 000000000..f9ad7aea8 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/default_sensors.json @@ -0,0 +1,14 @@ +[ + { + "sensor_type": "Camera", + "width": 960, + "height": 540, + "intrinsics": [420, 420, 480, 270], + "extrinsics": { + "pos": [0.4, 0.0, 2.2], + "eye": [-0.6, 0.0, 1.8], + "target": [0.0, 0.0, 0.75], + "up": [1.0, 0.0, 0.0] + } + } +] diff --git a/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json b/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json new file mode 100644 index 000000000..b5709f40d --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json @@ -0,0 +1,185 @@ +{ + "uid": "DualFrankaPanda", + "urdf_cfg": { + "fname": "dual_franka_panda_basket", + "name_case": { + "joint": "original", + "link": "original" + }, + "components": [ + { + "component_type": "left_arm", + "urdf_path": "Franka/Panda/Panda.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.25], + [0.0, 1.0, 0.0, 0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "left_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_arm", + "urdf_path": "Franka/Panda/Panda.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.25], + [0.0, 1.0, 0.0, -0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + } + ] + }, + "init_pos": [-0.7, 0.0, 0.0], + "init_rot": [0.0, 0.0, 180.0], + "init_qpos": [ + 0.0, + 0.0, + -0.569, + -0.569, + 0.0, + 0.0, + -2.81, + -2.81, + 0.0, + 0.0, + 3.037, + 3.037, + 0.0, + 0.0, + + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0 + ], + "drive_pros": { + "stiffness": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 50.0, + "right_eef": 50.0 + }, + "damping": { + "left_arm": 1000.0, + "right_arm": 1000.0, + "left_eef": 5.0, + "right_eef": 5.0 + }, + "max_effort": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 500.0, + "right_eef": 500.0 + } + }, + "control_parts": { + "left_arm": [ + "left_fr3_joint1", + "left_fr3_joint2", + "left_fr3_joint3", + "left_fr3_joint4", + "left_fr3_joint5", + "left_fr3_joint6", + "left_fr3_joint7" + ], + "left_eef": [ + "left_finger_joint", + "left_inner_knuckle_joint", + "left_inner_finger_joint", + "left_right_outer_knuckle_joint", + "left_right_inner_knuckle_joint", + "left_right_inner_finger_joint" + ], + "right_arm": [ + "right_fr3_joint1", + "right_fr3_joint2", + "right_fr3_joint3", + "right_fr3_joint4", + "right_fr3_joint5", + "right_fr3_joint6", + "right_fr3_joint7" + ], + "right_eef": [ + "right_finger_joint", + "right_left_inner_knuckle_joint", + "right_left_inner_finger_joint", + "right_outer_knuckle_joint", + "right_inner_knuckle_joint", + "right_inner_finger_joint" + ], + "dual_arm": [ + "left_fr3_joint1", + "left_fr3_joint2", + "left_fr3_joint3", + "left_fr3_joint4", + "left_fr3_joint5", + "left_fr3_joint6", + "left_fr3_joint7", + "right_fr3_joint1", + "right_fr3_joint2", + "right_fr3_joint3", + "right_fr3_joint4", + "right_fr3_joint5", + "right_fr3_joint6", + "right_fr3_joint7" + ] + }, + "observation_joint_parts": ["left_eef", "right_eef"], + "qpos_control_part_order": ["dual_arm", "left_eef", "right_eef"], + "solver_cfg": { + "left_arm": { + "class_type": "PytorchSolver", + "urdf_path": null, + "end_link_name": "left_fr3_link8", + "root_link_name": "left_base", + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ], + "num_samples": 15 + }, + "right_arm": { + "class_type": "PytorchSolver", + "urdf_path": null, + "end_link_name": "right_fr3_link8", + "root_link_name": "right_base", + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ], + "num_samples": 15 + } + } +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json b/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json new file mode 100644 index 000000000..8a7496547 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json @@ -0,0 +1,126 @@ +{ + "uid": "DualUR5", + "urdf_cfg": { + "fname": "dual_ur5_robotiq_arg2f_140_basket", + "name_case": {"joint": "lower", "link": "lower"}, + "components": [ + { + "component_type": "left_arm", + "urdf_path": "UniversalRobots/UR5/UR5.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.45], + [0.0, 1.0, 0.0, -0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "left_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_arm", + "urdf_path": "UniversalRobots/UR5/UR5.urdf", + "transform": [ + [1.0, 0.0, 0.0, -1.45], + [0.0, 1.0, 0.0, 0.3], + [0.0, 0.0, 1.0, 0.4], + [0.0, 0.0, 0.0, 1.0] + ] + }, + { + "component_type": "right_hand", + "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "transform": [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0] + ] + } + ] + }, + "init_pos": [2.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "init_qpos": [ + 0, 0, -1.57, -1.57, 1.57, 1.57, -1.57, -1.57, + -1.57, -1.57, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 + ], + "drive_pros": { + "stiffness": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 50.0, + "right_eef": 50.0 + }, + "damping": { + "left_arm": 1000.0, + "right_arm": 1000.0, + "left_eef": 5.0, + "right_eef": 5.0 + }, + "max_effort": { + "left_arm": 10000.0, + "right_arm": 10000.0, + "left_eef": 500.0, + "right_eef": 500.0 + } + }, + "control_parts": { + "left_arm": [ + "left_joint1", "left_joint2", "left_joint3", + "left_joint4", "left_joint5", "left_joint6" + ], + "left_eef": [ + "left_finger_joint", "left_inner_knuckle_joint", + "left_inner_finger_joint", "left_right_outer_knuckle_joint", + "left_right_inner_knuckle_joint", "left_right_inner_finger_joint" + ], + "right_arm": [ + "right_joint1", "right_joint2", "right_joint3", + "right_joint4", "right_joint5", "right_joint6" + ], + "right_eef": [ + "right_finger_joint", "right_left_inner_knuckle_joint", + "right_left_inner_finger_joint", "right_outer_knuckle_joint", + "right_inner_knuckle_joint", "right_inner_finger_joint" + ] + }, + "solver_cfg": { + "left_arm": { + "class_type": "URSolver", + "ur_type": "ur5", + "urdf_path": null, + "end_link_name": "left_ee_link", + "root_link_name": "left_base_link", + "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ] + }, + "right_arm": { + "class_type": "URSolver", + "ur_type": "ur5", + "urdf_path": null, + "end_link_name": "right_ee_link", + "root_link_name": "right_base_link", + "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], + "tcp": [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0] + ] + } + } +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json b/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json new file mode 100644 index 000000000..31084a497 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json @@ -0,0 +1,45 @@ +{ + "dual_franka": { + "aliases": ["franka", "panda", "dual_panda", "dual_franka_panda"], + "template": "dual_franka_robot.json", + "robot_family": "franka", + "tabletop_clearance": 0.05, + "arm_component_z": 0.4, + "arm_base_x": -1.45, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + }, + "dual_ur3": { + "aliases": ["ur3", "dual_ur3_dh_pgi", "dual_ur3_robotiq", "dual_ur3_robotiq_arg2f_140"], + "template": "dual_ur_robot.json", + "robot_family": "ur3", + "tabletop_clearance": 0.05, + "arm_component_z": 0.4, + "arm_base_x": -1.45, + "max_effort": 56.0, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + }, + "dual_ur5": { + "aliases": ["ur5", "dual_ur5_dh_pgi", "dual_ur5_robotiq", "dual_ur5_robotiq_arg2f_140"], + "template": "dual_ur_robot.json", + "robot_family": "ur5", + "tabletop_clearance": 0.05, + "arm_component_z": 0.4, + "arm_base_x": -1.45, + "max_effort": 10000.0, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + }, + "dual_ur10": { + "aliases": ["ur10", "dual_ur10_dh_pgi", "dual_ur10_robotiq", "dual_ur10_robotiq_arg2f_140"], + "template": "dual_ur_robot.json", + "robot_family": "ur10", + "tabletop_clearance": 0.05, + "arm_component_z": 0.3, + "arm_base_x": -1.1, + "max_effort": 330.0, + "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + } +} diff --git a/embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json b/embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json new file mode 100644 index 000000000..34e964569 --- /dev/null +++ b/embodichain/gen_sim/action_engine/generation/templates/vlm_sensors.json @@ -0,0 +1,58 @@ +[ + { + "uid": "vlm_front", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [-1.2, 0.0, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + }, + { + "uid": "vlm_left", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [0.0, 1.2, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + }, + { + "uid": "vlm_rear", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [1.2, 0.0, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + }, + { + "uid": "vlm_right", + "sensor_type": "Camera", + "width": 640, + "height": 480, + "intrinsics": [520, 520, 320, 240], + "enable_color": true, + "enable_depth": true, + "extrinsics": { + "eye": [0.0, -1.2, 1.65], + "target": [0.0, 0.0, 0.75], + "up": [0.0, 0.0, 1.0] + } + } +] diff --git a/embodichain/gen_sim/action_engine/graph_visualization.py b/embodichain/gen_sim/action_engine/graph_visualization.py new file mode 100644 index 000000000..4ad262c60 --- /dev/null +++ b/embodichain/gen_sim/action_engine/graph_visualization.py @@ -0,0 +1,938 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Headless PNG rendering for direct AtomicAction SeedGraphs. + +The renderer consumes the same validated coordinate-free v3 graph as runtime, +then builds an internal display view without grounding symbolic targets. E +TaskGroups remain the semantic grouping labels over the rendered action nodes. +Single chains use a folded timeline; DAGs use stable actor swimlanes. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import lru_cache +from io import BytesIO +from math import hypot +from typing import Any + +import matplotlib + +# Select the non-interactive backend before importing any canvas primitives. +matplotlib.use("Agg", force=True) + +from matplotlib import patheffects +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.font_manager import FontProperties, fontManager +from matplotlib.figure import Figure +from matplotlib.patches import Circle, FancyArrowPatch, FancyBboxPatch +import networkx as nx + +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + validate_execution_program, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + +__all__ = ["render_seed_task_graph_png", "render_task_graph_png"] + +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" +_EXECUTION_KEYS = frozenset( + { + "schema_version", + "task", + "goal_description", + "start", + "goal", + "nodes", + "edges", + "semantic_steps", + "allocation_groups", + "motion_policy_version", + } +) + +_BACKGROUND = "#F8FAFB" +_INK = "#17212B" +_MUTED = "#66727D" +_BORDER = "#CBD4DC" +_LEFT = "#168A78" +_RIGHT = "#D97706" +_AUTO = "#59636D" +_COORDINATED = "#7652A5" +_DEPENDENCY = "#8A94A0" + +# The figures are designed at this display width in inches; every type size +# below is chosen to stay readable when the PNG is shown at exactly this size. +_TARGET_WIDTH = 8.0 +_DPI = 300 +_LEVEL_STEP = 1.15 +_NODE_RADIUS = 0.16 +_SPECIAL_NODE_RADIUS = 0.20 +_SUCCESS = "#25834B" +_FAILED = "#C43E3E" +_SKIPPED = "#8B949C" +_LANE_COLORS = { + "left": _LEFT, + "auto": _AUTO, + "right": _RIGHT, + "coordinated": _COORDINATED, +} +_LANE_BACKGROUNDS = { + "left": "#EAF6F3", + "auto": "#F0F3F5", + "right": "#FFF4E6", +} +_LANE_LABELS = { + "left": "LEFT ARM [L]", + "auto": "WORLD / AUTO / COORDINATED", + "right": "RIGHT ARM [R]", +} +_STATUS_COLORS = { + "success": _SUCCESS, + "executed": _SUCCESS, + "failed": _FAILED, + "aborted": _FAILED, + "skipped": _SKIPPED, +} +_STATUS_BADGES = { + "success": "OK", + "executed": "OK", + "failed": "FAIL", + "aborted": "ABORT", + "skipped": "SKIP", +} + + +@dataclass(frozen=True) +class _RuntimeOverlay: + """Execution annotations kept separate from the immutable seed program.""" + + edge_status: Mapping[str, str] + edge_arm: Mapping[str, str] + step_status: Mapping[str, str] + graph_status: str | None = None + + +@dataclass(frozen=True) +class _GraphData: + """Validated program plus indices shared by both layout strategies.""" + + program: Mapping[str, Any] + graph: nx.MultiDiGraph + node_by_id: Mapping[str, Mapping[str, Any]] + edge_by_id: Mapping[str, Mapping[str, Any]] + step_by_id: Mapping[str, Mapping[str, Any]] + lane_override: Mapping[str, str] + runtime: _RuntimeOverlay + + +def render_seed_task_graph_png(seed_graph: Mapping[str, Any]) -> bytes: + """Render a v3 SeedGraph or package-owned legacy program through Agg.""" + program = _display_program(seed_graph) + return _render(program, _RuntimeOverlay({}, {}, {})) + + +def render_task_graph_png(task_graph: Mapping[str, Any]) -> bytes: + """Render an execution program with optional runtime event annotations. + + A bare program is accepted. Runtime events may be stored in its ``runtime`` + envelope, or beside a nested ``execution_program``, ``program``, or + ``seed_task_graph``. A record alone is rejected because it omits topology. + """ + program = _extract_execution_program(task_graph) + runtime = _extract_runtime_overlay(task_graph) + return _render(program, runtime) + + +def _render( + program: Mapping[str, Any], + runtime: _RuntimeOverlay, +) -> bytes: + data = _graph_data(program, runtime) + if _is_single_chain(data): + return _render_chain(data) + return _render_dag(data) + + +def _extract_execution_program(document: Mapping[str, Any]) -> dict[str, Any]: + """Find and validate the execution program embedded in a display document.""" + if not isinstance(document, Mapping): + raise ValueError("Task graph visualization input must be a mapping.") + + if document.get("schema_version") in {EXECUTION_PROGRAM_SCHEMA, SEED_GRAPH_SCHEMA}: + # A runtime artifact may preserve the program fields and add annotations. + if document.get("schema_version") == SEED_GRAPH_SCHEMA: + candidate = dict(document) + candidate.pop("runtime", None) + candidate.pop("runtime_record", None) + return _display_program(candidate) + candidate = {key: document[key] for key in _EXECUTION_KEYS if key in document} + return validate_execution_program(candidate) + + for key in ("execution_program", "program", "seed_task_graph"): + candidate = document.get(key) + if isinstance(candidate, Mapping): + return _display_program(candidate) + + # Supporting a full program plus a runtime schema at the top level keeps + # visualization useful for simple JSON joins without weakening validation. + if {"nodes", "edges", "semantic_steps"}.issubset(document): + candidate = {key: document[key] for key in _EXECUTION_KEYS if key in document} + return validate_execution_program(candidate) + + raise ValueError( + "Runtime records do not contain graph topology. Provide the matching " + "ExecutionProgram under 'execution_program', 'program', or " + "'seed_task_graph'." + ) + + +def _display_program(value: Mapping[str, Any]) -> dict[str, Any]: + if value.get("schema_version") == SEED_GRAPH_SCHEMA: + from embodichain.gen_sim.action_engine.compiler import ( + seed_graph_to_execution_program, + ) + + return seed_graph_to_execution_program(value, require_executable=False) + return validate_execution_program(value) + + +def _extract_runtime_overlay(document: Mapping[str, Any]) -> _RuntimeOverlay: + """Reduce a runtime record to the small set of display-only annotations.""" + record = document.get("runtime") + if record is None: + record = document.get("runtime_record", document) + if not isinstance(record, Mapping): + raise ValueError("runtime_record must be a mapping.") + raw_events = record.get("events", document.get("events", [])) + if not isinstance(raw_events, Sequence) or isinstance( + raw_events, (str, bytes, bytearray) + ): + raise ValueError("Runtime events must be a list.") + + edge_status: dict[str, str] = {} + edge_arm: dict[str, str] = {} + step_status: dict[str, str] = {} + for index, event in enumerate(raw_events): + if not isinstance(event, Mapping): + raise ValueError(f"Runtime events[{index}] must be a mapping.") + event_kind = event.get("event") + status = _optional_text(event.get("status")) + if event_kind == "edge": + edge_id = _optional_text(event.get("edge_id")) + if edge_id and status: + edge_status[edge_id] = status.lower() + arm = _optional_text(event.get("arm")) + if edge_id and arm: + edge_arm[edge_id] = arm + elif event_kind == "semantic_step": + step_id = _optional_text(event.get("semantic_step_id")) + if step_id and status: + step_status[step_id] = status.lower() + + graph_status = _optional_text(record.get("status")) + return _RuntimeOverlay( + edge_status=edge_status, + edge_arm=edge_arm, + step_status=step_status, + graph_status=graph_status.lower() if graph_status else None, + ) + + +def _graph_data( + program: Mapping[str, Any], + runtime: _RuntimeOverlay, +) -> _GraphData: + node_by_id = {str(node["id"]): node for node in program["nodes"]} + edge_by_id = {str(edge["id"]): edge for edge in program["edges"]} + step_by_id = {str(step["id"]): step for step in program["semantic_steps"]} + graph = nx.MultiDiGraph() + graph.add_nodes_from(node_by_id) + for edge in program["edges"]: + source = str(edge["source"]) + target = str(edge["target"]) + graph.add_edge(source, target, edge_id=str(edge["id"])) + if not nx.is_directed_acyclic_graph(graph): + raise ValueError("ExecutionProgram node topology must be a directed DAG.") + + return _GraphData( + program=program, + graph=graph, + node_by_id=node_by_id, + edge_by_id=edge_by_id, + step_by_id=step_by_id, + lane_override=_allocation_lane_overrides(program), + runtime=runtime, + ) + + +def _allocation_lane_overrides( + program: Mapping[str, Any], +) -> dict[str, str]: + """Give auto actors stable lanes when a distinct-arm group is declared.""" + result: dict[str, str] = {} + for group in program.get("allocation_groups", []): + if group.get("arm_constraint") != "distinct_arms": + continue + members = group.get("semantic_step_ids", []) + for index, step_id in enumerate(members): + result[str(step_id)] = "left" if index % 2 == 0 else "right" + return result + + +def _is_single_chain(data: _GraphData) -> bool: + graph = data.graph + if graph.number_of_edges() != graph.number_of_nodes() - 1: + return False + if any(graph.in_degree(node) > 1 for node in graph): + return False + if any(graph.out_degree(node) > 1 for node in graph): + return False + return ( + graph.in_degree(str(data.program["start"])) == 0 + and graph.out_degree(str(data.program["goal"])) == 0 + and nx.is_weakly_connected(graph) + ) + + +def _ordered_chain_edges(data: _GraphData) -> list[Mapping[str, Any]]: + current = str(data.program["start"]) + result: list[Mapping[str, Any]] = [] + while current != str(data.program["goal"]): + outgoing = list(data.graph.out_edges(current, data=True)) + if len(outgoing) != 1: + raise ValueError("ExecutionProgram chain has an incomplete path.") + _, target, attrs = outgoing[0] + result.append(data.edge_by_id[str(attrs["edge_id"])]) + current = str(target) + if len(result) != len(data.edge_by_id): + raise ValueError("ExecutionProgram chain does not cover every edge.") + return result + + +def _render_chain(data: _GraphData) -> bytes: + """Render a long linear program as a bounded, folded state timeline.""" + edges = _ordered_chain_edges(data) + nodes = [str(data.program["start"])] + nodes.extend(str(edge["target"]) for edge in edges) + + slots_per_row = 4 + row_count = (len(nodes) + slots_per_row - 1) // slots_per_row + width = _TARGET_WIDTH + height = max(3.4, 2.0 + row_count * 1.55) + figure, axis = _new_figure(width, height) + try: + _draw_header(axis, data, width) + left, right, first_y = 0.7, width - 0.7, 2.05 + spacing = (right - left) / (slots_per_row - 1) + positions: dict[str, tuple[float, float]] = {} + for index, node_id in enumerate(nodes): + row, column = divmod(index, slots_per_row) + visual_column = column if row % 2 == 0 else slots_per_row - 1 - column + positions[node_id] = ( + left + visual_column * spacing, + first_y + row * 1.55, + ) + + for edge in edges: + source = positions[str(edge["source"])] + target = positions[str(edge["target"])] + lane = _edge_lane(edge, data) + color = _edge_color(str(edge["id"]), lane, data.runtime) + label_position, label_align = _edge_label_position(source, target, width) + _draw_labeled_edge( + axis, + source, + target, + color=color, + label=_edge_label(edge, data), + label_position=label_position, + label_align=label_align, + ) + + for index, node_id in enumerate(nodes): + _draw_state_node( + axis, + positions[node_id], + start=node_id == str(data.program["start"]), + goal=node_id == str(data.program["goal"]), + fork=False, + join=False, + index=index, + ) + + _draw_legend(axis, width, height - 0.28) + return _figure_png_bytes(figure) + finally: + figure.clear() + + +def _render_dag(data: _GraphData) -> bytes: + """Render forks and joins against persistent actor swimlanes.""" + levels = _dag_levels(data.graph) + maximum_level = max(levels.values(), default=0) + width = _TARGET_WIDTH + height = max(5.2, 2.15 + maximum_level * _LEVEL_STEP + 1.35) + figure, axis = _new_figure(width, height) + try: + _draw_header(axis, data, width) + boundaries, lane_centers = _lane_geometry(width) + _draw_swimlanes(axis, height, boundaries, lane_centers) + positions = _dag_positions(data, levels, lane_centers) + + # Dependency arrows are drawn first and stay visually subordinate to + # physical state transitions; only constraints not already implied by + # the state topology are shown. + for source_id, target_id in _visible_dependencies(data): + _draw_dependency_arrow( + axis, + positions[source_id], + positions[target_id], + ) + + pair_groups: defaultdict[tuple[str, str], list[str]] = defaultdict(list) + for edge in data.edge_by_id.values(): + pair_groups[(str(edge["source"]), str(edge["target"]))].append( + str(edge["id"]) + ) + for edge in data.edge_by_id.values(): + edge_id = str(edge["id"]) + source_id = str(edge["source"]) + target_id = str(edge["target"]) + lane = _edge_lane(edge, data) + parallel_ids = pair_groups[(source_id, target_id)] + parallel_index = parallel_ids.index(edge_id) + curvature = (parallel_index - (len(parallel_ids) - 1) / 2.0) * 0.20 + label_position, label_align = _edge_label_position( + positions[source_id], + positions[target_id], + width, + ) + _draw_labeled_edge( + axis, + positions[source_id], + positions[target_id], + color=_edge_color(edge_id, lane, data.runtime), + label=_edge_label(edge, data), + label_position=label_position, + label_align=label_align, + curvature=curvature, + ) + + for index, node_id in enumerate(nx.topological_sort(data.graph)): + _draw_state_node( + axis, + positions[str(node_id)], + start=str(node_id) == str(data.program["start"]), + goal=str(node_id) == str(data.program["goal"]), + fork=data.graph.out_degree(node_id) > 1, + join=data.graph.in_degree(node_id) > 1, + index=index, + ) + + _draw_legend(axis, width, height - 0.28) + return _figure_png_bytes(figure) + finally: + figure.clear() + + +def _lane_geometry( + width: float, +) -> tuple[dict[str, tuple[float, float]], dict[str, float]]: + """Even thirds for lane boundaries with derived actor centers.""" + margin = 0.35 + area = width - 2 * margin + first = margin + area / 3.0 + second = margin + 2 * area / 3.0 + boundaries = { + "left": (margin, first), + "auto": (first, second), + "right": (second, width - margin), + } + centers = {lane: (left + right) / 2.0 for lane, (left, right) in boundaries.items()} + return boundaries, centers + + +def _edge_label_position( + source: tuple[float, float], + target: tuple[float, float], + width: float, +) -> tuple[tuple[float, float], str]: + """Place halo labels beside arrows instead of boxing them on the edge.""" + midpoint = _midpoint(source, target) + dx = target[0] - source[0] + dy = target[1] - source[1] + if abs(dx) < 0.3: + # Keep vertical-arrow labels inside the canvas: right side on the left + # half of the figure, left side on the right half. + if midpoint[0] > width / 2.0: + return (midpoint[0] - 0.14, midpoint[1]), "right" + return (midpoint[0] + 0.14, midpoint[1]), "left" + length = hypot(dx, dy) or 1.0 + normal_x, normal_y = dy / length, -dx / length + if normal_x < 0: + normal_x, normal_y = -normal_x, -normal_y + if abs(normal_x) < 0.2 and normal_y > 0: + # Horizontal arrows keep their label above the line in both directions. + normal_x, normal_y = -normal_x, -normal_y + return ( + (midpoint[0] + normal_x * 0.16, midpoint[1] + normal_y * 0.16), + "center", + ) + + +def _visible_dependencies(data: _GraphData) -> list[tuple[str, str]]: + """Node anchors for dependencies not implied by state continuity.""" + result: list[tuple[str, str]] = [] + for prerequisite_id, dependent_id in _dependency_pairs(data): + source = str(data.edge_by_id[prerequisite_id]["target"]) + target = str(data.edge_by_id[dependent_id]["source"]) + if source == target or nx.has_path(data.graph, source, target): + continue + result.append((source, target)) + return result + + +def _dag_levels(graph: nx.MultiDiGraph) -> dict[str, int]: + """Assign the longest-path depth so dependencies always flow downward.""" + levels: dict[str, int] = {} + for node in nx.topological_sort(graph): + predecessors = list(graph.predecessors(node)) + levels[str(node)] = ( + max(levels[str(parent)] for parent in predecessors) + 1 + if predecessors + else 0 + ) + return levels + + +def _dag_positions( + data: _GraphData, + levels: Mapping[str, int], + lane_centers: Mapping[str, float], +) -> dict[str, tuple[float, float]]: + """Place branch nodes in actor lanes and structural fork/join nodes centrally.""" + base: dict[str, tuple[str, int]] = {} + for node_id in data.node_by_id: + incoming = list(data.graph.in_edges(node_id, data=True)) + outgoing = list(data.graph.out_edges(node_id, data=True)) + if ( + node_id in {str(data.program["start"]), str(data.program["goal"])} + or len(incoming) > 1 + or len(outgoing) > 1 + ): + lane = "auto" + elif incoming: + edge = data.edge_by_id[str(incoming[0][2]["edge_id"])] + lane = _edge_lane(edge, data) + elif outgoing: + edge = data.edge_by_id[str(outgoing[0][2]["edge_id"])] + lane = _edge_lane(edge, data) + else: + lane = "auto" + if lane == "coordinated": + lane = "auto" + base[node_id] = (lane, levels[node_id]) + + groups: defaultdict[tuple[str, int], list[str]] = defaultdict(list) + for node_id, lane_level in base.items(): + groups[lane_level].append(node_id) + + result: dict[str, tuple[float, float]] = {} + for (lane, level), node_ids in groups.items(): + ordered = sorted(node_ids) + center = lane_centers[lane] + # Small symmetric offsets prevent same-level nodes from hiding each + # other while keeping every node visibly inside its actor lane. + offsets = [ + (index - (len(ordered) - 1) / 2.0) * 0.55 for index in range(len(ordered)) + ] + for node_id, offset in zip(ordered, offsets, strict=True): + result[node_id] = (center + offset, 2.15 + level * _LEVEL_STEP) + return result + + +def _dependency_pairs(data: _GraphData) -> list[tuple[str, str]]: + """Return explicit edge dependencies plus missing semantic dependencies.""" + result: list[tuple[str, str]] = [] + seen: set[tuple[str, str]] = set() + for edge in data.edge_by_id.values(): + dependent_id = str(edge["id"]) + for prerequisite_id in edge.get("depends_on", []): + pair = (str(prerequisite_id), dependent_id) + if pair not in seen: + seen.add(pair) + result.append(pair) + + for step in data.step_by_id.values(): + dependent_edges = step.get("edge_ids", []) + if not dependent_edges: + continue + for prerequisite_step_id in step.get("depends_on", []): + prerequisite = data.step_by_id[str(prerequisite_step_id)] + pair = ( + str(prerequisite["edge_ids"][-1]), + str(dependent_edges[0]), + ) + if pair not in seen: + seen.add(pair) + result.append(pair) + return result + + +def _edge_lane(edge: Mapping[str, Any], data: _GraphData) -> str: + edge_id = str(edge["id"]) + observed_arm = data.runtime.edge_arm.get(edge_id) + if observed_arm: + return _arm_lane(observed_arm) + + action_lanes = { + _actor_lane(action.get("actor", {})) for action in edge.get("actions", []) + } + action_lanes.discard("auto") + if action_lanes == {"left"}: + return "left" + if action_lanes == {"right"}: + return "right" + if "coordinated" in action_lanes or action_lanes == {"left", "right"}: + return "coordinated" + return data.lane_override.get(str(edge["semantic_step_id"]), "auto") + + +def _actor_lane(actor: Any) -> str: + if not isinstance(actor, Mapping): + return "auto" + mode = str(actor.get("mode", "auto")).lower() + if mode == "required": + return _arm_lane(str(actor.get("arm", ""))) + if mode == "coordinated": + return "coordinated" + return "auto" + + +def _arm_lane(arm: str) -> str: + normalized = arm.strip().lower() + if "left" in normalized: + return "left" + if "right" in normalized: + return "right" + if normalized in {"both", "coordinated", "dual_arm", "dual"}: + return "coordinated" + return "auto" + + +def _edge_color( + edge_id: str, + lane: str, + runtime: _RuntimeOverlay, +) -> str: + status = runtime.edge_status.get(edge_id) + return _STATUS_COLORS.get(status or "", _LANE_COLORS[lane]) + + +def _edge_label(edge: Mapping[str, Any], data: _GraphData) -> str: + """One-line semantic phrase; execution details live in the JSON artifacts.""" + edge_id = str(edge["id"]) + step = data.step_by_id[str(edge["semantic_step_id"])] + status = data.runtime.edge_status.get(edge_id) or data.runtime.step_status.get( + str(step["id"]) + ) + status_badge = f" [{_STATUS_BADGES.get(status, status.upper())}]" if status else "" + return _clip(f"{step['operator']}: {step['object']}", 40) + status_badge + + +def _draw_header(axis: Any, data: _GraphData, width: float) -> None: + status = data.runtime.graph_status + status_text = f" [{status.upper()}]" if status else "" + axis.text( + 0.4, + 0.42, + _clip(f"ACTION ENGINE / {data.program['task']}{status_text}", 84), + ha="left", + va="center", + color=_INK, + fontproperties=_font(10.0, "bold"), + zorder=20, + ) + axis.text( + 0.4, + 0.80, + _clip(str(data.program["goal_description"]), 115), + ha="left", + va="top", + color=_MUTED, + fontproperties=_font(7.0), + linespacing=1.25, + zorder=20, + ) + axis.plot( + [0.4, width - 0.4], + [1.28, 1.28], + color=_BORDER, + linewidth=0.7, + zorder=19, + ) + + +def _draw_swimlanes( + axis: Any, + height: float, + boundaries: Mapping[str, tuple[float, float]], + centers: Mapping[str, float], +) -> None: + for lane in ("left", "auto", "right"): + left, right = boundaries[lane] + axis.add_patch( + FancyBboxPatch( + (left, 1.50), + right - left, + height - 2.0, + boxstyle="round,pad=0.0,rounding_size=0.05", + facecolor=_LANE_BACKGROUNDS[lane], + edgecolor=_BORDER, + linewidth=0.6, + zorder=-10, + ) + ) + axis.plot( + [left, right], + [1.50, 1.50], + color=_LANE_COLORS[lane], + linewidth=1.1, + zorder=-9, + ) + axis.text( + centers[lane], + 1.74, + _LANE_LABELS[lane], + ha="center", + va="center", + color=_LANE_COLORS[lane], + fontproperties=_font(6.8, "bold"), + zorder=10, + ) + + +def _draw_legend(axis: Any, width: float, y: float) -> None: + """Single-row edge-type legend; START/GOAL labels are self-explanatory.""" + entries = ( + ("left", "left action", False), + ("right", "right action", False), + ("coordinated", "coordinated", False), + ("auto", "auto / world", False), + ("dependency", "dependency", True), + ) + slot = 1.32 + start = (width - slot * len(entries)) / 2.0 + for index, (key, label, dashed) in enumerate(entries): + x = start + index * slot + color = _DEPENDENCY if dashed else _LANE_COLORS[key] + axis.add_patch( + FancyArrowPatch( + (x, y), + (x + 0.3, y), + arrowstyle="-|>", + mutation_scale=7, + color=color, + linewidth=1.0, + linestyle=(0, (3.0, 2.6)) if dashed else "-", + zorder=20, + ) + ) + axis.text( + x + 0.38, + y, + label, + ha="left", + va="center", + color=_MUTED, + fontproperties=_font(6.2), + zorder=20, + ) + + +def _draw_labeled_edge( + axis: Any, + source: tuple[float, float], + target: tuple[float, float], + *, + color: str, + label: str, + label_position: tuple[float, float], + label_align: str = "center", + curvature: float = 0.0, +) -> None: + """Draw one solid state transition and its halo-backed one-line label.""" + axis.add_patch( + FancyArrowPatch( + source, + target, + arrowstyle="-|>", + mutation_scale=9, + color=color, + linewidth=1.15, + shrinkA=12, + shrinkB=12, + connectionstyle=f"arc3,rad={curvature}", + zorder=3, + ) + ) + axis.text( + *label_position, + label, + ha=label_align, + va="center", + color=_INK, + fontproperties=_font(6.5), + path_effects=[patheffects.withStroke(linewidth=1.7, foreground=_BACKGROUND)], + zorder=8, + ) + + +def _draw_dependency_arrow( + axis: Any, + source: tuple[float, float], + target: tuple[float, float], +) -> None: + if source == target: + return + axis.add_patch( + FancyArrowPatch( + source, + target, + arrowstyle="-|>", + mutation_scale=7, + color=_DEPENDENCY, + linewidth=0.9, + linestyle=(0, (3.0, 2.6)), + shrinkA=12, + shrinkB=12, + connectionstyle="arc3,rad=-0.2", + alpha=0.9, + zorder=1, + ) + ) + + +def _draw_state_node( + axis: Any, + center: tuple[float, float], + *, + start: bool, + goal: bool, + fork: bool, + join: bool, + index: int, +) -> None: + fill = "#DDEFEA" if start else ("#E7F2DD" if goal else "#FFFFFF") + edge = _SUCCESS if goal else (_LEFT if start else _INK) + radius = _SPECIAL_NODE_RADIUS if (start or goal or fork or join) else _NODE_RADIUS + axis.add_patch( + Circle( + center, + radius=radius, + facecolor=fill, + edgecolor=edge, + linewidth=1.1, + zorder=12, + ) + ) + axis.text( + center[0], + center[1], + str(index), + ha="center", + va="center", + color=_INK, + fontproperties=_font(6.5, "bold"), + zorder=13, + ) + role = ( + "START" + if start + else ("GOAL" if goal else ("FORK" if fork else "JOIN" if join else "")) + ) + if role: + axis.text( + center[0], + center[1] + radius + 0.12, + role, + ha="center", + va="top", + color=edge, + fontproperties=_font(6.0, "bold"), + zorder=13, + ) + + +def _new_figure(width: float, height: float) -> tuple[Figure, Any]: + figure = Figure(figsize=(width, height), dpi=_DPI, facecolor=_BACKGROUND) + axis = figure.subplots() + axis.set_facecolor(_BACKGROUND) + axis.set_axis_off() + axis.set_xlim(0.0, width) + axis.set_ylim(height, 0.0) + return figure, axis + + +def _figure_png_bytes(figure: Figure) -> bytes: + buffer = BytesIO() + FigureCanvasAgg(figure).print_png(buffer) + payload = buffer.getvalue() + if not payload.startswith(_PNG_SIGNATURE): + raise RuntimeError("Matplotlib did not produce a valid PNG payload.") + return payload + + +@lru_cache(maxsize=1) +def _font_family() -> str: + """Prefer a CJK-capable font while retaining a portable fallback.""" + available = {font.name for font in fontManager.ttflist} + for family in ( + "Noto Sans CJK SC", + "Noto Sans CJK JP", + "Source Han Sans CN", + "WenQuanYi Micro Hei", + "Microsoft YaHei", + "Arial Unicode MS", + "DejaVu Sans", + ): + if family in available: + return family + return "sans-serif" + + +def _font(size: float, weight: str = "normal") -> FontProperties: + return FontProperties(family=_font_family(), size=size, weight=weight) + + +def _optional_text(value: Any) -> str | None: + return value.strip() if isinstance(value, str) and value.strip() else None + + +def _clip(value: str, length: int) -> str: + return value if len(value) <= length else f"{value[: max(1, length - 3)]}..." + + +def _midpoint( + first: tuple[float, float], + second: tuple[float, float], +) -> tuple[float, float]: + return ((first[0] + second[0]) / 2.0, (first[1] + second[1]) / 2.0) diff --git a/embodichain/gen_sim/action_engine/orientation.py b/embodichain/gen_sim/action_engine/orientation.py new file mode 100644 index 000000000..3cac92e27 --- /dev/null +++ b/embodichain/gen_sim/action_engine/orientation.py @@ -0,0 +1,226 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Compile task-facing orientation goals into a small runtime contract.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import math +from typing import Any + +__all__ = [ + "AlignAxisConstraint", + "MatchRotationConstraint", + "OrientationConstraint", + "compile_orientation_constraint", +] + +_LONG_AXES = frozenset({"long", "long_axis", "longest"}) +_SCOPES = frozenset({"terminal"}) + + +@dataclass(frozen=True) +class AlignAxisConstraint: + """Require one local object axis to align with a target axis.""" + + local_axis: str + target_axis: str = "world_up" + directed: bool = True + tolerance: float | None = None + scope: str = "terminal" + + +@dataclass(frozen=True) +class MatchRotationConstraint: + """Require a complete object rotation relative to a captured reference.""" + + reference: str + equivalence: str = "none" + tolerance: float | None = None + scope: str = "terminal" + + +OrientationTerm = AlignAxisConstraint | MatchRotationConstraint + + +@dataclass(frozen=True) +class OrientationConstraint: + """Canonical hard constraints plus a separate planning preference.""" + + terms: tuple[OrientationTerm, ...] + planning_preference: str = "minimize_rotation_from_current" + + @property + def requires_reference(self) -> bool: + """Return whether execution must capture a step-start rotation.""" + return any( + isinstance(term, MatchRotationConstraint) and term.reference == "step_start" + for term in self.terms + ) + + @property + def allows_upright_yaw_search(self) -> bool: + """Return whether all hard terms leave world-up yaw unconstrained.""" + return bool(self.terms) and all( + isinstance(term, AlignAxisConstraint) and term.target_axis == "world_up" + for term in self.terms + ) + + +def compile_orientation_constraint( + goal: Mapping[str, Any], +) -> OrientationConstraint: + """Compile legacy goal enums or a composable serialized constraint. + + Existing persisted graphs continue to carry explicit ``orientation_goal`` + values. New tasks may omit the field, which intentionally means no hard + orientation constraint while retaining a minimum-rotation preference. + """ + serialized = goal.get("orientation_constraint") + if serialized is not None: + return _compile_serialized(serialized) + + orientation_goal = str(goal.get("orientation_goal", "none")) + if orientation_goal == "none": + terms: tuple[OrientationTerm, ...] = () + elif orientation_goal == "preserve": + terms = (MatchRotationConstraint(reference="step_start"),) + elif orientation_goal == "upright": + local_axis = str(goal.get("upright_local_axis", "long_axis")) + if local_axis == "auto": + local_axis = "long_axis" + directed = goal.get( + "orientation_directed", local_axis.lower() not in _LONG_AXES + ) + if not isinstance(directed, bool): + raise ValueError("orientation_directed must be a boolean.") + terms = ( + AlignAxisConstraint( + local_axis=local_axis, + target_axis="world_up", + directed=directed, + ), + ) + elif orientation_goal in {"lay_flat", "axis_align"}: + # These established modes materialize a full target rotation. Keep that + # contract until semantic face/axis metadata can express narrower terms. + terms = (MatchRotationConstraint(reference="target_pose"),) + else: + raise ValueError(f"Unsupported orientation_goal {orientation_goal!r}.") + return OrientationConstraint(terms=terms) + + +def _compile_serialized(value: Any) -> OrientationConstraint: + if not isinstance(value, Mapping): + raise ValueError("orientation_constraint must be a mapping.") + unknown = set(value) - {"terms", "planning_preference"} + if unknown: + raise ValueError( + "orientation_constraint contains unsupported fields: " f"{sorted(unknown)}." + ) + raw_terms = value.get("terms", ()) + if not isinstance(raw_terms, Sequence) or isinstance( + raw_terms, (str, bytes, bytearray) + ): + raise ValueError("orientation_constraint.terms must be a list.") + terms = tuple(_compile_term(item, index) for index, item in enumerate(raw_terms)) + preference = str(value.get("planning_preference", "minimize_rotation_from_current")) + if preference not in {"minimize_rotation_from_current", "none"}: + raise ValueError( + "orientation_constraint.planning_preference must be " + "'minimize_rotation_from_current' or 'none'." + ) + return OrientationConstraint(terms=terms, planning_preference=preference) + + +def _compile_term(value: Any, index: int) -> OrientationTerm: + context = f"orientation_constraint.terms[{index}]" + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + kind = str(value.get("type", "")) + scope = str(value.get("scope", "terminal")) + if scope not in _SCOPES: + raise ValueError( + f"{context}.scope {scope!r} is unsupported by the current runtime." + ) + if kind == "align_axis": + unknown = set(value) - { + "type", + "local_axis", + "target_axis", + "directed", + "tolerance", + "scope", + } + if unknown: + raise ValueError( + f"{context} contains unsupported fields: {sorted(unknown)}." + ) + local_axis = str(value.get("local_axis", "")) + if local_axis not in {"x", "y", "z", "long_axis"}: + raise ValueError(f"{context}.local_axis {local_axis!r} is unsupported.") + target_axis = str(value.get("target_axis", "world_up")) + if target_axis != "world_up": + raise ValueError(f"{context}.target_axis {target_axis!r} is unsupported.") + directed = value.get("directed", True) + if not isinstance(directed, bool): + raise ValueError(f"{context}.directed must be a boolean.") + return AlignAxisConstraint( + local_axis=local_axis, + target_axis=target_axis, + directed=directed, + tolerance=_optional_tolerance(value, context), + scope=scope, + ) + if kind == "match_rotation": + unknown = set(value) - { + "type", + "reference", + "equivalence", + "tolerance", + "scope", + } + if unknown: + raise ValueError( + f"{context} contains unsupported fields: {sorted(unknown)}." + ) + reference = str(value.get("reference", "")) + if reference not in {"step_start", "target_pose"}: + raise ValueError(f"{context}.reference {reference!r} is unsupported.") + equivalence = str(value.get("equivalence", "none")) + if equivalence != "none": + raise ValueError(f"{context}.equivalence {equivalence!r} is unsupported.") + return MatchRotationConstraint( + reference=reference, + equivalence=equivalence, + tolerance=_optional_tolerance(value, context), + scope=scope, + ) + raise ValueError(f"{context}.type {kind!r} is unsupported.") + + +def _optional_tolerance(value: Mapping[str, Any], context: str) -> float | None: + raw = value.get("tolerance") + if raw is None: + return None + if isinstance(raw, bool) or not isinstance(raw, (int, float)): + raise ValueError(f"{context}.tolerance must be a finite positive number.") + tolerance = float(raw) + if not math.isfinite(tolerance) or tolerance <= 0.0: + raise ValueError(f"{context}.tolerance must be a finite positive number.") + return tolerance diff --git a/embodichain/gen_sim/action_engine/planning/__init__.py b/embodichain/gen_sim/action_engine/planning/__init__.py new file mode 100644 index 000000000..a02f8ce3b --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/__init__.py @@ -0,0 +1,62 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stable route-free task planning API.""" + +from __future__ import annotations + +from .online import plan_online_seed_graph +from .dual import CandidatePair, plan_candidates_parallel +from .linker import ( + CONTRACT_LINKER_VERSION, + link_seed_graph, + link_task_dependencies, + validate_persisted_contracts, +) +from .planner import plan_task +from .selection import ( + CandidateEvaluation, + evaluate_candidate, + fuse_seed_graphs, + select_seed_graph, +) +from .vision import ( + CameraObservation, + SceneObservation, + analyze_visual_scene, + collect_scene_observation, + validate_visual_facts, +) + +__all__ = [ + "CONTRACT_LINKER_VERSION", + "CameraObservation", + "CandidatePair", + "CandidateEvaluation", + "SceneObservation", + "analyze_visual_scene", + "collect_scene_observation", + "evaluate_candidate", + "fuse_seed_graphs", + "link_seed_graph", + "link_task_dependencies", + "plan_online_seed_graph", + "plan_candidates_parallel", + "plan_task", + "select_seed_graph", + "validate_visual_facts", + "validate_persisted_contracts", +] diff --git a/embodichain/gen_sim/action_engine/planning/dual.py b/embodichain/gen_sim/action_engine/planning/dual.py new file mode 100644 index 000000000..d852464fd --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/dual.py @@ -0,0 +1,218 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Parallel offline/online candidate planning with isolated task views.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from dataclasses import dataclass +from time import perf_counter +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + public_task_spec, + seed_graph_hash, + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +from .linker import validate_persisted_contracts + +__all__ = ["CandidatePair", "plan_candidates_parallel"] + +CandidatePlanner = Callable[..., Mapping[str, Any]] + + +@dataclass(frozen=True) +class CandidatePair: + """Two independently planned graphs and branch-local planning metrics.""" + + offline: dict[str, Any] + online: dict[str, Any] + planning_metrics: dict[str, dict[str, Any]] + + +def plan_candidates_parallel( + task_spec: Mapping[str, Any], + *, + offline_planner: CandidatePlanner, + online_planner: CandidatePlanner, + known_objects: set[str] | None = None, + robot_profile: str = "dual_ur10", + registry: AtomicCapabilityRegistry | None = None, + require_executable: bool = False, +) -> CandidatePair: + """Plan both routes concurrently while hiding the oracle from online. + + Both returned graphs are validated against the same capability catalog and + motion-policy table before the pair is published. ``require_executable`` + is intentionally opt-in here: product planning may retain planning-only + candidates for inspection, while strict A/B execution enables the flag in + its final preflight. + """ + task = validate_task_spec(task_spec) + online_view = public_task_spec(task) + _reject_private_or_live_fields(online_view, "PublicTaskSpec") + capabilities = registry or build_atomic_capability_registry() + + def invoke(route: str) -> tuple[dict[str, Any], float]: + planner = offline_planner if route == "offline" else online_planner + # A planner is user/LLM supplied code. Give each route a detached + # copy so accidental mutation cannot change the other route's input or + # reintroduce private oracle fields after validation. + planner_input = deepcopy(task if route == "offline" else online_view) + started = perf_counter() + try: + result = planner(task_spec=planner_input) + except Exception as exc: + raise RuntimeError(f"{route} planner failed: {exc}") from exc + elapsed = perf_counter() - started + _reject_private_or_live_fields(result, f"{route} SeedGraph") + graph = validate_seed_graph( + result, + known_objects=known_objects, + known_actions=capabilities.names(), + executable_actions=capabilities.executable_names(), + require_executable=require_executable, + ) + if graph["planner_route"] != route: + raise ValueError( + f"{route} planner returned route {graph['planner_route']!r}." + ) + if graph["task_id"] != task["task_id"]: + raise ValueError(f"{route} planner returned a graph for another task.") + if graph["level"] != task["level"]: + raise ValueError(f"{route} planner returned a graph for another level.") + if graph["reasoning_type"] != task["reasoning_type"]: + raise ValueError( + f"{route} planner returned a graph with incompatible reasoning_type." + ) + if graph["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError( + f"{route} SeedGraph capability catalog does not match runtime." + ) + validate_persisted_contracts(graph, capabilities) + _validate_task_group_coverage(task, graph, route=route) + for node in graph["nodes"]: + capabilities.validate_binding(node) + if capabilities.get(str(node["atomic_action"])).runtime_available: + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + return graph, elapsed + + with ThreadPoolExecutor( + max_workers=2, thread_name_prefix="action-engine-plan" + ) as pool: + futures = {route: pool.submit(invoke, route) for route in ("offline", "online")} + results: dict[str, tuple[dict[str, Any], float]] = {} + for route, future in futures.items(): + try: + results[route] = future.result() + except Exception as exc: + # Do not expose a bare Future exception; callers need to know + # which route invalidated the pair before any environment is + # allowed to move. + for other_route, other in futures.items(): + if other_route != route: + other.cancel() + raise RuntimeError( + f"A/B {route} planning/preflight failed: {exc}" + ) from exc + + metrics = { + route: { + "planning_seconds": elapsed, + "vlm_call_count": int(graph.get("metadata", {}).get("vlm_call_count", 0)), + "seed_graph_hash": seed_graph_hash(graph), + "node_count": len(graph["nodes"]), + "task_group_count": len(graph["task_groups"]), + } + for route, (graph, elapsed) in results.items() + } + return CandidatePair( + offline=results["offline"][0], + online=results["online"][0], + planning_metrics=metrics, + ) + + +def _validate_task_group_coverage( + task: Mapping[str, Any], graph: Mapping[str, Any], *, route: str +) -> None: + """Ensure every explicit L1-L3 task instance has one complete group.""" + if task.get("level") == "L4": + # L4's reference instances are intentionally hidden from the online + # route; the graph validator still enforces non-empty, coherent groups. + return + expected = { + str(item["id"]) + for item in task.get("task_instances", ()) + if isinstance(item, Mapping) + } + actual = {str(group["id"]) for group in graph.get("task_groups", ())} + missing = expected - actual + unexpected = actual - expected + if missing or unexpected: + raise ValueError( + f"{route} SeedGraph TaskGroup coverage mismatch; " + f"missing={sorted(missing)}, unexpected={sorted(unexpected)}." + ) + + +_PRIVATE_OR_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "object_pose", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def _reject_private_or_live_fields(value: Any, context: str) -> None: + """Reject private-oracle and grounded state fields in online inputs/outputs.""" + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _PRIVATE_OR_LIVE_KEYS: + raise ValueError(f"{context} contains private/live field {key!r}.") + _reject_private_or_live_fields(child, f"{context}.{key}") + elif isinstance(value, (list, tuple)): + for index, child in enumerate(value): + _reject_private_or_live_fields(child, f"{context}[{index}]") diff --git a/embodichain/gen_sim/action_engine/planning/linker.py b/embodichain/gen_sim/action_engine/planning/linker.py new file mode 100644 index 000000000..c1db21e84 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/linker.py @@ -0,0 +1,1018 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Deterministic causal and resource linking for SeedGraph v3.""" + +from __future__ import annotations + +from collections.abc import Collection, Mapping, Sequence +from copy import deepcopy +import hashlib +import json +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + +__all__ = [ + "CONTRACT_LINKER_VERSION", + "link_seed_graph", + "link_task_dependencies", + "validate_persisted_contracts", +] + +CONTRACT_LINKER_VERSION = "action_contract_linker_v2" +_INITIAL_PREDICATES = frozenset({"arm_free", "object_free"}) +_REFERENCE_KEYS = frozenset( + { + "anchor", + "container", + "reference", + "reference_object", + "support", + "support_object", + "target", + "target_object", + } +) + + +def link_task_dependencies( + task_spec: Mapping[str, Any], + role_bindings: Mapping[str, str], + *, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Add the minimal stable TaskGroup dependencies implied by contracts.""" + del registry # Reserved for task-level capability specialization. + task = validate_task_spec(task_spec) + bindings = {str(key): str(value) for key, value in role_bindings.items()} + bindings_hash = hashlib.sha256( + json.dumps( + bindings, sort_keys=True, separators=(",", ":"), allow_nan=False + ).encode("utf-8") + ).hexdigest() + existing_metadata = task.get("metadata", {}) + existing_linker = ( + existing_metadata.get("action_contract_task_linker", {}) + if isinstance(existing_metadata, Mapping) + else {} + ) + if ( + isinstance(existing_linker, Mapping) + and existing_linker.get("version") == CONTRACT_LINKER_VERSION + and existing_linker.get("role_bindings_hash") == bindings_hash + ): + return task + instances = task["task_instances"] + order = [str(item["id"]) for item in instances] + dependencies = { + str(item["id"]): set(str(value) for value in item["depends_on"]) + for item in instances + } + dependency_order = { + str(item["id"]): [str(value) for value in item["depends_on"]] + for item in instances + } + claims = {str(item["id"]): _task_claims(item, bindings) for item in instances} + distinct_arm_pairs = _distinct_arm_pairs(task.get("metadata", {})) + linked: list[dict[str, str]] = [] + + latest_by_object: dict[str, tuple[str, str]] = {} + for instance in instances: + instance_id = str(instance["id"]) + task_type = str(instance["task_type"]) + primary = _task_primary_object(instance, bindings) + previous = latest_by_object.get(primary) + if ( + task_type == "E4" + and previous is not None + and previous[1] == "E2" + and previous[0] not in dependencies[instance_id] + and not _reaches(dependencies, previous[0], instance_id) + ): + dependencies[instance_id].add(previous[0]) + dependency_order[instance_id].append(previous[0]) + linked.append( + { + "from": previous[0], + "to": instance_id, + "reason": "causal", + "detail": f"object_held:{primary}", + } + ) + _assert_acyclic(dependencies, "TaskSpec causal linking") + latest_by_object[primary] = (instance_id, task_type) + + for later_index, later_id in enumerate(order): + for earlier_id in order[:later_index]: + if _reaches(dependencies, later_id, earlier_id) or _reaches( + dependencies, earlier_id, later_id + ): + continue + conflicts = _claim_conflicts(claims[earlier_id], claims[later_id]) + if frozenset({earlier_id, later_id}) in distinct_arm_pairs: + conflicts = [item for item in conflicts if item != "arm:auto"] + if not conflicts: + continue + dependencies[later_id].add(earlier_id) + dependency_order[later_id].append(earlier_id) + linked.append( + { + "from": earlier_id, + "to": later_id, + "reason": "resource", + "detail": ",".join(conflicts), + } + ) + _assert_acyclic(dependencies, "TaskSpec contract linking") + + for instance in instances: + instance_id = str(instance["id"]) + instance["depends_on"] = dependency_order[instance_id] + metadata = dict(task.get("metadata", {})) + metadata["action_contract_task_linker"] = { + "version": CONTRACT_LINKER_VERSION, + "role_bindings_hash": bindings_hash, + "linked_dependencies": linked, + } + task["metadata"] = metadata + return validate_task_spec(task) + + +def link_seed_graph( + draft: Mapping[str, Any], + *, + registry: AtomicCapabilityRegistry | None = None, + task_order: Sequence[str] = (), + completed_nodes: Collection[str] = (), + known_objects: Collection[str] | None = None, +) -> dict[str, Any]: + """Resolve contracts, link a draft graph, and return validated SeedGraph v3.""" + if not isinstance(draft, Mapping): + raise TypeError("SeedGraph draft must be a mapping.") + if draft.get("schema_version") != SEED_GRAPH_SCHEMA: + raise ValueError(f"Contract linker accepts only {SEED_GRAPH_SCHEMA!r} drafts.") + capabilities = registry or build_atomic_capability_registry() + graph = deepcopy(dict(draft)) + nodes = graph.get("nodes") + groups = graph.get("task_groups") + if not isinstance(nodes, list) or not nodes: + raise ValueError("SeedGraph draft nodes must be a non-empty list.") + if not isinstance(groups, list) or not groups: + raise ValueError("SeedGraph draft task_groups must be a non-empty list.") + + already_linked = _already_linked(graph) + for index, node in enumerate(nodes): + if not isinstance(node, dict): + raise TypeError(f"SeedGraph draft node {index} must be a mapping.") + action = str(node.get("atomic_action", "")) + expected = capabilities.get(action).resolve_contract(node).as_mapping() + persisted = node.get("contract") + if persisted is not None and persisted != expected: + raise ValueError( + f"SeedGraph node {node.get('id')!r} persisted Action Contract " + "does not match the current capability resolver." + ) + node["contract"] = expected + node.pop("resources", None) + if already_linked: + linked = validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + ) + validate_persisted_contracts(linked, capabilities) + return linked + + completed = {str(item) for item in completed_nodes} + node_by_id = _unique_by_id(nodes, "SeedGraph draft nodes") + group_by_id = _unique_by_id(groups, "SeedGraph draft task_groups") + ordered_groups = _ordered_group_ids(groups, task_order) + node_reasons: list[dict[str, str]] = [] + group_reasons: list[dict[str, str]] = [] + + for group in groups: + group_id = str(group.get("id", "")) + node_ids = [str(item) for item in group.get("node_ids", ())] + if not node_ids or any(node_id not in node_by_id for node_id in node_ids): + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} has missing or unknown node IDs." + ) + _link_internal_nodes( + node_ids, + node_by_id, + completed=completed, + reasons=node_reasons, + ) + _validate_internal_symbolic_state(node_ids, node_by_id) + group.pop("contract", None) + + group_dependencies = { + group_id: set(str(item) for item in group_by_id[group_id].get("depends_on", ())) + for group_id in ordered_groups + } + original_group_dependencies = { + group_id: [str(item) for item in group_by_id[group_id].get("depends_on", ())] + for group_id in ordered_groups + } + _assert_acyclic(group_dependencies, "SeedGraph TaskGroups") + summaries = { + group_id: _summarize_group(group_by_id[group_id], node_by_id) + for group_id in ordered_groups + } + distinct_arm_pairs = _distinct_arm_pairs(graph.get("metadata", {})) + + for later_index, later_id in enumerate(ordered_groups): + for earlier_id in ordered_groups[:later_index]: + if _reaches(group_dependencies, later_id, earlier_id) or _reaches( + group_dependencies, earlier_id, later_id + ): + continue + conflicts = _claim_conflicts( + summaries[earlier_id]["claims"], summaries[later_id]["claims"] + ) + if frozenset({earlier_id, later_id}) in distinct_arm_pairs: + conflicts = [item for item in conflicts if item != "arm:auto"] + if conflicts: + _add_group_dependency( + earlier_id, + later_id, + group_dependencies, + completed, + summaries, + group_reasons, + reason="resource", + detail=",".join(conflicts), + ) + + for later_index, group_id in enumerate(ordered_groups): + for requirement in summaries[group_id]["entry_requires"]: + if requirement["predicate"] in _INITIAL_PREDICATES: + continue + candidates = [ + candidate + for candidate in ordered_groups[:later_index] + if _adds_atom(summaries[candidate]["exit_effects"], requirement) + ] + if not candidates: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} has no producer for state " + f"{requirement}." + ) + maximal = [ + candidate + for candidate in candidates + if not any( + candidate != other + and _reaches(group_dependencies, other, candidate) + for other in candidates + ) + ] + if len(maximal) != 1: + raise ValueError( + f"SeedGraph TaskGroup {group_id!r} has multiple unordered " + f"producers for state {requirement}: {maximal}." + ) + producer = maximal[0] + if not _reaches(group_dependencies, group_id, producer): + _add_group_dependency( + producer, + group_id, + group_dependencies, + completed, + summaries, + group_reasons, + reason="causal", + detail=_atom_key(requirement), + ) + + _assert_acyclic(group_dependencies, "SeedGraph contract linking") + for group_id in ordered_groups: + group = group_by_id[group_id] + group["depends_on"] = original_group_dependencies[group_id] + [ + candidate + for candidate in ordered_groups + if candidate in group_dependencies[group_id] + and candidate not in original_group_dependencies[group_id] + ] + + _link_group_boundaries( + ordered_groups, + group_dependencies, + summaries, + node_by_id, + completed, + node_reasons, + ) + for group_id in ordered_groups: + summaries[group_id] = _summarize_group(group_by_id[group_id], node_by_id) + group_by_id[group_id]["contract"] = summaries[group_id] + + _validate_symbolic_state(ordered_groups, group_dependencies, summaries, group_by_id) + metadata = dict(graph.get("metadata", {})) + metadata["action_contract_linker"] = { + "version": CONTRACT_LINKER_VERSION, + "group_dependencies": _sorted_reasons(group_reasons), + "node_dependencies": _sorted_reasons(node_reasons), + } + graph["metadata"] = metadata + graph["schema_version"] = SEED_GRAPH_SCHEMA + graph["nodes"] = nodes + graph["task_groups"] = groups + return validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + ) + + +def validate_persisted_contracts( + graph: Mapping[str, Any], registry: AtomicCapabilityRegistry +) -> None: + """Reject persisted contracts that differ from the active capability catalog.""" + metadata = graph.get("metadata", {}) + linker = ( + metadata.get("action_contract_linker", {}) + if isinstance(metadata, Mapping) + else {} + ) + if ( + not isinstance(linker, Mapping) + or linker.get("version") != CONTRACT_LINKER_VERSION + ): + raise ValueError( + "SeedGraph was not produced by the current deterministic Contract Linker; " + "regenerate the configuration bundle." + ) + for node in graph.get("nodes", ()): + expected = ( + registry.get(str(node["atomic_action"])).resolve_contract(node).as_mapping() + ) + if node.get("contract") != expected: + raise ValueError( + f"SeedGraph node {node.get('id')!r} persisted Action Contract " + "does not match the current capability resolver." + ) + node_by_id = { + str(node["id"]): node + for node in graph.get("nodes", ()) + if isinstance(node, Mapping) and "id" in node + } + for group in graph.get("task_groups", ()): + expected = _summarize_group(group, node_by_id) + if group.get("contract") != expected: + raise ValueError( + f"SeedGraph TaskGroup {group.get('id')!r} persisted contract " + "does not match its linked AtomicAction topology." + ) + + +def _task_claims( + instance: Mapping[str, Any], bindings: Mapping[str, str] +) -> list[dict[str, str]]: + task_type = str(instance["task_type"]) + params = _resolve_roles(instance.get("params", {}), bindings) + primary_key = "source_role" if task_type == "E3" else "object_role" + primary = params.get(primary_key) + claims: list[dict[str, str]] = [] + if isinstance(primary, str) and primary: + claims.append(_claim(f"object:{primary}", "exclusive")) + target = params.get("target_role") + if isinstance(target, str) and target and target != primary: + claims.append(_claim(f"object:{target}", "shared_read")) + payloads = params.get("payload_roles", []) + if isinstance(payloads, Sequence) and not isinstance( + payloads, (str, bytes, bytearray) + ): + for payload in payloads: + if isinstance(payload, str) and payload and payload != primary: + claims.append(_claim(f"object:{payload}", "exclusive")) + if task_type == "E4": + transfer = str(params.get("transfer_arm", "")) + receive = str(params.get("receive_arm", "")) + if transfer not in {"left_arm", "right_arm"} or receive not in { + "left_arm", + "right_arm", + }: + raise ValueError( + "E4 contract linking requires explicit transfer/receive arms." + ) + if transfer == receive: + raise ValueError("E4 transfer_arm and receive_arm must be distinct.") + claims.extend((_claim(f"arm:{transfer}"), _claim(f"arm:{receive}"))) + elif task_type == "E5": + claims.extend((_claim("arm:left_arm"), _claim("arm:right_arm"))) + else: + required_arm = params.get("required_arm") + if required_arm in {"left_arm", "right_arm"}: + claims.append(_claim(f"arm:{required_arm}")) + else: + claims.append(_claim("arm:auto")) + return _merge_claims(claims) + + +def _task_primary_object( + instance: Mapping[str, Any], bindings: Mapping[str, str] +) -> str: + task_type = str(instance["task_type"]) + params = _resolve_roles(instance.get("params", {}), bindings) + key = "source_role" if task_type == "E3" else "object_role" + value = params.get(key) + if not isinstance(value, str) or not value: + raise ValueError( + f"TaskGroup {instance.get('id')!r} requires a resolved {key!r}." + ) + return value + + +def _link_internal_nodes( + node_ids: Sequence[str], + node_by_id: Mapping[str, dict[str, Any]], + *, + completed: set[str], + reasons: list[dict[str, str]], +) -> None: + positions = {node_id: index for index, node_id in enumerate(node_ids)} + for later_index, later_id in enumerate(node_ids): + later = node_by_id[later_id] + for requirement in later["contract"]["requires"]: + producers = [ + earlier_id + for earlier_id in node_ids[:later_index] + if node_by_id[earlier_id]["contract"]["failure_policy"] != "best_effort" + if _adds_atom( + node_by_id[earlier_id]["contract"]["effects"], requirement + ) + ] + if producers: + _add_node_dependency( + producers[-1], + later_id, + node_by_id, + completed, + reasons, + "causal", + _atom_key(requirement), + ) + for earlier_id in node_ids[:later_index]: + earlier = node_by_id[earlier_id] + if earlier.get("sync_group") is not None and earlier.get( + "sync_group" + ) == later.get("sync_group"): + continue + if _node_reaches(node_by_id, later_id, earlier_id) or _node_reaches( + node_by_id, earlier_id, later_id + ): + continue + conflicts = _claim_conflicts( + earlier["contract"]["claims"], later["contract"]["claims"] + ) + if conflicts: + _add_node_dependency( + earlier_id, + later_id, + node_by_id, + completed, + reasons, + "resource", + ",".join(conflicts), + ) + dependencies = { + node_id: { + str(parent) + for parent in node_by_id[node_id].get("depends_on", ()) + if str(parent) in positions + } + for node_id in node_ids + } + _assert_acyclic(dependencies, "AtomicAction contract linking") + + +def _summarize_group( + group: Mapping[str, Any], node_by_id: Mapping[str, Mapping[str, Any]] +) -> dict[str, Any]: + node_ids = [str(item) for item in group["node_ids"]] + node_set = set(node_ids) + entries = [ + node_id + for node_id in node_ids + if not any( + str(parent) in node_set for parent in node_by_id[node_id]["depends_on"] + ) + ] + depended = { + str(parent) + for node_id in node_ids + for parent in node_by_id[node_id]["depends_on"] + if str(parent) in node_set + } + terminals = [node_id for node_id in node_ids if node_id not in depended] + entry_requires: list[dict[str, str]] = [] + for node_id in node_ids: + node = node_by_id[node_id] + for requirement in node["contract"]["requires"]: + if any( + producer_id in node_set + and _node_reaches(node_by_id, node_id, producer_id) + and _adds_atom( + node_by_id[producer_id]["contract"]["effects"], requirement + ) + for producer_id in node_ids + ): + continue + if requirement not in entry_requires: + entry_requires.append(deepcopy(requirement)) + + last_effect: dict[str, dict[str, Any]] = {} + effect_order: list[str] = [] + for node_id in node_ids: + if node_by_id[node_id]["contract"]["failure_policy"] == "best_effort": + continue + for effect in node_by_id[node_id]["contract"]["effects"]: + key = _atom_key(effect["atom"]) + if key not in last_effect: + effect_order.append(key) + last_effect[key] = deepcopy(effect) + claims = [ + deepcopy(claim) + for node_id in node_ids + for claim in node_by_id[node_id]["contract"]["claims"] + ] + claims.extend(_goal_read_claims(group.get("goal", {}))) + merged_claims = _merge_claims(claims) + free_resources = { + ( + f"arm:{effect['atom']['arm']}" + if effect["atom"]["predicate"] == "arm_free" + else f"object:{effect['atom']['object_uid']}" + ) + for effect in last_effect.values() + if effect["op"] == "add" + and effect["atom"]["predicate"] in {"arm_free", "object_free"} + } + for claim in merged_claims: + if claim["resource"] in free_resources: + claim["lifetime"] = "action" + completion = ( + "terminal_barrier" + if terminals + and all( + node_by_id[node_id]["contract"]["completion"] == "terminal_barrier" + for node_id in terminals + ) + else "ordinary" + ) + return { + "entry_requires": entry_requires, + "exit_effects": [last_effect[key] for key in effect_order], + "claims": merged_claims, + "entry_node_ids": entries, + "terminal_node_ids": terminals, + "completion": completion, + } + + +def _validate_internal_symbolic_state( + node_ids: Sequence[str], node_by_id: Mapping[str, Mapping[str, Any]] +) -> None: + node_set = set(node_ids) + dependencies = { + node_id: { + str(parent) + for parent in node_by_id[node_id].get("depends_on", ()) + if str(parent) in node_set + } + for node_id in node_ids + } + entry_atoms = set() + for node_id in node_ids: + for requirement in node_by_id[node_id]["contract"]["requires"]: + has_prior_producer = any( + producer_id != node_id + and _node_reaches(node_by_id, node_id, producer_id) + and _adds_atom( + node_by_id[producer_id]["contract"]["effects"], requirement + ) + for producer_id in node_ids + ) + if not has_prior_producer: + entry_atoms.add(_atom_key(requirement)) + state = set(entry_atoms) + for node_id in _stable_topological(node_ids, dependencies): + contract = node_by_id[node_id]["contract"] + for requirement in contract["requires"]: + if _atom_key(requirement) not in state: + raise ValueError( + f"SeedGraph node {node_id!r} requires unavailable state " + f"{requirement}." + ) + for effect in contract["effects"]: + key = _atom_key(effect["atom"]) + if effect["op"] == "delete": + if key not in state: + raise ValueError( + f"SeedGraph node {node_id!r} deletes unavailable state " + f"{effect['atom']}." + ) + state.remove(key) + else: + state.add(key) + + +def _add_group_dependency( + parent: str, + child: str, + dependencies: dict[str, set[str]], + completed: set[str], + summaries: Mapping[str, Mapping[str, Any]], + reasons: list[dict[str, str]], + *, + reason: str, + detail: str, +) -> None: + if any(node_id in completed for node_id in summaries[child]["entry_node_ids"]): + raise ValueError( + f"Contract linking cannot add dependency into completed TaskGroup {child!r}." + ) + dependencies[child].add(parent) + _assert_acyclic(dependencies, "SeedGraph contract linking") + reasons.append({"from": parent, "to": child, "reason": reason, "detail": detail}) + + +def _link_group_boundaries( + ordered_groups: Sequence[str], + dependencies: Mapping[str, set[str]], + summaries: Mapping[str, Mapping[str, Any]], + node_by_id: Mapping[str, dict[str, Any]], + completed: set[str], + reasons: list[dict[str, str]], +) -> None: + for child in ordered_groups: + for parent in ordered_groups: + if parent not in dependencies[child]: + continue + for child_node in summaries[child]["entry_node_ids"]: + for parent_node in summaries[parent]["terminal_node_ids"]: + _add_node_dependency( + parent_node, + child_node, + node_by_id, + completed, + reasons, + "cleanup", + f"TaskGroup {parent} terminal barrier", + ) + + +def _add_node_dependency( + parent: str, + child: str, + node_by_id: Mapping[str, dict[str, Any]], + completed: set[str], + reasons: list[dict[str, str]], + reason: str, + detail: str, +) -> None: + if parent == child: + raise ValueError(f"Contract linker cannot add self-dependency {child!r}.") + dependencies = node_by_id[child].setdefault("depends_on", []) + if parent in dependencies or _node_reaches(node_by_id, child, parent): + return + if _node_reaches(node_by_id, parent, child): + raise ValueError( + f"Contract dependency {parent!r} -> {child!r} would create a cycle." + ) + if child in completed: + raise ValueError(f"Contract linker cannot modify completed node {child!r}.") + dependencies.append(parent) + reasons.append({"from": parent, "to": child, "reason": reason, "detail": detail}) + + +def _validate_symbolic_state( + ordered_groups: Sequence[str], + dependencies: Mapping[str, set[str]], + summaries: Mapping[str, Mapping[str, Any]], + groups: Mapping[str, Mapping[str, Any]], +) -> None: + atoms = [ + atom + for summary in summaries.values() + for atom in [ + *summary["entry_requires"], + *(effect["atom"] for effect in summary["exit_effects"]), + ] + ] + state = { + _atom_key({"predicate": "arm_free", "arm": str(atom["arm"])}) + for atom in atoms + if "arm" in atom + } + state.update( + _atom_key({"predicate": "object_free", "object_uid": str(atom["object_uid"])}) + for atom in atoms + if "object_uid" in atom + ) + for group_id in _stable_topological(ordered_groups, dependencies): + if groups[group_id].get("role") == "recovery": + for requirement in summaries[group_id]["entry_requires"]: + if requirement["predicate"] == "object_free": + object_uid = str(requirement["object_uid"]) + state = { + item + for item in state + if not ( + item.startswith("object_held|") + or item.startswith("object_coordinated_held|") + ) + or f"|{object_uid}|" not in f"|{item}|" + } + state.add(_atom_key(requirement)) + for requirement in summaries[group_id]["entry_requires"]: + key = _atom_key(requirement) + if key not in state: + raise ValueError( + _unavailable_group_state_message( + group_id, + requirement, + state, + groups[group_id], + ) + ) + for effect in summaries[group_id]["exit_effects"]: + key = _atom_key(effect["atom"]) + if effect["op"] == "add": + state.add(key) + else: + state.discard(key) + + +def _unavailable_group_state_message( + group_id: str, + requirement: Mapping[str, Any], + state: Collection[str], + group: Mapping[str, Any], +) -> str: + """Explain held-object conflicts without weakening symbolic validation.""" + if requirement.get("predicate") != "arm_free": + return ( + f"SeedGraph TaskGroup {group_id!r} requires unavailable state " + f"{dict(requirement)}." + ) + arm = str(requirement.get("arm", "")) + held_objects = sorted( + parts[1] + for item in state + if len(parts := item.split("|", maxsplit=2)) == 3 + and parts[0] == "object_held" + and parts[2] == arm + and parts[1] + ) + if not held_objects: + return ( + f"SeedGraph TaskGroup {group_id!r} requires unavailable state " + f"{dict(requirement)}." + ) + primary = str(group.get("object_uid", "")) + held = ", ".join(repr(item) for item in held_objects) + return ( + f"SeedGraph TaskGroup {group_id!r} requires arm {arm!r} to be free, " + f"but it currently holds {held}; the group's primary object is " + f"{primary!r}. A post-handover continuation must preserve object " + "identity and consume object_held instead of scheduling a fresh pickup." + ) + + +def _goal_read_claims(value: Any) -> list[dict[str, str]]: + claims: list[dict[str, str]] = [] + if isinstance(value, Mapping): + for key, child in value.items(): + if key in _REFERENCE_KEYS and isinstance(child, str): + if child not in {"table_center", "world"}: + claims.append(_claim(f"object:{child}", "shared_read")) + claims.extend(_goal_read_claims(child)) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for child in value: + claims.extend(_goal_read_claims(child)) + return claims + + +def _claim( + resource: str, access: str = "exclusive", lifetime: str = "action" +) -> dict[str, str]: + return {"resource": resource, "access": access, "lifetime": lifetime} + + +def _merge_claims(claims: Sequence[Mapping[str, Any]]) -> list[dict[str, str]]: + merged: dict[str, dict[str, str]] = {} + order: list[str] = [] + for claim in claims: + resource = str(claim["resource"]) + if resource not in merged: + order.append(resource) + merged[resource] = _claim( + resource, + str(claim.get("access", "exclusive")), + str(claim.get("lifetime", "action")), + ) + continue + current = merged[resource] + if claim.get("access") == "exclusive": + current["access"] = "exclusive" + if claim.get("lifetime") == "until_release": + current["lifetime"] = "until_release" + return [merged[resource] for resource in order] + + +def _claim_conflicts( + first: Sequence[Mapping[str, Any]], second: Sequence[Mapping[str, Any]] +) -> list[str]: + first_by_resource = {str(item["resource"]): str(item["access"]) for item in first} + second_by_resource = {str(item["resource"]): str(item["access"]) for item in second} + conflicts = { + resource + for resource in set(first_by_resource) & set(second_by_resource) + if "exclusive" in {first_by_resource[resource], second_by_resource[resource]} + } + first_arms = {item for item in first_by_resource if item.startswith("arm:")} + second_arms = {item for item in second_by_resource if item.startswith("arm:")} + if "arm:auto" in first_arms and second_arms: + conflicts.add("arm:auto") + if "arm:auto" in second_arms and first_arms: + conflicts.add("arm:auto") + return sorted(conflicts) + + +def _distinct_arm_pairs(value: Any) -> set[frozenset[str]]: + if not isinstance(value, Mapping): + return set() + groups = value.get("legacy_allocation_groups", value.get("allocation_groups", ())) + if not isinstance(groups, Sequence) or isinstance(groups, (str, bytes, bytearray)): + return set() + result: set[frozenset[str]] = set() + for group in groups: + if ( + not isinstance(group, Mapping) + or group.get("arm_constraint") != "distinct_arms" + ): + continue + members = group.get("semantic_step_ids", group.get("task_instance_ids", ())) + if not isinstance(members, Sequence) or isinstance( + members, (str, bytes, bytearray) + ): + continue + member_ids = [str(item) for item in members] + for index, first in enumerate(member_ids): + for second in member_ids[index + 1 :]: + result.add(frozenset({first, second})) + return result + + +def _resolve_roles(value: Any, bindings: Mapping[str, str]) -> Any: + if isinstance(value, Mapping): + return { + str(key): _resolve_roles(child, bindings) for key, child in value.items() + } + if isinstance(value, list): + return [_resolve_roles(child, bindings) for child in value] + if isinstance(value, tuple): + return tuple(_resolve_roles(child, bindings) for child in value) + if isinstance(value, str): + return bindings.get(value, value) + return value + + +def _adds_atom(effects: Sequence[Mapping[str, Any]], atom: Mapping[str, Any]) -> bool: + return any( + effect.get("op") == "add" and effect.get("atom") == atom for effect in effects + ) + + +def _atom_key(atom: Mapping[str, Any]) -> str: + return "|".join( + str(atom.get(key, "")) for key in ("predicate", "object_uid", "arm") + ) + + +def _unique_by_id(items: Sequence[Mapping[str, Any]], context: str) -> dict[str, Any]: + result: dict[str, Any] = {} + for item in items: + item_id = str(item.get("id", "")) + if not item_id: + raise ValueError(f"{context} require non-empty IDs.") + if item_id in result: + raise ValueError(f"{context} contain duplicate ID {item_id!r}.") + result[item_id] = item + return result + + +def _ordered_group_ids( + groups: Sequence[Mapping[str, Any]], task_order: Sequence[str] +) -> list[str]: + available = [str(group["id"]) for group in groups] + requested = [str(item) for item in task_order] + unknown = set(requested) - set(available) + if unknown: + raise ValueError( + f"task_order references unknown TaskGroups: {sorted(unknown)}." + ) + return requested + [item for item in available if item not in set(requested)] + + +def _reaches(dependencies: Mapping[str, set[str]], child: str, parent: str) -> bool: + pending = list(dependencies.get(child, ())) + visited: set[str] = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(dependencies.get(current, ())) + return False + + +def _node_reaches( + node_by_id: Mapping[str, Mapping[str, Any]], child: str, parent: str +) -> bool: + pending = [str(item) for item in node_by_id[child].get("depends_on", ())] + visited: set[str] = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited and current in node_by_id: + visited.add(current) + pending.extend( + str(item) for item in node_by_id[current].get("depends_on", ()) + ) + return False + + +def _assert_acyclic(dependencies: Mapping[str, set[str]], context: str) -> None: + for item_id in dependencies: + if _reaches(dependencies, item_id, item_id): + raise ValueError(f"{context} produced a dependency cycle at {item_id!r}.") + + +def _stable_topological( + order: Sequence[str], dependencies: Mapping[str, set[str]] +) -> list[str]: + remaining = set(order) + result: list[str] = [] + while remaining: + ready = [ + item + for item in order + if item in remaining and not (dependencies[item] & remaining) + ] + if not ready: + raise ValueError("SeedGraph TaskGroups contain a dependency cycle.") + result.extend(ready) + remaining.difference_update(ready) + return result + + +def _already_linked(graph: Mapping[str, Any]) -> bool: + metadata = graph.get("metadata", {}) + linker = ( + metadata.get("action_contract_linker", {}) + if isinstance(metadata, Mapping) + else {} + ) + return ( + isinstance(linker, Mapping) + and linker.get("version") == CONTRACT_LINKER_VERSION + and all("contract" in node for node in graph.get("nodes", ())) + and all("contract" in group for group in graph.get("task_groups", ())) + ) + + +def _sorted_reasons(reasons: Sequence[Mapping[str, str]]) -> list[dict[str, str]]: + unique = { + (item["from"], item["to"], item["reason"], item["detail"]) for item in reasons + } + return [ + {"from": source, "to": target, "reason": reason, "detail": detail} + for source, target, reason, detail in sorted(unique) + ] diff --git a/embodichain/gen_sim/action_engine/planning/online.py b/embodichain/gen_sim/action_engine/planning/online.py new file mode 100644 index 000000000..fc8b36080 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/online.py @@ -0,0 +1,371 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Online planner producing a complete direct AtomicAction SeedGraph.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +import json +from time import perf_counter +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + TASK_CONTRACTS, + public_task_spec, + requested_visual_task_predicates, + validate_public_task_spec, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +from .vision import ( + SceneObservation, + _reject_live_fields as _reject_visual_live_fields, + analyze_visual_scene, + validate_visual_facts, +) +from .linker import link_seed_graph + +__all__ = ["plan_online_seed_graph"] + +GraphCaller = Callable[..., Mapping[str, Any]] + +_GRAPH_OUTPUT_SCHEMA = { + "title": "ActionEngineOnlineSeedGraphBody", + "type": "object", + "additionalProperties": False, + "required": ["nodes", "task_groups", "success"], + "properties": { + "nodes": {"type": "array", "items": {"type": "object"}}, + "task_groups": {"type": "array", "items": {"type": "object"}}, + "success": {"type": "object"}, + }, +} + + +def plan_online_seed_graph( + task_spec: Mapping[str, Any], + observation: SceneObservation, + *, + visual_facts: Mapping[str, Any] | None = None, + vlm_model: str | None = None, + fact_caller: GraphCaller | None = None, + graph_caller: GraphCaller | None = None, + registry: AtomicCapabilityRegistry | None = None, + robot_profile: str = "dual_ur10", +) -> tuple[dict[str, Any], dict[str, Any]]: + """Extract visual facts and produce one validated online SeedGraph.""" + started = perf_counter() + task = ( + validate_public_task_spec(task_spec) + if "task_instances" not in task_spec and task_spec.get("level") == "L4" + else validate_task_spec(task_spec) + ) + _reject_private_or_live_fields(public_task_spec(task), "online TaskSpec") + capabilities = registry or build_atomic_capability_registry() + _reject_visual_live_fields(observation.entities, "SceneObservation.entities") + known_uids = {str(item["uid"]) for item in observation.entities} + if len(known_uids) != len(observation.entities): + raise ValueError("Online scene observation contains duplicate entity UIDs.") + if not known_uids: + raise ValueError("Online scene observation contains no simulator entities.") + visual_call_counter = [0] + allowed_task_predicates = requested_visual_task_predicates(task) + facts = ( + validate_visual_facts( + visual_facts, + known_uids=known_uids, + camera_uids={camera.uid for camera in observation.cameras}, + allowed_task_predicates=allowed_task_predicates, + ) + if visual_facts is not None + else analyze_visual_scene( + observation, + task, + model=vlm_model, + caller=fact_caller, + call_counter=visual_call_counter, + ) + ) + _validate_fact_information(facts) + prompt = _prompt(task, facts, capabilities, robot_profile=robot_profile) + if graph_caller is None: + # Facts remain the auditable planner input, but the production VLM also + # needs the same reset-time RGB/depth evidence to bind semantic TaskSpec + # roles (for example, "the purple can") to the known simulator UIDs. + # An injected graph caller keeps the compact facts-only contract used by + # deterministic tests and alternative planners. + def caller(**kwargs: Any) -> Mapping[str, Any]: + return _default_graph_caller(observation=observation, **kwargs) + + else: + caller = graph_caller + first_error: Exception | None = None + graph_call_count = 0 + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nThe previous graph was invalid. Correct only the JSON body. " + f"Validation error: {first_error}" + ) + graph_call_count += 1 + try: + response = caller( + prompt=current_prompt, + schema=_GRAPH_OUTPUT_SCHEMA, + model=vlm_model, + ) + graph = _wrap_graph(response, task, capabilities) + _reject_private_or_live_fields(graph, "online SeedGraph") + graph = link_seed_graph( + graph, + registry=capabilities, + task_order=[str(item["id"]) for item in task.get("task_instances", ())], + known_objects=known_uids, + ) + _validate_explicit_task_group_coverage(task, graph) + for node in graph["nodes"]: + capabilities.validate_binding(node) + if capabilities.get(str(node["atomic_action"])).runtime_available: + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + graph["metadata"].update( + { + "planning_latency_seconds": perf_counter() - started, + "vlm_call_count": graph_call_count + visual_call_counter[0], + "visual_fact_call_count": visual_call_counter[0], + "graph_call_count": graph_call_count, + } + ) + return graph, facts + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "Online SeedGraph failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def _prompt( + task: Mapping[str, Any], + facts: Mapping[str, Any], + capabilities: AtomicCapabilityRegistry, + *, + robot_profile: str, +) -> str: + from embodichain.gen_sim.action_engine.config import default_runtime_policy + + runtime_policy = default_runtime_policy(robot_profile) + motion_modifiers: dict[str, list[dict[str, str]]] = { + action: [] for action in runtime_policy.motion_defaults + } + for modifier_type, modes in runtime_policy.motion_modifiers.items(): + for mode, action_patches in modes.items(): + for action in action_patches: + motion_modifiers[action].append({"type": modifier_type, "mode": mode}) + grouping_instruction = ( + "Infer the necessary E TaskGroups from the abstract goal; the private " + "reference task instances are intentionally hidden." + if task["level"] == "L4" + else "Every public TaskSpec task instance must correspond to exactly one TaskGroup." + ) + return ( + "Produce the body of one coordinate-free direct AtomicAction SeedGraph. " + f"{grouping_instruction} " + "Nodes may contain only symbolic target bindings and scene UIDs; never " + "emit world coordinates, poses, qpos, trajectories, or grasp poses. " + "Do not emit Action Contracts or resource claims; the deterministic " + "Contract Linker owns those fields. " + "Use the supplied reset-time multi-view image evidence only to bind the " + "public task semantics to known UIDs; use normalized visual constraints " + "only when the facts justify them. " + "Do not output reasoning. Planning-only actions may appear but must not " + "be replaced with invented primitives.\n\n" + f"Public TaskSpec:\n{json.dumps(public_task_spec(task), ensure_ascii=False, sort_keys=True)}\n\n" + f"Visual facts:\n{json.dumps(facts, ensure_ascii=False, sort_keys=True)}\n\n" + f"E1-E9 task semantics:\n{json.dumps(_task_capability_catalog(), ensure_ascii=False, sort_keys=True)}\n\n" + f"Atomic capabilities:\n{json.dumps(capabilities.catalog(), ensure_ascii=False, sort_keys=True)}\n\n" + "Every node motion_policy must be an object with a modifiers list; " + "the AtomicAction selects its base policy implicitly. Use only the " + "typed modifiers supported by that action.\n" + f"Allowed motion modifiers by AtomicAction:\n" + f"{json.dumps(motion_modifiers, sort_keys=True)}" + ) + + +def _task_capability_catalog() -> dict[str, dict[str, Any]]: + """Return the Action Engine's runtime-aware E-task view.""" + executable = set(build_atomic_capability_registry().executable_names()) + return { + task_type: { + "semantics": contract.semantics, + "core_actions": list(contract.core_actions), + "runtime_available": set(contract.core_actions) <= executable, + } + for task_type, contract in TASK_CONTRACTS.items() + } + + +def _wrap_graph( + response: Mapping[str, Any], + task: Mapping[str, Any], + capabilities: AtomicCapabilityRegistry, +) -> dict[str, Any]: + if not isinstance(response, Mapping): + raise TypeError("Online planner output must be a mapping.") + if set(response) != {"nodes", "task_groups", "success"}: + raise ValueError( + "Online planner must return nodes, task_groups, and success only." + ) + for index, node in enumerate(response.get("nodes", ())): + if not isinstance(node, Mapping): + raise TypeError(f"Online planner node {index} must be a mapping.") + forbidden = sorted({"contract", "resources"} & set(node)) + if forbidden: + raise ValueError( + f"Online planner node {index} may not author linker-owned fields: " + f"{forbidden}." + ) + for index, group in enumerate(response.get("task_groups", ())): + if not isinstance(group, Mapping): + raise TypeError(f"Online planner TaskGroup {index} must be a mapping.") + if "contract" in group: + raise ValueError( + f"Online planner TaskGroup {index} may not author its contract." + ) + return { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": task["task_id"], + "instruction": task["instruction"], + "level": task["level"], + "reasoning_type": task["reasoning_type"], + "planner_route": "online", + "nodes": deepcopy(response["nodes"]), + "task_groups": deepcopy(response["task_groups"]), + "success": deepcopy(response["success"]), + "capability_catalog_hash": capabilities.catalog_hash(), + "metadata": { + "oracle_exposed": False, + "visual_facts_used": True, + "allocation_groups": deepcopy( + task.get("metadata", {}).get("allocation_groups", []) + ), + }, + } + + +def _default_graph_caller( + *, + prompt: str, + schema: Mapping[str, Any], + model: str | None, + observation: SceneObservation | None = None, +) -> Mapping[str, Any]: + from .vision import _camera_evidence, _default_structured_caller, _vlm_model + + images: list[str] = [] + if observation is not None: + _, images = _camera_evidence(observation) + + return _default_structured_caller( + prompt=prompt, + images=images, + schema=schema, + model=_vlm_model(model), + ) + + +def _validate_fact_information(facts: Mapping[str, Any]) -> None: + """Reject low-information visual outputs before graph planning.""" + confidence = facts.get("confidence", 0.0) + if float(confidence) < 0.5: + raise ValueError( + "VLM visual facts confidence is below the required 0.5 threshold." + ) + entities = facts.get("entities", ()) + if not any( + bool(item.get("visible", True)) and float(item.get("confidence", 0.0)) >= 0.5 + for item in entities + if isinstance(item, Mapping) + ): + raise ValueError("VLM visual facts contain no reliable visible entity.") + + +def _validate_explicit_task_group_coverage( + task: Mapping[str, Any], graph: Mapping[str, Any] +) -> None: + """Reject an online graph that drops or invents an explicit L1-L3 step.""" + if task.get("level") == "L4": + return + expected = { + str(item["id"]) + for item in task.get("task_instances", ()) + if isinstance(item, Mapping) + } + actual = {str(group["id"]) for group in graph.get("task_groups", ())} + if expected != actual: + raise ValueError( + "Online SeedGraph TaskGroup coverage mismatch; " + f"missing={sorted(expected - actual)}, " + f"unexpected={sorted(actual - expected)}." + ) + + +_PRIVATE_OR_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "object_pose", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "waypoints", + "xpos", + } +) + + +def _reject_private_or_live_fields(value: Any, context: str) -> None: + """Reject private oracle and grounded simulator fields recursively.""" + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _PRIVATE_OR_LIVE_KEYS: + raise ValueError(f"{context} contains private/live field {key!r}.") + _reject_private_or_live_fields(child, f"{context}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _reject_private_or_live_fields(child, f"{context}[{index}]") diff --git a/embodichain/gen_sim/action_engine/planning/planner.py b/embodichain/gen_sim/action_engine/planning/planner.py new file mode 100644 index 000000000..a3e71d532 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/planner.py @@ -0,0 +1,821 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Route-free LLM planning boundary for Action Engine.""" + +from __future__ import annotations + +import json +import os +import re +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from pathlib import Path +from string import Template +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import build_default_registry +from embodichain.gen_sim.action_engine.domain import ( + TASK_AGENT_SCHEMA, + validate_task_agent, +) +from embodichain.gen_sim.action_engine.orientation import ( + compile_orientation_constraint, +) + +from .task_planner_prompt import TASK_PLANNER_PROMPT + +__all__ = ["plan_task"] + +LLMCaller = Callable[..., Mapping[str, Any]] + +_GEN_CONFIG_PATH = ( + Path(__file__).resolve().parents[2] + / "simready_pipeline" + / "configs" + / "gen_config.json" +) +_GEN_SIM_ENV_PATH = Path(__file__).resolve().parents[2] / ".env" +_UNSAFE_ID_RE = re.compile(r"[^0-9a-z]+") +_MODEL_STEP_KEYS = frozenset( + {"id", "operator", "object", "objects", "actor", "goal", "depends_on"} +) + +_MODEL_OUTPUT_SCHEMA: dict[str, Any] = { + "title": "ActionEngineSemanticPlan", + "type": "object", + "additionalProperties": False, + "required": ["semantic_steps", "allocation_groups"], + "properties": { + "semantic_steps": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": False, + "required": ["operator"], + "properties": { + "id": {"type": "string"}, + "operator": {"type": "string"}, + "object": {"type": "string"}, + "objects": { + "type": "array", + "items": {"type": "string"}, + }, + "actor": {"type": "object"}, + "goal": {"type": "object"}, + "depends_on": { + "type": "array", + "items": {"type": "string"}, + }, + }, + }, + }, + "allocation_groups": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["id", "semantic_step_ids", "arm_constraint"], + "properties": { + "id": {"type": "string"}, + "semantic_step_ids": { + "type": "array", + "items": {"type": "string"}, + }, + "arm_constraint": {"const": "distinct_arms"}, + }, + }, + }, + }, +} + + +def plan_task( + task_description: str, + scene_objects: Sequence[Mapping[str, Any]], + *, + task_name: str = "task", + model: str | None = None, + llm_caller: LLMCaller | None = None, +) -> dict[str, Any]: + """Plan a natural-language task as route-free semantic steps. + + The model is intentionally prohibited from emitting atomic actions, graph + edges, resources, target coordinates, or motion-policy parameters. + ``compile_task_agent`` owns all of those deterministic decisions. + + Args: + task_description: User goal in natural language. + scene_objects: JSON-like scene inventory. ``runtime_uid`` is preferred + over ``uid`` and ``source_uid`` for all generated references. + task_name: Stable task identifier stored in the TaskAgent. + model: Optional model-name override for the default LLM caller. + llm_caller: Optional injected callable accepting ``prompt=`` and + ``model=`` keyword arguments. It must return a mapping whose only + top-level key is ``semantic_steps``. + Returns: + A validated ``action_engine_task_agent_v1`` mapping. + """ + task_name = _nonempty(task_name, "task_name") + task_description = _nonempty(task_description, "task_description") + scene = _normalize_scene_objects(scene_objects) + + prompt = _render_prompt( + task_name=task_name, + task_description=task_description, + scene_objects=scene, + ) + caller = llm_caller or _default_llm_caller + response = caller(prompt=prompt, model=model) + try: + return _task_agent_from_response( + response, + task_name=task_name, + task_description=task_description, + scene=scene, + ) + except (TypeError, ValueError) as first_error: + # One bounded repair gives the model the verifier's exact complaint + # without turning generation into an unbounded conversation. + repair_prompt = ( + f"{prompt}\n\n" + "Your previous JSON did not satisfy the TaskAgent contract.\n" + f"Validation error: {first_error}\n" + "Return one corrected JSON object. Do not explain the correction." + ) + repaired = caller(prompt=repair_prompt, model=model) + try: + return _task_agent_from_response( + repaired, + task_name=task_name, + task_description=task_description, + scene=scene, + ) + except (TypeError, ValueError) as second_error: + raise ValueError( + "Action Engine planner failed validation after one repair: " + f"{second_error}" + ) from second_error + + +def _task_agent_from_response( + response: Any, + *, + task_name: str, + task_description: str, + scene: Sequence[Mapping[str, Any]], +) -> dict[str, Any]: + """Normalize and validate one model response as a TaskAgent.""" + if not isinstance(response, Mapping): + raise ValueError("Action Engine planner output must be a JSON object.") + allowed_fields = {"semantic_steps", "allocation_groups"} + if not set(response) <= allowed_fields or "semantic_steps" not in response: + raise ValueError( + "Action Engine planner output may contain only 'semantic_steps' " + "and 'allocation_groups'; " + f"received fields {sorted(str(key) for key in response)}." + ) + raw_steps = response["semantic_steps"] + if not isinstance(raw_steps, Sequence) or isinstance( + raw_steps, (str, bytes, bytearray) + ): + raise ValueError("Planner semantic_steps must be a list.") + visible_operators = set(build_default_registry().operator_names()) + for index, step in enumerate(raw_steps): + operator = step.get("operator") if isinstance(step, Mapping) else None + if operator not in visible_operators: + raise ValueError( + f"Planner semantic_steps[{index}].operator must be one of " + f"{sorted(visible_operators)}; got {operator!r}." + ) + return _wrap_agent( + task_name, + task_description, + raw_steps, + scene, + allocation_groups=response.get("allocation_groups", []), + ) + + +def _wrap_agent( + task_name: str, + task_description: str, + raw_steps: Sequence[Any], + scene: Sequence[Mapping[str, Any]], + *, + allocation_groups: Any, +) -> dict[str, Any]: + steps = _normalize_semantic_steps(raw_steps, scene) + groups = deepcopy(allocation_groups) + task_agent = validate_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": task_name, + "goal": task_description, + "semantic_steps": steps, + "allocation_groups": groups, + }, + known_objects=[_scene_runtime_uid(item) for item in scene], + ) + _validate_operator_contracts(task_agent) + return task_agent + + +def _validate_operator_contracts(task_agent: Mapping[str, Any]) -> None: + """Validate capability-specific step shapes inside the planner repair loop.""" + registry = build_default_registry() + for step in task_agent["semantic_steps"]: + operator = str(step["operator"]) + try: + expanded = registry.operator(operator).expand(step) + except (TypeError, ValueError) as error: + raise ValueError( + f"Semantic step {step['id']!r} violates the {operator!r} " + f"operator contract: {error}" + ) from error + if not expanded: + raise ValueError( + f"Semantic step {step['id']!r} produced no executable " + f"{operator!r} operation." + ) + + +def _normalize_semantic_steps( + raw_steps: Sequence[Any], + scene: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + if not raw_steps: + raise ValueError("Planner semantic_steps must not be empty.") + aliases = _scene_uid_aliases(scene) + normalized: list[dict[str, Any]] = [] + known_ids: set[str] = set() + previous_id: str | None = None + + for index, raw_step in enumerate(raw_steps, start=1): + if not isinstance(raw_step, Mapping): + raise ValueError(f"Planner semantic_steps[{index - 1}] must be an object.") + step = deepcopy(dict(raw_step)) + unknown = sorted(set(step) - _MODEL_STEP_KEYS) + if unknown: + raise ValueError( + f"Planner semantic_steps[{index - 1}] contains unsupported " + f"fields: {unknown}." + ) + operator = _nonempty( + step.get("operator"), + f"semantic_steps[{index - 1}].operator", + ) + configured_id = str(step.get("id", "")).strip() + step_id = configured_id or f"s{index:02d}_{_slug(operator)}" + if step_id in known_ids: + raise ValueError( + f"Planner produced duplicate semantic step ID {step_id!r}." + ) + known_ids.add(step_id) + + result: dict[str, Any] = {"id": step_id, "operator": operator} + if "object" in step: + result["object"] = _resolve_scene_uid( + step["object"], + aliases, + f"semantic step {step_id!r} object", + ) + if "objects" in step: + objects = step["objects"] + if not isinstance(objects, Sequence) or isinstance( + objects, (str, bytes, bytearray) + ): + raise ValueError(f"Semantic step {step_id!r} objects must be a list.") + result["objects"] = [ + _resolve_scene_uid( + object_uid, + aliases, + f"semantic step {step_id!r} objects", + ) + for object_uid in objects + ] + + actor = step.get("actor", {"mode": "auto"}) + if not isinstance(actor, Mapping): + raise ValueError(f"Semantic step {step_id!r} actor must be an object.") + result["actor"] = deepcopy(dict(actor)) + raw_goal = step.get("goal", {}) + if not isinstance(raw_goal, Mapping): + raise ValueError(f"Semantic step {step_id!r} goal must be an object.") + goal = deepcopy(dict(raw_goal)) + for key in ( + "anchor", + "orientation_reference_object", + "reference_object", + "support_object", + ): + if key not in goal or goal[key] in {"table_center", "self"}: + continue + goal[key] = _resolve_scene_uid( + goal[key], + aliases, + f"semantic step {step_id!r} goal.{key}", + ) + result["goal"] = goal + + if "depends_on" in step: + depends_on = step["depends_on"] + if not isinstance(depends_on, Sequence) or isinstance( + depends_on, (str, bytes, bytearray) + ): + raise ValueError( + f"Semantic step {step_id!r} depends_on must be a list." + ) + result["depends_on"] = [str(value) for value in depends_on] + else: + # Sequential is the conservative default. The LLM must explicitly + # emit an empty list when two semantic operations are independent. + result["depends_on"] = [previous_id] if previous_id is not None else [] + normalized.append(result) + previous_id = step_id + return normalized + + +def _fuse_redundant_hold_place_steps( + steps: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + """Remove a preparatory hold that a complete placement would repeat. + + ``place_relative`` already owns the pickup, transport, release, retreat, + and home phases. A model may nevertheless emit ``hold_hover(object)`` + followed by ``place_relative(object)`` as if the operators were individual + motion commands. The runtime cannot safely transfer that implicit held + state between semantic steps, so normalize the unambiguous one-consumer + pattern before TaskAgent validation. + + A hold with multiple consumers is intentionally left intact because it may + reserve one arm while unrelated branches continue. Compilation rejects any + later reuse of the held object rather than guessing an implicit handover. + """ + result = [deepcopy(dict(step)) for step in steps] + by_id = {step["id"]: step for step in result} + dependents: dict[str, list[str]] = {step_id: [] for step_id in by_id} + for step in result: + for dependency in step["depends_on"]: + if dependency in dependents: + dependents[dependency].append(step["id"]) + + removable: set[str] = set() + claimed_places: set[str] = set() + for hold in result: + if hold["operator"] != "hold_hover": + continue + consumers = dependents[hold["id"]] + if len(consumers) != 1: + continue + place = by_id[consumers[0]] + if place["operator"] != "place_relative" or place.get("object") != hold.get( + "object" + ): + continue + if place["id"] in claimed_places: + raise ValueError( + f"Semantic step {place['id']!r} cannot consume more than one " + "hold_hover state." + ) + if not _is_default_hold_goal(hold): + raise ValueError( + f"Cannot fuse {hold['id']!r} into {place['id']!r}: a " + "non-default hold_hover goal would be discarded." + ) + + place["actor"] = _merge_fused_actors( + hold["actor"], + place["actor"], + hold_id=hold["id"], + place_id=place["id"], + ) + rewritten_dependencies: list[str] = [] + for dependency in place["depends_on"]: + replacements = ( + hold["depends_on"] if dependency == hold["id"] else [dependency] + ) + for replacement in replacements: + if replacement not in rewritten_dependencies: + rewritten_dependencies.append(replacement) + place["depends_on"] = rewritten_dependencies + removable.add(hold["id"]) + claimed_places.add(place["id"]) + + return [step for step in result if step["id"] not in removable] + + +def _is_default_hold_goal(hold: Mapping[str, Any]) -> bool: + """Return whether removing a preparatory hover loses no requested state.""" + goal = hold["goal"] + if set(goal) - { + "orientation_constraint", + "orientation_axis", + "orientation_directed", + "orientation_goal", + "reference_object", + "reference_state", + }: + return False + return ( + goal.get("orientation_axis", "none") == "none" + and not compile_orientation_constraint(goal).terms + and goal.get("reference_state", "initial") == "initial" + and goal.get("reference_object", "self") in ("self", hold.get("object")) + ) + + +def _merge_fused_actors( + hold_actor: Mapping[str, Any], + place_actor: Mapping[str, Any], + *, + hold_id: str, + place_id: str, +) -> dict[str, Any]: + """Preserve an explicit arm requirement while fusing semantic steps.""" + hold = deepcopy(dict(hold_actor)) + place = deepcopy(dict(place_actor)) + hold_mode = hold.get("mode") + place_mode = place.get("mode") + hold_group = hold.get("allocation_group") + place_group = place.get("allocation_group") + if hold_group is not None and place_group is not None and hold_group != place_group: + raise ValueError( + f"Cannot fuse {hold_id!r} into {place_id!r}: conflicting " + "allocation groups would lose explicit arm-allocation intent." + ) + if hold_mode == "required" and place_mode == "required": + if hold.get("arm") != place.get("arm"): + raise ValueError( + f"Cannot fuse {hold_id!r} into {place_id!r}: conflicting " + "required arms would require an unsupported handover." + ) + merged = place + elif hold_mode == "required" and place_mode == "auto": + merged = hold + else: + merged = place + allocation_group = hold_group if hold_group is not None else place_group + if allocation_group is not None: + merged["allocation_group"] = allocation_group + return merged + + +def _render_prompt( + *, + task_name: str, + task_description: str, + scene_objects: Sequence[Mapping[str, Any]], +) -> str: + capabilities = build_default_registry() + return Template(TASK_PLANNER_PROMPT).substitute( + task_name=task_name, + task_description=task_description, + scene_objects=json.dumps( + list(scene_objects), + ensure_ascii=False, + indent=2, + sort_keys=True, + ), + operator_catalog=json.dumps( + capabilities.operator_descriptions(), + ensure_ascii=False, + indent=2, + sort_keys=True, + ), + ) + + +def _default_llm_caller(*, prompt: str, model: str | None) -> Mapping[str, Any]: + """Invoke the configured OpenAI-compatible model with structured output.""" + # Heavy client imports remain lazy so validation and deterministic + # compilation work in minimal simulation test environments. + from langchain_core.messages import HumanMessage, SystemMessage + from langchain_openai import ChatOpenAI + + settings = _load_llm_settings(model=model) + kwargs: dict[str, Any] = { + "api_key": settings["api_key"], + "model": settings["model"], + "temperature": 0, + "http_socket_options": (), + } + if settings["base_url"]: + kwargs["base_url"] = settings["base_url"] + if settings["default_query"]: + kwargs["default_query"] = settings["default_query"] + if _is_mimo_compatible(settings): + # MiMo's OpenAI-compatible endpoint supports JSON mode but not the + # OpenAI ``json_schema`` response format. Disable hidden reasoning so + # the bounded semantic response is not truncated to a few fields. + kwargs.update( + { + "max_completion_tokens": _MIMO_MAX_COMPLETION_TOKENS, + "extra_body": {"thinking": {"type": "disabled"}}, + } + ) + client = ChatOpenAI(**kwargs) + structured = _structured_output_runnable( + client, _MODEL_OUTPUT_SCHEMA, settings=settings + ) + response = structured.invoke( + [ + SystemMessage( + content=( + "Return only the requested route-free semantic plan. " + "Never emit coordinates, atomic actions, or graph edges." + ) + ), + HumanMessage(content=prompt), + ] + ) + return _coerce_model_response(response) + + +_MIMO_MAX_COMPLETION_TOKENS = 4096 + + +def _is_mimo_compatible(settings: Mapping[str, Any]) -> bool: + """Identify MiMo models or regional compatible endpoints without secrets.""" + model = str(settings.get("model", "")).casefold() + base_url = str(settings.get("base_url", "")).casefold() + return "mimo" in model or "xiaomimimo.com" in base_url + + +def _structured_output_runnable( + client: Any, + schema: Mapping[str, Any], + *, + settings: Mapping[str, Any], +) -> Any: + """Bind a portable JSON contract while retaining local strict validation. + + OpenAI-compatible providers do not share the same structured-output + dialect. MiMo documents ``json_object`` JSON mode rather than + ``json_schema``; using the latter can return HTTP 200 with sparse nested + objects. The caller still validates the decoded object against its local + schema after this transport-level binding. + """ + if not hasattr(client, "with_structured_output"): + return client + method = "json_mode" if _is_mimo_compatible(settings) else "json_schema" + try: + return client.with_structured_output(schema, method=method) + except (TypeError, ValueError): + if method == "json_mode" and hasattr(client, "bind"): + # Compatibility with older LangChain adapters that do not expose + # the ``method`` keyword but do support response_format binding. + from langchain_core.output_parsers import JsonOutputParser + + return ( + client.bind(response_format={"type": "json_object"}) + | JsonOutputParser() + ) + # Preserve the historical adapter behavior for non-MiMo providers. + return client.with_structured_output(schema) + + +def _load_llm_settings(*, model: str | None) -> dict[str, Any]: + local_env = _load_env_file(_GEN_SIM_ENV_PATH) + config: dict[str, Any] = {} + if _GEN_CONFIG_PATH.exists(): + with _GEN_CONFIG_PATH.open("r", encoding="utf-8") as stream: + raw = json.load(stream) + if isinstance(raw, Mapping): + llm = raw.get("llm", {}) + if isinstance(llm, Mapping): + configured = llm.get("openai_compatible", {}) + if isinstance(configured, Mapping): + config = dict(configured) + + # A key and endpoint identify one provider transport and must not be mixed + # across process, dotenv, and JSON configuration sources. + api_key, base_url = _resolve_transport_settings(local_env, config) + selected_model = ( + (model.strip() if isinstance(model, str) else "") + or _first_env_value( + local_env, + "ACTION_ENGINE_LLM_MODEL", + "OPENAI_MODEL", + "LLM_MODEL", + ) + or str(config.get("model", "")).strip() + ) + default_query = config.get("default_query", {}) or {} + if not api_key: + raise ValueError( + "OPENAI_API_KEY is required for Action Engine planning. Set it in " + f"the process environment or {_GEN_SIM_ENV_PATH}." + ) + if not selected_model: + raise ValueError( + "An LLM model is required through model=, OPENAI_MODEL, LLM_MODEL, " + f"or {_GEN_CONFIG_PATH}." + ) + if not isinstance(default_query, Mapping): + raise ValueError("LLM default_query must be a mapping.") + return { + "api_key": api_key, + "model": selected_model, + "base_url": base_url, + "default_query": dict(default_query), + } + + +def _resolve_transport_settings( + local_env: Mapping[str, str], + config: Mapping[str, Any], +) -> tuple[str, str]: + """Resolve an API key and endpoint from one configuration source.""" + transports = ( + ( + _mapping_value(os.environ, "OPENAI_API_KEY"), + _mapping_value( + os.environ, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(local_env, "OPENAI_API_KEY"), + _mapping_value( + local_env, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(config, "api_key"), + _mapping_value(config, "base_url"), + ), + ) + for api_key, base_url in transports: + if api_key and base_url: + return api_key, base_url.rstrip("/") + for api_key, base_url in transports: + if api_key: + return api_key, base_url.rstrip("/") + return "", "" + + +def _mapping_value(source: Mapping[str, Any], *names: str) -> str: + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def _load_env_file(path: Path) -> dict[str, str]: + """Read a local dotenv file without exporting credentials process-wide.""" + if not path.is_file(): + return {} + values: dict[str, str] = {} + for line_number, raw_line in enumerate( + path.read_text(encoding="utf-8").splitlines(), + start=1, + ): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[len("export ") :].lstrip() + if "=" not in line: + continue + key, raw_value = line.split("=", 1) + key = key.strip() + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key): + raise ValueError(f"Invalid dotenv key at {path}:{line_number}.") + value = raw_value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + value = value[1:-1] + elif " #" in value: + value = value.split(" #", 1)[0].rstrip() + values[key] = value + return values + + +def _first_env_value(local_env: Mapping[str, str], *names: str) -> str | None: + """Resolve aliases while keeping every shell value above local dotenv.""" + for source in (os.environ, local_env): + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _coerce_model_response(response: Any) -> Mapping[str, Any]: + if isinstance(response, Mapping): + return dict(response) + model_dump = getattr(response, "model_dump", None) + if callable(model_dump): + dumped = model_dump() + if isinstance(dumped, Mapping): + return dict(dumped) + content = getattr(response, "content", response) + if isinstance(content, Mapping): + return dict(content) + if isinstance(content, list): + content = "\n".join( + str(item.get("text", "")) + for item in content + if isinstance(item, Mapping) and item.get("type") == "text" + ) + if not isinstance(content, str): + raise ValueError( + f"Planner model output has unsupported type {type(content).__name__}." + ) + text = content.strip() + if text.startswith("```"): + lines = text.splitlines() + lines = lines[1:] if lines else lines + lines = lines[:-1] if lines and lines[-1].startswith("```") else lines + text = "\n".join(lines).strip() + parsed = json.loads(text) + if not isinstance(parsed, Mapping): + raise ValueError("Planner model output must decode to a JSON object.") + return dict(parsed) + + +def _normalize_scene_objects( + scene_objects: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + if not isinstance(scene_objects, Sequence) or isinstance( + scene_objects, (str, bytes, bytearray) + ): + raise ValueError("scene_objects must be a list of mappings.") + normalized: list[dict[str, Any]] = [] + runtime_uids: set[str] = set() + for index, raw_object in enumerate(scene_objects): + if not isinstance(raw_object, Mapping): + raise ValueError(f"scene_objects[{index}] must be a mapping.") + item = deepcopy(dict(raw_object)) + runtime_uid = _scene_runtime_uid(item) + if runtime_uid in runtime_uids: + raise ValueError(f"Duplicate scene runtime UID {runtime_uid!r}.") + runtime_uids.add(runtime_uid) + item["runtime_uid"] = runtime_uid + normalized.append(item) + if not normalized: + raise ValueError("scene_objects must not be empty.") + return normalized + + +def _scene_uid_aliases( + scene_objects: Sequence[Mapping[str, Any]], +) -> dict[str, str]: + aliases: dict[str, str] = {} + for item in scene_objects: + runtime_uid = _scene_runtime_uid(item) + for key in ("runtime_uid", "uid", "source_uid"): + alias = item.get(key) + if isinstance(alias, str) and alias: + existing = aliases.get(alias) + if existing is not None and existing != runtime_uid: + raise ValueError(f"Ambiguous scene object alias {alias!r}.") + aliases[alias] = runtime_uid + return aliases + + +def _resolve_scene_uid(value: Any, aliases: Mapping[str, str], context: str) -> str: + uid = _nonempty(value, context) + try: + return aliases[uid] + except KeyError as exc: + raise ValueError(f"{context} references unknown scene object {uid!r}.") from exc + + +def _scene_runtime_uid(item: Mapping[str, Any]) -> str: + for key in ("runtime_uid", "uid", "source_uid"): + value = item.get(key) + if isinstance(value, str) and value.strip(): + return value + raise ValueError("Every scene object requires runtime_uid, uid, or source_uid.") + + +def _nonempty(value: Any, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{context} must be a non-empty string.") + return value.strip() + + +def _slug(value: str) -> str: + slug = _UNSAFE_ID_RE.sub("_", value.lower()).strip("_") + return slug[:48].rstrip("_") or "step" diff --git a/embodichain/gen_sim/action_engine/planning/selection.py b/embodichain/gen_sim/action_engine/planning/selection.py new file mode 100644 index 000000000..6ef92f089 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/selection.py @@ -0,0 +1,364 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Score, select, and conservatively fuse whole TaskGroups.""" + +from __future__ import annotations + +from collections import defaultdict, deque +from collections.abc import Collection, Mapping +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + +from .linker import link_seed_graph, validate_persisted_contracts + +__all__ = [ + "CandidateEvaluation", + "evaluate_candidate", + "fuse_seed_graphs", + "select_seed_graph", +] + + +@dataclass(frozen=True) +class CandidateEvaluation: + """Auditable static candidate score before any physical execution.""" + + route: str + valid: bool + executable: bool + coverage: float + visual_confidence: float + estimated_cost: float + score: float + errors: tuple[str, ...] = () + + +def evaluate_candidate( + graph: Mapping[str, Any], + task_spec: Mapping[str, Any], + *, + known_objects: Collection[str], + visual_confidence: float = 1.0, + exact_template_match: bool = False, + registry: AtomicCapabilityRegistry | None = None, + robot_profile: str = "dual_ur10", +) -> CandidateEvaluation: + """Apply schema, capabilities, object identity, coverage, and cost scoring.""" + task = validate_task_spec(task_spec) + capabilities = registry or build_atomic_capability_registry() + errors = [] + try: + seed = validate_seed_graph( + graph, + known_objects=known_objects, + known_actions=capabilities.names(), + ) + if seed["capability_catalog_hash"] != capabilities.catalog_hash(): + raise ValueError("SeedGraph capability catalog does not match runtime.") + validate_persisted_contracts(seed, capabilities) + for node in seed["nodes"]: + capabilities.validate_binding(node) + if capabilities.get(str(node["atomic_action"])).runtime_available: + resolve_motion_policy( + robot_profile, + node["atomic_action"], + node["motion_policy"], + ) + except (TypeError, ValueError) as error: + return CandidateEvaluation( + route=str(graph.get("planner_route", "unknown")), + valid=False, + executable=False, + coverage=0.0, + visual_confidence=0.0, + estimated_cost=float("inf"), + score=float("-inf"), + errors=(str(error),), + ) + + required = {str(item["id"]) for item in task["task_instances"]} + provided = {str(group["id"]) for group in seed["task_groups"]} + if task["level"] == "L4": + coverage = 1.0 if provided and seed["success"] else 0.0 + unexpected = set() + mismatched_types = {} + else: + coverage = len(required & provided) / max(len(required), 1) + unexpected = provided - required + if unexpected: + errors.append(f"unexpected task groups: {sorted(unexpected)}") + expected_types = { + str(item["id"]): str(item["task_type"]) for item in task["task_instances"] + } + mismatched_types = { + str(group["id"]): str(group["task_type"]) + for group in seed["task_groups"] + if group["id"] in expected_types + and group["task_type"] != expected_types[group["id"]] + } + if mismatched_types: + errors.append(f"task group type mismatches: {mismatched_types}") + unavailable = sorted( + { + str(node["atomic_action"]) + for node in seed["nodes"] + if not capabilities.get(str(node["atomic_action"])).runtime_available + } + ) + executable = not unavailable + if unavailable: + errors.append(f"planning-only actions: {unavailable}") + confidence = min(max(float(visual_confidence), 0.0), 1.0) + estimated_cost = float(len(seed["nodes"])) + score = coverage * 100.0 - estimated_cost + route = str(seed["planner_route"]) + if exact_template_match and route == "offline": + score += 15.0 + if task["level"] == "L4" and route == "online": + score += 20.0 * confidence + if not executable: + score -= 30.0 + return CandidateEvaluation( + route=route, + valid=not unexpected and not mismatched_types and coverage == 1.0, + executable=executable, + coverage=coverage, + visual_confidence=confidence, + estimated_cost=estimated_cost, + score=score, + errors=tuple(errors), + ) + + +def select_seed_graph( + offline: Mapping[str, Any], + online: Mapping[str, Any], + task_spec: Mapping[str, Any], + *, + known_objects: Collection[str], + visual_confidence: float = 1.0, + exact_template_match: bool = False, + registry: AtomicCapabilityRegistry | None = None, + robot_profile: str = "dual_ur10", +) -> tuple[dict[str, Any], dict[str, CandidateEvaluation]]: + """Choose one complete candidate; ties prefer mature offline templates.""" + evaluations = { + "offline": evaluate_candidate( + offline, + task_spec, + known_objects=known_objects, + visual_confidence=1.0, + exact_template_match=exact_template_match, + registry=registry, + robot_profile=robot_profile, + ), + "online": evaluate_candidate( + online, + task_spec, + known_objects=known_objects, + visual_confidence=visual_confidence, + registry=registry, + robot_profile=robot_profile, + ), + } + valid = [item for item in evaluations.items() if item[1].valid] + if not valid: + messages = {name: evaluation.errors for name, evaluation in evaluations.items()} + raise ValueError(f"Neither SeedGraph candidate is valid: {messages}.") + valid.sort( + key=lambda item: ( + item[1].score, + item[0] == "offline", + ), + reverse=True, + ) + selected = deepcopy(dict(offline if valid[0][0] == "offline" else online)) + selected["planner_route"] = "selected" + selected.setdefault("metadata", {})["selected_from"] = valid[0][0] + return selected, evaluations + + +def fuse_seed_graphs( + offline: Mapping[str, Any], + online: Mapping[str, Any], + group_routes: Mapping[str, str], + *, + registry: AtomicCapabilityRegistry | None = None, +) -> dict[str, Any]: + """Fuse candidates only at complete TaskGroup boundaries.""" + capabilities = registry or build_atomic_capability_registry() + if offline.get("task_id") != online.get("task_id"): + raise ValueError("Cannot fuse graphs for different tasks.") + for field in ("instruction", "level", "reasoning_type", "capability_catalog_hash"): + if offline.get(field) != online.get(field): + raise ValueError(f"Cannot fuse graphs with different {field} values.") + by_route = { + "offline": validate_seed_graph(offline, known_actions=capabilities.names()), + "online": validate_seed_graph(online, known_actions=capabilities.names()), + } + for graph in by_route.values(): + validate_persisted_contracts(graph, capabilities) + groups_by_route = { + route: {str(group["id"]): group for group in graph["task_groups"]} + for route, graph in by_route.items() + } + expected = set(groups_by_route["offline"]) + if set(groups_by_route["online"]) != expected or set(group_routes) != expected: + raise ValueError( + "Fusion requires the same complete TaskGroup set in both graphs." + ) + if set(group_routes.values()) - {"offline", "online"}: + raise ValueError("Every fused TaskGroup route must be offline or online.") + + selected_groups = { + group_id: deepcopy(groups_by_route[route][group_id]) + for group_id, route in group_routes.items() + } + _reject_state_conflicts(selected_groups) + source_nodes = { + route: {str(node["id"]): node for node in graph["nodes"]} + for route, graph in by_route.items() + } + selected_nodes_by_group: dict[str, list[dict[str, Any]]] = {} + id_map: dict[tuple[str, str], str] = {} + for group_id, route in group_routes.items(): + group = selected_groups[group_id] + group.pop("contract", None) + selected_nodes_by_group[group_id] = [] + for node_id in group["node_ids"]: + node = deepcopy(source_nodes[route][node_id]) + fused_id = f"{route}_{node_id}" + id_map[(route, node_id)] = fused_id + node["id"] = fused_id + selected_nodes_by_group[group_id].append(node) + + terminals = {} + for group_id, route in group_routes.items(): + original_ids = set(selected_groups[group_id]["node_ids"]) + referenced = { + dependency + for node_id in original_ids + for dependency in source_nodes[route][node_id]["depends_on"] + if dependency in original_ids + } + terminals[group_id] = [ + id_map[(route, node_id)] + for node_id in selected_groups[group_id]["node_ids"] + if node_id not in referenced + ] + nodes = [] + groups = [] + for group_id in _topological_groups(selected_groups): + route = group_routes[group_id] + group = selected_groups[group_id] + own_original_ids = set(group["node_ids"]) + group_nodes = selected_nodes_by_group[group_id] + for node in group_nodes: + original_id = node["id"][len(route) + 1 :] + original = source_nodes[route][original_id] + internal = [ + id_map[(route, dependency)] + for dependency in original["depends_on"] + if dependency in own_original_ids + ] + external = [ + terminal + for parent in group["depends_on"] + for terminal in terminals[parent] + ] + node["depends_on"] = list(dict.fromkeys([*internal, *external])) + nodes.append(node) + group["node_ids"] = [node["id"] for node in group_nodes] + groups.append(group) + + fused = deepcopy(by_route["offline"]) + fused["planner_route"] = "fused" + fused["nodes"] = nodes + fused["task_groups"] = groups + fused["success"] = {"op": "all", "terms": [group["success"] for group in groups]} + fused["metadata"] = { + "fusion_routes": dict(sorted(group_routes.items())), + "fusion_boundary": "task_group", + } + return link_seed_graph( + fused, + registry=capabilities, + task_order=[str(group["id"]) for group in groups], + ) + + +def _reject_state_conflicts(groups: Mapping[str, Mapping[str, Any]]) -> None: + by_object: dict[str, list[str]] = defaultdict(list) + for group_id, group in groups.items(): + by_object[str(group["object_uid"])].append(group_id) + dependencies = { + group_id: set(group["depends_on"]) for group_id, group in groups.items() + } + + def reaches(child: str, parent: str) -> bool: + pending = list(dependencies[child]) + visited = set() + while pending: + current = pending.pop() + if current == parent: + return True + if current not in visited: + visited.add(current) + pending.extend(dependencies[current]) + return False + + for object_uid, group_ids in by_object.items(): + for index, first in enumerate(group_ids): + for second in group_ids[index + 1 :]: + if not reaches(first, second) and not reaches(second, first): + raise ValueError( + f"Fusion has unordered state changes for object {object_uid!r}." + ) + + +def _topological_groups(groups: Mapping[str, Mapping[str, Any]]) -> list[str]: + outgoing = {group_id: [] for group_id in groups} + indegree = {group_id: 0 for group_id in groups} + for group_id, group in groups.items(): + for parent in group["depends_on"]: + outgoing[parent].append(group_id) + indegree[group_id] += 1 + ready = deque( + sorted(group_id for group_id, degree in indegree.items() if degree == 0) + ) + result = [] + while ready: + group_id = ready.popleft() + result.append(group_id) + for child in sorted(outgoing[group_id]): + indegree[child] -= 1 + if indegree[child] == 0: + ready.append(child) + return result diff --git a/embodichain/gen_sim/action_engine/planning/task_planner_prompt.py b/embodichain/gen_sim/action_engine/planning/task_planner_prompt.py new file mode 100644 index 000000000..9dce0b612 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/task_planner_prompt.py @@ -0,0 +1,159 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Prompt template for the Action Engine semantic planner.""" + +from __future__ import annotations + +__all__ = ["TASK_PLANNER_PROMPT"] + +TASK_PLANNER_PROMPT = """You are the semantic planner for a tabletop robot Action Engine. + +Return exactly one JSON object with exactly these two top-level fields: + +{ + "semantic_steps": [ + { + "id": "s01_short_stable_name", + "operator": "", + "object": "", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [] + } + ], + "allocation_groups": [] +} + +For collective operators, replace "object" with "objects": + +{ + "id": "s01_collective_goal", + "operator": "arrange_line", + "objects": ["object_a", "object_b"], + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [] +} + +Hard rules: + +- Plan a sequence or DAG of semantic operators. Do not select a task route. +- Emit semantic_steps and allocation_groups only. Do not emit explanations, + confidence, warnings, + atomic actions, graph nodes, graph edges, resources, motion policies, poses, + coordinates, offsets, distances, joint values, trajectories, or tolerances. +- Use runtime_uid values from the scene inventory. Never invent object IDs. +- Preserve every explicit before/after/then dependency with depends_on. +- Use depends_on=[] for genuinely independent operations that may run in + parallel. Otherwise depend on the preceding required semantic step. +- actor.mode is "auto" unless the user explicitly requires one arm. +- allocation_groups expresses an explicit distinct-arm constraint across + independent semantic steps. Use + {"id":"dual_arms_1","semantic_step_ids":["s01","s02"], + "arm_constraint":"distinct_arms"} only when the user explicitly requests + different arms. Merely independent steps must not receive a group. +- An explicitly required arm uses + {"mode": "required", "arm": "left_arm"} or "right_arm". +- Coordinated operators use + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}. +- Use named symbolic relations and policies only. Runtime observes geometry. +- Every operator is a complete skill, not an individual motion command. + place_relative already picks, transports, releases, retreats, and returns + home. Never emit individual robot motions. +- When the user asks both arms to handle two independent objects, emit two + direct object-level operators with + actor={"mode":"auto"} and depends_on=[], then reference their step IDs in one + allocation_groups entry. The deterministic compiler assigns distinct arms; + do not guess left/right from object positions. +- Spatial phrases that place objects on opposite sides describe object + locations, not an arm-allocation constraint. Emit an allocation group only + when the user explicitly requests both or distinct arms. + +Built-in operator shapes: + +1. arrange_line + - objects: at least two movable objects in requested order. + - goal fields: anchor="table_center"; axis="world_x"|"world_y"| + "table_long_axis"; order_constraint="free"|"ordered"; + order_by="explicit"|"size"|"color"; order_direction="given"| + "ascending"|"descending"; orientation_goal="none"|"preserve"|"upright"| + "lay_flat"|"axis_align"; orientation_axis="none"|"x"|"y"| + "long_axis"|"short_axis". + - In the rotated robot view, world_y is the horizontal left-to-right axis + and world_x is the front-to-back depth axis. For an unspecified line or + row direction, always use axis="world_y". Use axis="world_x" only when + the user explicitly requests a front-to-back, depth-wise, column, or + x-axis layout. Use table_long_axis only when the user explicitly names the + table's long axis; never infer it from a generic line request. + - Use order_constraint="free" when the user wants a line but does not care + which object occupies each slot. + - A line layout does not imply an orientation acceptance requirement. Use + orientation_goal="none" and orientation_axis="none" unless the task + explicitly asks to preserve orientation, make objects upright, lay them + flat, or align an axis. + +2. build_stack + - objects: bottom-to-top movable object order. + - goal fields: stack_mode="on_top"|"nested"; anchor="table_center" or a + passive support runtime_uid; orientation_goal and orientation_axis. + - A vertical stack chain is exactly one build_stack step. Always use the + plural "objects" list, never singular "object", and do not include the + passive anchor in that list. + - Repeated clauses such as "put A on anchor, then put B on top" describe one + chain: objects=[A,B], anchor=anchor. Use separate place_relative steps only + when every object should independently contact the same support. + +3. place_relative + - object: one movable object. + - goal fields: reference_object; relation="inside"|"on"|"left_of"| + "right_of"|"front_of"|"behind"|"front_left_of"|"front_right_of"| + "back_left_of"|"back_right_of"; reference_state="live"|"initial"; + orientation_goal; orientation_axis; optional + orientation_reference_object. + +4. orient_object + - object: one movable object. + - goal fields: orientation_goal="upright"|"lay_flat"|"axis_align"; + orientation_axis="none"|"x"|"y"|"long_axis"|"short_axis"; + support_object=; position_anchor="initial_xy"|"live_xy"; + upright_local_axis="auto"|"long_axis"|"x"|"y"|"z". + - Use orientation_goal="upright" only when the instruction explicitly asks + to make the object upright. + - Use support_object="table" and position_anchor="initial_xy" for an + in-place tabletop orientation request. Use upright_local_axis="auto" + unless the scene inventory explicitly supplies a local semantic axis; + never infer a mesh-local axis from an object name. + +5. coordinated_transport + - object: one shared object moved by both arms. + - goal fields: direction="none"|"world_x"|"world_y"|"front"|"back"| + "left"|"right"|"front_left"|"front_right"|"back_left"|"back_right"| + "up"|"down"; terminal_behavior="hold"|"place"; optional reference_object + and relation; orientation_goal and orientation_axis. + +Available operators: +$operator_catalog + +Task name: +$task_name + +Task description: +$task_description + +Scene inventory: +$scene_objects +""" diff --git a/embodichain/gen_sim/action_engine/planning/vision.py b/embodichain/gen_sim/action_engine/planning/vision.py new file mode 100644 index 000000000..cba17b9b7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/planning/vision.py @@ -0,0 +1,808 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Auditable multi-view observation and VLM fact extraction.""" + +from __future__ import annotations + +import base64 +from collections.abc import Callable, Collection, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from io import BytesIO +import json +import math +import os +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.domain import ( + VISUAL_RELATION_PARTICIPANTS, + public_task_spec, + requested_visual_task_predicates, +) + +__all__ = [ + "CameraObservation", + "SceneObservation", + "analyze_visual_scene", + "collect_scene_observation", + "validate_visual_facts", +] + +StructuredCaller = Callable[..., Mapping[str, Any]] + +_VISUAL_ENTITY_KEYS = frozenset( + { + "uid", + "camera_uid", + "bbox", + "keypoints", + "visible", + "confidence", + } +) +_VISUAL_RELATION_KEYS = frozenset({"type", "uids", "confidence"}) +_VISUAL_TASK_PREDICATE_KEYS = frozenset({"type", "confidence"}) + + +@dataclass(frozen=True) +class CameraObservation: + """One live camera sample with calibration for one vectorized env row.""" + + uid: str + rgb: torch.Tensor + depth: torch.Tensor | None + intrinsics: torch.Tensor | None + extrinsics: torch.Tensor | None + + +@dataclass(frozen=True) +class SceneObservation: + """Multi-view evidence and stable simulator entity IDs for online planning.""" + + cameras: tuple[CameraObservation, ...] + entities: tuple[dict[str, Any], ...] + env_id: int = 0 + + +_VISUAL_FACTS_SCHEMA = { + "title": "ActionEngineVisualFacts", + "type": "object", + "additionalProperties": False, + "required": ["entities", "relations", "task_predicates", "confidence"], + "properties": { + "entities": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["uid", "camera_uid", "confidence"], + "properties": { + "uid": {"type": "string"}, + "camera_uid": {"type": "string"}, + "bbox": { + "type": "array", + "minItems": 4, + "maxItems": 4, + "items": {"type": "number"}, + }, + "keypoints": { + "type": "object", + "additionalProperties": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": {"type": "number"}, + }, + }, + "visible": {"type": "boolean"}, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, + }, + }, + "relations": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["type", "uids", "confidence"], + "properties": { + "type": { + "type": "string", + "enum": sorted(VISUAL_RELATION_PARTICIPANTS), + }, + "uids": { + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": {"type": "string"}, + }, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, + }, + }, + "task_predicates": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": ["type", "confidence"], + "properties": { + "type": {"type": "string"}, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, + }, + }, + "confidence": {"type": "number", "minimum": 0.0, "maximum": 1.0}, + }, +} + +# Visual facts are deliberately a much smaller contract than a simulator +# snapshot. In particular, accepting arbitrary nested ``attributes`` would +# let a caller smuggle poses/qpos into the online planner while still passing +# the top-level schema. Keep the deny-list here (rather than relying only on +# the SeedGraph validator) because visual facts are persisted and may be +# consumed by an independent planner implementation. +_FORBIDDEN_LIVE_KEYS = frozenset( + { + "absolute_position", + "coordinates", + "extrinsics", + "grasp_pose", + "joint_positions", + "live_pose", + "live_transform", + "object_pose", + "oracle", + "pose", + "positions", + "qpos", + "target_pose", + "trajectory", + "transform", + "waypoints", + "xpos", + } +) + + +def collect_scene_observation( + env: Any, + *, + camera_uids: Sequence[str] | None = None, + env_id: int = 0, +) -> SceneObservation: + """Capture current RGB/depth/calibration and a simulator entity inventory.""" + if env_id < 0 or env_id >= int(env.num_envs): + raise ValueError("env_id is outside the vectorized environment range.") + sim = env.sim + uids = ( + list(camera_uids) + if camera_uids is not None + else list(sim.get_sensor_uid_list()) + ) + cameras = [] + for uid in uids: + sensor = sim.get_sensor(str(uid)) + if sensor is None: + raise ValueError(f"Unknown camera UID {uid!r}.") + update = getattr(sensor, "update", None) + if callable(update): + update() + data = sensor.get_data() + if not isinstance(data, Mapping): + raise TypeError(f"Camera {uid!r} returned non-mapping sensor data.") + rgb_data = data.get("color", data.get("rgb")) + if rgb_data is None: + raise ValueError(f"Camera {uid!r} does not provide RGB data.") + rgb = ( + _env_row( + rgb_data, + env_id, + num_envs=int(env.num_envs), + unbatched_ndim=3, + ) + .detach() + .cpu() + ) + depth = ( + _env_row( + data["depth"], + env_id, + num_envs=int(env.num_envs), + unbatched_ndim=2, + ) + .detach() + .cpu() + if data.get("depth") is not None + else None + ) + intrinsics = _optional_call( + sensor, "get_intrinsics", env_id, num_envs=int(env.num_envs) + ) + extrinsics = _optional_call( + sensor, + "get_arena_pose", + env_id, + num_envs=int(env.num_envs), + to_matrix=True, + ) + cameras.append( + CameraObservation( + uid=str(uid), + rgb=rgb, + depth=depth, + intrinsics=intrinsics, + extrinsics=extrinsics, + ) + ) + if not cameras: + raise ValueError("Online visual planning requires at least one camera.") + + entity_uids = list(sim.get_rigid_object_uid_list()) + articulation_uids = getattr(sim, "get_articulation_uid_list", lambda: [])() + entities = [] + for uid in [*entity_uids, *articulation_uids]: + item: dict[str, Any] = {"uid": str(uid)} + # Do not expose live simulator transforms to the online planner. The + # VLM receives RGB/depth evidence and stable UIDs only; JIT grounding + # resolves world-space targets inside the runtime immediately before + # each action. This also prevents an accidental pose oracle through + # the entity inventory prompt. + entities.append(item) + return SceneObservation(tuple(cameras), tuple(entities), env_id=env_id) + + +def analyze_visual_scene( + observation: SceneObservation, + task_spec: Mapping[str, Any], + *, + model: str | None = None, + caller: StructuredCaller | None = None, + call_counter: list[int] | None = None, +) -> dict[str, Any]: + """Ask a VLM for auditable facts, never hidden reasoning or an action plan.""" + _reject_live_fields(observation.entities, "SceneObservation.entities") + public = public_task_spec(task_spec) + allowed_task_predicates = requested_visual_task_predicates(public) + relation_contracts = { + name: list(participants) + for name, participants in VISUAL_RELATION_PARTICIPANTS.items() + } + _reject_live_fields(public, "PublicTaskSpec") + camera_manifest, images = _camera_evidence(observation) + prompt = ( + "Inspect every supplied camera view. Return only observable facts needed " + "for the task. Refer to simulator entities only by the supplied UID. " + "Use normalized [0,1] bbox/keypoint values, state uncertainty explicitly, " + "and do not provide reasoning or actions. The image blocks appear in the " + "camera_evidence order: each RGB image is followed by that camera's " + "normalized depth image when depth_image_index is present. Camera " + "calibration is input evidence only; never reproduce it in the facts. " + "Use only these canonical spatial relation contracts, whose values give " + "the ordered UID participants: " + f"{json.dumps(relation_contracts, sort_keys=True)}. Put task-level visual " + "judgments in task_predicates, never in relations; their allowed types " + f"are {json.dumps(sorted(allowed_task_predicates))}.\n\n" + f"TaskSpec:\n{json.dumps(public, ensure_ascii=False, sort_keys=True)}\n\n" + f"Entity inventory:\n{json.dumps(observation.entities, ensure_ascii=False, sort_keys=True)}\n\n" + f"Camera evidence:\n{json.dumps(camera_manifest, ensure_ascii=False, sort_keys=True)}" + ) + invoke = caller or _default_structured_caller + # Test/mocked callers own their transport and may intentionally receive no + # configured model. The production caller must resolve strictly through + # the visual-model priority rather than falling back to a text-only model. + selected_model = model if caller is not None else _vlm_model(model) + first_error: Exception | None = None + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nThe previous visual-facts JSON was invalid. Return corrected " + f"JSON only. Validation error: {first_error}" + ) + try: + if call_counter is not None: + call_counter[0] += 1 + response = invoke( + prompt=current_prompt, + images=images, + schema=_visual_facts_schema(allowed_task_predicates), + model=selected_model, + ) + facts = validate_visual_facts( + response, + known_uids={str(item["uid"]) for item in observation.entities}, + camera_uids={camera.uid for camera in observation.cameras}, + allowed_task_predicates=allowed_task_predicates, + ) + if facts["confidence"] < 0.5: + raise ValueError( + "VLM visual facts confidence is below the required 0.5 threshold." + ) + if not any( + item.get("visible", True) and item["confidence"] >= 0.5 + for item in facts["entities"] + ): + raise ValueError("VLM visual facts contain no reliable visible entity.") + return facts + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "VLM visual facts failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def validate_visual_facts( + value: Mapping[str, Any], + *, + known_uids: set[str], + camera_uids: set[str], + allowed_task_predicates: Collection[str] = (), +) -> dict[str, Any]: + """Validate entity identity and normalized image-space evidence.""" + if not isinstance(value, Mapping): + raise TypeError("VLM visual facts must be a mapping.") + required_fields = {"entities", "relations", "task_predicates", "confidence"} + if set(value) != required_fields: + raise ValueError( + "VLM visual facts require exactly fields " + f"{sorted(required_fields)}; received {sorted(value)}." + ) + confidence = _confidence(value.get("confidence"), "confidence") + entities = value.get("entities") + relations = value.get("relations") + task_predicates = value.get("task_predicates") + if not isinstance(entities, Sequence) or isinstance(entities, (str, bytes)): + raise ValueError("VLM visual facts entities must be a list.") + if not isinstance(relations, Sequence) or isinstance(relations, (str, bytes)): + raise ValueError("VLM visual facts relations must be a list.") + if not isinstance(task_predicates, Sequence) or isinstance( + task_predicates, (str, bytes) + ): + raise ValueError("VLM visual facts task_predicates must be a list.") + normalized_entities = [] + for index, item in enumerate(entities): + if not isinstance(item, Mapping): + raise ValueError(f"visual entities[{index}] must be a mapping.") + unsupported = set(item) - _VISUAL_ENTITY_KEYS + if unsupported: + raise ValueError( + f"visual entities[{index}] contains unsupported fields " + f"{sorted(unsupported)}." + ) + uid = item.get("uid") + camera_uid = item.get("camera_uid") + if not isinstance(uid, str) or not uid: + raise ValueError( + f"visual entities[{index}].uid must be a non-empty string." + ) + if not isinstance(camera_uid, str) or not camera_uid: + raise ValueError( + f"visual entities[{index}].camera_uid must be a non-empty string." + ) + if uid not in known_uids: + raise ValueError( + f"visual entities[{index}] references unknown UID {uid!r}." + ) + if camera_uid not in camera_uids: + raise ValueError( + f"visual entities[{index}] references unknown camera {camera_uid!r}." + ) + normalized = dict(item) + _reject_live_fields(normalized, f"visual entities[{index}]") + if "visible" in normalized and not isinstance(normalized["visible"], bool): + raise ValueError(f"visual entities[{index}].visible must be a boolean.") + if "bbox" in normalized: + normalized["bbox"] = _normalized_vector( + normalized["bbox"], 4, f"visual entities[{index}].bbox" + ) + x_min, y_min, x_max, y_max = normalized["bbox"] + if x_min >= x_max or y_min >= y_max: + raise ValueError( + f"visual entities[{index}].bbox must have non-zero ordered bounds." + ) + keypoints = normalized.get("keypoints", {}) + if not isinstance(keypoints, Mapping): + raise ValueError(f"visual entities[{index}].keypoints must be a mapping.") + normalized["keypoints"] = { + str(name): _normalized_vector(point, 2, f"keypoint {name!r}") + for name, point in keypoints.items() + } + if ( + normalized.get("visible", True) + and "bbox" not in normalized + and not normalized["keypoints"] + ): + raise ValueError( + f"visual entities[{index}] must include a bbox or keypoint evidence." + ) + normalized["confidence"] = _confidence( + normalized.get("confidence"), f"visual entities[{index}].confidence" + ) + normalized_entities.append(normalized) + normalized_relations = [] + for index, relation in enumerate(relations): + if not isinstance(relation, Mapping): + raise ValueError(f"visual relations[{index}] must be a mapping.") + unsupported = set(relation) - _VISUAL_RELATION_KEYS + if unsupported: + raise ValueError( + f"visual relations[{index}] contains unsupported fields " + f"{sorted(unsupported)}." + ) + relation_type = relation.get("type") + if ( + not isinstance(relation_type, str) + or relation_type not in VISUAL_RELATION_PARTICIPANTS + ): + raise ValueError( + f"visual relations[{index}] relation type must be one of " + f"{sorted(VISUAL_RELATION_PARTICIPANTS)}." + ) + participants = relation.get("uids", []) + if not isinstance(participants, Sequence) or isinstance( + participants, (str, bytes) + ): + raise ValueError(f"visual relations[{index}].uids must be a list.") + if any(not isinstance(uid, str) or not uid for uid in participants): + raise ValueError( + f"visual relations[{index}].uids must contain non-empty strings." + ) + expected_count = len(VISUAL_RELATION_PARTICIPANTS[relation_type]) + if len(participants) != expected_count: + raise ValueError( + f"visual relations[{index}].uids must contain exactly " + f"{expected_count} UIDs in canonical participant order." + ) + if len(set(participants)) != len(participants): + raise ValueError( + f"visual relations[{index}].uids must contain distinct UIDs." + ) + invalid = set(participants) - known_uids + if invalid: + raise ValueError( + f"visual relations[{index}] has unknown UIDs {sorted(invalid)}." + ) + normalized = dict(relation) + _reject_live_fields(normalized, f"visual relations[{index}]") + normalized["confidence"] = _confidence( + normalized.get("confidence"), f"visual relations[{index}].confidence" + ) + normalized_relations.append(normalized) + normalized_task_predicates = [] + allowed_predicates = {str(item) for item in allowed_task_predicates} + for index, predicate in enumerate(task_predicates): + if not isinstance(predicate, Mapping): + raise ValueError(f"visual task_predicates[{index}] must be a mapping.") + unsupported = set(predicate) - _VISUAL_TASK_PREDICATE_KEYS + if unsupported or set(predicate) != _VISUAL_TASK_PREDICATE_KEYS: + raise ValueError( + f"visual task_predicates[{index}] requires exactly fields " + f"{sorted(_VISUAL_TASK_PREDICATE_KEYS)}." + ) + predicate_type = predicate.get("type") + if ( + not isinstance(predicate_type, str) + or predicate_type not in allowed_predicates + ): + raise ValueError( + f"visual task_predicates[{index}].type must be one of " + f"{sorted(allowed_predicates)}." + ) + normalized = dict(predicate) + _reject_live_fields(normalized, f"visual task_predicates[{index}]") + normalized["confidence"] = _confidence( + normalized.get("confidence"), + f"visual task_predicates[{index}].confidence", + ) + normalized_task_predicates.append(normalized) + _reject_live_fields( + { + "entities": normalized_entities, + "relations": normalized_relations, + "task_predicates": normalized_task_predicates, + }, + "VLM visual facts", + ) + return { + "entities": normalized_entities, + "relations": normalized_relations, + "task_predicates": normalized_task_predicates, + "confidence": confidence, + } + + +def _visual_facts_schema( + allowed_task_predicates: Collection[str], +) -> dict[str, Any]: + """Return the visual-fact schema specialized for the current task.""" + schema = deepcopy(_VISUAL_FACTS_SCHEMA) + predicate_schema = schema["properties"]["task_predicates"] + allowed = sorted(str(item) for item in allowed_task_predicates) + if allowed: + predicate_schema["items"]["properties"]["type"]["enum"] = allowed + else: + predicate_schema["maxItems"] = 0 + return schema + + +def _default_structured_caller( + *, + prompt: str, + images: Sequence[str], + schema: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + from langchain_core.messages import HumanMessage, SystemMessage + from langchain_openai import ChatOpenAI + + from .planner import ( + _coerce_model_response, + _is_mimo_compatible, + _load_llm_settings, + _structured_output_runnable, + ) + + settings = _load_llm_settings(model=model) + kwargs: dict[str, Any] = { + "api_key": settings["api_key"], + "model": settings["model"], + "temperature": 0, + "http_socket_options": (), + } + for key in ("base_url", "default_query"): + if settings[key]: + kwargs[key] = settings[key] + if _is_mimo_compatible(settings): + kwargs.update( + { + "max_completion_tokens": 4096, + "extra_body": {"thinking": {"type": "disabled"}}, + } + ) + client = ChatOpenAI(**kwargs) + structured = _structured_output_runnable( + client, + schema, + settings=settings, + ) + schema_prompt = ( + f"{prompt}\n\nReturn one JSON object conforming exactly to this JSON " + f"Schema:\n{json.dumps(schema, ensure_ascii=False, sort_keys=True)}" + ) + content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + content[0]["text"] = schema_prompt + content.extend( + {"type": "image_url", "image_url": {"url": image}} for image in images + ) + response = structured.invoke( + [ + SystemMessage( + content="Report visual facts only. Never reveal chain-of-thought." + ), + HumanMessage(content=content), + ] + ) + return _coerce_model_response(response) + + +def _vlm_model(explicit: str | None) -> str: + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + from .planner import _GEN_SIM_ENV_PATH, _load_env_file + + local_env = _load_env_file(_GEN_SIM_ENV_PATH) + # A VLM-specific choice wins over the generic OpenAI default regardless of + # whether it comes from the shell or the project dotenv. Within each name, + # process variables retain their normal override behavior. + for key in ("ACTION_ENGINE_VLM_MODEL", "OPENAI_MODEL"): + for source in (os.environ, local_env): + value = source.get(key) + if isinstance(value, str) and value.strip(): + return value.strip() + raise ValueError( + "A VLM model is required through --vlm-model, agent_config.vlm_model, " + "ACTION_ENGINE_VLM_MODEL, or OPENAI_MODEL." + ) + + +def _rgb_data_url(value: torch.Tensor) -> str: + from PIL import Image + + image = value + if image.ndim != 3 or image.shape[-1] not in {3, 4}: + raise ValueError("Camera RGB must have shape (H, W, 3|4).") + if image.dtype != torch.uint8: + image = image.float() + if float(image.max()) <= 1.0: + image = image * 255.0 + image = image.clamp(0, 255).to(torch.uint8) + stream = BytesIO() + Image.fromarray(image.numpy()).convert("RGB").save(stream, format="PNG") + return "data:image/png;base64," + base64.b64encode(stream.getvalue()).decode( + "ascii" + ) + + +def _camera_evidence( + observation: SceneObservation, +) -> tuple[list[dict[str, Any]], list[str]]: + """Package calibrated RGB/depth evidence in a stable camera order.""" + manifest: list[dict[str, Any]] = [] + images: list[str] = [] + for camera in observation.cameras: + rgb_index = len(images) + images.append(_rgb_data_url(camera.rgb)) + item: dict[str, Any] = { + "uid": camera.uid, + "rgb_image_index": rgb_index, + "depth_available": camera.depth is not None, + "intrinsics": _calibration_list(camera.intrinsics), + "extrinsics": _calibration_list(camera.extrinsics), + } + if camera.depth is not None: + item["depth_image_index"] = len(images) + images.append(_depth_data_url(camera.depth)) + manifest.append(item) + return manifest, images + + +def _calibration_list(value: torch.Tensor | None) -> list[Any] | None: + """Serialize finite calibration tensors for the transient VLM prompt.""" + if value is None: + return None + tensor = torch.as_tensor(value).detach().cpu() + if not bool(torch.isfinite(tensor).all()): + raise ValueError("Camera calibration contains non-finite values.") + return tensor.tolist() + + +def _depth_data_url(value: torch.Tensor) -> str: + """Render one depth frame as a normalized grayscale VLM evidence image.""" + from PIL import Image + + depth = torch.as_tensor(value).detach().cpu().float() + if depth.ndim == 3 and depth.shape[-1] == 1: + depth = depth[..., 0] + elif depth.ndim == 3 and depth.shape[0] == 1: + depth = depth[0] + if depth.ndim != 2: + raise ValueError("Camera depth must have shape (H, W) or a singleton channel.") + finite = torch.isfinite(depth) + if not bool(finite.any()): + raise ValueError("Camera depth contains no finite values.") + minimum = depth[finite].min() + maximum = depth[finite].max() + normalized = torch.zeros_like(depth) + if float(maximum - minimum) > 0.0: + normalized[finite] = (depth[finite] - minimum) / (maximum - minimum) + image = (normalized.clamp(0.0, 1.0) * 255.0).to(torch.uint8).numpy() + stream = BytesIO() + Image.fromarray(image, mode="L").save(stream, format="PNG") + return "data:image/png;base64," + base64.b64encode(stream.getvalue()).decode( + "ascii" + ) + + +def _env_row( + value: Any, + env_id: int, + *, + num_envs: int | None = None, + unbatched_ndim: int | tuple[int, ...] | None = None, +) -> torch.Tensor: + """Select one vectorized environment row without slicing image dimensions. + + Sensor APIs return either ``(num_envs, ...)`` or an unbatched ``(...)`` + tensor. The old ``shape[0] > env_id`` heuristic sliced the first image row + for an unbatched ``(H, W, C)`` RGB tensor and similarly corrupted 4x4 poses. + Prefer the known environment count and only use the legacy heuristic when + no count is available. + """ + tensor = torch.as_tensor(value) + if unbatched_ndim is not None: + allowed_ndim = ( + (unbatched_ndim,) + if isinstance(unbatched_ndim, int) + else tuple(unbatched_ndim) + ) + if tensor.ndim in allowed_ndim: + return tensor + if tensor.ndim and num_envs is not None and tensor.shape[0] == int(num_envs): + if env_id >= tensor.shape[0]: + raise ValueError("env_id is outside the sensor batch dimension.") + return tensor[env_id] + if num_envs is None and tensor.ndim and tensor.shape[0] > env_id: + return tensor[env_id] + return tensor + + +def _optional_call( + sensor: Any, name: str, env_id: int, *, num_envs: int | None = None, **kwargs: Any +) -> torch.Tensor | None: + method = getattr(sensor, name, None) + if not callable(method): + return None + try: + value = method(env_id=env_id, **kwargs) + except TypeError: + try: + value = method(env_id, **kwargs) + except TypeError: + value = method(**kwargs) + value = torch.as_tensor(value) + # Calibration methods commonly return an unbatched matrix even for a + # vectorized simulator. Select a leading environment row only when the + # shape cannot itself be a canonical calibration matrix. This preserves + # 3x3/4x4 matrices while correctly handling batched compact vectors such as + # ``(num_envs, 4)``. + unbatched_matrix = value.ndim == 2 and tuple(value.shape) in { + (3, 3), + (4, 4), + } + if ( + num_envs is not None + and value.ndim >= 1 + and value.shape[0] == int(num_envs) + and not unbatched_matrix + ): + value = value[env_id] + return value.detach().cpu() + + +def _reject_live_fields(value: Any, context: str) -> None: + """Reject nested simulator state/geometry fields in VLM facts.""" + if isinstance(value, Mapping): + for key, child in value.items(): + if str(key).lower() in _FORBIDDEN_LIVE_KEYS: + raise ValueError( + f"{context} contains forbidden live-state field {key!r}." + ) + _reject_live_fields(child, f"{context}.{key}") + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)): + for index, child in enumerate(value): + _reject_live_fields(child, f"{context}[{index}]") + + +def _normalized_vector(value: Any, size: int, context: str) -> list[float]: + if ( + not isinstance(value, Sequence) + or isinstance(value, (str, bytes)) + or len(value) != size + ): + raise ValueError(f"{context} must contain {size} normalized values.") + if any( + not isinstance(item, (int, float)) or isinstance(item, bool) for item in value + ): + raise ValueError(f"{context} values must be numeric.") + result = [float(item) for item in value] + if any(not math.isfinite(item) or item < 0.0 or item > 1.0 for item in result): + raise ValueError(f"{context} values must lie in [0, 1].") + return result + + +def _confidence(value: Any, context: str) -> float: + if not isinstance(value, (int, float)) or isinstance(value, bool): + raise ValueError(f"{context} must be a number in [0, 1].") + result = float(value) + if not math.isfinite(result) or result < 0.0 or result > 1.0: + raise ValueError(f"{context} must lie in [0, 1].") + return result diff --git a/embodichain/gen_sim/action_engine/protocol.py b/embodichain/gen_sim/action_engine/protocol.py new file mode 100644 index 000000000..ef704952d --- /dev/null +++ b/embodichain/gen_sim/action_engine/protocol.py @@ -0,0 +1,60 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Cross-layer identifiers owned by Action Engine. + +These values are serialized into generated artifacts, so changing one is a +protocol migration rather than a local rename. +""" + +from __future__ import annotations + +from typing import Final + +__all__ = [ + "ACTION_ENGINE_CONFIG_SCHEMA", + "ACTION_ENGINE_ENV_ID", + "AGENT_CONFIG_FILENAME", + "COMPARISON_FILENAME", + "EXECUTION_PROGRAM_FILENAME", + "EXECUTION_PROGRAM_SCHEMA", + "FAST_GYM_CONFIG_FILENAME", + "SCENE_REQUIREMENTS_FILENAME", + "SCENE_REQUIREMENTS_SCHEMA", + "SEED_TASK_GRAPH_PNG_FILENAME", + "SEED_GRAPH_SCHEMA", + "TASK_SPEC_FILENAME", + "TASK_SPEC_SCHEMA", + "TASK_AGENT_FILENAME", + "TASK_AGENT_SCHEMA", +] + +ACTION_ENGINE_ENV_ID: Final = "ActionEngine-v1" +ACTION_ENGINE_CONFIG_SCHEMA: Final = "action_engine_config_v2" +TASK_AGENT_SCHEMA: Final = "action_engine_task_agent_v1" +EXECUTION_PROGRAM_SCHEMA: Final = "action_engine_execution_graph_v1" +SEED_GRAPH_SCHEMA: Final = "action_engine_seed_graph_v3" +TASK_SPEC_SCHEMA: Final = "action_engine_task_spec_v2" +SCENE_REQUIREMENTS_SCHEMA: Final = "action_engine_scene_requirements_v2" + +FAST_GYM_CONFIG_FILENAME: Final = "fast_gym_config.json" +AGENT_CONFIG_FILENAME: Final = "agent_config.json" +TASK_AGENT_FILENAME: Final = "task_agent.json" +EXECUTION_PROGRAM_FILENAME: Final = "seed_task_graph.json" +SEED_TASK_GRAPH_PNG_FILENAME: Final = "seed_task_graph.png" +TASK_SPEC_FILENAME: Final = "task_spec.json" +SCENE_REQUIREMENTS_FILENAME: Final = "scene_requirements.json" +COMPARISON_FILENAME: Final = "comparison.json" diff --git a/embodichain/gen_sim/action_engine/runtime/__init__.py b/embodichain/gen_sim/action_engine/runtime/__init__.py new file mode 100644 index 000000000..64774425f --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/__init__.py @@ -0,0 +1,66 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Runtime API for the compositional Action Engine.""" + +from __future__ import annotations + +from .executor import ProgramExecutor +from .dynamic import DynamicRecoveryController, RecoveryDirective +from .loader import load_agent_execution_program, load_execution_program +from .recovery import ( + FAILURE_TYPES, + GraphRevision, + RetryDecision, + RuntimeGraph, + build_upright_recovery, + classify_failure, +) +from .models import ExecutionProgram, ExecutionReport, ExecutionResult +from .reporting import ( + EXECUTION_REPORT_FILENAME, + EXECUTION_REPORT_SCHEMA, + build_execution_provenance, + validate_execution_report, + write_execution_report, +) +from .predicates import PREDICATE_TYPES, evaluate_predicate +from .state import ExecutionState + +__all__ = [ + "ExecutionProgram", + "ExecutionState", + "DynamicRecoveryController", + "PREDICATE_TYPES", + "ExecutionResult", + "ExecutionReport", + "EXECUTION_REPORT_FILENAME", + "EXECUTION_REPORT_SCHEMA", + "FAILURE_TYPES", + "GraphRevision", + "ProgramExecutor", + "RetryDecision", + "RecoveryDirective", + "RuntimeGraph", + "build_execution_provenance", + "build_upright_recovery", + "classify_failure", + "evaluate_predicate", + "load_agent_execution_program", + "load_execution_program", + "validate_execution_report", + "write_execution_report", +] diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py new file mode 100644 index 000000000..3a1411a9b --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -0,0 +1,1534 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Adapt Action Engine requests to the shared typed atomic-action planner.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import replace +import math +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapability, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.config import default_runtime_policy +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + ActionPlan, + AntipodalAffordance, + AtomicActionEngine, + ControlPartCommandProfile, + CoordinatedPickGoal, + DynamicCollisionMode, + EndEffectorPoseGoal, + EntityState, + ExecutionSession, + MotionPolicy, + ObjectSemantics, + PlanningContext, + RecoveryPolicy, + RobotObservation, + RigidObjectSceneProvider, + SceneProvider, + SceneSnapshot, + StateDelta, +) +from embodichain.lab.sim.planners import ( + CuroboPlannerCfg, + CuroboWorldCfg, + MotionGenCfg, + MotionGenerator, + ToppraPlannerCfg, +) +from embodichain.toolkits.graspkit.pg_grasp import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, + GripperCollisionCfg, +) +from embodichain.utils.logger import log_info + +from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache +from .models import ActionOutcome, GroundedAction +from .state import ExecutionState + +__all__ = ["AtomicActionAdapter"] + + +_DEFAULT_PLANNER_POLICY: dict[str, Any] = { + "backend": "curobo", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "fallback_strategy": "ik_interp", + "allow_fallback": True, + "dynamic_collision": False, + "static_obstacle_uids": [], + "dynamic_obstacle_uids": [], + "curobo": { + "log_level": "error", + "obstacle_representation": "cuboid", + "multi_env": False, + "use_cuda_graph": True, + "preserve_plan_samples": False, + "max_attempts": 5, + "collision_activation_distance": 0.01, + }, +} + +# Preserve cuRobo's fixed world shape while disabling intentional-contact objects. +_COLLISION_PARKING_Z_OFFSET = -100.0 + + +def _collision_cache_for_world( + representation: str, obstacle_count: int +) -> dict[str, int]: + """Size cuRobo's fixed collision cache for the generated scene.""" + cache = {"cuboid": 8, "mesh": 2} + if representation in cache: + cache[representation] = max(cache[representation], obstacle_count) + return cache + + +def _supported_kwargs(config_type: type, values: Mapping[str, Any]) -> dict[str, Any]: + names: set[str] = set() + for cls in reversed(config_type.__mro__): + names.update(getattr(cls, "__annotations__", {})) + return {key: value for key, value in values.items() if key in names} + + +def _as_hand_qpos(value: Any, dof: int, device: Any) -> torch.Tensor: + if dof == 0: + return torch.empty(0, dtype=torch.float32, device=device) + result = torch.as_tensor(value, dtype=torch.float32, device=device).flatten() + if result.numel() == 0: + return torch.zeros(dof, dtype=torch.float32, device=device) + if result.numel() == 1: + return result.repeat(dof) + if result.numel() >= dof: + return result[:dof] + repeats = (dof + result.numel() - 1) // result.numel() + return result.repeat(repeats)[:dof] + + +def _diagonal_approach_direction( + horizontal: torch.Tensor, + *, + vertical: float = -1.0, +) -> torch.Tensor: + """Combine one normalized horizontal role direction with a vertical component.""" + horizontal = horizontal.to(dtype=torch.float32) + norm = torch.linalg.vector_norm(horizontal) + if float(norm) <= 1.0e-6: + raise ValueError("Handover role direction must be non-zero.") + horizontal = horizontal / norm + direction = torch.stack( + (horizontal[0], horizontal[1], horizontal.new_tensor(float(vertical))) + ) + return direction / torch.linalg.vector_norm(direction) + + +class AtomicActionAdapter: + """Own the shared atomic engine and preserve Action Engine runtime contracts.""" + + def __init__( + self, + env: Any, + *, + grasp_policy: Mapping[str, Any] | None = None, + planner_policy: Mapping[str, Any] | None = None, + capability_registry: Any | None = None, + scene_provider: SceneProvider | None = None, + ) -> None: + self.env = env + self.num_envs = int(env.num_envs) + self.device = env.device + if grasp_policy is None: + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + grasp_policy = default_runtime_policy(profile).grasp + grasp_policy = { + **grasp_policy, + **(getattr(env, "agent_grasp_runtime_defaults", {}) or {}), + } + self.grasp_policy = deepcopy(dict(grasp_policy)) + self.planner_policy = deepcopy(_DEFAULT_PLANNER_POLICY) + if planner_policy is not None: + self._merge_planner_policy(self.planner_policy, planner_policy) + if not self.planner_policy.get("static_obstacle_uids"): + configured = getattr(env, "agent_static_obstacle_uids", ()) or () + if configured: + self.planner_policy["static_obstacle_uids"] = [ + str(uid) for uid in configured + ] + else: + get_rigid_object = getattr(env.sim, "get_rigid_object", None) + if callable(get_rigid_object) and get_rigid_object("table") is not None: + self.planner_policy["static_obstacle_uids"] = ["table"] + self.capabilities = capability_registry or build_atomic_capability_registry() + self._motion_generator: MotionGenerator | None = None + self._atomic_engine: AtomicActionEngine | None = None + self._semantics: dict[str, ObjectSemantics] = {} + self._scene_time = 0.0 + if scene_provider is not None and not isinstance(scene_provider, SceneProvider): + raise TypeError("scene_provider must implement SceneProvider.") + self.scene_provider = scene_provider or self._build_scene_provider() + + @staticmethod + def _merge_planner_policy( + target: dict[str, Any], + update: Mapping[str, Any], + ) -> None: + for key, value in update.items(): + if isinstance(value, Mapping) and isinstance(target.get(key), dict): + AtomicActionAdapter._merge_planner_policy(target[key], value) + else: + target[key] = deepcopy(value) + + def initial_state(self) -> ExecutionState: + """Capture the initial full-robot planning seed.""" + return ExecutionState(last_qpos=self.env.robot.get_qpos().clone()) + + def start_session( + self, + grounded: GroundedAction, + state: ExecutionState | None = None, + ) -> ExecutionSession: + """Start one closed-loop AtomicAction session from live scene state. + + ProgramExecutor may continue using its compatibility scheduler for + compound and per-arm merged trajectories. New callers can use this + boundary to adopt feedback-driven execution without constructing + private planning contexts. + """ + capability = self.capabilities.require_executable(grounded.action_class) + state = state or self.initial_state() + grounded = self._select_upright_transport_yaw(grounded, state) + context = self._planning_context(state, grounded) + invocation = self._invocation(grounded, capability) + return self._engine().start((invocation,), context) + + def _build_scene_provider(self) -> SceneProvider | None: + """Create the shared live rigid-object provider when entities are available.""" + sim = getattr(self.env, "sim", None) + if sim is None: + return None + dynamic_uids = tuple( + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ) + list_uids = getattr(sim, "get_rigid_object_uid_list", None) + uids = tuple(str(uid) for uid in list_uids()) if callable(list_uids) else () + if not uids: + uids = dynamic_uids + get_rigid_object = getattr(sim, "get_rigid_object", None) + if not callable(get_rigid_object): + return None + entities = { + uid: entity for uid in uids if (entity := get_rigid_object(uid)) is not None + } + if not entities: + return None + collision_uids = ( + dynamic_uids + if bool(self.planner_policy.get("dynamic_collision", False)) + else () + ) + return RigidObjectSceneProvider( + entities, + collision_entity_ids=collision_uids, + ) + + def semantics(self, uid: str) -> ObjectSemantics: + """Build object semantics once while retaining the live entity handle.""" + cached = self._semantics.get(uid) + if cached is not None: + return cached + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr( + self.env.sim, + "get_articulation", + lambda _uid: None, + )(uid) + if entity is None: + raise ValueError(f"Unknown grasp target {uid!r}.") + active_joint_ids = list(getattr(entity, "active_joint_ids", ())) + if len(active_joint_ids) != 1: + raise ValueError( + "Articulation semantics require exactly one active joint." + ) + backend_entities = getattr( + entity, + "_entities", + getattr(entity, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + joint_name = str(entity.joint_names[active_joint_ids[0]]) + joint_info = backend_entities[0].get_joint_info(joint_name) + child_link = str(getattr(joint_info, "child_link_name", "")) + vertices, triangles = entity.get_link_vert_face(child_link) + else: + vertices = entity.get_vertices(env_ids=[0], scale=True) + triangles = entity.get_triangles(env_ids=[0]) + if isinstance(vertices, (tuple, list)): + vertices = vertices[0] + if isinstance(triangles, (tuple, list)): + triangles = triangles[0] + vertices = torch.as_tensor(vertices, dtype=torch.float32) + triangles = torch.as_tensor(triangles, dtype=torch.int64) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if triangles.ndim == 3 and triangles.shape[0] == 1: + triangles = triangles[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Object {uid!r} has invalid mesh vertices.") + if triangles.ndim != 2 or triangles.shape[-1] != 3 or triangles.numel() == 0: + raise ValueError(f"Object {uid!r} has invalid mesh triangles.") + + grasp_options = self.grasp_policy + sampler = AntipodalSamplerCfg( + n_sample=int(grasp_options["antipodal_n_sample"]), + max_angle=float(grasp_options["antipodal_max_angle"]), + max_length=float(grasp_options["max_open_length"]), + min_length=float(grasp_options["min_open_length"]), + ) + generator = GraspGeneratorCfg( + viser_port=int(grasp_options["viser_port"]), + antipodal_sampler_cfg=sampler, + max_deviation_angle=float(grasp_options["max_deviation_angle"]), + n_deviated_approach_directions=int( + grasp_options["n_deviated_approach_directions"] + ), + ) + max_hulls = int(grasp_options["max_decomposition_hulls"]) + collision = GripperCollisionCfg( + max_open_length=float(grasp_options["max_open_length"]), + finger_length=float(grasp_options["finger_length"]), + point_sample_dense=float(grasp_options["point_sample_dense"]), + max_decomposition_hulls=max_hulls, + ) + cache_result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=max_hulls, + ) + if cache_result.status != "hit": + log_info(f"Prepared V-HACD grasp cache for {uid!r}: {cache_result.status}.") + + semantics = ObjectSemantics( + label=uid, + entity=entity, + geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + affordance=AntipodalAffordance( + object_label=uid, + mesh_vertices=vertices, + mesh_triangles=triangles, + generator_cfg=generator, + gripper_collision_cfg=collision, + force_reannotate=bool(grasp_options["force_grasp_reannotate"]), + ), + ) + self._semantics[uid] = semantics + return semantics + + def plan( + self, + grounded: GroundedAction, + state: ExecutionState | None = None, + ) -> ActionOutcome: + """Plan one grounded primitive through the mainline typed contract.""" + capability = self.capabilities.require_executable(grounded.action_class) + state = state or self.initial_state() + grounded = self._select_upright_transport_yaw(grounded, state) + context = self._planning_context(state, grounded) + invocation = self._invocation(grounded, capability) + plan = self._engine().plan(invocation, context) + selected_positions = self._positions_with_agent_holds( + plan, + grounded, + capability, + ) + primary_success = plan.plan_success.to(self.device) + reachability_search = None + if bool(grounded.motion_policy.get("retreat_reachability_search", False)): + ( + grounded, + selected_positions, + primary_success, + reachability_search, + ) = self._search_reachable_retreat( + grounded=grounded, + capability=capability, + state=state, + context=context, + invocation=invocation, + initial_positions=selected_positions, + initial_success=primary_success, + ) + invocation = replace(invocation, goal=grounded.target) + combined_success = primary_success.clone() + fallback_plan: ActionPlan | None = None + use_fallback = torch.zeros_like(combined_success) + fallback_attempted = torch.zeros_like(combined_success) + fallback_success = torch.zeros_like(combined_success) + + fallback_strategy = self.planner_policy.get("fallback_strategy") + collision_safety = str(grounded.motion_policy.get("collision_safety", "auto")) + fallback_allowed = bool(self.planner_policy.get("allow_fallback", True)) and ( + collision_safety != "required" + ) + if ( + fallback_allowed + and invocation.motion_policy.strategy == "motion_gen" + and fallback_strategy in {"ik_interp"} + and not bool(combined_success.all()) + ): + fallback_attempted = ~primary_success + fallback_policy = replace( + invocation.motion_policy, + strategy=str(fallback_strategy), + dynamic_collision_mode=DynamicCollisionMode.OFF, + plan_opts=None, + ) + fallback_plan = self._engine().plan( + replace(invocation, motion_policy=fallback_policy), + context, + ) + fallback_positions = self._positions_with_agent_holds( + fallback_plan, + grounded, + capability, + ) + fallback_success = fallback_plan.plan_success.to(self.device) + use_fallback = fallback_attempted & fallback_success + selected_positions = self._merge_plan_rows( + selected_positions, + fallback_positions, + use_fallback, + state.last_qpos, + ) + combined_success |= fallback_plan.plan_success.to(self.device) + + options = invocation.skill_options + if capability.config_materializer == "handover": + combined_success &= self._handover_receiver_hold_mask( + selected_positions, + grounded, + options, + tolerance=float( + grounded.motion_policy.get( + "receiver_hold_joint_tolerance", + 2.0e-3, + ) + ), + ) + + terminal_qpos = ( + selected_positions[:, -1] + if selected_positions.shape[1] + else state.last_qpos + ) + primary_rows = combined_success & primary_success + projected_task = plan.expected_effects.apply( + context.task, + primary_rows, + ) + held_keys = set(plan.expected_effects.held_object_updates) + if fallback_plan is not None: + fallback_rows = combined_success & use_fallback + projected_task = fallback_plan.expected_effects.apply( + projected_task, + fallback_rows, + ) + held_keys.update(fallback_plan.expected_effects.held_object_updates) + committed_effects = StateDelta( + held_object_updates={ + key: projected_task.held_objects.get(key) for key in held_keys + }, + ) + next_state = ExecutionState.from_task_state( + projected_task, + last_qpos=torch.where( + combined_success[:, None], terminal_qpos, state.last_qpos + ), + ) + return ActionOutcome( + trajectory=selected_positions, + success=combined_success, + next_state=next_state, + grounded=grounded, + prior_state=state, + expected_effects=committed_effects, + planner_trace={ + **self._planner_trace( + grounded=grounded, + invocation=invocation, + context=context, + state=state, + primary_success=primary_success, + fallback_allowed=fallback_allowed, + fallback_strategy=( + str(fallback_strategy) + if invocation.motion_policy.strategy == "motion_gen" + and fallback_strategy in {"ik_interp"} + else None + ), + fallback_attempted=fallback_attempted, + fallback_success=fallback_success, + fallback_used=use_fallback, + reachability_search=reachability_search, + ), + # Auditability takes precedence over compactness here: every + # selected planner route retains its complete joint path. + "planned_trajectory": selected_positions.detach().clone(), + }, + ) + + def _search_reachable_retreat( + self, + *, + grounded: GroundedAction, + capability: AtomicCapability, + state: ExecutionState, + context: PlanningContext, + invocation: ActionInvocation, + initial_positions: torch.Tensor, + initial_success: torch.Tensor, + ) -> tuple[GroundedAction, torch.Tensor, torch.Tensor, dict[str, Any]]: + """Select the highest row-local retreat accepted by the live planner.""" + candidates = self._retreat_search_targets(grounded) + target = getattr(grounded.target, "xpos", None) + if not isinstance(target, torch.Tensor) or len(candidates) <= 1: + return ( + grounded, + initial_positions, + initial_success, + { + "strategy": "bounded_motion_planner", + "attempts": [], + "selected_target_z": ( + None + if not isinstance(target, torch.Tensor) + else target[:, 2, 3] + ), + }, + ) + + selected_target = candidates[0][1].clone() + selected_positions = initial_positions + success = initial_success.clone() + attempts: list[dict[str, Any]] = [ + { + "candidate": candidates[0][0], + "target_z": candidates[0][1][:, 2, 3].detach().clone(), + "success": initial_success.detach().clone(), + } + ] + for label, candidate_target in candidates[1:]: + unresolved = ~success + if not bool(unresolved.any()): + break + row_target = torch.where( + unresolved[:, None, None], + candidate_target, + selected_target, + ) + candidate_grounded = replace( + grounded, + target=EndEffectorPoseGoal(xpos=row_target), + ) + candidate_invocation = replace( + invocation, + goal=candidate_grounded.target, + ) + candidate_plan = self._engine().plan(candidate_invocation, context) + candidate_positions = self._positions_with_agent_holds( + candidate_plan, + candidate_grounded, + capability, + ) + candidate_success = candidate_plan.plan_success.to(self.device) + selected_rows = unresolved & candidate_success + selected_positions = self._merge_plan_rows( + selected_positions, + candidate_positions, + selected_rows, + state.last_qpos, + ) + selected_target = torch.where( + selected_rows[:, None, None], + candidate_target, + selected_target, + ) + success |= candidate_success + attempts.append( + { + "candidate": label, + "target_z": candidate_target[:, 2, 3].detach().clone(), + "success": candidate_success.detach().clone(), + } + ) + + metadata = { + "retreat_selected_target_z": selected_target[:, 2, 3].detach().clone(), + "retreat_reachability_found": success.detach().clone(), + } + selected_grounded = replace( + grounded, + target=EndEffectorPoseGoal(xpos=selected_target), + cfg={**grounded.cfg, **metadata}, + motion_policy={**grounded.motion_policy, **metadata}, + ) + return ( + selected_grounded, + selected_positions, + success, + { + "strategy": "bounded_motion_planner", + "attempts": attempts, + "selected_target_z": selected_target[:, 2, 3].detach().clone(), + }, + ) + + def _retreat_search_targets( + self, + grounded: GroundedAction, + ) -> list[tuple[str, torch.Tensor]]: + """Build bounded height and baseward retreat candidates from live poses.""" + target = getattr(grounded.target, "xpos", None) + reference = grounded.motion_policy.get("retreat_reference_pose") + if not isinstance(target, torch.Tensor) or not isinstance( + reference, torch.Tensor + ): + return [] + target = target.to(device=self.device, dtype=torch.float32) + reference = reference.to(device=self.device, dtype=torch.float32) + if target.shape == (4, 4): + target = target.unsqueeze(0).repeat(self.num_envs, 1, 1) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + expected = (self.num_envs, 4, 4) + if target.shape != expected or reference.shape != expected: + return [] + + sample_count = int(grounded.cfg.get("retreat_search_samples", 6)) + if not 2 <= sample_count <= 16: + raise ValueError("retreat_search_samples must be in [2, 16].") + minimum_height = float(grounded.cfg.get("minimum_retreat_height", 0.05)) + if not math.isfinite(minimum_height) or minimum_height < 0.0: + raise ValueError("minimum_retreat_height must be finite and non-negative.") + desired_height = torch.clamp( + target[:, 2, 3] - reference[:, 2, 3], + min=0.0, + ) + minimum = torch.minimum( + desired_height, + torch.full_like(desired_height, minimum_height), + ) + fractions = torch.linspace( + 1.0, + 0.0, + sample_count, + dtype=target.dtype, + device=target.device, + ) + heights = ( + minimum[:, None] + (desired_height - minimum)[:, None] * fractions[None] + ) + candidates: list[tuple[str, torch.Tensor]] = [("requested", target.clone())] + for index in range(1, sample_count): + candidate = target.clone() + candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] + candidates.append((f"height_{index}", candidate)) + + from .frames import arm_base_poses + + left_base, right_base = arm_base_poses(self.env) + base = left_base if grounded.arm == "left_arm" else right_base + direction = base[:, :2, 3] - reference[:, :2, 3] + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + direction = torch.where( + norm > 1.0e-6, + direction / torch.clamp(norm, min=1.0e-6), + torch.zeros_like(direction), + ) + distance = float(grounded.cfg.get("retreat_distance", 0.10)) + if not math.isfinite(distance) or distance < 0.0: + raise ValueError("retreat_distance must be finite and non-negative.") + for index in range(sample_count): + candidate = target.clone() + candidate[:, :2, 3] = reference[:, :2, 3] + direction * distance + candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] + candidates.append((f"baseward_{index}", candidate)) + return candidates + + def _planner_trace( + self, + *, + grounded: GroundedAction, + invocation: ActionInvocation, + context: PlanningContext, + state: ExecutionState, + primary_success: torch.Tensor, + fallback_allowed: bool, + fallback_strategy: str | None, + fallback_attempted: torch.Tensor, + fallback_success: torch.Tensor, + fallback_used: torch.Tensor, + reachability_search: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """Build compact per-row evidence for the planner route actually used.""" + exclusions = self._collision_exclusion_masks(grounded, state) + obstacle_positions = { + uid: context.scene.entities[uid].pose[:, :3, 3].detach().clone() + for uid in context.scene.collision_entity_ids + } + revisions = torch.as_tensor( + context.scene.collision_world_revisions(self.num_envs), + dtype=torch.int64, + device=self.device, + ) + trace = { + "action_class": grounded.action_class, + "arm": grounded.arm, + "planner": str(self.planner_policy["backend"]), + "primary_strategy": invocation.motion_policy.strategy, + "dynamic_collision_mode": invocation.motion_policy.dynamic_collision_mode.value, + "primary_success": primary_success.detach().clone(), + "fallback_allowed": fallback_allowed, + "fallback_strategy": fallback_strategy, + "fallback_attempted": fallback_attempted.detach().clone(), + "fallback_success": fallback_success.detach().clone(), + "fallback_used": fallback_used.detach().clone(), + "search_budget": { + "primary_max_attempts": int( + self.planner_policy.get("curobo", {}).get("max_attempts", 1) + ), + "fallback_enabled": bool(fallback_allowed), + }, + "collision_world_revision": revisions, + "collision_obstacle_positions": obstacle_positions, + "collision_exclusions": { + uid: mask.detach().clone() for uid, mask in exclusions.items() + }, + } + if reachability_search is not None: + trace["reachability_search"] = deepcopy(dict(reachability_search)) + options = invocation.skill_options + object_part = getattr(options, "pick_object_part", None) + approach_direction = getattr(options, "approach_direction", None) + if object_part is None: + object_part = getattr(options, "receive_pick_object_part", None) + approach_direction = getattr( + options, + "receive_approach_direction", + approach_direction, + ) + if object_part is not None: + grasp_policy: dict[str, Any] = {"object_part": str(object_part)} + if isinstance(approach_direction, torch.Tensor): + direction = approach_direction.to(dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if bool(torch.isfinite(norm)) and float(norm) > 0.0: + grasp_policy["approach_direction"] = ( + (direction / norm).detach().cpu().tolist() + ) + trace["grasp_policy"] = grasp_policy + return trace + + def _select_upright_transport_yaw( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> GroundedAction: + """Choose the closest IK-feasible yaw for an upright object target.""" + sample_count = int(grounded.cfg.get("upright_yaw_samples", 1)) + capability = self.capabilities.get(grounded.action_class) + if ( + capability.target_materializer != "semantic_held_object" + or sample_count <= 1 + ): + return grounded + target_pose = getattr(grounded.target, "object_target_pose", None) + if not isinstance(target_pose, torch.Tensor): + return grounded + target_pose = target_pose.to(device=self.device, dtype=torch.float32) + if target_pose.shape == (4, 4): + target_pose = target_pose.unsqueeze(0).repeat(self.num_envs, 1, 1) + if target_pose.shape != (self.num_envs, 4, 4): + raise ValueError( + "Upright transport target must have shape (4, 4) or (N, 4, 4)." + ) + + arm_part, _, _ = self._parts(grounded.arm) + held = state.get_held_object(arm_part) + if held is None: + return grounded + object_to_eef = held.object_to_eef.to( + device=self.device, + dtype=target_pose.dtype, + ) + if object_to_eef.shape == (4, 4): + object_to_eef = object_to_eef.unsqueeze(0).repeat(self.num_envs, 1, 1) + variants = self._upright_yaw_variants(target_pose, sample_count) + eef_variants = torch.matmul(variants, object_to_eef[:, None]) + joint_ids = list(self.env.robot.get_joint_ids(name=arm_part)) + start_qpos = state.last_qpos[:, joint_ids] + seeds = start_qpos[:, None].expand(-1, sample_count, -1) + success, qpos = self.env.robot.compute_batch_ik( + pose=eef_variants, + name=arm_part, + joint_seed=seeds, + ) + success = torch.as_tensor( + success, + dtype=torch.bool, + device=self.device, + ).reshape(self.num_envs, sample_count) + qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) + success &= torch.isfinite(qpos).all(dim=-1) + distance = torch.linalg.vector_norm(qpos - seeds, dim=-1) + distance = torch.where( + success, + distance, + torch.full_like(distance, torch.inf), + ) + best = distance.argmin(dim=1) + env_ids = torch.arange(self.num_envs, device=self.device) + selected = variants[env_ids, best] + selected = torch.where( + success.any(dim=1)[:, None, None], + selected, + target_pose, + ) + return replace( + grounded, + target=replace(grounded.target, object_target_pose=selected), + target_object_pose=selected, + ) + + @staticmethod + def _upright_yaw_variants( + target_pose: torch.Tensor, + sample_count: int, + ) -> torch.Tensor: + signed_steps = [0] + for step in range(1, (sample_count + 1) // 2): + signed_steps.extend((step, -step)) + if sample_count % 2 == 0: + signed_steps.append(sample_count // 2) + angles = target_pose.new_tensor(signed_steps) * (2.0 * math.pi / sample_count) + yaw = target_pose.new_zeros((sample_count, 3, 3)) + yaw[:, 0, 0] = torch.cos(angles) + yaw[:, 0, 1] = -torch.sin(angles) + yaw[:, 1, 0] = torch.sin(angles) + yaw[:, 1, 1] = torch.cos(angles) + yaw[:, 2, 2] = 1.0 + variants = target_pose[:, None].repeat(1, sample_count, 1, 1) + variants[:, :, :3, :3] = torch.matmul(yaw[None], target_pose[:, None, :3, :3]) + return variants + + def _planning_context( + self, + state: ExecutionState, + grounded: GroundedAction, + ) -> PlanningContext: + qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) + get_qvel = getattr(self.env.robot, "get_qvel", None) + qvel = get_qvel() if callable(get_qvel) else None + if not isinstance(qvel, torch.Tensor) or qvel.shape != qpos.shape: + qvel = torch.zeros_like(qpos) + else: + qvel = qvel.to(device=self.device, dtype=qpos.dtype) + return PlanningContext( + robot=RobotObservation(timestamp=self._scene_time, qpos=qpos, qvel=qvel), + task=state.to_task_state(), + scene=self._scene_snapshot(grounded, state), + env_ids=torch.arange( + self.num_envs, + dtype=torch.long, + device=self.device, + ), + control_dt=float(getattr(self.env, "step_dt", 1.0 / 60.0)), + ) + + def _scene_snapshot( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> SceneSnapshot: + dynamic_uids = tuple( + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ) + env_ids = torch.arange( + self.num_envs, + dtype=torch.long, + device=self.device, + ) + if self.scene_provider is None: + base = SceneSnapshot(timestamp=self._scene_time, version=0) + else: + base = self.scene_provider.snapshot( + timestamp=self._scene_time, + env_ids=env_ids, + ) + if not bool(self.planner_policy.get("dynamic_collision", False)): + return base + exclusion_masks = self._collision_exclusion_masks(grounded, state) + entities = dict(base.entities) + for uid in dynamic_uids: + entity_state = entities.get(uid) + if entity_state is None: + raise ValueError( + f"SceneProvider omitted cuRobo dynamic obstacle {uid!r}." + ) + pose = entity_state.pose.to(dtype=torch.float32, device=self.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(self.num_envs, 1, 1) + if pose.shape != (self.num_envs, 4, 4): + raise ValueError( + f"Dynamic obstacle {uid!r} pose must have shape (4, 4) or " + f"({self.num_envs}, 4, 4), got {tuple(pose.shape)}." + ) + excluded = exclusion_masks.get(uid) + if excluded is not None and bool(excluded.any()): + pose = pose.clone() + pose[excluded, 2, 3] += _COLLISION_PARKING_Z_OFFSET + entities[uid] = EntityState( + pose=pose, + confidence=entity_state.confidence, + ) + return SceneSnapshot( + timestamp=base.timestamp, + version=base.version, + entities=entities, + collision_world_revision=base.collision_world_revision, + collision_entity_ids=dynamic_uids, + ) + + def _collision_exclusion_masks( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> dict[str, torch.Tensor]: + """Return per-environment masks for obstacles intentionally in contact.""" + dynamic_uids = { + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + } + masks: dict[str, torch.Tensor] = {} + + def include(uid: str | None, env_mask: torch.Tensor | None = None) -> None: + if uid is None or uid not in dynamic_uids: + return + mask = ( + torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + if env_mask is None + else torch.as_tensor( + env_mask, + dtype=torch.bool, + device=self.device, + ).reshape(-1) + ) + if mask.shape != (self.num_envs,): + raise ValueError( + f"Collision exclusion mask for {uid!r} must have shape " + f"({self.num_envs},), got {tuple(mask.shape)}." + ) + masks[uid] = masks.get(uid, torch.zeros_like(mask)) | mask + + if self.capabilities.get(grounded.action_class).allows_target_contact: + target_uid = grounded.object_uid + if target_uid is None: + target_uid = getattr( + getattr(grounded.target, "semantics", None), + "label", + None, + ) + include(target_uid) + + for held in state.held_objects.values(): + include(held.semantics.label, held.env_mask) + collision_exclusion_uids = grounded.motion_policy.get( + "collision_exclusion_uids", () + ) + if isinstance(collision_exclusion_uids, str): + collision_exclusion_uids = (collision_exclusion_uids,) + for uid in collision_exclusion_uids: + include(str(uid)) + return masks + + def _invocation( + self, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> ActionInvocation: + if capability.resource_mode == "coordinated_object": + strategy = str(self.planner_policy["coordinated_strategy"]) + elif grounded.control == "hand": + strategy = "ik_interp" + else: + strategy = str(self.planner_policy["single_arm_strategy"]) + sample_count = max(2, int(grounded.cfg.get("sample_interval", 50))) + dynamic_collision = bool(self.planner_policy.get("dynamic_collision", False)) + collision_required = ( + grounded.motion_policy.get("collision_safety") == "required" + ) + if dynamic_collision and strategy == "motion_gen": + dynamic_mode = ( + DynamicCollisionMode.REQUIRED + if collision_required + else DynamicCollisionMode.AUTO + ) + else: + dynamic_mode = DynamicCollisionMode.OFF + goal = ( + self._coordinated_pickment_goal(grounded) + if capability.config_materializer == "coordinated_pickment" + else grounded.target + ) + return ActionInvocation( + skill_id=str(capability.action_type.skill_id), + goal=goal, + binding=self._binding(grounded, capability), + motion_policy=MotionPolicy( + strategy=strategy, + sample_count=sample_count, + dynamic_collision_mode=dynamic_mode, + ), + recovery_policy=RecoveryPolicy(), + skill_options=self._build_config(grounded, capability), + ) + + @staticmethod + def _coordinated_pickment_goal(grounded: GroundedAction) -> CoordinatedPickGoal: + """Apply GenSim-only coordinated grasp filtering to an owned goal copy.""" + target = grounded.target + if not isinstance(target, CoordinatedPickGoal): + raise TypeError("CoordinatedPickment requires a CoordinatedPickGoal.") + requested = grounded.cfg.get("is_filter_ground_collision") + if requested is None: + return target + if not isinstance(requested, bool): + raise TypeError("is_filter_ground_collision must be a boolean.") + semantics = target.semantics + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise TypeError( + "CoordinatedPickment requires an AntipodalAffordance for GenSim " + "grasp filtering." + ) + generator_cfg = deepcopy(affordance.generator_cfg or GraspGeneratorCfg()) + generator_cfg.is_filter_ground_collision = requested + scoped_affordance = replace(affordance, generator_cfg=generator_cfg) + scoped_semantics = replace(semantics, affordance=scoped_affordance) + return replace(target, semantics=scoped_semantics) + + def _binding( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> ActionBinding: + engine = self._engine() + contract = getattr(capability.action_type, "binding_contract", None) + if contract is None: + return ActionBinding(owner_id=engine.binding_owner_id) + + slot_parts: dict[str, tuple[str, str | None]] = {} + if capability.config_materializer == "handover": + transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + transfer_arm, transfer_hand, _ = self._parts(transfer_side) + receive_arm, receive_hand, _ = self._parts(receive_side) + if transfer_hand is None or receive_hand is None: + raise ValueError("HandOver requires two configured end effectors.") + slot_parts = { + "source": (transfer_arm, transfer_hand), + "destination": (receive_arm, receive_hand), + } + elif capability.config_materializer == "coordinated_pickment": + left_arm, left_hand, _ = self._parts("left_arm") + right_arm, right_hand, _ = self._parts("right_arm") + if left_hand is None or right_hand is None: + raise ValueError("Coordinated pickup requires two end effectors.") + slot_parts = { + "left": (left_arm, left_hand), + "right": (right_arm, right_hand), + } + elif capability.config_materializer == "coordinated_placement": + placing_arm, placing_hand, _ = self._parts("left_arm") + support_arm, support_hand, _ = self._parts("right_arm") + if placing_hand is None or support_hand is None: + raise ValueError("Coordinated placement requires two end effectors.") + slot_parts = { + "placing": (placing_arm, placing_hand), + "support": (support_arm, support_hand), + } + else: + arm_part, hand_part, _ = self._parts(action.arm) + motion_part = hand_part if action.control == "hand" else arm_part + if motion_part is None: + raise ValueError( + f"{action.arm} has no configured {action.control} part." + ) + slot_parts = {"primary": (motion_part, hand_part)} + + endpoints: dict[str, dict[str, str]] = {} + for slot in contract.slots: + try: + motion_part, hand_part = slot_parts[slot.slot_id] + except KeyError as exc: + raise ValueError( + f"No GenSim binding is available for slot {slot.slot_id!r}." + ) from exc + selected: dict[str, str] = {} + for requirement in slot.endpoints: + if requirement.endpoint_id == "motion": + selected["motion"] = motion_part + elif requirement.endpoint_id == "grasp": + if hand_part is None: + raise ValueError( + f"{capability.name} requires a grasp endpoint for " + f"slot {slot.slot_id!r}." + ) + selected["grasp"] = hand_part + else: + raise ValueError( + f"Unsupported GenSim endpoint {slot.slot_id}." + f"{requirement.endpoint_id}." + ) + endpoints[slot.slot_id] = selected + return engine.bind_control_parts( + str(capability.action_type.skill_id), + endpoints, + ) + + def _build_config( + self, + action: GroundedAction, + capability: AtomicCapability | type, + ) -> Any: + """Build the mainline immutable ``ActionOptions`` value. + + The method name is retained as a narrow compatibility hook for existing + Action Engine tests and extensions; it no longer constructs legacy + hardware-bound ``ActionCfg`` objects. + """ + if isinstance(capability, type): + registered = self.capabilities.require_executable(action.action_class) + if registered.config_type is not capability: + raise ValueError( + f"Options type {capability.__name__!r} does not match " + f"AtomicAction {action.action_class!r}." + ) + capability = registered + if capability.config_materializer_hook is not None: + return capability.config_materializer_hook( + adapter=self, + action=action, + capability=capability, + ) + builder = getattr( + self, + f"_build_{capability.config_materializer}_config", + self._build_single_arm_config, + ) + return builder(action, capability) + + def _config_policy(self, action: GroundedAction) -> dict[str, Any]: + policy = dict(action.cfg) + for key in ( + "postcondition_tolerance", + "relation_distance", + "hover_height", + "staging_lift_height", + "transport_clearance", + "surface_clearance", + "receiver_hold_joint_tolerance", + "post_hold_steps", + ): + policy.pop(key, None) + return policy + + def _build_single_arm_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + config_type = capability.config_type + if capability.target_materializer == "semantic_held_object": + from .atomic_compat import ExactTargetMoveHeldObjectOptions + + config_type = ExactTargetMoveHeldObjectOptions + if int(action.cfg.get("upright_yaw_samples", 1)) > 1: + policy["allow_automatic_transport_rotation"] = False + if capability.target_materializer == "press": + press_depth = policy.pop("press_depth", None) + if press_depth is not None and "press_distance" not in policy: + policy["press_distance"] = press_depth + approach_mode = policy.pop("approach_direction_mode", None) + if approach_mode == "handover_transfer": + from .frames import robot_frame_axes + + _, lateral = robot_frame_axes(self.env) + outward = lateral[0] if action.arm == "left_arm" else -lateral[0] + policy["approach_direction"] = _diagonal_approach_direction( + -outward.to(device=self.device) + ) + elif approach_mode is not None: + raise ValueError(f"Unknown approach_direction_mode {approach_mode!r}.") + for name in ("approach_direction", "obj_upright_direction"): + if name in policy and not isinstance(policy[name], torch.Tensor): + policy[name] = torch.as_tensor( + policy[name], dtype=torch.float32, device=self.device + ) + return config_type(**_supported_kwargs(config_type, policy)) + + def _build_coordinated_pickment_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + from .frames import arm_base_poses + + policy = self._config_policy(action) + left_base, right_base = arm_base_poses(self.env) + direction = right_base[0, :3, 3] - left_base[0, :3, 3] + norm = torch.linalg.vector_norm(direction) + if not torch.isfinite(direction).all() or norm <= 1.0e-6: + raise ValueError( + "Coordinated pickup requires distinct finite left/right arm bases." + ) + policy["left_to_right_arm_direction"] = direction / norm + return capability.config_type( + **_supported_kwargs(capability.config_type, policy) + ) + + def _build_coordinated_placement_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + return self._build_single_arm_config(action, capability) + + def _build_handover_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + middle = action.cfg.get("middle_object_pose") + final = action.cfg.get("final_object_pose") + if middle is None or final is None: + raise ValueError("HandOver grounding must provide middle and final poses.") + transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + from .frames import robot_frame_axes + + _, lateral = robot_frame_axes(self.env) + receiver_outward = ( + lateral[0] if receive_side == "left_arm" else -lateral[0] + ).to(device=self.device) + receiver_inward_approach = -receiver_outward + policy.update( + { + "middle_object_pose": middle, + # Delivery is represented by a following MoveHeldObject node. + # Keep the receiver fixed while the source retreats here. + "final_object_pose": middle, + "receive_approach_direction": _diagonal_approach_direction( + receiver_inward_approach + ), + } + ) + return capability.config_type( + **_supported_kwargs(capability.config_type, policy) + ) + + def _positions_with_agent_holds( + self, + plan: ActionPlan, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> torch.Tensor: + trajectory = plan.joint_trajectory + if trajectory is None: + raise ValueError( + f"AtomicAction {plan.skill_id!r} did not retain a joint trajectory." + ) + positions = trajectory.positions.to( + device=self.device, + dtype=torch.float32, + ) + hold_steps = int(grounded.cfg.get("post_hold_steps", 0)) + if capability.state_effect != "release" or hold_steps <= 0: + return positions + release = next((item for item in plan.segments if item.name == "release"), None) + if release is None or release.stop <= 0 or release.stop > positions.shape[1]: + return positions + hold = positions[:, release.stop - 1 : release.stop].repeat(1, hold_steps, 1) + return torch.cat( + (positions[:, : release.stop], hold, positions[:, release.stop :]), + dim=1, + ) + + @staticmethod + def _merge_plan_rows( + primary: torch.Tensor, + fallback: torch.Tensor, + use_fallback: torch.Tensor, + hold_qpos: torch.Tensor, + ) -> torch.Tensor: + steps = max(primary.shape[1], fallback.shape[1], 1) + + def padded(value: torch.Tensor) -> torch.Tensor: + if value.shape[1] == 0: + return hold_qpos[:, None].repeat(1, steps, 1) + if value.shape[1] < steps: + value = torch.cat( + (value, value[:, -1:].repeat(1, steps - value.shape[1], 1)), + dim=1, + ) + return value + + primary = padded(primary) + fallback = padded(fallback) + return torch.where(use_fallback[:, None, None], fallback, primary) + + def _handover_receiver_hold_mask( + self, + trajectory: torch.Tensor, + grounded: GroundedAction, + options: Any, + *, + tolerance: float, + ) -> torch.Tensor: + if tolerance < 0.0: + raise ValueError("receiver_hold_joint_tolerance must be non-negative.") + retreat_steps = max(2, int(options.retreat_steps)) + if trajectory.shape[1] < retreat_steps: + return torch.zeros( + self.num_envs, dtype=torch.bool, device=trajectory.device + ) + transfer_side = str(grounded.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + receive_arm, _, _ = self._parts(receive_side) + receiver_ids = self.env.robot.get_joint_ids(name=receive_arm) + receiver = trajectory[:, -retreat_steps:, receiver_ids] + drift = torch.amax(torch.abs(receiver - receiver[:, :1]), dim=(1, 2)) + return torch.isfinite(drift) & (drift <= tolerance) + + def execute_trajectory( + self, + trajectory: torch.Tensor, + *, + active: torch.Tensor, + ) -> list[torch.Tensor]: + """Advance the environment while holding inactive vectorized rows.""" + if trajectory.ndim != 3 or trajectory.shape[0] != self.num_envs: + raise ValueError("Execution trajectory must have shape (N, T, robot_dof).") + active = active.to(device=trajectory.device, dtype=torch.bool) + current = self.env.robot.get_qpos().to( + device=trajectory.device, + dtype=trajectory.dtype, + ) + commands: list[torch.Tensor] = [] + for waypoint in trajectory.unbind(dim=1): + command = torch.where(active[:, None], waypoint, current) + self.env.step(command) + self._scene_time += self._scene_step_duration() + update = getattr(self.env, "update_obj_info", None) + if callable(update): + update() + commands.append(command.detach()) + current = command + sync = getattr(self.env, "sync_agent_state_from_qpos", None) + if callable(sync) and commands: + sync(commands[-1]) + return commands + + def _scene_step_duration(self) -> float: + """Return one positive logical waypoint duration for scene timestamps.""" + sim_config = getattr(getattr(self.env, "sim", None), "sim_config", None) + candidates = ( + getattr(self.env, "physics_dt", None), + getattr(sim_config, "physics_dt", None), + ) + for value in candidates: + if isinstance(value, (int, float)) and not isinstance(value, bool): + duration = float(value) + if math.isfinite(duration) and duration > 0.0: + return duration + return 1.0 + + def combine( + self, + outcomes: Mapping[str, ActionOutcome | None], + masks: Mapping[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Merge independently planned arm paths into one synchronized stream.""" + present = [item for item in outcomes.values() if item is not None] + if not present: + raise ValueError("At least one arm outcome is required.") + steps = max(int(item.trajectory.shape[1]) for item in present) + current = self.env.robot.get_qpos().to(self.device, dtype=torch.float32) + merged = current[:, None, :].repeat(1, max(steps, 1), 1) + success = torch.ones( + self.num_envs, + dtype=torch.bool, + device=self.device, + ) + for arm, outcome in outcomes.items(): + if outcome is None: + continue + mask = masks[arm].to(self.device, dtype=torch.bool) + success &= ~mask | outcome.success + trajectory = outcome.trajectory + if trajectory.shape[1] == 0: + continue + if trajectory.shape[1] < steps: + padding = trajectory[:, -1:].repeat(1, steps - trajectory.shape[1], 1) + trajectory = torch.cat((trajectory, padding), dim=1) + joint_ids = self.joint_ids(arm, include_hand=True) + if not joint_ids: + continue + selected = merged[:, :, joint_ids] + merged[:, :, joint_ids] = torch.where( + mask[:, None, None], trajectory[:, :, joint_ids], selected + ) + return merged, success + + def joint_ids(self, arm: str, *, include_hand: bool) -> list[int]: + if arm == "coordinated": + return list(range(int(self.env.robot.dof))) + side = "left" if arm == "left_arm" else "right" + result = list(getattr(self.env, f"{side}_arm_joints", ())) + if include_hand: + result.extend(getattr(self.env, f"{side}_eef_joints", ())) + return result + + def _engine(self) -> AtomicActionEngine: + if self._atomic_engine is None: + from .atomic_compat import ExactTargetMoveHeldObject + + engine = AtomicActionEngine( + self._generator(), + control_profiles=self._control_profiles(), + ) + engine.register(ExactTargetMoveHeldObject(), replace=True) + self._atomic_engine = engine + return self._atomic_engine + + def _generator(self) -> MotionGenerator: + if self._motion_generator is None: + backend = str(self.planner_policy.get("backend", "curobo")) + if backend == "curobo": + options = dict(self.planner_policy.get("curobo", {})) + obstacle_uids = tuple( + dict.fromkeys( + [ + *self.planner_policy.get("static_obstacle_uids", ()), + *self.planner_policy.get("dynamic_obstacle_uids", ()), + ] + ) + ) + rigid_objects: dict[str, Any] = {} + for uid in obstacle_uids: + obstacle_uid = str(uid) + entity = self.env.sim.get_rigid_object(obstacle_uid) + if entity is None: + raise ValueError(f"Unknown cuRobo obstacle {uid!r}.") + rigid_objects[obstacle_uid] = entity + obstacle_representation = str( + options.get("obstacle_representation", "cuboid") + ) + world = CuroboWorldCfg( + rigid_objects=rigid_objects or None, + obstacle_representation=obstacle_representation, + collision_cache=_collision_cache_for_world( + obstacle_representation, + len(rigid_objects), + ), + dynamic_obstacle_names=[ + str(uid) + for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ], + multi_env=bool(options.get("multi_env", False)), + ) + planner_cfg = CuroboPlannerCfg( + robot_uid=self.env.robot.uid, + log_level=str(options.get("log_level", "error")), + world=world, + use_cuda_graph=bool(options.get("use_cuda_graph", True)), + preserve_plan_samples=bool( + options.get("preserve_plan_samples", False) + ), + max_attempts=int(options.get("max_attempts", 5)), + collision_activation_distance=float( + options.get("collision_activation_distance", 0.01) + ), + ) + elif backend == "toppra": + planner_cfg = ToppraPlannerCfg(robot_uid=self.env.robot.uid) + else: + raise ValueError( + f"Unsupported Action Engine planner backend {backend!r}." + ) + self._motion_generator = MotionGenerator( + cfg=MotionGenCfg(planner_cfg=planner_cfg) + ) + return self._motion_generator + + def _control_profiles(self) -> dict[str, ControlPartCommandProfile]: + profiles: dict[str, ControlPartCommandProfile] = {} + for side in ("left_arm", "right_arm"): + try: + _, hand_part, hand_dof = self._parts(side) + except ValueError: + continue + if hand_part is None or hand_dof == 0 or hand_part in profiles: + continue + profiles[hand_part] = ControlPartCommandProfile.joint_positions( + open=_as_hand_qpos(self.env.open_state, hand_dof, self.device), + grasp=_as_hand_qpos(self.env.close_state, hand_dof, self.device), + ) + return profiles + + def _parts(self, arm: str) -> tuple[str, str | None, int]: + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a physical arm, got {arm!r}.") + is_left = arm == "left_arm" + if hasattr(self.env, "get_agent_arm_control_part"): + arm_part = self.env.get_agent_arm_control_part(is_left) + hand_part = self.env.get_agent_eef_control_part(is_left) + else: + arm_part = arm + hand_part = "left_eef" if is_left else "right_eef" + hand_ids = ( + [] + if hand_part is None + else list(self.env.robot.get_joint_ids(name=hand_part)) + ) + return ( + str(arm_part), + None if hand_part is None else str(hand_part), + len(hand_ids), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py new file mode 100644 index 000000000..8c6a0d63b --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py @@ -0,0 +1,85 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Action Engine-specific adapters for mainline atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionPlan, + HeldObjectPoseGoal, + MoveHeldObject, + MoveHeldObjectOptions, + PlanningContext, + ResolvedActionRequest, +) + +__all__ = ["ExactTargetMoveHeldObject", "ExactTargetMoveHeldObjectOptions"] + + +@dataclass(frozen=True, slots=True, eq=False) +class ExactTargetMoveHeldObjectOptions(MoveHeldObjectOptions): + """Action Engine transport options with an exact-orientation switch.""" + + allow_automatic_transport_rotation: bool = True + """Whether the mainline transport heuristic may replace target rotation.""" + + +class ExactTargetMoveHeldObject(MoveHeldObject): + """Preserve a selected semantic orientation when explicitly requested.""" + + OptionsType = ExactTargetMoveHeldObjectOptions + binding_contract = MoveHeldObject.binding_contract + + def __init__( + self, + default_options: ExactTargetMoveHeldObjectOptions | None = None, + ) -> None: + super().__init__(default_options) + self._allow_automatic_transport_rotation = True + + def _plan( + self, + request: ResolvedActionRequest[ + HeldObjectPoseGoal, + ExactTargetMoveHeldObjectOptions, + ], + context: PlanningContext, + ) -> ActionPlan: + previous = self._allow_automatic_transport_rotation + self._allow_automatic_transport_rotation = ( + request.skill_options.allow_automatic_transport_rotation + ) + try: + return super()._plan(request, context) + finally: + self._allow_automatic_transport_rotation = previous + + def _apply_automatic_transport_rotation( + self, + move_eef_xpos: torch.Tensor, + end_arm_xpos: torch.Tensor, + ) -> None: + """Apply the heuristic unless semantic grounding selected exact yaw.""" + if self._allow_automatic_transport_rotation: + super()._apply_automatic_transport_rotation( + move_eef_xpos, + end_arm_xpos, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/dynamic.py b/embodichain/gen_sim/action_engine/runtime/dynamic.py new file mode 100644 index 000000000..616e14c80 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/dynamic.py @@ -0,0 +1,164 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Route-explicit recovery and suffix-replanning coordinator.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from .recovery import RuntimeGraph + +__all__ = ["DynamicRecoveryController", "RecoveryDirective"] + +Replanner = Callable[..., Mapping[str, Any]] + + +@dataclass(frozen=True) +class RecoveryDirective: + """A graph revision that must execute before suffix replanning.""" + + failure_type: str + failed_node_id: str + recovery_group_id: str | None + graph: dict[str, Any] + requires_recovery_execution: bool + active_env_ids: tuple[int, ...] + + +class DynamicRecoveryController: + """Keep offline and online dynamic replanning as separately testable modes.""" + + def __init__( + self, + runtime_graph: RuntimeGraph, + *, + mode: str, + offline_replanner: Replanner | None = None, + online_replanner: Replanner | None = None, + ) -> None: + if mode not in {"offline_dynamic", "online_dynamic"}: + raise ValueError("Dynamic mode must be offline_dynamic or online_dynamic.") + selected = offline_replanner if mode == "offline_dynamic" else online_replanner + if not callable(selected): + raise ValueError(f"{mode} requires its matching replanner callback.") + self.runtime_graph = runtime_graph + self.mode = mode + self._replanner = selected + + def handle_failure( + self, + *, + failed_node_id: str, + failure_type: str, + active_env_ids: Sequence[int] | None = None, + ) -> RecoveryDirective: + """Insert recovery when known; otherwise request immediate full replanning.""" + env_ids = tuple( + sorted( + set( + range(self.runtime_graph.num_envs) + if active_env_ids is None + else (int(env_id) for env_id in active_env_ids) + ) + ) + ) + if not env_ids or env_ids[0] < 0 or env_ids[-1] >= self.runtime_graph.num_envs: + raise ValueError( + "Recovery active_env_ids are outside the environment range." + ) + if failure_type == "object_fallen": + graph = self.runtime_graph.insert_default_recovery( + failed_node_id=failed_node_id, + failure_type=failure_type, + active_env_ids=env_ids, + ) + group_id = self.runtime_graph.revisions[-1].inserted_group_ids[0] + return RecoveryDirective( + failure_type, + failed_node_id, + group_id, + graph, + True, + self.runtime_graph.revisions[-1].active_env_ids, + ) + return RecoveryDirective( + failure_type, + failed_node_id, + None, + self.runtime_graph.graph, + False, + env_ids, + ) + + def handle_execution_result(self, result: Any) -> RecoveryDirective: + """Create a directive from the first actionable runtime failure event.""" + events = getattr(result, "failure_events", None) + if not isinstance(events, Sequence) or not events: + raise ValueError("Execution result contains no recoverable failure event.") + event = next( + ( + item + for item in events + if isinstance(item, Mapping) and bool(item.get("fatal", True)) + ), + None, + ) + if event is None: + raise ValueError( + "Execution result contains no fatal recoverable failure event." + ) + if not isinstance(event, Mapping): + raise ValueError("Execution failure events must be mappings.") + node_id = event.get("node_id") + failure_type = event.get("failure_type") + env_ids = event.get("env_ids", ()) + if not isinstance(node_id, str) or not node_id: + raise ValueError("Dynamic recovery requires a v3 SeedGraph node_id.") + if not isinstance(failure_type, str): + raise ValueError("Execution failure event requires failure_type.") + return self.handle_failure( + failed_node_id=node_id, + failure_type=failure_type, + active_env_ids=env_ids, + ) + + def replan( + self, + directive: RecoveryDirective, + *, + completed_group_ids: Sequence[str], + recovery_succeeded: bool, + observations: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """Replace only the unfinished suffix after recovery or escalation.""" + if directive.requires_recovery_execution and not recovery_succeeded: + reason = f"{directive.failure_type}:recovery_failed" + else: + reason = f"{directive.failure_type}:state_restored" + replacement = self._replanner( + graph=self.runtime_graph.graph, + completed_group_ids=tuple(completed_group_ids), + failure_type=directive.failure_type, + observations=dict(observations or {}), + ) + return self.runtime_graph.replace_unfinished_suffix( + replacement, + completed_group_ids=completed_group_ids, + reason=reason, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py new file mode 100644 index 000000000..31d26a847 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -0,0 +1,4359 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Closed-loop executor for action-engine execution-program DAGs.""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping, Sequence +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass, field, replace +import logging +from threading import RLock +from typing import Any + +import numpy as np +import torch + +from embodichain.gen_sim.action_engine.config import ( + ArmSelectionPolicyCfg, + RuntimePolicyCfg, + default_runtime_policy, + runtime_policy_hash, +) +from embodichain.gen_sim.action_engine.domain import normalize_placement_relation +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + compile_orientation_constraint, +) +from embodichain.lab.sim.atomic_actions import ( + HeldObjectState, + SceneProvider, + StateDelta, +) +from embodichain.utils import logger as project_logger +from embodichain.utils.logger import log_info, log_warning + +from .actions import AtomicActionAdapter +from .frames import DIRECTIONAL_RELATIONS, robot_frame_axes +from .grounding import ActionGrounder, LiveArrangementPlan, LivePlacementPlan +from .models import ( + ActionOutcome, + ExecutionEdge, + ExecutionProgram, + ExecutionResult, + GroundedAction, + SemanticStep, +) +from .predicates import evaluate_predicate +from .recording import RuntimeRecorder +from .recovery import RuntimeGraph +from .robot_parts import arm_control_part +from .state import ExecutionState + +__all__ = ["ProgramExecutor"] + + +@dataclass +class _Candidate: + feasible: torch.Tensor + cost: torch.Tensor + plans: dict[str, tuple[GroundedAction, ActionOutcome]] + score_components: dict[str, torch.Tensor] = field(default_factory=dict) + warnings: tuple[str, ...] = () + blockers: tuple[dict[str, Any], ...] = () + + +@dataclass +class _EdgeResult: + actions: list[torch.Tensor] + failed: torch.Tensor + grounded: list[GroundedAction] + planner_traces: list[dict[str, Any]] = field(default_factory=list) + executed: torch.Tensor | None = None + + +@dataclass(frozen=True) +class _SupportRelation: + support_uid: str + semantic_step_id: str + + +@dataclass +class _PlacementRecoveryResult: + failed: torch.Tensor + succeeded: torch.Tensor + observed: torch.Tensor + actions: list[torch.Tensor] + failure_events: list[dict[str, Any]] = field(default_factory=list) + covered_failures: torch.Tensor | None = None + + +def _score_arm_candidate( + *, + arm: str, + motion_cost: torch.Tensor, + source_pose: torch.Tensor, + target_pose: torch.Tensor | None, + workspace_center_xy: torch.Tensor, + workspace_half_width: torch.Tensor, + robot_lateral_axis: torch.Tensor, + policy: ArmSelectionPolicyCfg, +) -> dict[str, torch.Tensor]: + """Combine motion length with soft, table-normalized cross-zone costs.""" + arm_sign = 1.0 if arm == "left_arm" else -1.0 + deadband = workspace_half_width * float(policy.crossing_deadband_ratio) + + def crossing(pose: torch.Tensor | None, weight: float) -> torch.Tensor: + if pose is None: + return torch.zeros_like(motion_cost) + lateral = torch.sum( + (pose[:, :2, 3] - workspace_center_xy) * robot_lateral_axis, + dim=1, + ) + wrong_side_depth = torch.clamp( + -arm_sign * lateral - deadband, + min=0.0, + ) + return weight * torch.square(wrong_side_depth / workspace_half_width) + + normalized_motion = motion_cost / float(policy.motion_cost_scale) + pickup_penalty = crossing(source_pose, float(policy.pickup_crossing_weight)) + placement_penalty = crossing( + target_pose, + float(policy.placement_crossing_weight), + ) + return { + "motion_cost": motion_cost, + "normalized_motion_cost": normalized_motion, + "pickup_crossing_penalty": pickup_penalty, + "placement_crossing_penalty": placement_penalty, + "total_cost": normalized_motion + pickup_penalty + placement_penalty, + } + + +_SPECULATIVE_LOG_LOCK = RLock() + + +@contextmanager +def _capture_speculative_warnings() -> Iterator[list[str]]: + """Temporarily capture project warnings without changing its log level.""" + messages: list[str] = [] + collector = logging.Handler(level=logging.WARNING) + collector.emit = lambda record: messages.append(record.getMessage()) + logger = project_logger.logger + with _SPECULATIVE_LOG_LOCK: + handlers = list(logger.handlers) + propagate = logger.propagate + try: + logger.handlers[:] = [collector] + logger.propagate = False + yield messages + finally: + logger.handlers[:] = handlers + logger.propagate = propagate + + +class ProgramExecutor: + """Schedule, ground, plan, execute, and verify one immutable program.""" + + def __init__( + self, + program: ExecutionProgram, + env: Any, + *, + max_transitions: int | None = None, + settle_steps: int | None = None, + record_runtime: bool = True, + record_root: str | None = None, + runtime_policy: RuntimePolicyCfg | None = None, + capability_registry: Any | None = None, + scene_provider: SceneProvider | None = None, + ) -> None: + self.program = program + self.env = env + self.record_runtime = bool(record_runtime) + self.record_root = record_root + if runtime_policy is None: + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + runtime_policy = default_runtime_policy(profile) + if not isinstance(runtime_policy, RuntimePolicyCfg): + raise TypeError("ProgramExecutor runtime_policy must be RuntimePolicyCfg.") + self.runtime_policy = runtime_policy + self.capability_registry = capability_registry + self.env.runtime_policy = runtime_policy + execution = runtime_policy.execution + self.max_transitions = int( + execution["max_transitions"] if max_transitions is None else max_transitions + ) + self.settle_steps = int( + execution["semantic_step_settle_steps"] + if settle_steps is None + else settle_steps + ) + self.max_retries_per_action = int(execution["max_retries_per_action"]) + self.support_stability_samples = int(execution["support_stability_samples"]) + self.support_stability_interval_steps = int( + execution["support_stability_interval_steps"] + ) + self.support_linear_velocity_tolerance = float( + execution["support_linear_velocity_tolerance"] + ) + self.support_angular_velocity_tolerance = float( + execution["support_angular_velocity_tolerance"] + ) + self.placement_recovery_attempts = int( + runtime_policy.grounding["placement"]["recovery_attempts"] + ) + self.runtime_graph = ( + RuntimeGraph( + program.seed_graph, + num_envs=int(env.num_envs), + max_retries=self.max_retries_per_action, + max_revisions=int(execution["max_graph_revisions"]), + max_recovery_actions=int(execution["max_recovery_actions"]), + registry=capability_registry, + ) + if program.seed_graph is not None + else None + ) + self.retry_count = 0 + self.edges = {edge.id: edge for edge in program.edges} + self.steps = {step.id: step for step in program.semantic_steps} + self.step_by_edge = { + edge_id: step + for step in program.semantic_steps + for edge_id in step.edge_ids + } + missing = set(self.edges) - set(self.step_by_edge) + if missing: + raise ValueError( + "Every execution edge must belong to one semantic step; missing " + f"{sorted(missing)}." + ) + self._completion_only_dependencies = self._completion_only_dependency_edges() + self.group_by_step = { + str(step_id): group + for group in program.allocation_groups + for step_id in group.get("semantic_step_ids", ()) + } + arrangement_steps = [ + step + for step in program.semantic_steps + if step.operator in {"arrange_line", "place_in_line"} + ] + arrangement_groups: dict[str, list[SemanticStep]] = {} + for step in arrangement_steps: + arrangement_groups.setdefault(step.parent_step_id, []).append(step) + arrangement_policy = runtime_policy.grounding["arrangement"] + plans = [ + LiveArrangementPlan( + env, + steps, + slot_margin=float(arrangement_policy["slot_margin"]), + minimum_spacing=float(arrangement_policy["minimum_spacing"]), + clearance=float(arrangement_policy["layout_clearance"]), + row_search_step=float(arrangement_policy["row_search_step"]), + row_search_radius=float(arrangement_policy["row_search_radius"]), + ) + for steps in arrangement_groups.values() + ] + self.arrangements = {step.id: plan for plan in plans for step in plan.steps} + # Retain the singular attribute as a convenient introspection hook for + # the common one-arrangement case. + self.arrangement = plans[0] if len(plans) == 1 else None + placement_groups: dict[str, list[SemanticStep]] = {} + for step in program.semantic_steps: + if ( + step.operator == "place_relative" + and step.goal.get("relation") == "inside" + and isinstance(step.goal.get("reference_object"), str) + ): + placement_groups.setdefault( + str(step.goal["reference_object"]), + [], + ).append(step) + placement_plans = [ + LivePlacementPlan( + env, + steps, + clearance=float(runtime_policy.grounding["placement"]["clearance"]), + ) + for steps in placement_groups.values() + if len(steps) > 1 + ] + self.placements = { + step.id: plan for plan in placement_plans for step in plan.steps + } + self.adapter = AtomicActionAdapter( + env, + grasp_policy=runtime_policy.grasp, + planner_policy=runtime_policy.planner, + capability_registry=capability_registry, + scene_provider=scene_provider, + ) + self.grounder = ActionGrounder( + program, + env, + self.adapter.semantics, + self.arrangements, + self.placements, + runtime_policy=runtime_policy, + capability_registry=capability_registry, + ) + self._step_states: dict[tuple[str, str], ExecutionState] = {} + self._object_states: dict[tuple[str, str], ExecutionState] = {} + self._object_owners: dict[str, list[str | None]] = {} + self._arm_owners: dict[str, list[str | None]] = { + "left_arm": [None] * int(env.num_envs), + "right_arm": [None] * int(env.num_envs), + } + self._assignments: dict[str, list[str | None]] = {} + self._candidate_cache: dict[tuple[str, str], _Candidate] = {} + self._candidate_failures: dict[tuple[str, str], str] = {} + self._candidate_diagnostics: dict[str, tuple[str, ...]] = {} + self._candidate_blockers: dict[str, tuple[dict[str, Any], ...]] = {} + self._reported_candidates: set[str] = set() + self._pickup_retry_exclusions: dict[tuple[str, int], set[str]] = {} + self._targets: dict[str, torch.Tensor] = {} + self._target_poses: dict[str, torch.Tensor] = {} + self._orientation_references: dict[str, torch.Tensor] = {} + self._orientation_errors: dict[str, torch.Tensor] = {} + self._policies: dict[str, dict[str, Any]] = {} + self._payload_initial: dict[str, dict[str, torch.Tensor]] = {} + self._support_relations: dict[str, list[_SupportRelation | None]] = {} + self._placement_candidate_history: dict[tuple[str, str], set[int]] = {} + self._robot_lateral_axis_cache: torch.Tensor | None = None + self._transition_count = 0 + self._retry_counts = [0] * int(self.env.num_envs) + + def run( + self, + *, + run_id: str | None = None, + episode_index: int = 0, + ) -> ExecutionResult: + """Execute ready edges until the DAG completes or raises a structural error.""" + self._reset_runtime_state() + recorder = RuntimeRecorder( + self.program, + num_envs=int(self.env.num_envs), + run_id=run_id, + episode_index=episode_index, + output_root=self.record_root, + enabled=self.record_runtime, + runtime_policy=self.runtime_policy.as_mapping(), + runtime_policy_hash=runtime_policy_hash(self.runtime_policy), + ) + aggregate_failed = torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + edge_failures: dict[str, torch.Tensor] = {} + semantic_success: dict[str, torch.Tensor] = {} + failure_events: list[dict[str, Any]] = [] + completed: set[str] = set() + remaining = [edge.id for edge in self.program.edges] + executed_actions: list[torch.Tensor] = [] + error_message = None + try: + while remaining: + ready = [ + self.edges[edge_id] + for edge_id in remaining + if set(self.edges[edge_id].depends_on) <= completed + ] + if not ready: + raise RuntimeError( + "Execution program is deadlocked: no remaining edge is ready." + ) + ready_blocked = { + edge.id: self._dependency_failures(edge, edge_failures) + for edge in ready + } + batch = self._pack_ready_edges( + ready, + inactive=ready_blocked, + completed=completed, + ) + blocked = {edge.id: ready_blocked[edge.id] for edge in batch} + # A synchronized pair needs the same active rows. Execute a + # healthy independent branch separately when its peer is blocked. + if len(batch) == 2 and not torch.equal( + blocked[batch[0].id], blocked[batch[1].id] + ): + batch = (batch[0],) + self._consume_transitions(len(batch)) + + if len(batch) == 2: + posture_before = { + edge.id: self._object_not_fallen(self.step_by_edge[edge.id]) + for edge in batch + } + edge_results, _ = self._execute_parallel_pickups( + batch, + failed=blocked[batch[0].id], + ) + for edge in batch: + result = edge_results[edge.id] + step = self.step_by_edge[edge.id] + active = ~blocked[edge.id] + recorder.edge( + edge.id, + step, + assignments=self._assignments[step.id], + grounded=result.grounded, + active=active, + failed=result.failed, + action_steps=len(result.actions), + planner_traces=getattr(result, "planner_traces", ()), + diagnostics=self._edge_diagnostics( + step, + edge, + result.failed, + ), + ) + failure_events.extend( + self._failure_events( + edge, + step, + result.failed & ~blocked[edge.id], + postcondition=False, + executed=result.executed, + fallen_transition=self._fallen_transition( + step, + posture_before[edge.id], + result, + ), + planner_traces=result.planner_traces, + ) + ) + # Both edge records describe the same synchronized command + # stream. Store it once in the returned execution trace. + executed_actions.extend(edge_results[batch[0].id].actions) + for edge in batch: + edge_failures[edge.id] = edge_results[edge.id].failed.clone() + else: + edge = batch[0] + step = self.step_by_edge[edge.id] + branch_failed = blocked[edge.id] + self._ensure_assignment(step, branch_failed) + active = ~branch_failed + posture_before = self._object_not_fallen(step) + failure_policy = self._edge_failure_policy(edge) + try: + primary_result = self._execute_edge_with_retries( + edge, + step, + failed=branch_failed, + ) + except Exception as exc: + if failure_policy != "best_effort": + raise + primary_result = self._edge_exception_result( + edge, + step, + branch_failed, + exc, + ) + newly_failed = primary_result.failed & ~branch_failed + fallen_transition = self._fallen_transition( + step, + posture_before, + primary_result, + ) + recorder.edge( + edge.id, + step, + assignments=self._assignments[step.id], + grounded=primary_result.grounded, + active=active, + failed=primary_result.failed, + action_steps=len(primary_result.actions), + planner_traces=getattr(primary_result, "planner_traces", ()), + diagnostics=self._edge_diagnostics( + step, + edge, + primary_result.failed, + ), + phase="primary", + ) + edge_result = primary_result + if failure_policy == "task_required": + edge_result = self._recover_object_fallen( + edge, + step, + edge_result, + inherited_failed=branch_failed, + fallen_transition=fallen_transition, + recorder=recorder, + ) + if failure_policy == "best_effort": + # Best-effort parking is observable but cannot invalidate + # an already verified task or safety condition. + next_failed = branch_failed + else: + next_failed = edge_result.failed + executed_actions.extend(edge_result.actions) + edge_failures[edge.id] = next_failed + failure_events.extend( + self._failure_events( + edge, + step, + newly_failed & edge_result.failed, + postcondition=False, + executed=getattr(primary_result, "executed", None), + fallen_transition=fallen_transition, + planner_traces=getattr( + primary_result, "planner_traces", () + ), + ) + ) + + for edge in batch: + completed.add(edge.id) + remaining.remove(edge.id) + step = self.step_by_edge[edge.id] + if edge.id != step.edge_ids[-1]: + continue + prior_failed = edge_failures[edge.id] + verified_failed, step_success, observed = self._verify_step( + step, prior_failed + ) + postcondition_failed = verified_failed & ~prior_failed + recovery_covered = torch.zeros_like(verified_failed) + primary_step_recorded = False + if ( + self.placement_recovery_attempts + and bool(postcondition_failed.any()) + and step.operator == "place_relative" + and normalize_placement_relation( + step.goal.get("relation", "on") + ) + == "on" + ): + recorder.step( + step, + step_success, + observed=observed, + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) + if self.record_runtime + else None + ), + phase="primary", + ) + primary_step_recorded = True + recovery = self._recover_unstable_placement( + step, + postcondition_failed, + recorder=recorder, + ) + executed_actions.extend(recovery.actions) + failure_events.extend(recovery.failure_events) + recovery_covered = ( + torch.zeros_like(verified_failed) + if recovery.covered_failures is None + else recovery.covered_failures + ) + verified_failed = ( + verified_failed & ~postcondition_failed + ) | recovery.failed + step_success |= recovery.succeeded + observed = recovery.observed + failure_events.extend( + self._failure_events( + edge, + step, + verified_failed & ~prior_failed & ~recovery_covered, + postcondition=True, + executed=~prior_failed, + fallen_transition=None, + ) + ) + edge_failures[edge.id] = verified_failed + aggregate_failed |= ~step_success + semantic_success[step.id] = step_success + arrangement = self.arrangements.get(step.id) + if arrangement is not None: + arrangement.mark_completed(step.id, step_success) + if not primary_step_recorded: + recorder.step( + step, + step_success, + observed=observed, + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) + if self.record_runtime + else None + ), + ) + revalidation_failures = self._revalidate_support_relations() + for step_id, lost in revalidation_failures.items(): + step = self.steps[step_id] + edge = self.edges[step.edge_ids[-1]] + aggregate_failed |= lost + semantic_success[step_id] = semantic_success[step_id] & ~lost + edge_failures[edge.id] |= lost + failure_events.extend( + self._failure_events( + edge, + step, + lost, + postcondition=True, + executed=torch.ones_like(lost), + fallen_transition=None, + ) + ) + recorder.step( + step, + semantic_success[step_id], + observed=self._entity_pose(step.object_uid)[:, :3, 3], + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) + if self.record_runtime + else None + ), + phase="final_revalidation", + ) + record_dir = recorder.finalize(~aggregate_failed) + except BaseException as exc: + error_message = f"{type(exc).__name__}: {exc}" + recorder.finalize(~aggregate_failed, error=error_message) + raise + finally: + if error_message is not None: + log_warning(f"Action Engine execution aborted: {error_message}") + + return ExecutionResult( + actions=executed_actions, + success=~aggregate_failed, + semantic_success=semantic_success, + record_dir=record_dir, + retry_count=self.retry_count, + retry_counts=list(self._retry_counts), + recovery_count=( + 0 + if self.runtime_graph is None + else sum( + revision.kind == "insert_recovery" + for revision in self.runtime_graph.revisions + ) + ), + revision_count=( + 0 if self.runtime_graph is None else len(self.runtime_graph.revisions) + ), + failure_events=failure_events, + runtime_revisions=( + [] + if self.runtime_graph is None + else [ + { + "revision": revision.revision, + "kind": revision.kind, + "reason": revision.reason, + "failed_node_id": revision.failed_node_id, + "inserted_group_ids": list(revision.inserted_group_ids), + "replaced_group_ids": list(revision.replaced_group_ids), + "active_env_ids": list(revision.active_env_ids), + } + for revision in self.runtime_graph.revisions + ] + ), + ) + + def _failure_events( + self, + edge: ExecutionEdge, + step: SemanticStep, + failed: torch.Tensor, + *, + postcondition: bool, + executed: torch.Tensor | None, + fallen_transition: torch.Tensor | None, + planner_traces: Sequence[Mapping[str, Any]] = (), + ) -> list[dict[str, Any]]: + if not bool(failed.any()): + return [] + action = edge.actions[-1] + action_name = str(action["atomic_action_class"]) + capability = self.adapter.capabilities.get(action_name) + executed_mask = ( + torch.zeros_like(failed) + if executed is None + else torch.as_tensor( + executed, + dtype=torch.bool, + device=failed.device, + ).reshape(-1) + ) + if executed_mask.shape != failed.shape: + raise ValueError("Failure provenance mask must match failed rows.") + transitioned = ( + torch.zeros_like(failed) + if fallen_transition is None + else torch.as_tensor( + fallen_transition, + dtype=torch.bool, + device=failed.device, + ).reshape(-1) + ) + if transitioned.shape != failed.shape: + raise ValueError("Fallen-transition mask must match failed rows.") + failure_policy = ( + "task_required" if postcondition else self._edge_failure_policy(edge) + ) + fatal = failure_policy != "best_effort" + if postcondition: + classified = (("postcondition_failed", failed),) + else: + fallen = failed & executed_mask & transitioned + planning = failed & ~executed_mask + execution = failed & executed_mask & ~fallen + if capability.failure_classifier == "grasp": + execution_type = "grasp_missed" + elif capability.state_effect in {"preserve_hold", "transfer_hold"}: + execution_type = "object_dropped" + else: + execution_type = "plan_failed" + classified = ( + ("object_fallen", fallen), + ("search_exhausted", planning), + (execution_type, execution), + ) + result: list[dict[str, Any]] = [] + for failure_type, mask in classified: + env_ids = torch.nonzero(mask, as_tuple=False).flatten().tolist() + if not env_ids: + continue + if failure_type == "search_exhausted": + covered: set[int] = set() + for blocker in getattr(self, "_candidate_blockers", {}).get( + step.id, () + ): + env_id = int(blocker["env_id"]) + if env_id not in env_ids: + continue + assignment = self._assignments.get(step.id, [None] * len(failed))[ + env_id + ] + if assignment is not None and blocker.get("arm") != assignment: + continue + blocker_policy = str(blocker.get("failure_policy", failure_policy)) + result.append( + { + "node_id": blocker.get("node_id"), + "edge_id": edge.id, + "origin_edge_id": edge.id, + "blocking_edge_id": blocker["blocking_edge_id"], + "task_instance_id": step.id, + "atomic_action": blocker["atomic_action"], + "object_uid": step.object_uid, + "arm": blocker.get("arm"), + "failure_type": "search_exhausted", + "failure_policy": blocker_policy, + "fatal": blocker_policy != "best_effort", + "planning_stage": blocker["planning_stage"], + "search_strategy": blocker["search_strategy"], + "search_budget": deepcopy(blocker["search_budget"]), + "reason": ( + "Bounded candidate search exhausted without a " + "valid plan; this is not a geometric proof of " + "unreachability." + ), + "evidence": deepcopy(blocker["evidence"]), + "env_ids": [env_id], + } + ) + covered.add(env_id) + for env_id in (item for item in env_ids if item not in covered): + trace = next( + ( + item + for item in planner_traces + if str(item.get("arm", "")) + == str(self._assignments.get(step.id, [None])[env_id]) + ), + planner_traces[0] if planner_traces else {}, + ) + result.append( + { + "node_id": action.get("seed_node_id"), + "edge_id": edge.id, + "blocking_edge_id": edge.id, + "task_instance_id": step.id, + "atomic_action": action_name, + "arm": self._assignments.get(step.id, [None])[env_id], + "failure_type": "search_exhausted", + "failure_policy": failure_policy, + "fatal": fatal, + "planning_stage": "runtime_planning", + **self._planner_failure_details(trace, env_id), + "reason": ( + "Bounded runtime search exhausted without a valid " + "plan; this is not a geometric proof of " + "unreachability." + ), + "env_ids": [env_id], + } + ) + continue + result.append( + { + "node_id": action.get("seed_node_id"), + "edge_id": edge.id, + "blocking_edge_id": edge.id, + "task_instance_id": step.id, + "atomic_action": action_name, + "object_uid": step.object_uid, + "failure_type": failure_type, + "failure_policy": failure_policy, + "fatal": fatal, + "planning_stage": ( + "postcondition" if postcondition else "execution" + ), + "env_ids": env_ids, + } + ) + return result + + def _edge_exception_result( + self, + edge: ExecutionEdge, + step: SemanticStep, + inherited_failed: torch.Tensor, + exc: Exception, + ) -> _EdgeResult: + """Convert a planning exception into an auditable failed edge result.""" + action = edge.actions[0] + assignments = self._assignments.get(step.id, [None] * int(self.env.num_envs)) + arm = next((item for item in assignments if item is not None), None) + trace = { + "action_class": str(action.get("atomic_action_class")), + "arm": arm, + "primary_strategy": "planner_exception", + "primary_success": torch.zeros_like(inherited_failed), + "fallback_attempted": torch.zeros_like(inherited_failed), + "fallback_success": torch.zeros_like(inherited_failed), + "search_budget": self._planner_search_budget(), + "exception": f"{type(exc).__name__}: {exc}", + } + return _EdgeResult( + [], + torch.ones_like(inherited_failed), + [], + [trace], + torch.zeros_like(inherited_failed), + ) + + def _object_not_fallen(self, step: SemanticStep) -> torch.Tensor | None: + """Return the live posture predicate when the object supports it.""" + try: + return evaluate_predicate( + self.env, + {"type": "object_not_fallen", "object": step.object_uid}, + ) + except (TypeError, ValueError): + return None + + def _fallen_transition( + self, + step: SemanticStep, + before: torch.Tensor | None, + result: _EdgeResult, + ) -> torch.Tensor: + """Identify rows where an executed action changed upright to fallen.""" + result_executed = getattr(result, "executed", None) + if before is None or result_executed is None: + return torch.zeros_like(result.failed) + after = self._object_not_fallen(step) + if after is None: + return torch.zeros_like(result.failed) + executed = torch.as_tensor( + result_executed, + dtype=torch.bool, + device=result.failed.device, + ).reshape(-1) + return ( + executed & before.to(result.failed.device) & ~after.to(result.failed.device) + ) + + def _execute_edge_with_retries( + self, + edge: ExecutionEdge, + step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + """Retry a complete AtomicAction with fresh Grounding on failed rows.""" + result = self._execute_edge(edge, step, failed=failed) + if self.runtime_graph is None or len(edge.actions) != 1: + return result + action = edge.actions[0] + node_id = action.get("seed_node_id") + if not isinstance(node_id, str): + return result + aggregate_actions = list(result.actions) + grounded = list(result.grounded) + planner_traces = list(getattr(result, "planner_traces", ())) + executed = ( + torch.zeros_like(result.failed) + if getattr(result, "executed", None) is None + else result.executed.clone() + ) + current_failed = result.failed.clone() + attempted_failure = current_failed & ~failed + while bool(attempted_failure.any()): + precondition = self._retry_precondition(node_id, attempted_failure) + decision = self.runtime_graph.record_failure( + node_id, + attempted_failure, + precondition_holds=precondition, + ) + if not bool(decision.retry.any()): + break + self.retry_count += int(decision.retry.sum()) + for env_id in ( + torch.nonzero(decision.retry, as_tuple=False).flatten().tolist() + ): + self._retry_counts[env_id] += 1 + self._consume_transitions(1) + for arm in ("left_arm", "right_arm"): + self._candidate_cache.pop((step.id, arm), None) + self._candidate_failures.pop((step.id, arm), None) + capability = self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ) + if capability.state_effect == "hold": + previous = list(self._assignments[step.id]) + for env_id in ( + torch.nonzero(decision.retry, as_tuple=False).flatten().tolist() + ): + arm = previous[env_id] + if step.actor.get("mode") == "auto" and arm in { + "left_arm", + "right_arm", + }: + self._pickup_retry_exclusions.setdefault( + (step.id, env_id), set() + ).add(str(arm)) + for arm in ("left_arm", "right_arm"): + self._step_states.pop((step.id, arm), None) + if step.actor.get("mode") == "auto": + self._assignments.pop(step.id, None) + self._ensure_assignment(step, ~decision.retry) + refreshed = self._assignments[step.id] + self._assignments[step.id] = [ + refreshed[index] if bool(decision.retry[index]) else assignment + for index, assignment in enumerate(previous) + ] + retry_result = self._execute_edge( + edge, + step, + failed=~decision.retry, + ) + aggregate_actions.extend(retry_result.actions) + grounded.extend(retry_result.grounded) + planner_traces.extend(getattr(retry_result, "planner_traces", ())) + if getattr(retry_result, "executed", None) is not None: + executed |= retry_result.executed + succeeded = decision.retry & ~retry_result.failed + current_failed &= ~succeeded + attempted_failure = decision.retry & retry_result.failed + return _EdgeResult( + aggregate_actions, + current_failed, + grounded, + planner_traces, + executed, + ) + + def _recover_object_fallen( + self, + edge: ExecutionEdge, + step: SemanticStep, + result: _EdgeResult, + *, + inherited_failed: torch.Tensor, + fallen_transition: torch.Tensor, + recorder: RuntimeRecorder, + ) -> _EdgeResult: + """Run the bounded E2 repair and replay only the failed vector rows.""" + if self.runtime_graph is None or len(edge.actions) != 1: + return result + node_id = edge.actions[0].get("seed_node_id") + if not isinstance(node_id, str) or not node_id: + return result + newly_failed = result.failed & ~inherited_failed + if not bool(newly_failed.any()): + return result + executed = ( + torch.zeros_like(result.failed) + if getattr(result, "executed", None) is None + else torch.as_tensor( + result.executed, + dtype=torch.bool, + device=result.failed.device, + ).reshape(-1) + ) + transition = torch.as_tensor( + fallen_transition, + dtype=torch.bool, + device=result.failed.device, + ).reshape(-1) + if ( + executed.shape != result.failed.shape + or transition.shape != result.failed.shape + ): + raise ValueError("Recovery provenance masks must match failed rows.") + fallen = newly_failed & executed & transition + if not bool(fallen.any()): + return result + + env_ids = torch.nonzero(fallen, as_tuple=False).flatten().tolist() + original_assignment = list( + self._assignments.get(step.id, [None] * int(self.env.num_envs)) + ) + try: + patched = self.runtime_graph.insert_default_recovery( + failed_node_id=node_id, + failure_type="object_fallen", + active_env_ids=env_ids, + resume_failed_group=True, + ) + revision = self.runtime_graph.revisions[-1] + recovery_group_id = revision.inserted_group_ids[0] + from .loader import load_execution_program + + recovery_program = load_execution_program( + patched, + registry=self.capability_registry, + require_executable=True, + ) + recovery_step = next( + item + for item in recovery_program.semantic_steps + if item.id == recovery_group_id + ) + recovery_spec = next( + item + for item in recovery_program.raw["semantic_steps"] + if str(item["id"]) == recovery_group_id + ) + recorder.register_step(recovery_step, recovery_spec) + recovery_edges = { + item.id: item + for item in recovery_program.edges + if item.id in set(recovery_step.edge_ids) + } + if set(recovery_edges) != set(recovery_step.edge_ids): + raise RuntimeError("Compiled recovery group is incomplete.") + except Exception as exc: + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="rejected", + error=f"{type(exc).__name__}: {exc}", + semantic_step_id=step.id, + ) + return result + + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="started", + recovery_group_id=recovery_group_id, + semantic_step_id=step.id, + ) + aggregate_actions = list(result.actions) + grounded = list(result.grounded) + planner_traces = list(getattr(result, "planner_traces", ())) + self._clear_recovery_rows(step, fallen) + + # Recovery edges are compiled from the revised graph but execute through + # this executor so live ownership, recorder, and simulator state remain + # continuous. They are removed from the scheduling maps afterwards. + installed_edge_ids: list[str] = [] + self.steps[recovery_step.id] = recovery_step + for recovery_edge in recovery_edges.values(): + self.edges[recovery_edge.id] = recovery_edge + self.step_by_edge[recovery_edge.id] = recovery_step + installed_edge_ids.append(recovery_edge.id) + recovery_failed = ~fallen + try: + self._assignments.pop(recovery_step.id, None) + self._ensure_assignment(recovery_step, recovery_failed) + for recovery_edge_id in recovery_step.edge_ids: + self._consume_transitions(1) + recovery_edge = recovery_edges[recovery_edge_id] + recovery_result = self._execute_edge_with_retries( + recovery_edge, + recovery_step, + failed=recovery_failed, + ) + recorder.edge( + recovery_edge.id, + recovery_step, + assignments=self._assignments[recovery_step.id], + grounded=recovery_result.grounded, + active=~recovery_failed, + failed=recovery_result.failed, + action_steps=len(recovery_result.actions), + planner_traces=recovery_result.planner_traces, + phase="recovery", + ) + aggregate_actions.extend(recovery_result.actions) + grounded.extend(recovery_result.grounded) + planner_traces.extend(recovery_result.planner_traces) + recovery_failed = recovery_result.failed + _, recovery_success, observed = self._verify_step( + recovery_step, + recovery_failed, + ) + recorder.step( + recovery_step, + recovery_success, + observed=observed, + target=self._targets.get(recovery_step.id), + metadata=( + self._step_runtime_metadata(recovery_step) + if self.record_runtime + else None + ), + phase="recovery", + ) + except Exception as exc: + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="failed", + recovery_group_id=recovery_group_id, + error=f"{type(exc).__name__}: {exc}", + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + result.failed, + grounded, + planner_traces, + result.executed, + ) + finally: + for recovery_edge_id in installed_edge_ids: + self.edges.pop(recovery_edge_id, None) + self.step_by_edge.pop(recovery_edge_id, None) + self.steps.pop(recovery_step.id, None) + + recovered = fallen & recovery_success + if not bool(recovered.any()): + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="failed", + recovery_group_id=recovery_group_id, + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + result.failed, + grounded, + planner_traces, + result.executed, + ) + + # Recompute this TaskGroup's assignment for recovered rows, retaining + # the untouched assignments of healthy vector rows. Replay the prefix + # through the failed edge; the ordinary main loop will then continue at + # the next edge and verify the TaskGroup exactly once. + try: + self._assignments.pop(step.id, None) + for arm in ("left_arm", "right_arm"): + self._candidate_cache.pop((step.id, arm), None) + self._candidate_failures.pop((step.id, arm), None) + self._ensure_assignment(step, ~recovered) + replay_assignment = self._assignments[step.id] + self._assignments[step.id] = [ + ( + replay_assignment[index] + if bool(recovered[index]) + else original_assignment[index] + ) + for index in range(int(self.env.num_envs)) + ] + replay_failed = ~recovered + for prefix_edge_id in step.edge_ids: + self._consume_transitions(1) + prefix_edge = self.edges[prefix_edge_id] + replay_active = ~replay_failed + prefix_result = self._execute_edge_with_retries( + prefix_edge, + step, + failed=replay_failed, + ) + recorder.edge( + prefix_edge.id, + step, + assignments=self._assignments[step.id], + grounded=prefix_result.grounded, + active=replay_active, + failed=prefix_result.failed, + action_steps=len(prefix_result.actions), + planner_traces=prefix_result.planner_traces, + diagnostics=self._edge_diagnostics( + step, + prefix_edge, + prefix_result.failed, + ), + phase="replay", + ) + aggregate_actions.extend(prefix_result.actions) + grounded.extend(prefix_result.grounded) + planner_traces.extend(prefix_result.planner_traces) + replay_failed = prefix_result.failed + if prefix_edge_id == edge.id: + break + except Exception as exc: + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status="failed", + recovery_group_id=recovery_group_id, + error=f"{type(exc).__name__}: {exc}", + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + result.failed, + grounded, + planner_traces, + result.executed, + ) + final_failed = result.failed.clone() + final_failed[fallen] = replay_failed[fallen] + recorder.recovery( + failure_type="object_fallen", + failed_node_id=node_id, + active=fallen, + status=("succeeded" if not bool(final_failed[fallen].any()) else "failed"), + recovery_group_id=recovery_group_id, + semantic_step_id=step.id, + ) + return _EdgeResult( + aggregate_actions, + final_failed, + grounded, + planner_traces, + result.executed, + ) + + def _clear_recovery_rows( + self, + step: SemanticStep, + mask: torch.Tensor, + ) -> None: + """Discard stale hold projections only for rows entering recovery.""" + owners = self._object_owners.setdefault( + step.object_uid, [None] * int(self.env.num_envs) + ) + for env_id in torch.nonzero(mask, as_tuple=False).flatten().tolist(): + owner = owners[env_id] + owners[env_id] = None + if owner in self._arm_owners and ( + self._arm_owners[str(owner)][env_id] == step.object_uid + ): + self._arm_owners[str(owner)][env_id] = None + for arm in ("left_arm", "right_arm"): + if self._arm_owners[arm][env_id] == step.object_uid: + self._arm_owners[arm][env_id] = None + + candidate_keys = [ + key for key in self._object_states if key[0] == step.object_uid + ] + step_keys = [key for key in self._step_states if key[0] == step.id] + for cache, keys in ( + (self._object_states, candidate_keys), + (self._step_states, step_keys), + ): + for key in keys: + state = cache[key] + delta = StateDelta( + held_object_updates={name: None for name in state.held_objects}, + ) + if delta.is_empty: + continue + cache[key] = ExecutionState.from_task_state( + delta.apply(state.to_task_state(), mask), + last_qpos=self.env.robot.get_qpos().clone(), + ) + + def _recover_unstable_placement( + self, + step: SemanticStep, + failed: torch.Tensor, + *, + recorder: RuntimeRecorder, + ) -> _PlacementRecoveryResult: + """Regrasp after release, then retry unused placement poses only.""" + pending = failed.clone() + recovered = torch.zeros_like(failed) + observed = self._entity_pose(step.object_uid)[:, :3, 3] + actions: list[torch.Tensor] = [] + blocking_failures: list[tuple[ExecutionEdge, _EdgeResult, torch.Tensor]] = [] + terminal_edge = self.edges[step.edge_ids[-1]] + failed_node_id = str( + terminal_edge.actions[-1].get("seed_node_id", terminal_edge.id) + ) + recorder.recovery( + failure_type="placement_unstable", + failed_node_id=failed_node_id, + active=failed, + status="started", + semantic_step_id=step.id, + ) + for _attempt in range(self.placement_recovery_attempts): + if not bool(pending.any()): + break + self._consume_transitions(len(step.edge_ids)) + attempt_active = pending.clone() + self._clear_recovery_rows(step, attempt_active) + self._assignments.pop(step.id, None) + for arm in ("left_arm", "right_arm"): + self._candidate_cache.pop((step.id, arm), None) + self._candidate_failures.pop((step.id, arm), None) + try: + self._ensure_assignment(step, ~attempt_active) + except Exception as exc: + blocking_edge = self.edges[step.edge_ids[0]] + blocking_result = self._edge_exception_result( + blocking_edge, + step, + ~attempt_active, + exc, + ) + blocking_failures.append( + (blocking_edge, blocking_result, attempt_active) + ) + recorder.edge( + blocking_edge.id, + step, + assignments=self._assignments.get( + step.id, + [None] * int(self.env.num_envs), + ), + grounded=(), + active=attempt_active, + failed=blocking_result.failed, + action_steps=0, + planner_traces=blocking_result.planner_traces, + diagnostics=self._edge_diagnostics( + step, + blocking_edge, + blocking_result.failed, + ), + phase="recovery", + ) + break + + replay_failed = ~attempt_active + for edge_id in step.edge_ids: + edge = self.edges[edge_id] + edge_active = ~replay_failed + try: + result = self._execute_edge_with_retries( + edge, + step, + failed=replay_failed, + ) + except Exception as exc: + result = self._edge_exception_result( + edge, + step, + replay_failed, + exc, + ) + actions.extend(result.actions) + recorder.edge( + edge.id, + step, + assignments=self._assignments[step.id], + grounded=result.grounded, + active=edge_active, + failed=result.failed, + action_steps=len(result.actions), + planner_traces=result.planner_traces, + diagnostics=self._edge_diagnostics(step, edge, result.failed), + phase="recovery", + ) + newly_failed = edge_active & result.failed + if bool(newly_failed.any()): + blocking_failures.append((edge, result, newly_failed)) + replay_failed = result.failed + if not bool((attempt_active & ~replay_failed).any()): + break + + execution_succeeded = attempt_active & ~replay_failed + if not bool(execution_succeeded.any()): + break + verified_failed, verified_success, observed = self._verify_step( + step, + ~execution_succeeded, + ) + del verified_failed + recovered_now = attempt_active & verified_success + recovered |= recovered_now + recorder.step( + step, + verified_success, + observed=observed, + target=self._targets.get(step.id), + metadata=( + self._step_runtime_metadata(step) if self.record_runtime else None + ), + phase="recovery", + ) + action_failed = attempt_active & replay_failed + pending &= ~recovered_now + if bool(action_failed.any()): + break + + final_failed = failed & ~recovered + recovery_events: list[dict[str, Any]] = [] + covered_failures = torch.zeros_like(failed) + for blocking_edge, blocking_result, blocking_rows in blocking_failures: + event_rows = final_failed & blocking_rows & ~covered_failures + if not bool(event_rows.any()): + continue + events = self._failure_events( + blocking_edge, + step, + event_rows, + postcondition=False, + executed=blocking_result.executed, + fallen_transition=None, + planner_traces=blocking_result.planner_traces, + ) + for event in events: + event["phase"] = "recovery" + event["origin_edge_id"] = terminal_edge.id + recovery_events.extend(events) + covered_failures |= event_rows + recorder.recovery( + failure_type="placement_unstable", + failed_node_id=failed_node_id, + active=failed, + status="failed" if bool(final_failed.any()) else "succeeded", + semantic_step_id=step.id, + ) + return _PlacementRecoveryResult( + failed=final_failed, + succeeded=recovered, + observed=observed, + actions=actions, + failure_events=recovery_events, + covered_failures=covered_failures, + ) + + def _retry_precondition( + self, + node_id: str, + failed: torch.Tensor, + ) -> torch.Tensor: + assert self.runtime_graph is not None + node = next( + item for item in self.runtime_graph.graph["nodes"] if item["id"] == node_id + ) + predicate = node.get("precondition", {}) + if not predicate: + return failed.clone() + try: + return failed & evaluate_predicate( + self.env, + predicate, + held_owners=self._object_owners, + held_states=self._object_states, + coordinated_state=self._step_states.get( + (str(node.get("task_instance_id", "")), "coordinated") + ), + ) + except (TypeError, ValueError): + return torch.zeros_like(failed) + + def _dependency_failures( + self, + edge: ExecutionEdge, + failures: Mapping[str, torch.Tensor], + ) -> torch.Tensor: + """Return success-required failures that can reach this edge.""" + result = torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + for dependency in edge.depends_on: + if (dependency, edge.id) in self._completion_only_dependencies: + continue + result |= failures[dependency] + return result + + def _completion_only_dependency_edges(self) -> frozenset[tuple[str, str]]: + """Resolve linker-added resource ordering to executable edge pairs.""" + graph = self.program.seed_graph + if not isinstance(graph, Mapping): + return frozenset() + metadata = graph.get("metadata", {}) + if not isinstance(metadata, Mapping): + return frozenset() + + reasons_by_pair: dict[tuple[str, str], set[str]] = {} + sources = ( + ("action_contract_task_linker", "linked_dependencies"), + ("action_contract_linker", "group_dependencies"), + ) + for metadata_key, dependency_key in sources: + provenance = metadata.get(metadata_key, {}) + if not isinstance(provenance, Mapping): + continue + dependencies = provenance.get(dependency_key, ()) + if not isinstance(dependencies, Sequence) or isinstance( + dependencies, (str, bytes, bytearray) + ): + continue + for dependency in dependencies: + if not isinstance(dependency, Mapping): + continue + parent = dependency.get("from") + child = dependency.get("to") + reason = dependency.get("reason") + if not all( + isinstance(value, str) and value for value in (parent, child) + ): + continue + if reason not in {"causal", "resource"}: + continue + reasons_by_pair.setdefault((parent, child), set()).add(reason) + + completion_only_steps = { + pair for pair, reasons in reasons_by_pair.items() if reasons == {"resource"} + } + return frozenset( + (dependency, edge.id) + for edge in self.program.edges + for dependency in edge.depends_on + if ( + self.step_by_edge[dependency].id, + self.step_by_edge[edge.id].id, + ) + in completion_only_steps + ) + + def _reset_runtime_state(self) -> None: + self.retry_count = 0 + if self.program.seed_graph is not None: + execution = self.runtime_policy.execution + self.runtime_graph = RuntimeGraph( + self.program.seed_graph, + num_envs=int(self.env.num_envs), + max_retries=self.max_retries_per_action, + max_revisions=int(execution["max_graph_revisions"]), + max_recovery_actions=int(execution["max_recovery_actions"]), + registry=self.capability_registry, + ) + self._step_states.clear() + self._object_states.clear() + self._object_owners.clear() + for owners in self._arm_owners.values(): + owners[:] = [None] * int(self.env.num_envs) + self._assignments.clear() + self._candidate_cache.clear() + self._candidate_failures.clear() + self._candidate_diagnostics.clear() + self._candidate_blockers.clear() + self._reported_candidates.clear() + self._pickup_retry_exclusions.clear() + self._targets.clear() + self._target_poses.clear() + self._orientation_references.clear() + self._orientation_errors.clear() + self._policies.clear() + self._payload_initial.clear() + self._support_relations.clear() + self._placement_candidate_history.clear() + self._robot_lateral_axis_cache = None + self._transition_count = 0 + self._retry_counts = [0] * int(self.env.num_envs) + + def _consume_transitions(self, count: int) -> None: + """Charge ordinary, retry, and recovery edges to one runtime budget.""" + self._transition_count += int(count) + if self._transition_count > self.max_transitions: + raise RuntimeError("Execution exceeded max_transitions.") + + def _pack_ready_edges( + self, + ready: Sequence[ExecutionEdge], + *, + inactive: Mapping[str, torch.Tensor] | None = None, + completed: set[str] | None = None, + ) -> tuple[ExecutionEdge, ...]: + """Prefer progress on held payloads and pack only resource-safe pickups.""" + inactive = inactive or {} + completed = completed or set() + schedulable = [ + edge + for edge in ready + if not self._temporarily_resource_blocked( + edge, + inactive.get(edge.id), + ) + ] + started = [ + edge + for edge in schedulable + if any( + edge_id in completed for edge_id in self.step_by_edge[edge.id].edge_ids + ) + ] + candidates = started or schedulable or list(ready) + first = candidates[0] + if not self._parallel_pickup_candidate(first): + return (first,) + if not self._two_arms_available(inactive.get(first.id)): + return (first,) + first_step = self.step_by_edge[first.id] + for second in candidates[1:]: + if not self._parallel_pickup_candidate(second): + continue + second_step = self.step_by_edge[second.id] + if first_step.object_uid == second_step.object_uid: + continue + shared = set(first.resources) & set(second.resources) + same_group = self.group_by_step.get(first_step.id) is not None and ( + self.group_by_step.get(first_step.id) + is self.group_by_step.get(second_step.id) + ) + if same_group: + # A shared destination workspace constrains transport/place, + # not two independent pickups declared by this group. + conflicts = { + item + for item in shared + if item != "arm:auto" and not item.startswith("workspace:") + } + else: + conflicts = shared - {"arm:auto"} + if conflicts: + continue + required_opposite = ( + first_step.actor.get("mode") == "required" + and second_step.actor.get("mode") == "required" + and first_step.actor.get("arm") != second_step.actor.get("arm") + ) + if same_group or required_opposite: + return first, second + return (first,) + + def _temporarily_resource_blocked( + self, + edge: ExecutionEdge, + inactive: torch.Tensor | None, + ) -> bool: + """Defer a new pickup while its arm is carrying another payload.""" + if not self._parallel_pickup_candidate(edge): + return False + step = self.step_by_edge[edge.id] + mode = str(step.actor.get("mode", "auto")) + inactive_mask = ( + torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + if inactive is None + else inactive + ) + for env_id in range(int(self.env.num_envs)): + if bool(inactive_mask[env_id]): + continue + if mode == "required": + arms = (str(step.actor["arm"]),) + else: + arms = ("left_arm", "right_arm") + if not any( + self._arm_owners[arm][env_id] in {None, step.object_uid} for arm in arms + ): + return True + return False + + def _two_arms_available(self, inactive: torch.Tensor | None) -> bool: + inactive_mask = ( + torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + if inactive is None + else inactive + ) + for env_id in range(int(self.env.num_envs)): + if bool(inactive_mask[env_id]): + continue + free = sum( + self._arm_owners[arm][env_id] is None + for arm in ("left_arm", "right_arm") + ) + if free < 2: + return False + return True + + def _parallel_pickup_candidate(self, edge: ExecutionEdge) -> bool: + if len(edge.actions) != 1: + return False + step = self.step_by_edge[edge.id] + capability = self.adapter.capabilities.get( + str(edge.actions[0].get("atomic_action_class")) + ) + return ( + capability.state_effect == "hold" + and capability.resource_mode == "single_arm_object" + and step.actor.get("mode") in {"auto", "required"} + and step.operator != "orient_object" + ) + + def _preferred_in_place_arm( + self, + step: SemanticStep, + env_id: int, + ) -> str | None: + """Map a clearly sided in-place object to the robot-view arm slot.""" + if step.operator != "orient_object": + return None + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + if initial is None: + entity = self.env.sim.get_rigid_object(step.object_uid) + if entity is None: + return None + initial = entity.get_local_pose(to_matrix=True) + pose = torch.as_tensor(initial, device=self.env.device) + if pose.ndim == 2: + pose = pose.unsqueeze(0) + center, _, lateral_axis = self._arm_selection_workspace(step) + index = min(env_id, pose.shape[0] - 1) + lateral = float( + torch.sum((pose[index, :2, 3] - center[index]) * lateral_axis[index]) + ) + if ( + abs(lateral) + <= self.runtime_policy.arm_selection.orient_object_preferred_arm_deadband + ): + return None + return "left_arm" if lateral > 0.0 else "right_arm" + + def _preferred_live_pickup_arm( + self, + step: SemanticStep, + env_id: int, + ) -> str | None: + """Return the object's same-side arm outside the central deadband.""" + pose = self._entity_pose(step.object_uid) + center, half_width, lateral_axis = self._arm_selection_workspace(step) + index = min(env_id, pose.shape[0] - 1) + lateral = float( + torch.sum((pose[index, :2, 3] - center[index]) * lateral_axis[index]) + ) + deadband = float(half_width[index]) * float( + self.runtime_policy.arm_selection.crossing_deadband_ratio + ) + if abs(lateral) <= deadband: + return None + return "left_arm" if lateral > 0.0 else "right_arm" + + def _auto_arm_is_allowed( + self, + step: SemanticStep, + arm: str, + env_id: int, + ) -> bool: + """Apply the same-side constraint to automatic arm allocation.""" + if step.actor.get("mode") != "auto": + return True + preferred = self._preferred_live_pickup_arm(step, env_id) + if preferred is None or arm == preferred: + return True + excluded = self._pickup_retry_exclusions.get((step.id, env_id), set()) + return bool( + self.runtime_policy.arm_selection.allow_cross_side_fallback + and preferred in excluded + ) + + def _ensure_assignment( + self, + step: SemanticStep, + failed: torch.Tensor, + *, + allow_rematch: bool = True, + ) -> None: + self._capture_orientation_reference(step) + if step.id in self._assignments: + return + if bool(failed.all()): + self._assignments[step.id] = [None] * int(self.env.num_envs) + return + mode = str(step.actor.get("mode", "auto")) + group = self.group_by_step.get(step.id) + if mode == "auto" and group is not None: + self._ensure_serial_group_assignments(group, failed) + if step.id in self._assignments: + return + if mode == "coordinated": + self._assignments[step.id] = [ + ( + None + if bool(failed[index]) + or self._arm_owners["left_arm"][index] is not None + or self._arm_owners["right_arm"][index] is not None + else "coordinated" + ) + for index in range(len(failed)) + ] + return + if mode == "required": + arm = str(step.actor["arm"]) + first_action = self.edges[step.edge_ids[0]].actions[0] + first_capability = self.adapter.capabilities.get( + str(first_action.get("atomic_action_class")) + ) + if step.operator == "handover" and first_capability.state_effect != "hold": + # A coordinated handover has an internal, multi-arm planner. + # Do not let a speculative single-arm suffix plan veto the + # real execution (or create a misleading downstream pickup + # error) once a predecessor already established the transfer + # hold. A standalone E4 starts with PickUp and still needs its + # cached candidate plan for that first action. + source_state = self._state_for(step, arm) + has_hold = ( + source_state.get_held_object(arm_control_part(self.env, arm)) + is not None + ) + self._assignments[step.id] = [ + arm if has_hold and not bool(failed[index]) else None + for index in range(len(failed)) + ] + return + candidate = self._candidate(step, arm, failed) + conflicts = self._resource_conflicts(step, arm) + self._assignments[step.id] = [ + ( + arm + if not bool(failed[index]) and not bool(conflicts[index]) + else None + ) + for index in range(len(failed)) + ] + self._report_candidates(step, (candidate,)) + return + + left = self._candidate(step, "left_arm", failed) + right = self._candidate(step, "right_arm", failed) + candidates = {"left_arm": left, "right_arm": right} + conflicts = { + arm: self._resource_conflicts(step, arm) + for arm in ("left_arm", "right_arm") + } + owners = self._object_owners.get(step.object_uid, [None] * len(failed)) + assignments: list[str | None] = [] + selection_failed = torch.zeros_like(failed) + for env_id in range(len(failed)): + if bool(failed[env_id]): + assignments.append(None) + continue + if owners[env_id] is not None: + owner = str(owners[env_id]) + excluded = self._pickup_retry_exclusions.get((step.id, env_id), set()) + if owner not in excluded and not bool(conflicts[owner][env_id]): + assignments.append(owner) + else: + assignments.append(None) + selection_failed[env_id] = True + continue + excluded = self._pickup_retry_exclusions.get((step.id, env_id), set()) + available = [ + arm + for arm in ("left_arm", "right_arm") + if arm not in excluded + and not bool(conflicts[arm][env_id]) + and self._auto_arm_is_allowed(step, arm, env_id) + ] + if not available: + assignments.append(None) + selection_failed[env_id] = True + continue + preferred = self._preferred_in_place_arm(step, env_id) + feasible = [ + arm for arm in available if bool(candidates[arm].feasible[env_id]) + ] + if preferred in feasible: + assignments.append(preferred) + elif feasible: + assignments.append( + min(feasible, key=lambda arm: float(candidates[arm].cost[env_id])) + ) + else: + live_preferred = self._preferred_live_pickup_arm(step, env_id) + assignments.append( + live_preferred if live_preferred in available else available[0] + ) + + if ( + allow_rematch + and bool(selection_failed.any()) + and step.id in self.arrangements + and step.goal.get("slot_constraint") == "free_reassignable" + and bool(self._rematch_arrangement(step, selection_failed, failed).any()) + ): + self._assignments.pop(step.id, None) + self._ensure_assignment(step, failed, allow_rematch=False) + return + self._assignments[step.id] = assignments + self._report_candidates(step, (left, right)) + + def _ensure_serial_group_assignments( + self, + group: Mapping[str, Any], + failed: torch.Tensor, + ) -> None: + """Bind a distinct-arm pair even when its operators execute serially.""" + step_ids = [str(value) for value in group.get("semantic_step_ids", ())] + if len(step_ids) != 2 or any( + step_id in self._assignments for step_id in step_ids + ): + return + steps = [self.steps[step_id] for step_id in step_ids] + for candidate_step in steps: + self._capture_orientation_reference(candidate_step) + candidates = { + (candidate_step.id, arm): self._candidate(candidate_step, arm, failed) + for candidate_step in steps + for arm in ("left_arm", "right_arm") + } + for candidate_step in steps: + self._report_candidates( + candidate_step, + ( + candidates[(candidate_step.id, "left_arm")], + candidates[(candidate_step.id, "right_arm")], + ), + ) + assignments = { + candidate_step.id: [None] * len(failed) for candidate_step in steps + } + permutations = (("left_arm", "right_arm"), ("right_arm", "left_arm")) + for env_id in range(len(failed)): + if bool(failed[env_id]): + continue + ranked: list[tuple[bool, bool, float, float, str, str]] = [] + for first_arm, second_arm in permutations: + first = candidates[(steps[0].id, first_arm)] + second = candidates[(steps[1].id, second_arm)] + feasible = bool(first.feasible[env_id] and second.feasible[env_id]) + required_match = all( + candidate_step.actor.get("mode") != "required" + or str(candidate_step.actor.get("arm")) == candidate_arm + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + available = ( + required_match + and not bool( + self._resource_conflicts(steps[0], first_arm)[env_id] + or self._resource_conflicts(steps[1], second_arm)[env_id] + ) + and all( + self._auto_arm_is_allowed(candidate_step, candidate_arm, env_id) + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + ) + preferred = ( + self._preferred_in_place_arm(steps[0], env_id), + self._preferred_in_place_arm(steps[1], env_id), + ) + side_penalty = float(first_arm != preferred[0]) if preferred[0] else 0.0 + side_penalty += ( + float(second_arm != preferred[1]) if preferred[1] else 0.0 + ) + ranked.append( + ( + not available, + not feasible, + side_penalty, + float(first.cost[env_id] + second.cost[env_id]), + first_arm, + second_arm, + ) + ) + ranked.sort() + unavailable, _, _, _, first_arm, second_arm = ranked[0] + if unavailable: + continue + assignments[steps[0].id][env_id] = first_arm + assignments[steps[1].id][env_id] = second_arm + self._assignments.update(assignments) + + def _candidate( + self, + step: SemanticStep, + arm: str, + failed: torch.Tensor, + ) -> _Candidate: + """Plan the complete semantic suffix before fixing an arm.""" + if step.actor.get("mode") == "required" and str(step.actor.get("arm")) != arm: + return _Candidate( + feasible=torch.zeros_like(failed), + cost=torch.full( + failed.shape, + torch.inf, + dtype=torch.float32, + device=self.env.device, + ), + plans={}, + ) + cached = self._candidate_cache.get((step.id, arm)) + if cached is not None: + return _Candidate( + feasible=cached.feasible & ~failed, + cost=cached.cost, + plans=cached.plans, + score_components=cached.score_components, + warnings=cached.warnings, + blockers=cached.blockers, + ) + feasible = ~failed.clone() & ~self._resource_conflicts(step, arm) + motion_cost = torch.zeros( + int(self.env.num_envs), + dtype=torch.float32, + device=self.env.device, + ) + source_pose = self._entity_pose(step.object_uid) + target_pose = None + state = self._state_for(step, arm) + reference_eef_pose = None + plans: dict[str, tuple[GroundedAction, ActionOutcome]] = {} + warnings: list[str] = [] + blockers: list[dict[str, Any]] = [] + try: + with _capture_speculative_warnings() as captured: + for edge_id in step.edge_ids: + edge = self.edges[edge_id] + if len(edge.actions) != 1: + raise ValueError( + "Auto/required arm candidates require one action per edge." + ) + action = edge.actions[0] + capability = self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ) + if ( + step.operator == "handover" + and capability.resource_mode == "coordinated_object" + ): + # A standalone E4 needs a speculative PickUp/staging + # prefix to choose and cache its transfer arm. The + # actual HandOver is coordinated, however, and must + # only be planned from the live post-staging state. + break + failure_policy = self._edge_failure_policy(edge) + try: + if capability.state_effect == "hold": + grounded = self.grounder.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=self._orientation_references.get( + step.id + ), + ) + grounded = self._with_downstream_targets( + step, edge_id, arm, state, grounded + ) + outcome = self.adapter.plan(grounded, state) + else: + grounded, outcome = self._ground_and_plan_candidates( + action, + step, + arm=arm, + state=state, + active=feasible & ~failed, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=self._orientation_references.get( + step.id + ), + ) + except Exception as exc: + if failure_policy != "best_effort": + blockers.extend( + self._candidate_exception_blockers( + step, + edge, + arm, + failed, + exc, + ) + ) + raise + warnings.append( + f"{arm} best-effort action could not be planned at " + f"{edge_id} ({capability.name}): " + f"{type(exc).__name__}: {exc}" + ) + continue + plans[edge_id] = (grounded, outcome) + if failure_policy != "best_effort": + feasible &= outcome.success + motion_cost += outcome.cost + blockers.extend( + self._candidate_outcome_blockers( + step, + edge, + arm, + failed, + outcome, + ) + ) + elif not bool(outcome.success.all()): + warnings.append( + f"{arm} best-effort action degraded at {edge_id} " + f"({capability.name}); required suffix remains feasible." + ) + state = outcome.next_state + target = outcome.grounded.target_object_pose + if isinstance(target, torch.Tensor): + reference_eef_pose = self._eef_target(outcome) + binding = edge.actions[0].get("target_binding", {}) + if ( + binding.get("kind") + in { + "semantic_goal", + "coordinated_goal", + } + and binding.get("phase", "final") != "staging" + ): + target_pose = target + if not bool((feasible & ~failed).any()): + target = getattr(grounded.target, "xpos", None) + target_detail = "" + if isinstance(target, torch.Tensor) and target.shape[-2:] == ( + 4, + 4, + ): + target_z = target[..., 2, 3] + target_detail = ( + f" target_z=[{float(target_z.min()):.3f}, " + f"{float(target_z.max()):.3f}]" + ) + warnings.append( + f"{arm} candidate became infeasible at {edge_id} " + f"({capability.name}).{target_detail}" + ) + break + warnings.extend(captured) + except Exception as exc: + self._candidate_failures[(step.id, arm)] = f"{type(exc).__name__}: {exc}" + feasible = torch.zeros_like(failed) + motion_cost[:] = torch.inf + center_xy, half_width, lateral_axis = self._arm_selection_workspace(step) + score_components = _score_arm_candidate( + arm=arm, + motion_cost=motion_cost, + source_pose=source_pose, + target_pose=target_pose, + workspace_center_xy=center_xy, + workspace_half_width=half_width, + robot_lateral_axis=lateral_axis, + policy=self.runtime_policy.arm_selection, + ) + cost = score_components["total_cost"] + candidate = _Candidate( + feasible=feasible, + cost=cost, + plans=plans, + score_components=score_components, + warnings=tuple(warnings), + blockers=tuple(blockers), + ) + self._candidate_cache[(step.id, arm)] = candidate + return _Candidate( + feasible=feasible & ~failed, + cost=cost, + plans=plans, + score_components=score_components, + warnings=tuple(warnings), + blockers=tuple(blockers), + ) + + def _ground_and_plan_candidates( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + active: torch.Tensor, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[GroundedAction, ActionOutcome]: + """Plan live grounding candidates and retain the best bounded attempt.""" + groundings = self.grounder.ground_candidates( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + used = self._placement_candidate_history.get((step.id, arm), set()) + selected: tuple[GroundedAction, ActionOutcome] | None = None + selected_rank: tuple[int, float, int] | None = None + attempts: list[dict[str, Any]] = [] + last_error: Exception | None = None + for ordinal, grounded in enumerate(groundings): + candidate_index = int( + grounded.motion_policy.get("placement_candidate_index", ordinal) + ) + is_placement = "placement_candidate_index" in grounded.motion_policy + if is_placement and candidate_index in used: + attempts.append( + { + "candidate_index": candidate_index, + "status": "previously_released", + } + ) + continue + try: + outcome = self.adapter.plan(grounded, state) + except Exception as exc: + last_error = exc + attempts.append( + { + "candidate_index": candidate_index, + "status": "planning_error", + "error": f"{type(exc).__name__}: {exc}", + } + ) + continue + failed_count = int((active & ~outcome.success).sum()) + active_cost = ( + float(outcome.cost[active].sum()) if bool(active.any()) else 0.0 + ) + rank = (failed_count, active_cost, candidate_index) + attempts.append( + { + "candidate_index": candidate_index, + "status": "planned", + "failed_rows": failed_count, + "cost": active_cost, + } + ) + if selected is None or rank < selected_rank: + selected = (grounded, outcome) + selected_rank = rank + if failed_count == 0: + break + if selected is None: + if last_error is not None: + detail = f"{type(last_error).__name__}: {last_error}" + if last_error.__cause__ is not None: + cause = last_error.__cause__ + detail += f"; caused by {type(cause).__name__}: {cause}" + raise RuntimeError( + "All grounding candidates raised during planning. " + f"Last error: {detail}" + ) from last_error + raise RuntimeError("No unused grounding candidate remains.") + grounded, outcome = selected + outcome = replace( + outcome, + planner_trace={ + **outcome.planner_trace, + "grounding_candidates": attempts, + "selected_grounding_candidate": int( + grounded.motion_policy.get("placement_candidate_index", 0) + ), + }, + ) + return grounded, outcome + + def _arm_selection_workspace( + self, + step: SemanticStep, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return workspace geometry along the robot's live lateral axis.""" + lateral_axis = self._robot_view_lateral_axis() + arrangement = self.arrangements.get(step.id) + if arrangement is not None: + minimum = arrangement.table_bounds[:, 0, :2] + maximum = arrangement.table_bounds[:, 1, :2] + center = (minimum + maximum) * 0.5 + half_extents = (maximum - minimum) * 0.5 + half_width = torch.sum(torch.abs(lateral_axis) * half_extents, dim=1) + return center, half_width, lateral_axis + count = int(self.env.num_envs) + centers = torch.zeros((count, 2), dtype=torch.float32, device=self.env.device) + half_widths = torch.full( + (count,), + float(self.runtime_policy.arm_selection.fallback_workspace_half_width), + dtype=torch.float32, + device=self.env.device, + ) + table = self.env.sim.get_rigid_object("table") + if table is None or not hasattr(table, "get_vertices"): + return centers, half_widths, lateral_axis + table_pose = self._entity_pose("table") + for env_id in range(count): + value = table.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (list, tuple)): + value = value[0] + vertices = torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3: + continue + world = ( + vertices @ table_pose[env_id, :3, :3].transpose(0, 1) + + table_pose[env_id, :3, 3] + ) + minimum = world[:, :2].min(dim=0).values + maximum = world[:, :2].max(dim=0).values + center = (minimum + maximum) * 0.5 + lateral = torch.sum((world[:, :2] - center) * lateral_axis[env_id], dim=1) + half_width = torch.max(torch.abs(lateral)) + if float(half_width) > 1.0e-6: + centers[env_id] = center + half_widths[env_id] = half_width + return centers, half_widths, lateral_axis + + def _robot_view_lateral_axis(self) -> torch.Tensor: + """Return the normalized world-space axis pointing right-arm to left-arm.""" + if self._robot_lateral_axis_cache is not None: + return self._robot_lateral_axis_cache + _, self._robot_lateral_axis_cache = robot_frame_axes(self.env) + return self._robot_lateral_axis_cache + + def _report_candidates( + self, + step: SemanticStep, + candidates: Sequence[_Candidate], + ) -> None: + if step.id in self._reported_candidates: + return + warning_count = sum(len(item.warnings) for item in candidates) + failures = [ + message + for (step_id, _), message in self._candidate_failures.items() + if step_id == step.id + ] + diagnostics = tuple( + dict.fromkeys(message for item in candidates for message in item.warnings) + ) + tuple(dict.fromkeys(failures)) + diagnostics = tuple(dict.fromkeys(diagnostics)) + if diagnostics: + self._candidate_diagnostics[step.id] = diagnostics + blockers = tuple( + deepcopy(item) + for candidate in candidates + for item in getattr(candidate, "blockers", ()) + ) + if blockers: + self._candidate_blockers[step.id] = blockers + if warning_count or failures: + feasible = ", ".join( + f"{int(item.feasible.sum())}/{len(item.feasible)}" + for item in candidates + ) + log_info( + f"Speculative arm candidates for {step.id}: feasible=[{feasible}], " + f"suppressed_warnings={warning_count}, exceptions={len(failures)}." + ) + edge_failures = tuple( + message + for message in diagnostics + if "candidate became infeasible" in message + ) + prioritized = tuple( + dict.fromkeys((*failures, *edge_failures, *diagnostics)) + ) + for message in prioritized[:3]: + log_warning(f"Candidate planning for {step.id}: {message}") + self._reported_candidates.add(step.id) + + def _candidate_outcome_blockers( + self, + step: SemanticStep, + edge: ExecutionEdge, + arm: str, + inherited_failed: torch.Tensor, + outcome: ActionOutcome, + ) -> list[dict[str, Any]]: + """Capture the real suffix edge that exhausted bounded planning.""" + failed = ~outcome.success & ~inherited_failed + action = edge.actions[0] + return [ + { + "env_id": int(env_id), + "node_id": action.get("seed_node_id"), + "blocking_edge_id": edge.id, + "atomic_action": str(action.get("atomic_action_class")), + "arm": arm, + "failure_policy": self._edge_failure_policy(edge), + "planning_stage": "candidate_suffix", + **self._planner_failure_details(outcome.planner_trace, env_id), + } + for env_id in torch.nonzero(failed, as_tuple=False).flatten().tolist() + ] + + def _candidate_exception_blockers( + self, + step: SemanticStep, + edge: ExecutionEdge, + arm: str, + inherited_failed: torch.Tensor, + exc: Exception, + ) -> list[dict[str, Any]]: + """Record a bounded candidate-planning exception without claiming proof.""" + del step + action = edge.actions[0] + budget = self._planner_search_budget() + return [ + { + "env_id": int(env_id), + "node_id": action.get("seed_node_id"), + "blocking_edge_id": edge.id, + "atomic_action": str(action.get("atomic_action_class")), + "arm": arm, + "failure_policy": self._edge_failure_policy(edge), + "planning_stage": "candidate_suffix", + "search_strategy": "planner_exception", + "search_budget": budget, + "evidence": {"exception": f"{type(exc).__name__}: {exc}"}, + } + for env_id in torch.nonzero(~inherited_failed, as_tuple=False) + .flatten() + .tolist() + ] + + def _planner_search_budget(self) -> dict[str, Any]: + """Return the configured finite search budget used by motion planning.""" + runtime_policy = getattr(self, "runtime_policy", None) + planner = getattr(runtime_policy, "planner", {}) + curobo = planner.get("curobo", {}) if isinstance(planner, Mapping) else {} + return { + "primary_max_attempts": int(curobo.get("max_attempts", 1)), + "fallback_enabled": bool(planner.get("allow_fallback", False)), + } + + def _planner_failure_details( + self, + trace: Mapping[str, Any], + env_id: int, + ) -> dict[str, Any]: + """Extract compact row-local evidence from one planner trace.""" + reachability = trace.get("reachability_search") + reachability = reachability if isinstance(reachability, Mapping) else {} + strategy = str( + reachability.get("strategy") or trace.get("primary_strategy") or "unknown" + ) + budget = deepcopy( + dict(trace.get("search_budget", self._planner_search_budget())) + ) + attempts = reachability.get("attempts", ()) + evidence: dict[str, Any] = { + "primary_success": bool( + self._row_trace_value(trace.get("primary_success", False), env_id) + ), + "fallback_attempted": bool( + self._row_trace_value(trace.get("fallback_attempted", False), env_id) + ), + "fallback_success": bool( + self._row_trace_value(trace.get("fallback_success", False), env_id) + ), + } + if trace.get("exception") is not None: + evidence["exception"] = str(trace["exception"]) + if isinstance(attempts, Sequence) and not isinstance( + attempts, (str, bytes, bytearray) + ): + evidence["reachability_attempts"] = [ + { + "candidate": str(item.get("candidate", "")), + "target_z": self._row_trace_value(item.get("target_z"), env_id), + "success": bool( + self._row_trace_value(item.get("success", False), env_id) + ), + } + for item in attempts + if isinstance(item, Mapping) + ] + budget["reachability_candidate_count"] = len( + evidence["reachability_attempts"] + ) + return { + "search_strategy": strategy, + "search_budget": budget, + "evidence": evidence, + } + + @staticmethod + def _row_trace_value(value: Any, env_id: int) -> Any: + """Detach one environment row from JSON-like or tensor trace data.""" + if isinstance(value, torch.Tensor): + detached = value.detach().cpu() + if detached.ndim == 0: + return detached.item() + row = detached[min(env_id, detached.shape[0] - 1)] + return row.item() if row.ndim == 0 else row.tolist() + if isinstance(value, Sequence) and not isinstance( + value, (str, bytes, bytearray) + ): + if not value: + return None + return deepcopy(value[min(env_id, len(value) - 1)]) + return deepcopy(value) + + def _edge_diagnostics( + self, + step: SemanticStep, + edge: ExecutionEdge, + failed: torch.Tensor, + ) -> tuple[str, ...]: + if edge.id != step.edge_ids[0] or not bool(failed.any()): + return () + return self._candidate_diagnostics.get(step.id, ()) + + def _with_downstream_targets( + self, + step: SemanticStep, + pickup_edge_id: str, + arm: str, + state: ExecutionState, + grounded: GroundedAction, + ) -> GroundedAction: + """Screen grasp poses against every later held-object target. + + A handover is split across semantic steps: its staging ``MoveHeldObject`` + edge is not part of the pickup step's local edge suffix. Include that + first exchange pose here so ``PickUp`` can reject a grasp whose + ``object_to_eef`` transform makes the later transfer arm unreachable. + This keeps the screening speculative and bounded; no simulator steps + are sent while a candidate is being built. + """ + targets: list[torch.Tensor] = [] + start = step.edge_ids.index(pickup_edge_id) + 1 + for edge_id in step.edge_ids[start:]: + edge = self.edges[edge_id] + if len(edge.actions) != 1: + continue + action = edge.actions[0] + if ( + self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ).target_materializer + != "semantic_held_object" + ): + continue + future = self.grounder.ground( + action, + step, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + if future.target_object_pose is not None: + targets.append(future.target_object_pose) + targets.extend(self._handover_successor_targets(step, arm, state)) + if not targets: + return grounded + existing = tuple(grounded.cfg.get("downstream_object_target_poses", ())) + return replace( + grounded, + cfg={ + **grounded.cfg, + "downstream_object_target_poses": existing + tuple(targets), + }, + ) + + def _handover_successor_targets( + self, + step: SemanticStep, + arm: str, + state: ExecutionState, + ) -> list[torch.Tensor]: + """Return staging poses for handovers downstream of a pickup. + + ``SemanticStep.depends_on`` contains semantic IDs rather than edge IDs, + so walk the small dependency graph instead of assuming the handover is + an immediate child. Only a handover that transfers this object from + the selected pickup arm is relevant to the grasp screen. + """ + reachable = {step.id} + changed = True + while changed: + changed = False + for candidate in self.steps.values(): + if candidate.id in reachable: + continue + if any(dependency in reachable for dependency in candidate.depends_on): + reachable.add(candidate.id) + changed = True + + targets: list[torch.Tensor] = [] + for successor in self.steps.values(): + if ( + successor.id not in reachable + or successor.id == step.id + or successor.operator != "handover" + or successor.object_uid != step.object_uid + ): + continue + for edge_id in successor.edge_ids: + edge = self.edges[edge_id] + if len(edge.actions) != 1: + continue + action = edge.actions[0] + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + continue + if binding.get("kind") != "handover_staging": + continue + transfer_arm = str( + binding.get( + "transfer_arm", + successor.goal.get("transfer_arm", ""), + ) + ) + if transfer_arm != arm: + break + try: + grounded = self.grounder.ground( + action, + successor, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get( + successor.id, + self._orientation_references.get(step.id), + ), + ) + except (AttributeError, KeyError, ValueError): + # A malformed/incomplete successor must not make an + # otherwise valid pickup candidate disappear. The normal + # successor execution will report that grounding error. + break + if grounded.target_object_pose is not None: + targets.append(grounded.target_object_pose) + break + return targets + + def _eef_target(self, outcome: ActionOutcome) -> torch.Tensor | None: + state = outcome.next_state + held_object = state.get_held_object( + arm_control_part(self.env, outcome.grounded.arm) + ) + object_target = outcome.grounded.target_object_pose + if object_target is not None and held_object is not None: + object_to_eef = held_object.object_to_eef.to( + device=object_target.device, + dtype=object_target.dtype, + ) + return torch.bmm(object_target, object_to_eef) + if held_object is not None: + return held_object.grasp_xpos + target = outcome.grounded.target + return getattr(target, "xpos", None) + + def _state_for(self, step: SemanticStep, arm: str) -> ExecutionState: + """Refresh qpos while retaining holds across TaskGroup boundaries.""" + cached = self._step_states.get((step.id, arm)) + if cached is None: + cached = self._object_states.get((step.object_uid, arm)) + live_qpos = self.env.robot.get_qpos().clone() + if cached is None: + return ExecutionState(last_qpos=live_qpos) + return cached.with_updates(last_qpos=live_qpos) + + def _resource_conflicts( + self, + step: SemanticStep, + arm: str, + ) -> torch.Tensor: + object_owners = self._object_owners.get( + step.object_uid, [None] * int(self.env.num_envs) + ) + arm_owners = self._arm_owners[arm] + return torch.tensor( + [ + (object_owner not in {None, arm}) + or (arm_owner not in {None, step.object_uid}) + for object_owner, arm_owner in zip(object_owners, arm_owners) + ], + dtype=torch.bool, + device=self.env.device, + ) + + def _update_ownership( + self, + step: SemanticStep, + arm: str, + action_class: str, + state: ExecutionState, + successful: torch.Tensor, + ) -> None: + capability = self.adapter.capabilities.get(action_class) + owners = self._object_owners.setdefault( + step.object_uid, [None] * int(self.env.num_envs) + ) + if capability.state_effect == "release": + for env_id in torch.nonzero(successful, as_tuple=False).flatten().tolist(): + if owners[env_id] == arm: + owners[env_id] = None + if self._arm_owners[arm][env_id] == step.object_uid: + self._arm_owners[arm][env_id] = None + if arm not in owners: + self._object_states.pop((step.object_uid, arm), None) + return + held_object = state.get_held_object(arm_control_part(self.env, arm)) + if held_object is None or not bool(successful.any()): + return + self._object_states[(step.object_uid, arm)] = state + if capability.state_effect == "hold": + self._clear_support_relation(step.object_uid, successful) + for env_id in torch.nonzero(successful, as_tuple=False).flatten().tolist(): + owners[env_id] = arm + self._arm_owners[arm][env_id] = step.object_uid + + def _rematch_arrangement( + self, + trigger_step: SemanticStep, + trigger: torch.Tensor, + failed: torch.Tensor, + ) -> torch.Tensor: + """Globally rematch unfinished objects to feasible free slots.""" + from scipy.optimize import linear_sum_assignment + + arrangement = self.arrangements[trigger_step.id] + changed = torch.zeros_like(trigger) + for env_id in torch.nonzero(trigger, as_tuple=False).flatten().tolist(): + step_ids = arrangement.remaining(env_id) + slots = arrangement.available_slots(env_id) + if len(step_ids) != len(slots): + continue + original = { + step_id: int(arrangement.assignments[step_id][env_id]) + for step_id in step_ids + } + costs = np.full((len(step_ids), len(slots)), np.inf, dtype=np.float64) + isolate = torch.ones_like(failed) + isolate[env_id] = failed[env_id] + for row, step_id in enumerate(step_ids): + for column, slot_id in enumerate(slots): + arrangement.assignments[step_id][env_id] = slot_id + self._candidate_cache.pop((step_id, "left_arm"), None) + self._candidate_cache.pop((step_id, "right_arm"), None) + arm_costs = [] + for arm in ("left_arm", "right_arm"): + candidate = self._candidate(self.steps[step_id], arm, isolate) + if bool(candidate.feasible[env_id]): + arm_costs.append(float(candidate.cost[env_id])) + if arm_costs: + costs[row, column] = min(arm_costs) + arrangement.assignments[step_id][env_id] = original[step_id] + self._candidate_cache.pop((step_id, "left_arm"), None) + self._candidate_cache.pop((step_id, "right_arm"), None) + if not np.isfinite(costs).any(axis=1).all(): + continue + rows, columns = linear_sum_assignment( + np.where(np.isfinite(costs), costs, 1.0e12) + ) + if not np.isfinite(costs[rows, columns]).all(): + continue + arrangement.assign( + env_id, + { + step_ids[int(row)]: slots[int(column)] + for row, column in zip(rows, columns) + }, + ) + changed[env_id] = True + return changed + + def _execute_edge( + self, + edge: ExecutionEdge, + step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + if step.goal.get("payloads"): + # Capture before the first physical action, including an ordinary + # single-arm pickup. Verification then measures whether every + # direct payload stayed fixed relative to its carrier. + self._capture_payloads(step) + if ( + len(edge.actions) == 1 + and self.adapter.capabilities.get( + str(edge.actions[0].get("atomic_action_class")) + ).resource_mode + == "coordinated_object" + ): + return self._execute_coordinated(edge, step, failed) + if len(edge.actions) == 2: + return self._execute_explicit_dual(edge, step, failed) + if len(edge.actions) != 1: + raise ValueError( + f"Edge {edge.id!r} must contain one action or an explicit dual pair." + ) + assignments = self._assignments[step.id] + outcomes: dict[str, ActionOutcome | None] = { + "left_arm": None, + "right_arm": None, + } + masks = { + arm: torch.tensor( + [assignment == arm for assignment in assignments], + dtype=torch.bool, + device=self.env.device, + ) + & ~failed + for arm in outcomes + } + grounded_items: list[GroundedAction] = [] + planner_traces: list[dict[str, Any]] = [] + planning_failed = torch.zeros_like(failed) + action_class = str(edge.actions[0]["atomic_action_class"]) + capability = self.adapter.capabilities.get(action_class) + for arm in outcomes: + if not bool(masks[arm].any()): + continue + state = self._state_for(step, arm) + if capability.state_effect == "hold": + try: + grounded, outcome = self._plan_live_hold(edge, step, arm) + except Exception as exc: + planning_failed |= masks[arm] + planner_traces.append( + self._live_hold_failure_trace(edge, step, arm, exc) + ) + continue + else: + # Re-ground transport and placement from live simulator state. + grounded, outcome = self._ground_and_plan_candidates( + edge.actions[0], + step, + arm=arm, + state=state, + active=masks[arm], + orientation_reference_pose=self._orientation_references.get( + step.id + ), + ) + grounded = outcome.grounded + outcomes[arm] = outcome + grounded_items.append(grounded) + planner_traces.append(outcome.planner_trace) + self._remember_target(step, grounded) + placement_index = grounded.motion_policy.get("placement_candidate_index") + if placement_index is not None and bool( + (masks[arm] & outcome.success).any() + ): + self._placement_candidate_history.setdefault((step.id, arm), set()).add( + int(placement_index) + ) + assigned = masks["left_arm"] | masks["right_arm"] + if not grounded_items: + return _EdgeResult( + [], + failed | (~failed & ~assigned) | planning_failed, + [], + planner_traces, + executed=torch.zeros_like(failed), + ) + trajectory, action_success = self.adapter.combine(outcomes, masks) + active = assigned & action_success & ~failed & ~planning_failed + actions = self.adapter.execute_trajectory(trajectory, active=active) + physical_failed = torch.zeros_like(failed) + for arm, outcome in outcomes.items(): + if outcome is not None: + successful = masks[arm] & outcome.success & active + if capability.state_effect == "hold": + physical = self._physical_pickup( + step.object_uid, arm, outcome.next_state, successful + ) + physical_failed |= successful & ~physical + successful = physical + elif capability.state_effect == "preserve_hold": + physical = self._physical_hold( + step.object_uid, arm, outcome.next_state, successful + ) + lost = successful & ~physical + physical_failed |= lost + self._release_ownership(step.object_uid, arm, lost) + successful = physical + if capability.verifier_hook is not None: + verified = torch.as_tensor( + capability.verifier_hook( + executor=self, + step=step, + arm=arm, + outcome=outcome, + attempted=successful, + ), + dtype=torch.bool, + device=self.env.device, + ).reshape(-1) + if verified.numel() != int(self.env.num_envs): + raise ValueError( + f"AtomicAction {action_class!r} verifier returned " + "an invalid vectorized mask." + ) + physical_failed |= successful & ~verified + successful &= verified + committed_state = outcome.state_after(successful) + if capability.state_effect in {"hold", "preserve_hold"}: + committed_state = self._rebase_held_state( + step.object_uid, + arm, + committed_state, + successful, + from_planned_qpos=capability.state_effect == "preserve_hold", + ) + self._step_states[(step.id, arm)] = committed_state + self._update_ownership( + step, + arm, + action_class, + committed_state, + successful, + ) + edge_failed = ( + failed + | (~failed & ~assigned) + | (assigned & ~action_success) + | planning_failed + | physical_failed + ) + return _EdgeResult( + actions, + edge_failed, + grounded_items, + planner_traces, + active, + ) + + def _plan_live_hold( + self, + edge: ExecutionEdge, + step: SemanticStep, + arm: str, + ) -> tuple[GroundedAction, ActionOutcome]: + """Replace a speculative hold plan with one grounded at execution time.""" + candidate = self._candidate_cache.get((step.id, arm)) + cached_plan_available = candidate is not None and edge.id in candidate.plans + update_obj_info = getattr(self.env, "update_obj_info", None) + if callable(update_obj_info): + update_obj_info() + object_pose = self._entity_pose(step.object_uid).detach().clone() + state = self._state_for(step, arm) + grounded = self.grounder.ground( + edge.actions[0], + step, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + grounded = self._with_downstream_targets(step, edge.id, arm, state, grounded) + outcome = self.adapter.plan(grounded, state) + return grounded, replace( + outcome, + planner_trace={ + **outcome.planner_trace, + "execution_replanned_from_live_state": True, + "speculative_candidate_available": cached_plan_available, + "speculative_candidate_replaced": cached_plan_available, + "execution_object_pose": object_pose, + }, + ) + + def _live_hold_failure_trace( + self, + edge: ExecutionEdge, + step: SemanticStep, + arm: str, + exc: Exception, + ) -> dict[str, Any]: + """Describe a live PickUp planning exception without aborting the task.""" + candidate = self._candidate_cache.get((step.id, arm)) + return { + "action_class": str(edge.actions[0].get("atomic_action_class")), + "arm": arm, + "primary_strategy": "live_pickup_replan", + "primary_success": torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ), + "execution_replanned_from_live_state": True, + "speculative_candidate_available": ( + candidate is not None and edge.id in candidate.plans + ), + "speculative_candidate_replaced": False, + "execution_object_pose": self._entity_pose(step.object_uid) + .detach() + .clone(), + "exception": f"{type(exc).__name__}: {exc}", + } + + def _physical_pickup( + self, + uid: str, + arm: str, + state: ExecutionState, + attempted: torch.Tensor, + ) -> torch.Tensor: + owners = list(self._object_owners.get(uid, [None] * int(self.env.num_envs))) + for env_id in torch.nonzero(attempted, as_tuple=False).flatten().tolist(): + owners[env_id] = arm + states = dict(self._object_states) + states[(uid, arm)] = state + physical = attempted & evaluate_predicate( + self.env, + { + "type": "object_held", + "object": uid, + "position_tolerance": self.runtime_policy.predicate_fallbacks[ + "held_position_tolerance" + ], + }, + held_owners={**self._object_owners, uid: owners}, + held_states=states, + ) + return physical + + def _physical_hold( + self, + uid: str, + arm: str, + state: ExecutionState, + attempted: torch.Tensor, + *, + owners: Mapping[str, Sequence[str | None]] | None = None, + states: Mapping[tuple[str, str], ExecutionState] | None = None, + position_tolerance: float | None = None, + ) -> torch.Tensor: + candidate_states = dict(self._object_states if states is None else states) + candidate_states[(uid, arm)] = state + return attempted & evaluate_predicate( + self.env, + { + "type": "object_held", + "object": uid, + "position_tolerance": ( + self.runtime_policy.predicate_fallbacks["held_position_tolerance"] + if position_tolerance is None + else float(position_tolerance) + ), + "arm": arm, + }, + held_owners=self._object_owners if owners is None else owners, + held_states=candidate_states, + ) + + def _release_ownership( + self, + uid: str, + arm: str, + lost: torch.Tensor, + ) -> None: + owners = self._object_owners.get(uid) + if owners is None: + return + for env_id in torch.nonzero(lost, as_tuple=False).flatten().tolist(): + if owners[env_id] == arm: + owners[env_id] = None + if self._arm_owners[arm][env_id] == uid: + self._arm_owners[arm][env_id] = None + if arm not in owners: + self._object_states.pop((uid, arm), None) + + def _execute_coordinated( + self, + edge: ExecutionEdge, + step: SemanticStep, + failed: torch.Tensor, + ) -> _EdgeResult: + action = edge.actions[0] + action_name = str(action["atomic_action_class"]) + capability = self.adapter.capabilities.get(action_name) + binding = action.get("target_binding", {}) + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + accepted_assignments = ( + {"coordinated", transfer_arm} + if capability.state_effect == "transfer_hold" + else {"coordinated"} + ) + assigned = torch.tensor( + [item in accepted_assignments for item in self._assignments[step.id]], + dtype=torch.bool, + device=self.env.device, + ) + receiver_arm = str(binding.get("receive_arm", "right_arm")) + receiver_conflict = torch.tensor( + [ + owner not in {None, step.object_uid} + for owner in self._arm_owners[receiver_arm] + ], + dtype=torch.bool, + device=self.env.device, + ) + active = assigned & ~failed & ~receiver_conflict + if not bool(active.any()): + return _EdgeResult( + [], + failed | (~failed & ~assigned) | receiver_conflict, + [], + executed=torch.zeros_like(failed), + ) + state_key = ( + transfer_arm + if capability.state_effect == "transfer_hold" + else "coordinated" + ) + state = self._state_for(step, state_key) + if capability.state_effect == "coordinated_release": + held_objects = dict(state.held_objects) + for arm in ("left_arm", "right_arm"): + arm_state = self._step_states.get((step.id, arm)) + if arm_state is None: + continue + control_part = arm_control_part(self.env, arm) + held_object = arm_state.get_held_object(control_part) + if held_object is not None: + held_objects[control_part] = held_object + state = state.with_updates(held_objects=held_objects) + update_obj_info = getattr(self.env, "update_obj_info", None) + if callable(update_obj_info): + update_obj_info() + groundings = self.grounder.ground_candidates( + action, + step, + arm="coordinated", + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + selected: tuple[GroundedAction, ActionOutcome] | None = None + selected_warnings: tuple[str, ...] = () + best_failure_count = int(active.sum()) + 1 + rejected_warning_count = 0 + for candidate in groundings: + with _capture_speculative_warnings() as captured: + candidate_outcome = self.adapter.plan(candidate, state) + failure_count = int((active & ~candidate_outcome.success).sum()) + if selected is None or failure_count < best_failure_count: + selected = (candidate, candidate_outcome) + selected_warnings = tuple(captured) + best_failure_count = failure_count + if failure_count == 0: + if rejected_warning_count: + log_info( + "Selected a feasible coordinated grounding after " + f"suppressing {rejected_warning_count} warnings from " + "rejected candidates." + ) + break + rejected_warning_count += len(captured) + if selected is None: + raise RuntimeError("Coordinated action grounding produced no candidates.") + if best_failure_count: + for message in dict.fromkeys(selected_warnings): + log_warning(message) + grounded, outcome = selected + if capability.state_effect == "transfer_hold": + outcome = replace( + outcome, + planner_trace={ + **outcome.planner_trace, + "execution_replanned_from_live_state": True, + "execution_object_pose": self._entity_pose(step.object_uid) + .detach() + .clone(), + }, + ) + self._remember_target(step, grounded) + successful = active & outcome.success + actions = self.adapter.execute_trajectory( + outcome.trajectory, + active=successful, + ) + physical_failed = torch.zeros_like(failed) + committed_state = outcome.state_after(successful) + if capability.state_effect == "coordinated_hold": + self._clear_support_relation(step.object_uid, successful) + if capability.state_effect == "transfer_hold": + if bool(successful.any()): + current_owners = list( + self._object_owners.get( + step.object_uid, + [None] * int(self.env.num_envs), + ) + ) + tentative_owners = list(current_owners) + for env_id in ( + torch.nonzero(successful, as_tuple=False).flatten().tolist() + ): + tentative_owners[env_id] = receiver_arm + tentative_states = dict(self._object_states) + tentative_states[(step.object_uid, receiver_arm)] = outcome.next_state + physical = self._physical_hold( + step.object_uid, + receiver_arm, + outcome.next_state, + successful, + owners={ + **self._object_owners, + step.object_uid: tentative_owners, + }, + states=tentative_states, + position_tolerance=min( + float( + self.runtime_policy.predicate_fallbacks[ + "held_position_tolerance" + ] + ), + float( + grounded.motion_policy.get( + "held_position_tolerance", + self.runtime_policy.predicate_fallbacks[ + "held_position_tolerance" + ], + ) + ), + ), + ) + lost = successful & ~physical + physical_failed |= lost + successful = physical + committed_state = outcome.state_after(successful) + committed_state = self._rebase_held_state( + step.object_uid, + receiver_arm, + committed_state, + successful, + from_planned_qpos=True, + ) + committed_owners = list(current_owners) + for env_id in ( + torch.nonzero( + active & outcome.success, + as_tuple=False, + ) + .flatten() + .tolist() + ): + if bool(physical[env_id]): + committed_owners[env_id] = receiver_arm + self._arm_owners[receiver_arm][env_id] = step.object_uid + else: + committed_owners[env_id] = None + self._arm_owners[transfer_arm][env_id] = None + self._object_owners[step.object_uid] = committed_owners + if any(owner == receiver_arm for owner in committed_owners): + self._step_states[(step.id, receiver_arm)] = committed_state + self._object_states[(step.object_uid, receiver_arm)] = ( + committed_state + ) + else: + self._object_states.pop((step.object_uid, receiver_arm), None) + if not any(owner == transfer_arm for owner in committed_owners): + self._object_states.pop((step.object_uid, transfer_arm), None) + self._step_states[(step.id, "coordinated")] = committed_state + return _EdgeResult( + actions, + failed + | (~failed & ~assigned) + | (active & ~outcome.success) + | physical_failed, + [grounded], + [outcome.planner_trace], + active & outcome.success, + ) + + def _rebase_held_state( + self, + uid: str, + arm: str, + state: ExecutionState, + mask: torch.Tensor, + *, + from_planned_qpos: bool = True, + ) -> ExecutionState: + """Refresh a held object's object-to-EEF transform after execution.""" + control_part = arm_control_part(self.env, arm) + held = state.get_held_object(control_part) + entity = self.env.sim.get_rigid_object(uid) + if held is None or entity is None or not bool(mask.any()): + return state + if from_planned_qpos: + # Preserve-hold planning must stay in its terminal qpos/FK frame; + # get_current_xpos_agent() may still expose the previous command. + joint_ids = self.env.robot.get_joint_ids(name=control_part) + eef_pose = self.env.robot.compute_fk( + state.last_qpos[:, joint_ids], + name=control_part, + to_matrix=True, + ) + else: + eef_poses = self.env.get_current_xpos_agent() + eef_pose = eef_poses[0 if arm == "left_arm" else 1] + eef_pose = torch.as_tensor( + eef_pose, + dtype=held.object_to_eef.dtype, + device=held.object_to_eef.device, + ) + object_pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=eef_pose.dtype, + device=eef_pose.device, + ) + if eef_pose.ndim == 2: + eef_pose = eef_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + if object_pose.ndim == 2: + object_pose = object_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + selector = mask[:, None, None] + live_object_to_eef = torch.bmm(torch.linalg.inv(object_pose), eef_pose) + rebased = HeldObjectState( + semantics=held.semantics, + object_to_eef=torch.where( + selector, + live_object_to_eef, + held.object_to_eef, + ), + grasp_xpos=torch.where(selector, eef_pose, held.grasp_xpos), + env_mask=held.env_mask, + ) + held_objects = dict(state.held_objects) + held_objects[control_part] = rebased + return state.with_updates(held_objects=held_objects) + + def _execute_explicit_dual( + self, + edge: ExecutionEdge, + step: SemanticStep, + failed: torch.Tensor, + ) -> _EdgeResult: + assigned = torch.tensor( + [item == "coordinated" for item in self._assignments[step.id]], + dtype=torch.bool, + device=self.env.device, + ) + if not bool((assigned & ~failed).any()): + return _EdgeResult( + [], + failed | (~failed & ~assigned), + [], + executed=torch.zeros_like(failed), + ) + outcomes: dict[str, ActionOutcome | None] = { + "left_arm": None, + "right_arm": None, + } + masks = {arm: assigned & ~failed for arm in outcomes} + grounded_items = [] + coordinated_state = self._state_for(step, "coordinated") + for action in edge.actions: + actor = action.get("actor", {}) + arm = str(actor.get("arm", "")) + if arm not in outcomes or outcomes[arm] is not None: + raise ValueError( + f"Explicit dual edge {edge.id!r} must bind each arm once." + ) + state = self._step_states.get((step.id, arm)) + if state is None: + state = coordinated_state + else: + state = self._state_for(step, arm) + grounded = self.grounder.ground( + action, + step, + arm=arm, + state=state, + orientation_reference_pose=self._orientation_references.get(step.id), + ) + outcome = self.adapter.plan(grounded, state) + outcomes[arm] = outcome + grounded_items.append(grounded) + trajectory, action_success = self.adapter.combine(outcomes, masks) + active = assigned & ~failed & action_success + actions = self.adapter.execute_trajectory(trajectory, active=active) + is_coordinated_release = { + str( + action.get("target_binding", {}).get( + "coordinated_release_role", + "", + ) + ) + for action in edge.actions + } == {"participant", "commit"} and all( + action.get("control") == "hand" + and action.get("target_binding", {}).get("kind") == "joint_state" + and action.get("target_binding", {}).get("source") == "gripper_open" + for action in edge.actions + ) + physical_failed = torch.zeros_like(failed) + if is_coordinated_release: + opened = evaluate_predicate(self.env, {"type": "both_grippers_open"}) + released = active & opened + physical_failed = active & ~opened + control_parts = ( + arm_control_part(self.env, "left_arm"), + arm_control_part(self.env, "right_arm"), + ) + released_task = StateDelta( + held_object_updates={name: None for name in control_parts} + ).apply(coordinated_state.to_task_state(), released) + released_state = ExecutionState.from_task_state( + released_task, + last_qpos=self.env.robot.get_qpos().clone(), + ) + for key in ("coordinated", "left_arm", "right_arm"): + self._step_states[(step.id, key)] = released_state + else: + for arm, outcome in outcomes.items(): + if outcome is not None: + self._step_states[(step.id, arm)] = outcome.state_after( + active & outcome.success + ) + return _EdgeResult( + actions, + failed + | (~failed & ~assigned) + | (assigned & ~action_success) + | physical_failed, + grounded_items, + [ + outcome.planner_trace + for outcome in outcomes.values() + if outcome is not None + ], + active, + ) + + def _execute_parallel_pickups( + self, + edges: Sequence[ExecutionEdge], + *, + failed: torch.Tensor, + ) -> tuple[dict[str, _EdgeResult], torch.Tensor]: + steps = [self.step_by_edge[edge.id] for edge in edges] + for step in steps: + self._capture_orientation_reference(step) + candidates = { + (step.id, arm): self._candidate(step, arm, failed) + for step in steps + for arm in ("left_arm", "right_arm") + } + for step in steps: + self._report_candidates( + step, + ( + candidates[(step.id, "left_arm")], + candidates[(step.id, "right_arm")], + ), + ) + assignments = {step.id: [None] * len(failed) for step in steps} + selection_failed = torch.zeros_like(failed) + permutations = ( + ("left_arm", "right_arm"), + ("right_arm", "left_arm"), + ) + for env_id in range(len(failed)): + if bool(failed[env_id]): + continue + ranked: list[tuple[bool, bool, float, float, str, str]] = [] + for first_arm, second_arm in permutations: + first = candidates[(steps[0].id, first_arm)] + second = candidates[(steps[1].id, second_arm)] + feasible = bool(first.feasible[env_id] and second.feasible[env_id]) + required_match = all( + candidate_step.actor.get("mode") != "required" + or str(candidate_step.actor.get("arm")) == candidate_arm + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + available = ( + required_match + and not bool( + self._resource_conflicts(steps[0], first_arm)[env_id] + or self._resource_conflicts(steps[1], second_arm)[env_id] + ) + and all( + self._auto_arm_is_allowed(candidate_step, candidate_arm, env_id) + for candidate_step, candidate_arm in ( + (steps[0], first_arm), + (steps[1], second_arm), + ) + ) + ) + first_preferred = self._preferred_in_place_arm(steps[0], env_id) + second_preferred = self._preferred_in_place_arm(steps[1], env_id) + side_penalty = ( + float(first_arm != first_preferred) if first_preferred else 0.0 + ) + side_penalty += ( + float(second_arm != second_preferred) if second_preferred else 0.0 + ) + cost = float(first.cost[env_id] + second.cost[env_id]) + ranked.append( + ( + not available, + not feasible, + side_penalty, + cost, + first_arm, + second_arm, + ) + ) + ranked.sort() + unavailable, _, _, _, first_arm, second_arm = ranked[0] + if unavailable: + selection_failed[env_id] = True + continue + assignments[steps[0].id][env_id] = first_arm + assignments[steps[1].id][env_id] = second_arm + self._assignments.update(assignments) + + base_failed = failed | selection_failed + results = { + edge.id: _EdgeResult( + [], + base_failed.clone(), + [], + executed=torch.zeros_like(failed), + ) + for edge in edges + } + for first_arm, second_arm in permutations: + partition = torch.tensor( + [ + assignments[steps[0].id][env_id] == first_arm + and assignments[steps[1].id][env_id] == second_arm + for env_id in range(len(failed)) + ], + dtype=torch.bool, + device=self.env.device, + ) + if not bool(partition.any()): + continue + outcomes: dict[str, ActionOutcome | None] = { + "left_arm": None, + "right_arm": None, + } + masks = { + "left_arm": partition, + "right_arm": partition, + } + edge_by_arm = {first_arm: edges[0], second_arm: edges[1]} + parallel_planning_failed = False + for arm, edge in edge_by_arm.items(): + step = self.step_by_edge[edge.id] + try: + grounded, outcome = self._plan_live_hold(edge, step, arm) + except Exception as exc: + parallel_planning_failed = True + results[edge.id].planner_traces.append( + self._live_hold_failure_trace(edge, step, arm, exc) + ) + continue + outcomes[arm] = outcome + results[edge.id].grounded.append(outcome.grounded) + results[edge.id].planner_traces.append(outcome.planner_trace) + if bool((partition & ~outcome.success).any()): + parallel_planning_failed = True + if parallel_planning_failed: + serial_actions: list[torch.Tensor] = [] + for edge in edges: + step = self.step_by_edge[edge.id] + serial = self._execute_edge_with_retries( + edge, + step, + failed=~partition, + ) + serial_actions.extend(serial.actions) + results[edge.id].grounded.extend(serial.grounded) + results[edge.id].planner_traces.extend(serial.planner_traces) + results[edge.id].failed = torch.where( + partition, + serial.failed, + results[edge.id].failed, + ) + assert results[edge.id].executed is not None + if serial.executed is not None: + results[edge.id].executed |= serial.executed + for edge in edges: + results[edge.id].actions.extend(serial_actions) + continue + trajectory, action_success = self.adapter.combine(outcomes, masks) + active = partition & ~base_failed & action_success + commands = self.adapter.execute_trajectory(trajectory, active=active) + for edge in edges: + assert results[edge.id].executed is not None + results[edge.id].executed |= active + for arm, edge in edge_by_arm.items(): + step = self.step_by_edge[edge.id] + outcome = outcomes[arm] + assert outcome is not None + attempted = active & outcome.success + physical = self._physical_pickup( + step.object_uid, arm, outcome.next_state, attempted + ) + committed_state = outcome.state_after(physical) + committed_state = self._rebase_held_state( + step.object_uid, + arm, + committed_state, + physical, + from_planned_qpos=False, + ) + self._step_states[(step.id, arm)] = committed_state + results[edge.id].failed |= partition & ~physical + self._update_ownership( + step, + arm, + str(edge.actions[0]["atomic_action_class"]), + committed_state, + physical, + ) + for edge in edges: + # Both edge records refer to this one synchronized stream. The + # run loop adds only the first copy to its returned trace. + results[edge.id].actions.extend(commands) + aggregate_failed = torch.zeros_like(failed) + for result in results.values(): + aggregate_failed |= result.failed + return results, aggregate_failed + + def _remember_target( + self, + step: SemanticStep, + grounded: GroundedAction, + ) -> None: + self._policies[step.id] = grounded.motion_policy + target = grounded.target_object_pose + if target is not None: + self._targets[step.id] = target[:, :3, 3].clone() + self._target_poses[step.id] = target.clone() + + def _capture_orientation_reference(self, step: SemanticStep) -> None: + """Freeze preserve orientation before speculative pickup can disturb it.""" + if ( + compile_orientation_constraint(step.goal).requires_reference + and step.id not in self._orientation_references + ): + predecessor_references = [ + self._orientation_references[predecessor.id] + for dependency in step.depends_on + if (predecessor := self.steps.get(dependency)) is not None + and predecessor.object_uid == step.object_uid + and predecessor.id in self._orientation_references + ] + if predecessor_references: + self._orientation_references[step.id] = predecessor_references[ + 0 + ].clone() + return + self._orientation_references[step.id] = self._entity_pose( + step.object_uid + ).clone() + + def _step_runtime_metadata(self, step: SemanticStep) -> list[dict[str, Any]]: + """Expose the live grounding and allocation decisions for diagnosis.""" + observed_pose = self._entity_pose(step.object_uid) + assignments = self._assignments.get( + step.id, + [None] * int(self.env.num_envs), + ) + target_pose = self._target_poses.get(step.id) + orientation_reference = self._orientation_references.get(step.id) + orientation_error = self._orientation_errors.get(step.id) + arrangement = self.arrangements.get(step.id) + policy = self._policies.get(step.id, {}) + articulation_state: dict[str, Any] | None = None + joint_name = policy.get("articulation_joint_name") + if isinstance(joint_name, str) and joint_name: + articulation = getattr( + self.env.sim, + "get_articulation", + lambda _uid: None, + )(step.object_uid) + if articulation is not None and joint_name in articulation.joint_names: + joint_id = articulation.joint_names.index(joint_name) + articulation_state = { + "joint_name": joint_name, + "initial_qpos": policy.get("articulation_initial_qpos"), + "target_qpos": policy.get("articulation_target_qpos"), + "observed_qpos": articulation.get_qpos()[:, joint_id], + } + result = [] + for env_id, assignment in enumerate(assignments): + same_side_arm = self._preferred_live_pickup_arm(step, env_id) + physical_part = assignment + if assignment in {"left_arm", "right_arm"}: + physical_part = arm_control_part(self.env, assignment) + candidate_scores = {} + for arm in ("left_arm", "right_arm"): + candidate = self._candidate_cache.get((step.id, arm)) + if candidate is None: + candidate_scores[arm] = None + continue + scores = { + name: float(values[env_id]) + for name, values in candidate.score_components.items() + } + candidate_scores[arm] = { + "feasible": bool(candidate.feasible[env_id]), + **scores, + "failure": self._candidate_failures.get((step.id, arm)), + } + item: dict[str, Any] = { + "assigned_arm": assignment, + "physical_control_part": physical_part, + "same_side_arm": same_side_arm, + "inside_arm_deadband": same_side_arm is None, + "cross_side_fallback_allowed": bool( + self.runtime_policy.arm_selection.allow_cross_side_fallback + ), + "cross_side_fallback_used": bool( + same_side_arm is not None + and assignment in {"left_arm", "right_arm"} + and assignment != same_side_arm + ), + "observed_object_pose": observed_pose[env_id], + "final_target_pose": ( + None if target_pose is None else target_pose[env_id] + ), + "orientation_reference_pose": ( + None + if orientation_reference is None + else orientation_reference[env_id] + ), + "orientation_error": ( + None + if orientation_error is None + else float(orientation_error[env_id]) + ), + "candidate_scores": candidate_scores, + } + if arrangement is not None: + item["arrangement"] = arrangement.metadata(step, env_id) + if articulation_state is not None: + item["articulation_state"] = articulation_state + result.append(item) + return result + + def _verify_step( + self, + step: SemanticStep, + failed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self.settle_steps < 0: + raise ValueError("settle_steps must be non-negative.") + if self.settle_steps and bool((~failed).any()): + self.env.sim.update(step=self.settle_steps) + entity = self.env.sim.get_rigid_object(step.object_uid) + if entity is None: + entity = getattr(self.env.sim, "get_articulation", lambda _uid: None)( + step.object_uid + ) + if entity is None: + raise ValueError(f"Unknown semantic scene entity {step.object_uid!r}.") + observed_pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + observed = observed_pose[:, :3, 3] + active = ~failed + if not bool(active.any()): + success = torch.zeros_like(failed) + log_info(f"Skipped verification for {step.id}: no active environments.") + return failed, success, observed + relation = ( + normalize_placement_relation(step.goal.get("relation", "on")) + if step.operator == "place_relative" + else str(step.goal.get("relation", "")) + ) + reference = self._support_reference_uid(step) + postcondition_type = step.postcondition.get("type") + if postcondition_type in {"object_held", "handover_complete"}: + # A planned hover target is not evidence that the object remains + # grasped. Verify live TCP/object geometry and gripper closure. + satisfied = evaluate_predicate( + self.env, + step.postcondition, + held_owners=self._object_owners, + held_states=self._object_states, + ) + satisfied &= self._placement_orientation_satisfied(step, observed_pose) + elif postcondition_type in { + "held_by_both_grippers", + "object_held_by_both_grippers", + }: + satisfied = evaluate_predicate( + self.env, + step.postcondition, + coordinated_state=self._step_states.get((step.id, "coordinated")), + ) + target = self._targets.get(step.id) + if target is not None: + policy = self._policies.get(step.id, {}) + tolerance = float( + policy.get( + "postcondition_tolerance", + self.runtime_policy.predicate_fallbacks["position_tolerance"], + ) + ) + target = target.to(device=observed.device, dtype=observed.dtype) + satisfied &= ( + torch.linalg.vector_norm(observed - target, dim=1) <= tolerance + ) + elif postcondition_type == "pressed": + satisfied = evaluate_predicate(self.env, step.postcondition) + elif postcondition_type == "articulation_joint_near": + policy = self._policies.get(step.id, {}) + predicate = dict(step.postcondition) + if "articulation_joint_name" in policy: + predicate["joint_name"] = policy["articulation_joint_name"] + if "articulation_target_qpos" in policy: + predicate["target_qpos"] = policy["articulation_target_qpos"] + satisfied = evaluate_predicate(self.env, predicate) + elif relation == "inside" and isinstance(reference, str): + satisfied = evaluate_predicate( + self.env, + { + "type": "object_in_container", + "object": step.object_uid, + "container": reference, + }, + ) + elif relation in {"on", "on_top", "on_top_of"} and isinstance(reference, str): + satisfied = self._support_stable_for(step, reference, active) + satisfied &= self._support_cycle_free( + step.object_uid, + reference, + active, + ) + elif step.operator == "orient_object": + position_anchor = str(step.goal.get("position_anchor", "initial_xy")) + anchor_pose = None + if position_anchor == "initial_xy": + anchor_pose = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + if anchor_pose is None: + anchor_pose = self._targets.get(step.id) + if anchor_pose is None: + raise ValueError( + f"orient_object step {step.id!r} has no {position_anchor} anchor." + ) + anchor_pose = torch.as_tensor( + anchor_pose, + dtype=observed.dtype, + device=observed.device, + ) + if anchor_pose.ndim == 2 and anchor_pose.shape == (4, 4): + anchor_pose = anchor_pose.unsqueeze(0).repeat( + int(self.env.num_envs), 1, 1 + ) + target_xy = ( + anchor_pose[:, :2, 3] if anchor_pose.ndim == 3 else anchor_pose[:, :2] + ) + policy = self._policies.get(step.id, {}) + fallbacks = self.runtime_policy.predicate_fallbacks + upright = evaluate_predicate( + self.env, + { + "type": "object_upright", + "object": step.object_uid, + "local_axis": policy.get("upright_local_axis", "long_axis"), + "max_tilt": float( + policy.get("upright_max_tilt", fallbacks["upright_max_tilt"]) + ), + }, + ) + xy_near_initial = evaluate_predicate( + self.env, + { + "type": "object_xy_near", + "object": step.object_uid, + "target_xy": target_xy, + "tolerance": float( + policy.get("upright_xy_tolerance", fallbacks["xy_tolerance"]) + ), + }, + ) + satisfied = upright & xy_near_initial + elif step.id in self._targets: + target = self._targets[step.id].to( + device=observed.device, + dtype=observed.dtype, + ) + policy = self._policies.get(step.id, {}) + fallbacks = self.runtime_policy.predicate_fallbacks + tolerance = float( + policy.get("postcondition_tolerance", fallbacks["position_tolerance"]) + ) + arrangement = self.arrangements.get(step.id) + if arrangement is not None: + # Line membership is a planar relation. Height changes after + # release (for example a can settling onto another stable face) + # must not invalidate an otherwise correct row placement. + delta = torch.abs(observed - target) + axis_tolerance = float( + policy.get( + "line_axis_tolerance", + fallbacks["line_axis_tolerance"], + ) + ) + perpendicular_tolerance = float( + policy.get( + "line_perpendicular_tolerance", + fallbacks["line_perpendicular_tolerance"], + ) + ) + satisfied = (delta[:, arrangement.axis_index] <= axis_tolerance) & ( + delta[:, arrangement.perpendicular_index] <= perpendicular_tolerance + ) + elif relation in DIRECTIONAL_RELATIONS: + # Left/right/front/behind constrain the support plane. The + # grounded release height is a transport target and may differ + # from the stable height after the object settles. + satisfied = ( + torch.linalg.vector_norm(observed[:, :2] - target[:, :2], dim=-1) + <= tolerance + ) + else: + satisfied = ( + torch.linalg.vector_norm(observed - target, dim=-1) <= tolerance + ) + else: + satisfied = evaluate_predicate(self.env, step.postcondition) + if relation in DIRECTIONAL_RELATIONS and isinstance(reference, str): + policy = self._policies.get(step.id, {}) + satisfied &= evaluate_predicate( + self.env, + { + "type": "object_relative_position", + "object": step.object_uid, + "reference_object": reference, + "relation": relation, + "relation_frame": step.goal.get("relation_frame", "world"), + "minimum_distance": float(policy.get("relation_clearance", 0.01)), + }, + ) + verifies_placement_orientation = bool( + compile_orientation_constraint(step.goal).terms + ) and ( + postcondition_type == "semantic_goal" + or self.arrangements.get(step.id) is not None + ) + if verifies_placement_orientation: + satisfied &= self._placement_orientation_satisfied(step, observed_pose) + if step.goal.get("payloads"): + satisfied &= self._verify_payloads(step) + success = active & satisfied + failed = failed | (active & ~satisfied) + if relation in {"on", "on_top", "on_top_of"} and isinstance(reference, str): + self._commit_support_relation(step, reference, success) + log_info( + f"Verified {step.id}: {int(success.sum())}/{len(success)} envs succeeded." + ) + return failed, success, observed + + def _capture_payloads(self, step: SemanticStep) -> None: + if step.id in self._payload_initial: + return + carrier = self._entity_pose(step.object_uid) + record = {"carrier_rotation": carrier[:, :3, :3].clone()} + for payload in step.goal.get("payloads", []): + uid = str(payload["object"]) + record[uid] = torch.bmm(torch.linalg.inv(carrier), self._entity_pose(uid)) + self._payload_initial[step.id] = record + + def _verify_payloads(self, step: SemanticStep) -> torch.Tensor: + record = self._payload_initial.get(step.id) + if record is None: + return torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + carrier = self._entity_pose(step.object_uid) + initial_up = record["carrier_rotation"][:, :3, 2] + live_up = carrier[:, :3, 2] + fallbacks = self.runtime_policy.predicate_fallbacks + tilt_ok = torch.sum(initial_up * live_up, dim=-1) >= float( + fallbacks["payload_minimum_upright_cosine"] + ) + result = tilt_ok + carrier_entity = self.env.sim.get_rigid_object(step.object_uid) + for payload in step.goal["payloads"]: + uid = str(payload["object"]) + expected = torch.bmm(carrier, record[uid]) + observed = self._entity_pose(uid) + drift_ok = torch.linalg.vector_norm( + observed[:, :3, 3] - expected[:, :3, 3], + dim=-1, + ) <= float(fallbacks["payload_position_tolerance"]) + support_ok = torch.ones_like(drift_ok) + for env_id in range(int(self.env.num_envs)): + vertices = carrier_entity.get_vertices( + env_ids=[env_id], + scale=True, + ) + if isinstance(vertices, (list, tuple)): + vertices = vertices[0] + vertices = torch.as_tensor( + vertices, + dtype=carrier.dtype, + device=carrier.device, + ) + if vertices.ndim == 3: + vertices = vertices[0] + world = ( + vertices @ carrier[env_id, :3, :3].transpose(0, 1) + + carrier[env_id, :3, 3] + ) + position = observed[env_id, :2, 3] + margin = float(fallbacks["payload_support_margin"]) + lower = world[:, :2].min(dim=0).values - margin + upper = world[:, :2].max(dim=0).values + margin + support_ok[env_id] = bool( + torch.all(position >= lower) and torch.all(position <= upper) + ) + result &= drift_ok & support_ok + return result + + def _entity_pose(self, uid: str) -> torch.Tensor: + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr(self.env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") + pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + if pose.ndim == 2: + pose = pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + return pose + + def _entity_motion_stable(self, uid: str) -> torch.Tensor: + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + + def velocity(value: Any, name: str) -> torch.Tensor | None: + if callable(value): + value = value() + if value is None: + return None + tensor = torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0).repeat(int(self.env.num_envs), 1) + if tensor.shape != (int(self.env.num_envs), 3): + raise ValueError( + f"Rigid object {uid!r} {name} must have shape " + f"({int(self.env.num_envs)}, 3)." + ) + return tensor + + linear = velocity(getattr(entity, "lin_vel", None), "lin_vel") + angular = velocity(getattr(entity, "ang_vel", None), "ang_vel") + if linear is None or angular is None: + body_state = getattr(entity, "body_state", None) + if callable(body_state): + body_state = body_state() + if body_state is not None: + state = torch.as_tensor( + body_state, + dtype=torch.float32, + device=self.env.device, + ) + if state.ndim == 1: + state = state.unsqueeze(0).repeat(int(self.env.num_envs), 1) + if state.shape == (int(self.env.num_envs), 13): + linear = state[:, 7:10] + angular = state[:, 10:13] + if linear is None or angular is None: + body_data = getattr(entity, "body_data", None) + if body_data is not None: + if linear is None: + linear = velocity(getattr(body_data, "lin_vel", None), "lin_vel") + if angular is None: + angular = velocity(getattr(body_data, "ang_vel", None), "ang_vel") + if linear is None or angular is None: + return torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + return ( + torch.linalg.vector_norm(linear, dim=1) + <= self.support_linear_velocity_tolerance + ) & ( + torch.linalg.vector_norm(angular, dim=1) + <= self.support_angular_velocity_tolerance + ) + + def _support_stable_for( + self, + step: SemanticStep, + support_uid: str, + active: torch.Tensor, + ) -> torch.Tensor: + """Require the support relation and low motion across a time window.""" + stable = active.clone() + for sample_index in range(self.support_stability_samples): + supported = evaluate_predicate( + self.env, + { + "type": "object_supported_by", + "object": step.object_uid, + "support": support_uid, + }, + ) + stable &= ( + supported + & self._entity_motion_stable(step.object_uid) + & self._entity_motion_stable(support_uid) + ) + if ( + sample_index + 1 < self.support_stability_samples + and self.support_stability_interval_steps + and bool(active.any()) + ): + self.env.sim.update(step=self.support_stability_interval_steps) + return stable + + def _clear_support_relation(self, object_uid: str, mask: torch.Tensor) -> None: + relations = self._support_relations.get(object_uid) + if relations is None: + return + for env_id in torch.nonzero(mask, as_tuple=False).flatten().tolist(): + relations[env_id] = None + if not any(relation is not None for relation in relations): + self._support_relations.pop(object_uid, None) + + def _support_cycle_free( + self, + object_uid: str, + support_uid: str, + active: torch.Tensor, + ) -> torch.Tensor: + result = active.clone() + for env_id in torch.nonzero(active, as_tuple=False).flatten().tolist(): + current = support_uid + visited: set[str] = set() + while current and current not in visited: + if current == object_uid: + result[env_id] = False + break + visited.add(current) + relations = self._support_relations.get(current) + relation = None if relations is None else relations[env_id] + current = "" if relation is None else relation.support_uid + return result + + def _commit_support_relation( + self, + step: SemanticStep, + support_uid: str, + successful: torch.Tensor, + ) -> None: + relations = self._support_relations.setdefault( + step.object_uid, + [None] * int(self.env.num_envs), + ) + relation = _SupportRelation( + support_uid=support_uid, + semantic_step_id=step.id, + ) + for env_id in torch.nonzero(successful, as_tuple=False).flatten().tolist(): + relations[env_id] = relation + + def _placement_orientation_satisfied( + self, + step: SemanticStep, + observed_pose: torch.Tensor, + ) -> torch.Tensor: + constraint = compile_orientation_constraint(step.goal) + satisfied = torch.ones( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ) + if not constraint.terms or ( + step.goal.get("orientation_goal") == "preserve" + and step.goal.get("relation") == "inside" + ): + return satisfied + policy = self._policies.get(step.id, {}) + fallbacks = self.runtime_policy.predicate_fallbacks + errors = [] + for term in constraint.terms: + if isinstance(term, AlignAxisConstraint): + if term.target_axis != "world_up": + raise ValueError( + f"Unsupported orientation target axis {term.target_axis!r}." + ) + satisfied &= evaluate_predicate( + self.env, + { + "type": "object_upright", + "object": step.object_uid, + "local_axis": term.local_axis, + "directed": term.directed, + "max_tilt": float( + term.tolerance + if term.tolerance is not None + else policy.get( + "upright_max_tilt", fallbacks["upright_max_tilt"] + ) + ), + }, + ) + continue + if not isinstance(term, MatchRotationConstraint): + raise TypeError(f"Unsupported orientation term {type(term)!r}.") + reference_pose = ( + self._orientation_references.get(step.id) + if term.reference == "step_start" + else self._target_poses.get(step.id) + ) + if reference_pose is None: + satisfied &= False + continue + reference_rotation = reference_pose[:, :3, :3].to( + device=observed_pose.device, + dtype=observed_pose.dtype, + ) + relative = torch.bmm( + reference_rotation.transpose(1, 2), + observed_pose[:, :3, :3], + ) + cosine = (relative.diagonal(dim1=-2, dim2=-1).sum(dim=-1) - 1.0) * 0.5 + error = torch.acos(cosine.clamp(-1.0, 1.0)) + errors.append(error) + satisfied &= error <= float( + term.tolerance + if term.tolerance is not None + else policy.get( + "preserve_orientation_tolerance", + fallbacks["preserve_orientation_tolerance"], + ) + ) + if errors: + self._orientation_errors[step.id] = torch.stack(errors).amax(dim=0) + return satisfied + + def _revalidate_support_relations(self) -> dict[str, torch.Tensor]: + active_by_step: dict[str, torch.Tensor] = {} + for relations in self._support_relations.values(): + for env_id, relation in enumerate(relations): + if relation is None: + continue + active = active_by_step.setdefault( + relation.semantic_step_id, + torch.zeros( + int(self.env.num_envs), + dtype=torch.bool, + device=self.env.device, + ), + ) + active[env_id] = True + failures: dict[str, torch.Tensor] = {} + for step_id, active in active_by_step.items(): + step = self.steps[step_id] + support_uid = self._support_reference_uid(step) + if support_uid is None: + failures[step_id] = active + continue + observed_pose = self._entity_pose(step.object_uid) + valid = self._support_stable_for(step, support_uid, active) + valid &= self._placement_orientation_satisfied(step, observed_pose) + lost = active & ~valid + if bool(lost.any()): + failures[step_id] = lost + return failures + + @staticmethod + def _support_reference_uid(step: SemanticStep) -> str | None: + value = step.goal.get("reference_object", step.goal.get("support_object")) + if isinstance(value, str) and value: + return value + if ( + step.postcondition.get("type") == "stack_layer_supported" + and int(step.goal.get("layer_index", -1)) == 0 + ): + return "table" + return None + + @staticmethod + def _edge_failure_policy(edge: ExecutionEdge) -> str: + """Return the persisted node policy for one synchronized edge.""" + policies = { + str(action.get("failure_policy", "task_required")) + for action in edge.actions + } + if not policies <= {"task_required", "safety_required", "best_effort"}: + raise ValueError( + f"Edge {edge.id!r} contains unknown failure policies {policies}." + ) + if len(policies) != 1: + raise ValueError( + f"Edge {edge.id!r} mixes incompatible failure policies {policies}." + ) + return next(iter(policies)) diff --git a/embodichain/gen_sim/action_engine/runtime/frames.py b/embodichain/gen_sim/action_engine/runtime/frames.py new file mode 100644 index 000000000..513308484 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/frames.py @@ -0,0 +1,165 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Resolve directional relations in live robot and world frames.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from .robot_parts import arm_control_part + +__all__ = [ + "DIRECTIONAL_RELATIONS", + "arm_base_poses", + "relation_axes", + "relation_offset", + "robot_frame_axes", +] + + +_RELATION_COMPONENTS = { + "left": ("left",), + "left_of": ("left",), + "right": ("right",), + "right_of": ("right",), + "front": ("front",), + "front_of": ("front",), + "in_front_of": ("front",), + "behind": ("back",), + "back": ("back",), + "front_left": ("front", "left"), + "front_left_of": ("front", "left"), + "front_right": ("front", "right"), + "front_right_of": ("front", "right"), + "back_left": ("back", "left"), + "back_left_of": ("back", "left"), + "back_right": ("back", "right"), + "back_right_of": ("back", "right"), +} +DIRECTIONAL_RELATIONS = frozenset(_RELATION_COMPONENTS) + + +def arm_base_poses(env: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Return live world poses of the left and right arm bases.""" + left_part = arm_control_part(env, "left_arm") + right_part = arm_control_part(env, "right_arm") + robot = env.robot + if hasattr(robot, "get_solver") and hasattr(robot, "get_link_pose"): + left_solver = robot.get_solver(name=left_part) + right_solver = robot.get_solver(name=right_part) + left_root = getattr(left_solver, "root_link_name", None) + right_root = getattr(right_solver, "root_link_name", None) + if left_root is None or right_root is None: + raise ValueError("Directional grounding requires both arm root links.") + left = robot.get_link_pose(link_name=left_root, to_matrix=True) + right = robot.get_link_pose(link_name=right_root, to_matrix=True) + elif hasattr(robot, "get_control_part_base_pose"): + left = robot.get_control_part_base_pose(name=left_part, to_matrix=True) + right = robot.get_control_part_base_pose(name=right_part, to_matrix=True) + elif hasattr(env, "get_current_xpos_agent"): + left, right = env.get_current_xpos_agent() + else: + raise ValueError( + "Directional grounding requires live left/right arm-base or TCP poses." + ) + + left = _batched_pose(left, env) + right = _batched_pose(right, env) + return left, right + + +def robot_frame_axes(env: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Return normalized world-space forward and left axes for a dual-arm robot.""" + left, right = arm_base_poses(env) + lateral = left[:, :2, 3] - right[:, :2, 3] + norm = torch.linalg.vector_norm(lateral, dim=1, keepdim=True) + if bool((norm <= 1.0e-6).any()): + raise ValueError("Left and right arm bases must have distinct XY positions.") + lateral = lateral / norm + forward = torch.stack((lateral[:, 1], -lateral[:, 0]), dim=1) + return forward, lateral + + +def relation_axes( + env: Any, + relation: str, + *, + frame: str, +) -> tuple[torch.Tensor, ...]: + """Return signed world-space axes whose projections define a relation.""" + relation = str(relation) + if relation not in DIRECTIONAL_RELATIONS: + return () + if frame == "robot": + forward, lateral = robot_frame_axes(env) + elif frame == "world": + count = int(env.num_envs) + forward = torch.tensor( + [1.0, 0.0], dtype=torch.float32, device=env.device + ).repeat(count, 1) + lateral = torch.tensor( + [0.0, 1.0], dtype=torch.float32, device=env.device + ).repeat(count, 1) + else: + raise ValueError(f"Unsupported directional relation frame {frame!r}.") + + component_axes = { + "front": forward, + "back": -forward, + "left": lateral, + "right": -lateral, + } + return tuple(component_axes[item] for item in _RELATION_COMPONENTS[relation]) + + +def relation_offset( + env: Any, + relation: str, + *, + frame: str, + forward_distance: float, + lateral_distance: float, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + """Resolve one directional relation into a batched world-space offset.""" + axes = relation_axes(env, relation, frame=frame) + if not axes: + return None + offset = torch.zeros((int(env.num_envs), 3), dtype=dtype, device=device) + components = _RELATION_COMPONENTS[relation] + for component, axis in zip(components, axes): + axis = axis.to(dtype=dtype, device=device) + distance = ( + forward_distance if component in {"front", "back"} else lateral_distance + ) + offset[:, :2] += axis * float(distance) + return offset + + +def _batched_pose(value: Any, env: Any) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + if pose.shape != (int(env.num_envs), 4, 4): + raise ValueError( + "Frame pose must have shape (4, 4) or " + f"({int(env.num_envs)}, 4, 4), got {tuple(pose.shape)}." + ) + return pose diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py b/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py new file mode 100644 index 000000000..252c388d6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py @@ -0,0 +1,330 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Prepare checksummed V-HACD caches for the shared grasp collision checker. + +The sidecar identifies the backend without changing Main's cache key or pickle +payload, so an unlabelled CoACD cache is never silently reused as V-HACD. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import hashlib +import io +import json +import operator +import os +from pathlib import Path +import pickle +import stat +import tempfile +from typing import Literal + +import numpy as np +import torch + +__all__ = [ + "GraspCollisionCacheError", + "GraspCollisionCacheResult", + "ensure_vhacd_grasp_collision_cache", + "grasp_collision_cache_path", +] + +_CACHE_SCHEMA_VERSION = 1 +_METADATA_SUFFIX = ".action_engine.json" +_DEFAULT_CACHE_DIR = ( + Path.home() / ".cache" / "embodichain_cache" / "convex_decomposition" +) + +CacheStatus = Literal["hit", "generated", "replaced"] + + +class GraspCollisionCacheError(RuntimeError): + """Raised when a safe, Main-compatible V-HACD cache cannot be prepared.""" + + +@dataclass(frozen=True) +class GraspCollisionCacheResult: + """Describe the prepared cache files and whether decomposition ran.""" + + status: CacheStatus + cache_path: Path + metadata_path: Path + + +def grasp_collision_cache_path( + mesh_vertices: torch.Tensor | np.ndarray, + mesh_triangles: torch.Tensor | np.ndarray, + max_decomposition_hulls: int, + *, + cache_dir: str | Path | None = None, +) -> Path: + """Return Main's exact ``_.pkl`` cache path.""" + vertices, triangles = _validate_mesh(mesh_vertices, mesh_triangles) + hull_limit = _validate_hull_limit(max_decomposition_hulls) + mesh_hash = hashlib.md5(vertices.tobytes() + triangles.tobytes()).hexdigest() + return _resolve_cache_dir(cache_dir) / f"{mesh_hash}_{hull_limit}.pkl" + + +def ensure_vhacd_grasp_collision_cache( + *, + mesh_vertices: torch.Tensor | np.ndarray, + mesh_triangles: torch.Tensor | np.ndarray, + max_decomposition_hulls: int, + cache_dir: str | Path | None = None, +) -> GraspCollisionCacheResult: + """Create or validate a V-HACD cache and its checksummed backend sidecar.""" + vertices, triangles = _validate_mesh(mesh_vertices, mesh_triangles) + hull_limit = _validate_hull_limit(max_decomposition_hulls) + mesh_hash = hashlib.md5(vertices.tobytes() + triangles.tobytes()).hexdigest() + cache_path = _resolve_cache_dir(cache_dir) / f"{mesh_hash}_{hull_limit}.pkl" + metadata_path = cache_path.with_name(f"{cache_path.name}{_METADATA_SUFFIX}") + expected_metadata: dict[str, object] = { + "schema_version": _CACHE_SCHEMA_VERSION, + "backend": "vhacd", + "mesh_hash": mesh_hash, + "max_decomposition_hulls": hull_limit, + } + + _prepare_private_directory(cache_path.parent) + _refuse_symlink(cache_path) + _refuse_symlink(metadata_path) + if _cache_matches_metadata(cache_path, metadata_path, expected_metadata): + return GraspCollisionCacheResult("hit", cache_path, metadata_path) + + exists = cache_path.exists() or metadata_path.exists() + status: CacheStatus = "replaced" if exists else "generated" + try: + plane_equations = _compute_vhacd_plane_equations( + vertices, + triangles, + hull_limit, + ) + cache_bytes = _serialize_checker_payload(plane_equations) + metadata = { + **expected_metadata, + "cache_sha256": hashlib.sha256(cache_bytes).hexdigest(), + } + + # Publish the complete pickle before its sidecar. A crash between the + # two replaces leaves a cache miss on retry, never a partial pickle. + _write_bytes_atomic(cache_path, cache_bytes) + metadata_bytes = (json.dumps(metadata, sort_keys=True) + "\n").encode() + _write_bytes_atomic(metadata_path, metadata_bytes) + except GraspCollisionCacheError: + raise + except Exception as exc: + raise GraspCollisionCacheError( + f"Failed to prepare V-HACD grasp collision cache {cache_path}: {exc}" + ) from exc + + return GraspCollisionCacheResult(status, cache_path, metadata_path) + + +def _compute_vhacd_plane_equations( + vertices: np.ndarray, + triangles: np.ndarray, + max_decomposition_hulls: int, +) -> list[tuple[np.ndarray, np.ndarray]]: + """Run DexSim V-HACD and convert its hulls to checker plane equations.""" + import open3d as o3d + from dexsim.kit.meshproc import convex_decomposition_vhacd + + from embodichain.toolkits.graspkit.pg_grasp.collision_checker import ( + extract_plane_equations, + ) + + mesh = o3d.t.geometry.TriangleMesh() + mesh.vertex.positions = o3d.core.Tensor(vertices.astype(np.float32, copy=False)) + mesh.triangle.indices = o3d.core.Tensor(triangles.astype(np.int32, copy=False)) + is_success, hull_meshes = convex_decomposition_vhacd( + mesh, + max_convex_hull_num=max_decomposition_hulls, + ) + if not is_success or not hull_meshes: + raise GraspCollisionCacheError( + "V-HACD returned no convex hulls for the grasp collision mesh." + ) + + convex_parts = [ + ( + np.asarray(hull.vertex.positions.numpy()), + np.asarray(hull.triangle.indices.numpy()), + ) + for hull in hull_meshes + ] + plane_equations = extract_plane_equations(convex_parts) + if not plane_equations: + raise GraspCollisionCacheError( + "V-HACD hulls produced no grasp collision plane equations." + ) + return plane_equations + + +def _serialize_checker_payload( + plane_equations: list[tuple[np.ndarray, np.ndarray]], +) -> bytes: + """Pack plane equations in the exact tensor dictionary Main unpickles.""" + if not plane_equations: + raise ValueError("V-HACD must produce at least one convex hull.") + + normalized: list[tuple[np.ndarray, np.ndarray]] = [] + for normals_value, offsets_value in plane_equations: + normals = np.asarray(normals_value, dtype=np.float32) + offsets = np.asarray(offsets_value, dtype=np.float32) + if normals.ndim != 2 or normals.shape[1:] != (3,) or not len(normals): + raise ValueError("Each V-HACD hull must have normals shaped [K, 3].") + if offsets.shape != (len(normals),): + raise ValueError("Each hull needs one offset per plane normal.") + if not np.isfinite(normals).all() or not np.isfinite(offsets).all(): + raise ValueError("V-HACD plane equations must contain finite values.") + normalized.append((normals, offsets)) + + max_plane_count = max(normals.shape[0] for normals, _ in normalized) + equations = torch.zeros((len(normalized), max_plane_count, 4)) + counts = torch.zeros(len(normalized), dtype=torch.int32) + for index, (normals, offsets) in enumerate(normalized): + plane_count = normals.shape[0] + equations[index, :plane_count, :3] = torch.from_numpy(normals) + equations[index, :plane_count, 3] = torch.from_numpy(offsets) + counts[index] = plane_count + + stream = io.BytesIO() + payload = {"plane_equations": equations, "plane_equation_counts": counts} + pickle.dump(payload, stream, protocol=pickle.HIGHEST_PROTOCOL) + return stream.getvalue() + + +def _cache_matches_metadata( + cache_path: Path, + metadata_path: Path, + expected_metadata: dict[str, object], +) -> bool: + if not cache_path.is_file() or not metadata_path.is_file(): + return False + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + checksum = metadata.get("cache_sha256") + expected_checksum = hashlib.sha256(cache_path.read_bytes()).hexdigest() + return ( + isinstance(metadata, dict) + and all( + metadata.get(key) == value for key, value in expected_metadata.items() + ) + and isinstance(checksum, str) + and checksum == expected_checksum + ) + except (AttributeError, OSError, UnicodeDecodeError, json.JSONDecodeError): + return False + + +def _validate_mesh( + mesh_vertices: torch.Tensor | np.ndarray, + mesh_triangles: torch.Tensor | np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + if isinstance(mesh_vertices, torch.Tensor): + mesh_vertices = mesh_vertices.detach().cpu().numpy() + if isinstance(mesh_triangles, torch.Tensor): + mesh_triangles = mesh_triangles.detach().cpu().numpy() + if not isinstance(mesh_vertices, np.ndarray): + raise TypeError("mesh_vertices must be a torch.Tensor or numpy.ndarray.") + if not isinstance(mesh_triangles, np.ndarray): + raise TypeError("mesh_triangles must be a torch.Tensor or numpy.ndarray.") + vertices = np.ascontiguousarray(mesh_vertices) + triangles = np.ascontiguousarray(mesh_triangles) + if vertices.ndim != 2 or vertices.shape[1:] != (3,) or len(vertices) == 0: + raise ValueError("mesh_vertices must have non-empty shape [N, 3].") + if triangles.ndim != 2 or triangles.shape[1:] != (3,) or len(triangles) == 0: + raise ValueError("mesh_triangles must have non-empty shape [M, 3].") + if not np.issubdtype(vertices.dtype, np.number): + raise TypeError("mesh_vertices must contain numeric values.") + if not np.isfinite(vertices).all(): + raise ValueError("mesh_vertices must contain only finite values.") + if not np.issubdtype(triangles.dtype, np.integer): + raise TypeError("mesh_triangles must contain integer indices.") + if triangles.min() < 0 or triangles.max() >= len(vertices): + raise ValueError("mesh_triangles contains out-of-range vertex indices.") + return vertices, triangles + + +def _validate_hull_limit(value: int) -> int: + if isinstance(value, (bool, np.bool_)): + raise TypeError("max_decomposition_hulls must be an integer.") + try: + hull_limit = operator.index(value) + except TypeError as exc: + raise TypeError("max_decomposition_hulls must be an integer.") from exc + if hull_limit <= 0: + raise ValueError("max_decomposition_hulls must be positive.") + return hull_limit + + +def _resolve_cache_dir(cache_dir: str | Path | None) -> Path: + if cache_dir is not None: + return Path(cache_dir).expanduser().resolve() + try: + from embodichain.lab.sim import CONVEX_DECOMP_DIR + except Exception: + return _DEFAULT_CACHE_DIR + return Path(CONVEX_DECOMP_DIR).expanduser().resolve() + + +def _prepare_private_directory(path: Path) -> None: + try: + path.mkdir(parents=True, exist_ok=True, mode=0o700) + path.chmod(0o700) + except OSError as exc: + raise GraspCollisionCacheError( + f"Cannot secure grasp collision cache directory: {path}" + ) from exc + if path.stat().st_mode & (stat.S_IWGRP | stat.S_IWOTH): + raise GraspCollisionCacheError(f"Refusing writable cache directory: {path}") + + +def _refuse_symlink(path: Path) -> None: + if path.is_symlink(): + raise GraspCollisionCacheError( + f"Refusing symlinked grasp collision cache path: {path}" + ) + + +def _write_bytes_atomic(path: Path, payload: bytes) -> None: + """Publish one complete file with a same-directory atomic replacement.""" + _refuse_symlink(path) + file_descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=path.parent, + ) + temporary_path = Path(temporary_name) + try: + os.fchmod(file_descriptor, 0o600) + with os.fdopen(file_descriptor, "wb") as output: + file_descriptor = -1 + output.write(payload) + output.flush() + os.fsync(output.fileno()) + _refuse_symlink(path) + os.replace(temporary_path, path) + path.chmod(0o600) + finally: + if file_descriptor >= 0: + os.close(file_descriptor) + try: + temporary_path.unlink(missing_ok=True) + except OSError: + pass diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py new file mode 100644 index 000000000..8687a12b7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -0,0 +1,3155 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Resolve symbolic bindings from live simulator state.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, replace +import math +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + default_runtime_policy, +) +from embodichain.gen_sim.action_engine.domain import normalize_placement_relation +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + OrientationConstraint, + compile_orientation_constraint, +) +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + AxisAlignAffordance, + AxisAlignGoal, + CoordinatedPickGoal, + CoordinatedPlacementGoal, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectPoseGoal, + JointPositionGoal, + ObjectSemantics, + PlaceGoal, + PourGoal, + PressAffordance, + PressGoal, + SlideAffordance, + SlideGoal, + TwistAffordance, + TwistGoal, +) +from embodichain.toolkits.graspkit.pg_grasp import ( + AntipodalSamplerCfg, + GraspGeneratorCfg, + GripperCollisionCfg, +) +from embodichain.utils.logger import log_info +from .frames import arm_base_poses, relation_offset, robot_frame_axes +from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache +from .models import ExecutionProgram, GroundedAction, SemanticStep +from .motion_policy import resolve_motion_policy, with_motion_modifiers +from .robot_parts import arm_control_part +from .state import ExecutionState + +__all__ = ["ActionGrounder", "LiveArrangementPlan", "LivePlacementPlan"] + + +def _batched_pose(value: Any, env: Any) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + if pose.shape != (int(env.num_envs), 4, 4): + raise ValueError( + "Live pose must have shape (4, 4) or " + f"({int(env.num_envs)}, 4, 4), got {tuple(pose.shape)}." + ) + return pose + + +def _object(env: Any, uid: str) -> Any: + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + return entity + + +def _live_pose(env: Any, uid: str) -> torch.Tensor: + entity = env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") + return _batched_pose(entity.get_local_pose(to_matrix=True), env) + + +def _local_vertices(entity: Any, env: Any, env_id: int = 0) -> torch.Tensor: + value = entity.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (list, tuple)): + value = value[0] + vertices = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError("Rigid-object mesh vertices must have shape (N, 3).") + return vertices + + +def _world_vertices(entity: Any, env: Any, env_id: int) -> torch.Tensor: + vertices = _local_vertices(entity, env, env_id) + pose = _batched_pose(entity.get_local_pose(to_matrix=True), env)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +@dataclass(frozen=True) +class _Geometry: + radius: torch.Tensor + half_height: torch.Tensor + + +class LiveArrangementPlan: + """Materialize collision-aware line slots independently in every env.""" + + def __init__( + self, + env: Any, + steps: Sequence[SemanticStep], + *, + slot_margin: float | None = None, + minimum_spacing: float | None = None, + clearance: float | None = None, + row_search_step: float | None = None, + row_search_radius: float | None = None, + ) -> None: + if not steps: + raise ValueError("An arrangement plan requires at least one step.") + self.env = env + self.steps = tuple(steps) + self.step_by_id = {step.id: step for step in steps} + self.num_envs = int(env.num_envs) + self.device = env.device + self.slot_count = len(steps) + self.axis = str(steps[0].goal.get("axis", "world_x")) + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + defaults = default_runtime_policy(profile).grounding["arrangement"] + slot_margin = defaults["slot_margin"] if slot_margin is None else slot_margin + minimum_spacing = ( + defaults["minimum_spacing"] if minimum_spacing is None else minimum_spacing + ) + self.clearance = float( + defaults["layout_clearance"] if clearance is None else clearance + ) + self.row_search_step = float( + defaults["row_search_step"] if row_search_step is None else row_search_step + ) + self.row_search_radius = float( + defaults["row_search_radius"] + if row_search_radius is None + else row_search_radius + ) + + table = _object(env, "table") + bounds = [] + for env_id in range(self.num_envs): + vertices = _world_vertices(table, env, env_id) + bounds.append( + torch.stack((vertices.min(dim=0).values, vertices.max(dim=0).values)) + ) + self.table_bounds = torch.stack(bounds) + self.table_center = self.table_bounds.mean(dim=1) + self.table_top = self.table_bounds[:, 1, 2] + if self.axis == "table_long_axis": + mean_extent = ( + self.table_bounds[:, 1, :2] - self.table_bounds[:, 0, :2] + ).mean(dim=0) + self.axis_index = int(torch.argmax(mean_extent).item()) + else: + self.axis_index = 0 if self.axis in {"x", "world_x"} else 1 + self.perpendicular_index = 1 - self.axis_index + self.geometry = {step.id: self._geometry(step) for step in self.steps} + diameters = torch.stack( + [self.geometry[step.id].radius * 2.0 for step in self.steps], + dim=1, + ) + self.spacing = torch.maximum( + diameters.max(dim=1).values + float(slot_margin), + torch.full( + (self.num_envs,), + float(minimum_spacing), + dtype=torch.float32, + device=self.device, + ), + ) + self.positions = self._make_slots() + self.reassignment_reason: list[str | None] = [None] * self.num_envs + self.reassignment_cost = torch.full( + (self.num_envs,), + float("nan"), + dtype=torch.float32, + device=self.device, + ) + self.assignments = self._initial_slot_assignments() + order_by = str(self.steps[0].goal.get("order_by", "explicit")) + direction = str(self.steps[0].goal.get("order_direction", "given")) + if order_by == "size" and not any( + step.goal.get("slot_constraint") == "free_reassignable" + for step in self.steps + ): + for env_id in range(self.num_envs): + ordered = sorted( + self.steps, + key=lambda step: float(self.geometry[step.id].radius[env_id]), + reverse=direction != "ascending", + ) + for slot_id, step in enumerate(ordered): + self.assignments[step.id][env_id] = slot_id + self.completed = { + step.id: torch.zeros( + self.num_envs, + dtype=torch.bool, + device=self.device, + ) + for step in self.steps + } + + def _initial_slot_assignments(self) -> dict[str, torch.Tensor]: + """Match free-order objects to slots in their current spatial order.""" + assignments = { + step.id: torch.full( + (self.num_envs,), + int(step.goal.get("nominal_slot_index", index)), + dtype=torch.long, + device=self.device, + ) + for index, step in enumerate(self.steps) + } + free_steps = [ + step + for step in self.steps + if step.goal.get("slot_constraint") == "free_reassignable" + ] + if not free_steps: + return assignments + required_slots = { + int(step.goal.get("nominal_slot_index", index)) + for index, step in enumerate(self.steps) + if step.goal.get("slot_constraint") != "free_reassignable" + } + available_slots = [ + slot_id + for slot_id in range(self.slot_count) + if slot_id not in required_slots + ] + if len(available_slots) != len(free_steps): + raise ValueError( + "Arrangement slot constraints do not define a one-to-one assignment." + ) + axis_positions = { + step.id: _live_pose(self.env, step.object_uid)[:, self.axis_index, 3] + for step in free_steps + } + for env_id in range(self.num_envs): + ordered_steps = sorted( + free_steps, + key=lambda step: ( + float(axis_positions[step.id][env_id]), + int(step.goal.get("nominal_slot_index", 0)), + step.id, + ), + ) + ordered_slots = sorted( + available_slots, + key=lambda slot_id: ( + float(self.positions[env_id, slot_id, self.axis_index]), + slot_id, + ), + ) + matching_cost = 0.0 + changed = False + for step, slot_id in zip(ordered_steps, ordered_slots): + nominal = int(step.goal.get("nominal_slot_index", 0)) + assignments[step.id][env_id] = slot_id + changed |= slot_id != nominal + matching_cost += abs( + float(axis_positions[step.id][env_id]) + - float(self.positions[env_id, slot_id, self.axis_index]) + ) + if changed: + self.reassignment_reason[env_id] = ( + "free arrangement initialized from live spatial order" + ) + self.reassignment_cost[env_id] = matching_cost + return assignments + + def _geometry(self, step: SemanticStep) -> _Geometry: + entity = _object(self.env, step.object_uid) + radii = [] + heights = [] + for env_id in range(self.num_envs): + vertices = _local_vertices(entity, self.env, env_id) + half_extent = ( + vertices.max(dim=0).values - vertices.min(dim=0).values + ) * 0.5 + if step.goal.get("orientation_goal", "none") in {"none", "preserve"}: + rotation = _live_pose(self.env, step.object_uid)[env_id, :3, :3] + rotated = vertices @ rotation.transpose(0, 1) + radii.append(torch.linalg.vector_norm(rotated[:, :2], dim=-1).max()) + else: + # A non-preserve target may rotate the longest local dimension + # into the table plane, so retain the conservative bound. + radii.append( + torch.linalg.vector_norm(torch.topk(half_extent, k=2).values) + ) + heights.append((vertices[:, 2].max() - vertices[:, 2].min()) * 0.5) + return _Geometry(torch.stack(radii), torch.stack(heights)) + + def _make_slots(self) -> torch.Tensor: + offsets = ( + torch.arange(self.slot_count, device=self.device, dtype=torch.float32) + - (self.slot_count - 1) / 2.0 + ) + slots = torch.empty( + self.num_envs, + self.slot_count, + 3, + dtype=torch.float32, + device=self.device, + ) + radii = torch.stack( + [self.geometry[step.id].radius for step in self.steps], + dim=1, + ) + # Free slot rematching allows any remaining object to occupy any slot. + # Size every slot for the largest member in that environment rather + # than accidentally baking the nominal object order into geometry. + slot_radii = radii.max(dim=1).values[:, None].repeat(1, self.slot_count) + obstacles = self._obstacle_bounds() + search_offsets = [0.0] + steps = int(self.row_search_radius / self.row_search_step) + for index in range(1, steps + 1): + offset = self.row_search_step * index + search_offsets.extend((offset, -offset)) + for env_id in range(self.num_envs): + chosen = None + for perpendicular in search_offsets: + candidate = self.table_center[env_id].repeat(self.slot_count, 1) + candidate[:, self.axis_index] += self.spacing[env_id] * offsets + candidate[:, self.perpendicular_index] += perpendicular + candidate[:, 2] = self.table_top[env_id] + if self._safe( + candidate, + slot_radii[env_id], + self.table_bounds[env_id], + obstacles[env_id], + ): + chosen = candidate + break + if chosen is None: + raise ValueError( + f"Environment {env_id} has no collision-free arrangement row." + ) + slots[env_id] = chosen + return slots + + def _obstacle_bounds( + self, + ) -> list[list[tuple[torch.Tensor, torch.Tensor]]]: + result: list[list[tuple[torch.Tensor, torch.Tensor]]] = [ + [] for _ in range(self.num_envs) + ] + getter = getattr(self.env.sim, "get_rigid_object_uid_list", None) + if not callable(getter): + return result + movable = {step.object_uid for step in self.steps} + for uid in getter(): + if uid == "table" or uid in movable: + continue + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + continue + for env_id in range(self.num_envs): + vertices = _world_vertices(entity, self.env, env_id) + if float(vertices[:, 2].max()) < float( + self.table_top[env_id] - self.clearance + ): + continue + result[env_id].append( + ( + vertices[:, :2].min(dim=0).values, + vertices[:, :2].max(dim=0).values, + ) + ) + return result + + def _safe( + self, + slots: torch.Tensor, + radii: torch.Tensor, + table_bounds: torch.Tensor, + obstacles: Sequence[tuple[torch.Tensor, torch.Tensor]], + ) -> bool: + lower = table_bounds[0, :2] + radii[:, None] + self.clearance + upper = table_bounds[1, :2] - radii[:, None] - self.clearance + if bool(((slots[:, :2] < lower) | (slots[:, :2] > upper)).any()): + return False + for center, radius in zip(slots[:, :2], radii): + for obstacle_lower, obstacle_upper in obstacles: + closest = torch.maximum( + obstacle_lower, + torch.minimum(center, obstacle_upper), + ) + if float(torch.linalg.vector_norm(center - closest)) <= float( + radius + self.clearance + ): + return False + return True + + def target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + phase: str, + policy: Mapping[str, Any], + ) -> torch.Tensor: + """Return a live final or collision-clear staging object pose.""" + if phase not in {"staging", "final"}: + raise ValueError(f"Unsupported arrangement phase {phase!r}.") + target = object_pose.clone() + env_ids = torch.arange(self.num_envs, device=self.device) + slot_ids = self.assignments[step.id] + target[:, :2, 3] = self.positions[env_ids, slot_ids, :2] + final_z = ( + self.table_top + + self.geometry[step.id].half_height + + float(policy["surface_clearance"]) + ) + target[:, 2, 3] = final_z + if phase == "staging": + target[:, 2, 3] = final_z + float(policy["transport_clearance"]) + return target + + def mark_completed(self, step_id: str, success: torch.Tensor) -> None: + self.completed[step_id] |= success.to(self.device, dtype=torch.bool) + + def remaining(self, env_id: int) -> list[str]: + return [ + step.id for step in self.steps if not bool(self.completed[step.id][env_id]) + ] + + def available_slots(self, env_id: int) -> list[int]: + occupied = { + int(self.assignments[step.id][env_id]) + for step in self.steps + if bool(self.completed[step.id][env_id]) + } + return [index for index in range(self.slot_count) if index not in occupied] + + def assign(self, env_id: int, assignment: Mapping[str, int]) -> None: + for step_id, slot_id in assignment.items(): + self.assignments[step_id][env_id] = int(slot_id) + + def metadata(self, step: SemanticStep, env_id: int) -> dict[str, Any]: + """Describe the live slot resolution used by one environment.""" + nominal = int(step.goal.get("nominal_slot_index", 0)) + resolved = int(self.assignments[step.id][env_id]) + return { + "nominal_slot_index": nominal, + "resolved_slot_index": resolved, + "slot_constraint": str(step.goal.get("slot_constraint", "required")), + "slot_reassigned": resolved != nominal, + "reassignment_reason": self.reassignment_reason[env_id], + "matching_cost": ( + float(self.reassignment_cost[env_id]) + if torch.isfinite(self.reassignment_cost[env_id]) + else None + ), + "spacing": float(self.spacing[env_id]), + "resolved_slot_position": self.positions[env_id, resolved].tolist(), + } + + +class LivePlacementPlan: + """Allocate non-overlapping live slots for one shared container.""" + + def __init__( + self, + env: Any, + steps: Sequence[SemanticStep], + *, + clearance: float | None = None, + ) -> None: + if not steps: + raise ValueError("A placement plan requires at least one step.") + references = {step.goal.get("reference_object") for step in steps} + if len(references) != 1 or not isinstance(next(iter(references)), str): + raise ValueError("Placement-plan steps must share one reference object.") + self.env = env + self.steps = tuple(steps) + self.reference_uid = str(next(iter(references))) + self.num_envs = int(env.num_envs) + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + default_clearance = default_runtime_policy(profile).grounding["placement"][ + "clearance" + ] + self.clearance = float(default_clearance if clearance is None else clearance) + self.positions = self._make_slots() + + def _make_slots(self) -> dict[str, torch.Tensor]: + container = _object(self.env, self.reference_uid) + positions = { + step.id: torch.empty( + self.num_envs, + 3, + dtype=torch.float32, + device=self.env.device, + ) + for step in self.steps + } + named_slots = [str(step.goal.get("slot", "auto")) for step in self.steps] + for slot in named_slots: + if slot not in {"auto", "left", "center", "right"}: + raise ValueError(f"Unsupported container slot {slot!r}.") + + for env_id in range(self.num_envs): + vertices = _world_vertices(container, self.env, env_id) + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + center = (lower + upper) * 0.5 + extent = upper[:2] - lower[:2] + axis = int(torch.argmax(extent).item()) + radii = [] + for step in self.steps: + moved_vertices = _local_vertices( + _object(self.env, step.object_uid), + self.env, + env_id, + ) + half = ( + moved_vertices.max(dim=0).values - moved_vertices.min(dim=0).values + )[:2] * 0.5 + radii.append(float(torch.linalg.vector_norm(half))) + radius = max(radii) + usable_span = float(extent[axis]) - 2.0 * (radius + self.clearance) + required_span = 2.0 * radius * max(len(self.steps) - 1, 0) + if usable_span + 1.0e-6 < required_span: + raise ValueError( + f"Environment {env_id} container {self.reference_uid!r} " + "has no non-overlapping slot plan." + ) + offsets = torch.linspace( + -required_span * 0.5, + required_span * 0.5, + len(self.steps), + device=self.env.device, + ) + named_offsets = { + "left": required_span * 0.5, + "center": 0.0, + "right": -required_span * 0.5, + } + used: list[float] = [] + for index, step in enumerate(self.steps): + slot = named_slots[index] + offset = ( + float(offsets[index]) if slot == "auto" else named_offsets[slot] + ) + if any(abs(offset - item) < 2.0 * radius for item in used): + raise ValueError( + f"Container slot {slot!r} overlaps another requested slot." + ) + used.append(offset) + target = center.clone() + target[axis] += offset + target[2] = lower[2] + positions[step.id][env_id] = target + return positions + + def target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + rotation: torch.Tensor, + *, + surface_clearance: float, + ) -> torch.Tensor: + """Return a slot pose corrected for the rotated object mesh bottom.""" + target = object_pose.clone() + target[:, :3, :3] = rotation + target[:, :2, 3] = self.positions[step.id][:, :2] + entity = _object(self.env, step.object_uid) + for env_id in range(self.num_envs): + bottom = ( + _local_vertices(entity, self.env, env_id) + @ rotation[env_id].transpose(0, 1) + )[:, 2].min() + target[env_id, 2, 3] = ( + self.positions[step.id][env_id, 2] + surface_clearance - bottom + ) + return target + + +class ActionGrounder: + """Translate one symbolic action into a public typed atomic-action target.""" + + def __init__( + self, + program: ExecutionProgram, + env: Any, + semantics_factory: Callable[[str], ObjectSemantics], + arrangement: ( + LiveArrangementPlan | Mapping[str, LiveArrangementPlan] | None + ) = None, + placements: Mapping[str, LivePlacementPlan] | None = None, + runtime_policy: RuntimePolicyCfg | None = None, + capability_registry: Any | None = None, + ) -> None: + self.program = program + self.env = env + self.semantics_factory = semantics_factory + self.capabilities = capability_registry or build_atomic_capability_registry() + self.robot_profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + self.runtime_policy = runtime_policy or default_runtime_policy( + self.robot_profile + ) + if isinstance(arrangement, Mapping): + self.arrangements = dict(arrangement) + elif arrangement is None: + self.arrangements = {} + else: + self.arrangements = { + step.id: arrangement + for step in program.semantic_steps + if step.operator in {"arrange_line", "place_in_line"} + } + self.placements = dict(placements or {}) + + def policy( + self, + action: Mapping[str, Any], + *, + extra_modifiers: tuple[tuple[str, str], ...] = (), + ) -> dict[str, Any]: + action_class = str(action.get("atomic_action_class", "")) + capability = self.capabilities.get(action_class) + motion_base = capability.motion_base or capability.name + policy_spec = action.get("motion_policy", {"modifiers": []}) + if extra_modifiers: + policy_spec = with_motion_modifiers(policy_spec, *extra_modifiers) + inline = action.get("motion_policy_config", action.get("cfg")) + return resolve_motion_policy( + self.robot_profile, + motion_base, + policy_spec, + motion_defaults=self.runtime_policy.motion_defaults, + motion_modifiers=self.runtime_policy.motion_modifiers, + inline_overrides=inline if isinstance(inline, Mapping) else None, + ) + + def _policy_value(self, policy: Mapping[str, Any], key: str) -> Any: + defaults = self.runtime_policy.grounding["semantic_defaults"] + return policy[key] if key in policy else defaults[key] + + def ground( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + _handover_workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> GroundedAction: + action_class = str(action["atomic_action_class"]) + capability = self.capabilities.require_executable(action_class) + self.capabilities.validate_binding(action) + control = str(action.get("control", "arm")) + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("target_binding must be a mapping.") + kind = str(binding.get("kind", "")) + orientation = compile_orientation_constraint(step.goal) + is_handover_continuation = self._is_handover_continuation(step) + uses_handover_staging = ( + kind == "handover_staging" + and capability.target_materializer == "semantic_held_object" + ) + use_upright_yaw_search = ( + is_handover_continuation or uses_handover_staging + ) and self._uses_upright_yaw_search( + step, + orientation, + ) + extra_modifiers: tuple[tuple[str, str], ...] = () + if ( + is_handover_continuation + and use_upright_yaw_search + and capability.target_materializer + in { + "semantic_held_object", + "current_held_pose", + "eef_pose", + } + ): + extra_modifiers = (("orientation", "upright"),) + policy = self.policy(action, extra_modifiers=extra_modifiers) + if kind == "joint_state": + joint_defaults = self.runtime_policy.grounding["joint_state"] + source = binding.get("source") + if source == "gripper_closed": + policy["sample_interval"] = int( + joint_defaults["hand_close_sample_interval"] + ) + elif source == "gripper_open": + policy["sample_interval"] = int( + joint_defaults["hand_open_sample_interval"] + ) + elif source == "initial" and control == "arm": + # Returning home after release is a safety motion. If the + # collision-aware planner cannot find a route, do not silently + # replace it with collision-unaware joint interpolation. + policy["collision_safety"] = "required" + if uses_handover_staging and use_upright_yaw_search: + # Handover consumes the live payload pose immediately after this + # move. Use the existing upright-yaw feasibility search instead + # of the generic transport orientation heuristic, which can tilt + # a payload while moving it to the exchange point. + policy["upright_yaw_samples"] = max( + int(policy.get("upright_yaw_samples", 1)), + 8, + ) + object_pose = _live_pose(self.env, step.object_uid) + if step.operator == "orient_object": + policy["upright_local_axis"] = self._upright_local_axis(step) + if capability.target_materializer == "object_grasp": + policy["obj_upright_direction"] = self._upright_local_direction(step) + reference_pose = self._reference_pose(step) + target_object_pose = None + + if capability.target_materializer_hook is not None: + grounded = capability.target_materializer_hook( + grounder=self, + action=action, + step=step, + arm=arm, + state=state, + binding=binding, + policy=policy, + object_pose=object_pose, + reference_pose=reference_pose, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + if not isinstance(grounded, GroundedAction): + raise TypeError( + f"AtomicAction {action_class!r} target materializer must " + "return GroundedAction." + ) + return grounded + + if kind == "object": + semantics = self.semantics_factory( + str(binding.get("object", step.object_uid)) + ) + if capability.target_materializer == "object_grasp": + if step.operator == "pour": + semantics = self._pour_source_semantics( + step, + semantics, + object_pose, + reference_pose, + ) + target: Any = GraspGoal(semantics=semantics) + elif capability.target_materializer == "axis_align": + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise ValueError( + "AxisAlign requires an AntipodalAffordance with mesh geometry." + ) + semantics = replace( + semantics, + affordance=AxisAlignAffordance( + object_label=affordance.object_label, + custom_config=dict(affordance.custom_config), + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + generator_cfg=affordance.generator_cfg, + gripper_collision_cfg=affordance.gripper_collision_cfg, + force_reannotate=affordance.force_reannotate, + internal_axis=self._upright_local_direction(step), + ), + ) + policy.setdefault( + "surface_clearance", + float(self._policy_value(policy, "surface_clearance")), + ) + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = AxisAlignGoal( + semantics=semantics, + object_target_pose=target_object_pose, + ) + elif capability.target_materializer == "coordinated_pickment": + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = CoordinatedPickGoal( + object_target_pose=target_object_pose, + semantics=semantics, + object_initial_pose=object_pose, + ) + elif capability.target_materializer == "press": + target_object_pose = object_pose.clone() + target, press_policy = self._press_goal( + step.object_uid, + object_pose, + semantics=semantics, + terminal_state=str( + step.postcondition.get("terminal_state", "activated") + ), + ) + policy.update(press_policy) + else: + raise ValueError( + f"{action_class} does not support object target bindings." + ) + elif kind == "pour_goal": + if capability.target_materializer != "pour": + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve a pour_goal." + ) + contents = step.goal.get("contents", ()) + if not isinstance(contents, Sequence) or isinstance( + contents, (str, bytes, bytearray) + ): + raise ValueError("Pour contents must be a list of object bindings.") + content_uids = [ + item.get("object") if isinstance(item, Mapping) else item + for item in contents + ] + if not content_uids or any( + not isinstance(uid, str) or not uid for uid in content_uids + ): + raise ValueError( + "Pour requires independently observable content objects; " + "a texture or contents baked into the source mesh is not " + "physical transfer evidence." + ) + for uid in content_uids: + _object(self.env, str(uid)) + policy.setdefault("rotate_angle", math.pi / 2.0) + target = PourGoal() + elif kind == "articulation_goal": + if capability.target_materializer == "slide": + target, policy = self._slide_target( + step, + arm, + policy, + ) + elif capability.target_materializer == "twist": + target, policy = self._twist_target(step, arm, policy) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve an articulation_goal." + ) + elif kind in {"semantic_goal", "coordinated_goal"}: + phase = str(binding.get("phase", "final")) + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase=phase, + orientation_reference_pose=orientation_reference_pose, + ) + if capability.target_materializer == "coordinated_pickment": + semantics = self.semantics_factory(step.object_uid) + target = CoordinatedPickGoal( + object_target_pose=target_object_pose, + semantics=semantics, + object_initial_pose=object_pose, + ) + elif capability.target_materializer == "press": + # Press moves the TCP, not the target object. Keep the object's + # live pose as the postcondition reference while grounding a + # downward contact point from its current surface geometry. + target_object_pose = object_pose.clone() + target, press_policy = self._press_goal( + step.object_uid, + object_pose, + terminal_state=str( + step.postcondition.get("terminal_state", "activated") + ), + ) + policy.update(press_policy) + elif capability.target_materializer == "semantic_held_object": + target = HeldObjectPoseGoal(object_target_pose=target_object_pose) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} cannot " + f"resolve {kind!r}." + ) + elif kind == "coordinated_placement_goal": + support_uid = binding.get( + "support_object", + step.goal.get("support_object"), + ) + placing_uid = binding.get("placing_object", step.object_uid) + if not isinstance(placing_uid, str) or not placing_uid: + raise ValueError("coordinated_placement_goal requires placing_object.") + if not isinstance(support_uid, str) or not support_uid: + raise ValueError("coordinated_placement_goal requires support_object.") + support_pose = _live_pose(self.env, support_uid) + target_object_pose = self._semantic_target( + step, + object_pose, + support_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = CoordinatedPlacementGoal( + placing_object_target_pose=target_object_pose, + support_object_target_pose=support_pose, + release=bool(step.goal.get("release", True)), + ) + elif kind == "current_held_pose": + if state.get_held_object(arm_control_part(self.env, arm)) is None: + raise ValueError("Place requires a held object from a prior PickUp.") + target = PlaceGoal( + xpos=( + reference_eef_pose + if reference_eef_pose is not None + else self._current_eef_pose(arm) + ) + ) + elif kind == "policy_pose": + source = binding.get("source") + retreat_reference = self._retreat_reference_pose( + arm, + reference_eef_pose, + ) + if binding.get("operation") == "retreat": + policy["retreat_reachability_search"] = True + policy["retreat_reference_pose"] = retreat_reference.clone() + if source in {"release", "handover"}: + policy["clearance_object_uid"] = step.object_uid + policy["collision_safety"] = "required" + contact_uids = [step.object_uid] + reference_uid = step.goal.get("reference_object") + if isinstance(reference_uid, str) and reference_uid: + contact_uids.append(reference_uid) + policy["collision_exclusion_uids"] = list(dict.fromkeys(contact_uids)) + if source == "handover": + policy.update(self.runtime_policy.grounding["handover"]) + policy["transfer_arm"] = arm + policy["transfer_role_axis"] = self._handover_role_axis( + arm, + dtype=object_pose.dtype, + device=object_pose.device, + ) + target = EndEffectorPoseGoal( + xpos=self._retreat_pose( + arm, + policy, + retreat_reference, + clear_exchange=source == "handover", + ) + ) + elif kind == "visual_constraint": + visual_pose = self._visual_target(binding, arm) + if capability.target_materializer == "semantic_held_object": + target_object_pose = object_pose.clone() + target_object_pose[:, :3, 3] = visual_pose[:, :3, 3] + target = HeldObjectPoseGoal(object_target_pose=target_object_pose) + elif capability.target_materializer == "eef_pose": + target = EndEffectorPoseGoal(xpos=visual_pose) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve a visual_constraint." + ) + elif kind == "joint_state": + target = JointPositionGoal( + target=self._joint_target( + arm, + control, + str(binding.get("source", "initial")), + binding, + ) + ) + elif kind in {"eef_pose", "pose"}: + target = EndEffectorPoseGoal(xpos=self._explicit_pose(binding, object_pose)) + elif kind == "handover_goal": + target, target_object_pose, policy = self._handover_target( + step, + binding, + object_pose, + reference_pose, + policy, + state, + orientation_reference_pose=orientation_reference_pose, + workspace=_handover_workspace, + ) + elif kind == "handover_staging": + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + receive_arm = str(binding.get("receive_arm", "right_arm")) + middle, _ = self._handover_workspace_poses( + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + step=step, + orientation_reference_pose=orientation_reference_pose, + ) + middle[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + target_object_pose = middle + target = HeldObjectPoseGoal(object_target_pose=middle) + else: + raise ValueError(f"Unsupported target binding kind {kind!r}.") + return GroundedAction( + action_class=action_class, + arm=arm, + control=control, + target=target, + cfg=policy, + object_pose=object_pose, + reference_pose=reference_pose, + target_object_pose=target_object_pose, + motion_policy=policy, + object_uid=step.object_uid, + ) + + def _handover_role_axis( + self, + transfer_arm: str, + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Return the world-space axis from the receiver base to transfer base.""" + if transfer_arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Unknown handover arm {transfer_arm!r}.") + _, lateral = robot_frame_axes(self.env) + horizontal = lateral if transfer_arm == "left_arm" else -lateral + return torch.cat( + ( + horizontal.to(dtype=dtype, device=device), + torch.zeros( + (int(self.env.num_envs), 1), + dtype=dtype, + device=device, + ), + ), + dim=1, + ) + + def ground_candidates( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[GroundedAction, ...]: + """Return deterministic grounding candidates for an opt-in capability.""" + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + return ( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ), + ) + placement_support_uid = self._placement_support_uid(step) + placement_relation = ( + normalize_placement_relation(step.goal.get("relation", "on")) + if step.operator == "place_relative" + else str(step.goal.get("relation", "none")) + ) + is_on_placement = ( + binding.get("kind") == "semantic_goal" + and binding.get("phase", "final") != "staging" + and placement_relation in {"on", "on_top", "on_top_of"} + and placement_support_uid is not None + ) + if is_on_placement: + base = self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + return self._placement_grounding_candidates( + base, + step, + support_uid=placement_support_uid, + ) + if binding.get("kind") != "handover_goal": + return ( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ), + ) + policy = self.policy(action) + object_pose = _live_pose(self.env, step.object_uid) + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + workspaces = self._handover_workspace_candidates( + step, + object_pose, + transfer_arm=str(binding.get("transfer_arm", "left_arm")), + receive_arm=str(binding.get("receive_arm", "right_arm")), + policy=policy, + rotation=rotation, + ) + return tuple( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + _handover_workspace=workspace, + ) + for workspace in workspaces + ) + + def _placement_grounding_candidates( + self, + base: GroundedAction, + step: SemanticStep, + *, + support_uid: str, + ) -> tuple[GroundedAction, ...]: + """Sample bounded support-relative poses from live object geometry.""" + if base.target_object_pose is None or not isinstance( + base.target, HeldObjectPoseGoal + ): + return (base,) + support = _object(self.env, support_uid) + moved = _object(self.env, step.object_uid) + placement = self.runtime_policy.grounding["placement"] + count = int(placement["candidate_count"]) + fraction = float(placement["candidate_offset_fraction"]) + margin = float(placement["support_margin"]) + patterns = ( + (0.0, 0.0), + (1.0, 0.0), + (-1.0, 0.0), + (0.0, 1.0), + (0.0, -1.0), + (1.0, 1.0), + (1.0, -1.0), + (-1.0, 1.0), + (-1.0, -1.0), + )[:count] + candidates: list[GroundedAction] = [] + seen_offsets: list[torch.Tensor] = [] + for candidate_index, pattern in enumerate(patterns): + target_pose = base.target_object_pose.clone() + offsets = target_pose.new_zeros((int(self.env.num_envs), 2)) + for env_id in range(int(self.env.num_envs)): + support_vertices = _world_vertices(support, self.env, env_id) + moved_local = _local_vertices(moved, self.env, env_id) + rotated = moved_local @ target_pose[env_id, :3, :3].transpose(0, 1) + support_lower = support_vertices[:, :2].min(dim=0).values + support_upper = support_vertices[:, :2].max(dim=0).values + moved_lower = rotated[:, :2].min(dim=0).values + moved_upper = rotated[:, :2].max(dim=0).values + allowed_lower = support_lower + margin - moved_lower + allowed_upper = support_upper - margin - moved_upper + if bool(torch.all(allowed_lower <= allowed_upper)): + base_xy = target_pose[env_id, :2, 3].clone() + center = torch.minimum( + torch.maximum(base_xy, allowed_lower), + allowed_upper, + ) + direction = target_pose.new_tensor(pattern) + room = torch.where( + direction >= 0.0, + allowed_upper - center, + center - allowed_lower, + ) + candidate_xy = center + direction * room * fraction + offsets[env_id] = candidate_xy - base_xy + target_pose[env_id, :2, 3] = candidate_xy + + footprint_lower = target_pose[env_id, :2, 3] + moved_lower + footprint_upper = target_pose[env_id, :2, 3] + moved_upper + local_mask = torch.all( + (support_vertices[:, :2] >= footprint_lower - margin) + & (support_vertices[:, :2] <= footprint_upper + margin), + dim=1, + ) + if bool(local_mask.any()): + support_height = support_vertices[local_mask, 2].max() + else: + distances = torch.linalg.vector_norm( + support_vertices[:, :2] - target_pose[env_id, :2, 3], + dim=1, + ) + nearest_count = min(8, int(support_vertices.shape[0])) + nearest = torch.topk( + distances, + nearest_count, + largest=False, + ).indices + support_height = support_vertices[nearest, 2].max() + target_pose[env_id, 2, 3] = ( + support_height + + float(self._policy_value(base.motion_policy, "surface_clearance")) + - rotated[:, 2].min() + ) + if any(torch.allclose(offsets, prior) for prior in seen_offsets): + continue + seen_offsets.append(offsets) + candidates.append( + replace( + base, + target=replace(base.target, object_target_pose=target_pose), + target_object_pose=target_pose, + motion_policy={ + **base.motion_policy, + "placement_candidate_index": candidate_index, + "placement_xy_offset": offsets, + }, + ) + ) + return tuple(candidates) or (base,) + + @staticmethod + def _placement_support_uid(step: SemanticStep) -> str | None: + value = step.goal.get("reference_object", step.goal.get("support_object")) + if isinstance(value, str) and value: + return value + if ( + step.postcondition.get("type") == "stack_layer_supported" + and int(step.goal.get("layer_index", -1)) == 0 + ): + return "table" + return None + + def _is_handover_continuation(self, step: SemanticStep) -> bool: + if step.operator != "place_relative": + return False + predecessors = { + candidate.id: candidate for candidate in self.program.semantic_steps + } + return any( + (predecessor := predecessors.get(dependency)) is not None + and predecessor.operator == "handover" + and predecessor.object_uid == step.object_uid + for dependency in step.depends_on + ) + + def _visual_target( + self, + binding: Mapping[str, Any], + arm: str, + ) -> torch.Tensor: + """Unproject one normalized image keypoint using live camera depth.""" + camera_uid = str(binding.get("camera_uid", "")) + sensor = self.env.sim.get_sensor(camera_uid) + if sensor is None: + raise ValueError(f"Unknown visual-constraint camera {camera_uid!r}.") + keypoint_value = binding.get("normalized_keypoint") + if keypoint_value is None: + bbox = binding.get("normalized_bbox") + if isinstance(bbox, Sequence) and len(bbox) == 4: + keypoint_value = [ + (float(bbox[0]) + float(bbox[2])) * 0.5, + (float(bbox[1]) + float(bbox[3])) * 0.5, + ] + if keypoint_value is None: + raise ValueError( + "visual_constraint requires a normalized keypoint or bbox in [0, 1]." + ) + keypoint = torch.as_tensor( + keypoint_value, + dtype=torch.float32, + device=self.env.device, + ).flatten() + if keypoint.numel() != 2 or bool( + ((~torch.isfinite(keypoint)) | (keypoint < 0.0) | (keypoint > 1.0)).any() + ): + raise ValueError( + "visual_constraint requires a normalized keypoint or bbox in [0, 1]." + ) + data = sensor.get_data() + if "depth" not in data: + raise ValueError( + f"Camera {camera_uid!r} must provide depth for visual Grounding." + ) + depth = torch.as_tensor(data["depth"], device=self.env.device).squeeze(-1) + if depth.ndim == 2: + depth = depth.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + if depth.ndim != 3 or depth.shape[0] != int(self.env.num_envs): + raise ValueError("Camera depth must have shape (N, H, W) or (N, H, W, 1).") + height, width = depth.shape[-2:] + pixel_x = min(max(int(round(float(keypoint[0]) * (width - 1))), 0), width - 1) + pixel_y = min(max(int(round(float(keypoint[1]) * (height - 1))), 0), height - 1) + distance = depth[:, pixel_y, pixel_x].to(torch.float32) + if bool((~torch.isfinite(distance) | (distance <= 0.0)).any()): + raise ValueError("visual_constraint keypoint has no valid live depth.") + intrinsics = torch.as_tensor( + sensor.get_intrinsics(), + dtype=torch.float32, + device=self.env.device, + ) + if intrinsics.ndim == 2: + intrinsics = intrinsics.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + camera_pose = torch.as_tensor( + sensor.get_arena_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + if camera_pose.ndim == 2: + camera_pose = camera_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + fx = intrinsics[:, 0, 0] + fy = intrinsics[:, 1, 1] + cx = intrinsics[:, 0, 2] + cy = intrinsics[:, 1, 2] + point = torch.stack( + ( + (float(pixel_x) - cx) * distance / fx, + (float(pixel_y) - cy) * distance / fy, + distance, + torch.ones_like(distance), + ), + dim=1, + ) + world = torch.bmm(camera_pose, point.unsqueeze(-1)).squeeze(-1) + target = self._current_eef_pose(arm).clone() + target[:, :3, 3] = world[:, :3] + return target + + def _handover_target( + self, + step: SemanticStep, + binding: Mapping[str, Any], + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + policy: Mapping[str, Any], + state: ExecutionState, + *, + orientation_reference_pose: torch.Tensor | None, + workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[GraspGoal, torch.Tensor, dict[str, Any]]: + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + receive_arm = str( + binding.get( + "receive_arm", + "right_arm" if transfer_arm == "left_arm" else "left_arm", + ) + ) + if transfer_arm == receive_arm or {transfer_arm, receive_arm} != { + "left_arm", + "right_arm", + }: + raise ValueError("HandOver requires distinct left_arm/right_arm roles.") + transfer_part = arm_control_part(self.env, transfer_arm) + held = state.get_held_object(transfer_part) + if held is None: + raise ValueError( + f"HandOver requires {transfer_arm} to hold {step.object_uid!r}." + ) + + del reference_pose + if workspace is None: + middle, final = self._handover_workspace_poses( + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + step=step, + orientation_reference_pose=orientation_reference_pose, + ) + else: + middle, final = (item.clone() for item in workspace) + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + middle[:, :3, :3] = rotation + final[:, :3, :3] = rotation + semantics = self.semantics_factory(step.object_uid) + grounded_policy = dict(policy) + grounded_policy.update( + { + "transfer_arm": transfer_arm, + "receive_arm": receive_arm, + "middle_object_pose": middle, + "final_object_pose": final, + } + ) + return ( + GraspGoal(semantics=semantics), + middle, + grounded_policy, + ) + + def _handover_workspace_poses( + self, + object_pose: torch.Tensor, + *, + transfer_arm: str, + receive_arm: str, + policy: Mapping[str, Any], + step: SemanticStep, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Choose the highest-ranked collision-aware handover workspace.""" + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + candidates = self._handover_workspace_candidates( + step, + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + rotation=rotation, + ) + return candidates[0] + + def _handover_workspace_candidates( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + transfer_arm: str, + receive_arm: str, + policy: Mapping[str, Any], + rotation: torch.Tensor, + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + """Rank exchange poses inside the two arm workspaces and above obstacles.""" + if transfer_arm == receive_arm or {transfer_arm, receive_arm} != { + "left_arm", + "right_arm", + }: + raise ValueError("Handover workspace requires distinct arm roles.") + table = self.env.sim.get_rigid_object("table") + if table is not None and hasattr(table, "get_vertices"): + centers = [] + tops = [] + bounds = [] + for env_id in range(int(self.env.num_envs)): + vertices = _world_vertices(table, self.env, env_id) + lower = vertices[:, :2].min(dim=0).values + upper = vertices[:, :2].max(dim=0).values + centers.append((lower + upper) * 0.5) + tops.append(vertices[:, 2].max()) + bounds.append(torch.stack((lower, upper))) + center = torch.stack(centers) + table_top = torch.stack(tops) + table_bounds = torch.stack(bounds) + else: + left = self._current_eef_pose("left_arm") + right = self._current_eef_pose("right_arm") + center = (left[:, :2, 3] + right[:, :2, 3]) * 0.5 + table_top = object_pose[:, 2, 3] + extent = float(policy.get("exchange_candidate_offset", 0.16)) * 2.0 + table_bounds = torch.stack((center - extent, center + extent), dim=1) + + forward, lateral = robot_frame_axes(self.env) + left_base, right_base = arm_base_poses(self.env) + transfer_base = left_base if transfer_arm == "left_arm" else right_base + receive_base = right_base if receive_arm == "right_arm" else left_base + base_midpoint = (transfer_base[:, :2, 3] + receive_base[:, :2, 3]) * 0.5 + table_forward = torch.sum((center - base_midpoint) * forward, dim=1) + shared_center = base_midpoint + forward * table_forward[:, None] + offset = float(policy.get("exchange_candidate_offset", 0.16)) + obstacle_clearance = float(policy.get("exchange_obstacle_clearance", 0.04)) + tool_horizontal_envelope = float( + policy.get("exchange_gripper_horizontal_envelope", 0.035) + ) + float(policy.get("exchange_wrist_horizontal_envelope", 0.055)) + tool_vertical_envelope = float( + policy.get("exchange_gripper_vertical_envelope", 0.025) + ) + float(policy.get("exchange_wrist_vertical_envelope", 0.04)) + minimum_reach = float(policy.get("exchange_minimum_reach", 0.10)) + maximum_reach = float(policy.get("exchange_maximum_reach", 1.00)) + if not 0.0 <= minimum_reach < maximum_reach: + raise ValueError("Handover reach bounds require 0 <= minimum < maximum.") + requested_count = max(1, int(policy.get("exchange_candidate_count", 4))) + object_clearance = float(policy.get("exchange_clearance", 0.06)) + if ( + min( + obstacle_clearance, + tool_horizontal_envelope, + tool_vertical_envelope, + object_clearance, + ) + < 0.0 + ): + raise ValueError("Handover geometry clearances must be non-negative.") + xy_coefficients = ( + (0.0, 0.0), + (1.0, 0.0), + (-1.0, 0.0), + (2.0, 0.0), + (-2.0, 0.0), + (0.0, 0.5), + (0.0, -0.5), + ) + ranked_by_env: list[list[tuple[float, torch.Tensor]]] = [] + moved = _object(self.env, step.object_uid) + obstacle_uids = ( + self.env.sim.get_rigid_object_uid_list() + if hasattr(self.env.sim, "get_rigid_object_uid_list") + else [] + ) + for env_id in range(int(self.env.num_envs)): + local_vertices = _local_vertices(moved, self.env, env_id) + rotated = local_vertices @ rotation[env_id].transpose(0, 1) + half_xy = ( + rotated[:, :2].max(dim=0).values - rotated[:, :2].min(dim=0).values + ) * 0.5 + bottom = rotated[:, 2].min() + margin = half_xy + obstacle_clearance + tool_horizontal_envelope + lower_limit = table_bounds[env_id, 0] + margin + upper_limit = table_bounds[env_id, 1] - margin + options: list[tuple[float, torch.Tensor]] = [] + for forward_scale, lateral_scale in xy_coefficients: + xy = ( + shared_center[env_id] + + forward[env_id] * (offset * forward_scale) + + lateral[env_id] * (offset * lateral_scale) + ) + if bool(((xy < lower_limit) | (xy > upper_limit)).any()): + continue + transfer_distance = torch.linalg.vector_norm( + xy - transfer_base[env_id, :2, 3] + ) + receive_distance = torch.linalg.vector_norm( + xy - receive_base[env_id, :2, 3] + ) + if not ( + minimum_reach <= float(transfer_distance) <= maximum_reach + and minimum_reach <= float(receive_distance) <= maximum_reach + ): + continue + obstacle_score, nearby_obstacle_top = self._handover_obstacle_metrics( + xy, + env_id=env_id, + object_uid=step.object_uid, + obstacle_uids=obstacle_uids, + half_xy=half_xy, + clearance=obstacle_clearance + tool_horizontal_envelope, + ) + center_cost = float( + torch.linalg.vector_norm(xy - shared_center[env_id]) + ) + pose = object_pose[env_id].clone() + pose[:3, :3] = rotation[env_id] + pose[:2, 3] = xy + safety_floor = torch.maximum(table_top[env_id], nearby_obstacle_top) + safe_z = ( + safety_floor + object_clearance + tool_vertical_envelope - bottom + ) + pose[2, 3] = torch.maximum(object_pose[env_id, 2, 3], safe_z) + lift_cost = max( + 0.0, + float(pose[2, 3] - object_pose[env_id, 2, 3]), + ) + options.append( + (obstacle_score + center_cost * 0.25 + lift_cost * 0.1, pose) + ) + if not options: + raise ValueError( + "No handover exchange pose lies inside the table bounds and " + "the reachable intersection of both arm bases." + ) + options.sort(key=lambda item: item[0]) + ranked_by_env.append(options[:requested_count]) + + candidate_count = min( + requested_count, + max(len(options) for options in ranked_by_env), + ) + candidates = [] + for candidate_index in range(candidate_count): + middle = object_pose.clone() + for env_id, options in enumerate(ranked_by_env): + middle[env_id] = options[min(candidate_index, len(options) - 1)][1] + # The built-in HandOver primitive plans its final transfer/receiver + # phase concurrently. An exchange-to-exchange target makes that + # receiver path stationary; graph-level retreat/home nodes then + # clear the transfer arm before any receiver-side continuation. + final = middle.clone() + candidates.append((middle, final)) + return tuple(candidates) + + def _handover_obstacle_metrics( + self, + xy: torch.Tensor, + *, + env_id: int, + object_uid: str, + obstacle_uids: Sequence[str], + half_xy: torch.Tensor, + clearance: float, + ) -> tuple[float, torch.Tensor]: + score = 0.0 + highest_top = torch.tensor( + -torch.inf, + dtype=xy.dtype, + device=xy.device, + ) + for uid in obstacle_uids: + if uid in {"table", object_uid}: + continue + obstacle = self.env.sim.get_rigid_object(uid) + if obstacle is None or not hasattr(obstacle, "get_vertices"): + continue + vertices = _world_vertices(obstacle, self.env, env_id) + lower = vertices[:, :2].min(dim=0).values - half_xy - clearance + upper = vertices[:, :2].max(dim=0).values + half_xy + clearance + outside = torch.maximum( + torch.maximum(lower - xy, xy - upper), + torch.zeros_like(xy), + ) + if bool((outside > 0.0).any()): + distance = float(torch.linalg.vector_norm(outside)) + score += 1.0 / max(distance, 1.0e-3) + else: + score += 1.0e3 + highest_top = torch.maximum(highest_top, vertices[:, 2].max()) + return score, highest_top + + def _handover_receiver_exit( + self, + middle: torch.Tensor, + receive_arm: str, + policy: Mapping[str, Any], + ) -> torch.Tensor: + final = middle.clone() + receive_pose = self._current_eef_pose(receive_arm) + direction = receive_pose[:, :2, 3] - middle[:, :2, 3] + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + fallback = direction.new_zeros(direction.shape) + fallback[:, 1] = -1.0 if receive_arm == "right_arm" else 1.0 + direction = torch.where( + norm > 1.0e-6, direction / norm.clamp_min(1.0e-6), fallback + ) + final[:, :2, 3] += direction * min( + 0.12, + float(self._policy_value(policy, "relation_distance")) * 0.5, + ) + return final + + def _reference_pose(self, step: SemanticStep) -> torch.Tensor | None: + uid = step.goal.get("reference_object", step.goal.get("support_object")) + if not isinstance(uid, str) or not uid: + return None + if step.goal.get("reference_state") == "initial": + initial = getattr(self.env, "agent_initial_object_poses", {}).get(uid) + if initial is None: + raise ValueError(f"Initial pose for {uid!r} is unavailable.") + return _batched_pose(initial, self.env) + return _live_pose(self.env, uid) + + def _semantic_target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + policy: Mapping[str, Any], + *, + phase: str, + orientation_reference_pose: torch.Tensor | None = None, + ) -> torch.Tensor: + if step.operator in {"arrange_line", "place_in_line"}: + arrangement = self.arrangements.get(step.id) + if arrangement is None: + raise ValueError("arrange_line requires a live arrangement plan.") + target = arrangement.target( + step, + object_pose, + phase=phase, + policy=policy, + ) + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + arrangement.table_top[env_id] + + float(policy["surface_clearance"]) + - bottom + ) + if phase == "staging": + target[:, 2, 3] += float(policy["transport_clearance"]) + return target + placement = self.placements.get(step.id) + if placement is not None: + target = placement.target( + step, + object_pose, + self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ), + surface_clearance=float(policy["surface_clearance"]), + ) + if phase == "staging": + target[:, 2, 3] += float(policy["transport_clearance"]) + return target + if step.operator == "orient_object": + initial = None + if step.goal.get("position_anchor", "initial_xy") == "initial_xy": + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + target = ( + _batched_pose(initial, self.env).clone() + if initial is not None + else object_pose.clone() + ) + target[:, :3, :3] = self._target_rotation( + step, + target, + orientation_reference_pose=orientation_reference_pose, + ) + support_uid = str(step.goal.get("support_object", "table")) + support = _object(self.env, support_uid) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + float(policy["surface_clearance"]) - bottom + ) + if phase == "staging": + target[:, 2, 3] += float(policy["staging_lift_height"]) + return target + if step.operator == "pour": + if reference_pose is None: + raise ValueError("Pour requires a live target-container pose.") + target = object_pose.clone() + source = _object(self.env, step.object_uid) + receiver_uid = str(step.goal.get("reference_object", "")) + receiver = _object(self.env, receiver_uid) + clearance = float(self._policy_value(policy, "transport_clearance")) + for env_id in range(int(self.env.num_envs)): + receiver_vertices = _world_vertices(receiver, self.env, env_id) + target[env_id, :2, 3] = ( + receiver_vertices[:, :2].min(dim=0).values + + receiver_vertices[:, :2].max(dim=0).values + ) * 0.5 + source_vertices = _local_vertices(source, self.env, env_id) + rotation_radius = torch.linalg.vector_norm(source_vertices, dim=1).max() + target[env_id, 2, 3] = ( + receiver_vertices[:, 2].max() + clearance + rotation_radius + ) + return target + target = object_pose.clone() + if reference_pose is not None: + target[:, :3, 3] = reference_pose[:, :3, 3] + # Operators without a relational goal (for example press or a + # direction-only coordinated transport) must preserve the live origin + # instead of being silently projected onto a synthetic table support. + relation = ( + normalize_placement_relation(step.goal.get("relation", "on")) + if step.operator == "place_relative" + else str(step.goal.get("relation", "none")) + ) + distance = float(self._policy_value(policy, "relation_distance")) + relation_frame = str(step.goal.get("relation_frame", "world")) + forward_distance = distance + lateral_distance = distance + if relation_frame == "robot" and reference_pose is not None: + nominal = float(policy.get("robot_relative_distance", 0.10)) + clearance = float(policy.get("relation_clearance", 0.02)) + reference_uid = str(step.goal.get("reference_object", "")) + if reference_uid: + forward_axis, lateral_axis = robot_frame_axes(self.env) + forward_distance = self._relative_object_spacing( + step.object_uid, + reference_uid, + axis=forward_axis, + nominal=nominal, + clearance=clearance, + ) + lateral_distance = self._relative_object_spacing( + step.object_uid, + reference_uid, + axis=lateral_axis, + nominal=nominal, + clearance=clearance, + ) + directional_offset = relation_offset( + self.env, + relation, + frame=relation_frame, + forward_distance=forward_distance, + lateral_distance=lateral_distance, + dtype=target.dtype, + device=target.device, + ) + offsets = { + "above": (0.0, 0.0, float(self._policy_value(policy, "hover_height"))), + "held_above_initial": ( + 0.0, + 0.0, + float(self._policy_value(policy, "hover_height")), + ), + } + if directional_offset is not None: + target[:, :3, 3] += directional_offset + elif (offset := offsets.get(relation)) is not None: + target[:, :3, 3] += torch.tensor( + offset, + dtype=target.dtype, + device=target.device, + ) + slot = str(step.goal.get("slot", "auto")) + if relation in {"on", "on_top", "on_top_of", "inside"} and slot in { + "left", + "right", + }: + slot_offset = relation_offset( + self.env, + slot, + frame=relation_frame, + forward_distance=forward_distance, + lateral_distance=lateral_distance, + dtype=target.dtype, + device=target.device, + ) + if slot_offset is not None: + target[:, :3, 3] += slot_offset + direction = str(step.goal.get("direction", "none")) + direction_offsets = { + "world_x": (distance, 0.0, 0.0), + "world_y": (0.0, distance, 0.0), + "up": (0.0, 0.0, distance), + "down": (0.0, 0.0, -distance), + } + planar_direction_offset = relation_offset( + self.env, + direction, + frame=relation_frame, + forward_distance=distance, + lateral_distance=distance, + dtype=target.dtype, + device=target.device, + ) + if planar_direction_offset is not None: + target[:, :3, 3] += planar_direction_offset + elif direction in direction_offsets: + target[:, :3, 3] += torch.tensor( + direction_offsets[direction], + dtype=target.dtype, + device=target.device, + ) + + root_stack_layer = ( + step.operator == "build_stack" + and int(step.goal.get("layer_index", 0)) == 0 + and reference_pose is None + ) + if root_stack_layer: + table = _object(self.env, "table") + for env_id in range(int(self.env.num_envs)): + vertices = _world_vertices(table, self.env, env_id) + target[env_id, :2, 3] = ( + vertices[:, :2].min(dim=0).values + + vertices[:, :2].max(dim=0).values + ) * 0.5 + + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + if ( + step.operator == "coordinated_transport" + and relation not in {"on", "on_top", "on_top_of", "inside"} + and direction not in {"up", "down"} + ): + release = str(step.goal.get("terminal_behavior", "hold")) == "place" + if not release: + target[:, 2, 3] = object_pose[:, 2, 3] + float( + self._policy_value(policy, "transport_clearance") + ) + else: + table = _object(self.env, "table") + moved = _object(self.env, step.object_uid) + clearance = float(self._policy_value(policy, "surface_clearance")) + for env_id in range(int(self.env.num_envs)): + table_top = _world_vertices(table, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = table_top + clearance - bottom + if relation in {"on", "on_top", "on_top_of"} or root_stack_layer: + support_uid = ( + step.goal.get("reference_object") + or step.goal.get("support_object") + or "table" + ) + support = _object(self.env, str(support_uid)) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + + float(self._policy_value(policy, "surface_clearance")) + - bottom + ) + elif relation == "inside" and reference_pose is not None: + # Grounding the final move happens after the staging lift. Preserve + # the pre-pick supported height rather than the lifted live height. + supported_pose = orientation_reference_pose + if supported_pose is None: + supported_pose = object_pose + supported_pose = _batched_pose(supported_pose, self.env) + target[:, 2, 3] = supported_pose[:, 2, 3] + if phase == "staging": + # Staging is a runtime waypoint, not a persisted coordinate. This + # keeps in-place orientation robust to the object's live height. + target[:, 2, 3] += float(self._policy_value(policy, "transport_clearance")) + elif self._is_handover_continuation(step) and relation not in { + "on", + "on_top", + "on_top_of", + "inside", + }: + # A handover can leave the live rigid-body center a few centimetres + # below the original table-supported height. Reusing that drifted + # height for the lateral placement target makes the can intersect + # the table during release and it may tip or slide. Preserve the + # predecessor's supported height for the final held-object pose. + supported_pose = orientation_reference_pose + if supported_pose is None: + supported_pose = object_pose + supported_pose = _batched_pose(supported_pose, self.env) + target[:, 2, 3] = torch.maximum( + target[:, 2, 3], + supported_pose[:, 2, 3], + ) + return target + + def _pour_source_semantics( + self, + step: SemanticStep, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + ) -> ObjectSemantics: + """Attach the target-directed local rotation axis used by Pour.""" + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise ValueError("Pour requires an AntipodalAffordance for pickup.") + if reference_pose is None: + raise ValueError("Pour requires a live target-container pose.") + direction = reference_pose[:, :3, 3] - object_pose[:, :3, 3] + direction[:, 2] = 0.0 + norm = torch.linalg.vector_norm(direction, dim=1) + if torch.any(norm <= 1.0e-6): + raise ValueError( + "Pour source and target must define a non-zero horizontal direction." + ) + direction = direction / norm.unsqueeze(1) + world_up = direction.new_tensor([0.0, 0.0, 1.0]).expand_as(direction) + world_axis = torch.linalg.cross(world_up, direction, dim=1) + local_axes = torch.bmm( + object_pose[:, :3, :3].transpose(1, 2), + world_axis.unsqueeze(2), + ).squeeze(2) + local_axes = torch.nn.functional.normalize(local_axes, dim=1) + if not torch.allclose( + local_axes, + local_axes[:1].expand_as(local_axes), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched Pour environments require one shared object-local " + "target-directed rotation axis." + ) + return replace( + semantics, + affordance=AxisAlignAffordance( + object_label=affordance.object_label, + custom_config=dict(affordance.custom_config), + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + generator_cfg=affordance.generator_cfg, + gripper_collision_cfg=affordance.gripper_collision_cfg, + force_reannotate=affordance.force_reannotate, + internal_axis=local_axes[0], + ), + ) + + def _slide_target( + self, + step: SemanticStep, + arm: str, + policy: Mapping[str, Any], + ) -> tuple[SlideGoal, dict[str, Any]]: + """Build one Slide goal from live prismatic-joint metadata.""" + articulation = getattr(self.env.sim, "get_articulation", lambda _uid: None)( + step.object_uid + ) + if articulation is None: + raise ValueError( + f"Articulation action requires live articulation {step.object_uid!r}." + ) + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + backend = backend_entities[0] + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == "prismatic": + candidates.append((int(joint_id), joint_name, info)) + if len(candidates) != 1: + raise ValueError( + "Slide grounding requires exactly one active prismatic joint; " + f"found {[name for _, name, _ in candidates]}." + ) + joint_id, joint_name, joint_info = candidates[0] + child_link = str(getattr(joint_info, "child_link_name", "")) + parent_link = str(getattr(joint_info, "parent_link_name", "")) + if ( + child_link not in articulation.link_names + or parent_link not in articulation.link_names + ): + raise ValueError( + "Prismatic joint metadata must identify live parent and child links." + ) + contact_links = [] + for candidate_name in getattr(articulation, "all_joint_names", ()): + candidate_info = backend.get_joint_info(str(candidate_name)) + candidate_type = ( + str( + getattr( + getattr(candidate_info, "joint_type", None), + "name", + candidate_info.joint_type, + ) + ) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + candidate_child = str(getattr(candidate_info, "child_link_name", "")) + if ( + candidate_type == "fixed" + and str(getattr(candidate_info, "parent_link_name", "")) == child_link + and candidate_child in articulation.link_names + ): + contact_links.append(candidate_child) + if len(contact_links) != 1: + raise ValueError( + "Slide grounding requires exactly one fixed contact endpoint " + f"on prismatic child link {child_link!r}; found {contact_links}." + ) + grasp_pose = _batched_pose( + articulation.get_link_pose(contact_links[0], to_matrix=True), self.env + ) + # The bundled drawer contact frame owns TCP +Z as its approach axis. + # A local quarter turn aligns the parallel-gripper closing direction + # with the narrow handle, matching scripts/tutorials/sim/open_drawer.py. + grasp_roll = torch.eye( + 4, + dtype=grasp_pose.dtype, + device=grasp_pose.device, + ) + grasp_roll[0, 0] = 0.0 + grasp_roll[0, 1] = -1.0 + grasp_roll[1, 0] = 1.0 + grasp_roll[1, 1] = 0.0 + grasp_pose = torch.matmul(grasp_pose, grasp_roll) + + vertices, triangles = articulation.get_link_vert_face(child_link) + vertices = torch.as_tensor( + vertices, dtype=torch.float32, device=self.env.device + ) + triangles = torch.as_tensor( + triangles, dtype=torch.int64, device=self.env.device + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): + raise ValueError("Slide child link has no valid grasp geometry.") + if triangles.ndim != 2 or triangles.shape[-1] != 3 or not triangles.numel(): + raise ValueError("Slide child link has no valid triangle geometry.") + + child_pose = _batched_pose( + articulation.get_link_pose(child_link, to_matrix=True), self.env + ) + parent_pose = _batched_pose( + articulation.get_link_pose(parent_link, to_matrix=True), self.env + ) + raw_axis = joint_info.axis + raw_origin = joint_info.origin_pose + axis = torch.tensor( + raw_axis.tolist() if hasattr(raw_axis, "tolist") else raw_axis, + dtype=torch.float32, + device=self.env.device, + ).reshape(3) + origin = torch.tensor( + raw_origin.tolist() if hasattr(raw_origin, "tolist") else raw_origin, + dtype=torch.float32, + device=self.env.device, + ).reshape(4, 4) + opening_world = torch.matmul( + torch.matmul(parent_pose[:, :3, :3], origin[:3, :3]), axis + ) + push_world = -torch.nn.functional.normalize(opening_world, dim=1) + push_local = torch.bmm( + child_pose[:, :3, :3].transpose(1, 2), + push_world.unsqueeze(2), + ).squeeze(2) + push_local = torch.nn.functional.normalize(push_local, dim=1) + if not torch.allclose( + push_local, + push_local[:1].expand_as(push_local), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched Slide environments require one shared child-link axis." + ) + + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + qpos = articulation.get_qpos()[:, joint_id] + if limits.shape != (int(self.env.num_envs), 2): + raise ValueError("Prismatic joint limits have an invalid batch shape.") + if not torch.isfinite(limits).all() or torch.any(limits[:, 0] >= limits[:, 1]): + raise ValueError("Slide requires finite ordered prismatic joint limits.") + if step.operator not in {"pull_articulated_part", "push_articulated_part"}: + raise ValueError("Slide grounding requires a pull or push operator.") + direction = "pull" if step.operator == "pull_articulated_part" else "push" + target_qpos = limits[:, 1] if direction == "pull" else limits[:, 0] + distances = torch.abs(target_qpos - qpos) + if torch.any(distances <= 1.0e-5): + raise ValueError("Articulation joint is already at the requested target.") + if not torch.allclose( + distances, + distances[:1].expand_as(distances), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched Slide environments require one shared translation distance." + ) + scoped_policy = dict(policy) + scoped_policy.update( + { + "direction": direction, + "translation_distance": float(distances[0]), + "articulation_joint_name": joint_name, + "articulation_joint_id": joint_id, + "articulation_initial_qpos": qpos, + "articulation_target_qpos": target_qpos, + "articulation_grasp_position": grasp_pose[:, :3, 3], + "articulation_push_axis_world": push_world, + } + ) + left_base, right_base = arm_base_poses(self.env) + arm_base = left_base if arm == "left_arm" else right_base + current_eef = self._current_eef_pose(arm) + log_info( + f"Slide grounding {step.id}/{arm}: grasp_position=" + f"{grasp_pose[0, :3, 3].detach().cpu().tolist()}, " + f"push_axis={push_world[0].detach().cpu().tolist()}, " + f"eef_position={current_eef[0, :3, 3].detach().cpu().tolist()}, " + f"arm_base_position={arm_base[0, :3, 3].detach().cpu().tolist()}." + ) + grasp_options = self.runtime_policy.grasp + sampler = AntipodalSamplerCfg( + n_sample=int(grasp_options["antipodal_n_sample"]), + max_angle=float(grasp_options["antipodal_max_angle"]), + max_length=float(grasp_options["max_open_length"]), + min_length=float(grasp_options["min_open_length"]), + ) + generator = GraspGeneratorCfg( + viser_port=int(grasp_options["viser_port"]), + antipodal_sampler_cfg=sampler, + max_deviation_angle=float(grasp_options["max_deviation_angle"]), + n_deviated_approach_directions=int( + grasp_options["n_deviated_approach_directions"] + ), + ) + max_hulls = int(grasp_options["max_decomposition_hulls"]) + collision = GripperCollisionCfg( + max_open_length=float(grasp_options["max_open_length"]), + finger_length=float(grasp_options["finger_length"]), + point_sample_dense=float(grasp_options["point_sample_dense"]), + max_decomposition_hulls=max_hulls, + ) + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=max_hulls, + ) + affordance = SlideAffordance( + object_label=f"{step.object_uid}:{child_link}", + mesh_vertices=vertices, + mesh_triangles=triangles, + generator_cfg=generator, + gripper_collision_cfg=collision, + force_reannotate=bool(grasp_options["force_grasp_reannotate"]), + translation_axis=push_local[0], + joint_name=joint_name, + joint_limits=(float(limits[0, 0]), float(limits[0, 1])), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + label=f"{step.object_uid}:{child_link}", + ) + return ( + SlideGoal( + semantics=semantics, + target_pose=child_pose, + grasp_xpos=grasp_pose, + ), + scoped_policy, + ) + + def _twist_target( + self, + step: SemanticStep, + arm: str, + policy: Mapping[str, Any], + ) -> tuple[TwistGoal, dict[str, Any]]: + """Build one Twist goal from a live revolute joint and setting map.""" + articulation = getattr(self.env.sim, "get_articulation", lambda _uid: None)( + step.object_uid + ) + if articulation is None: + raise ValueError( + f"TurnKnob requires live articulation {step.object_uid!r}." + ) + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + backend = backend_entities[0] + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == "revolute": + candidates.append((int(joint_id), joint_name, info)) + if len(candidates) != 1: + raise ValueError( + "TurnKnob grounding requires exactly one active revolute joint; " + f"found {[name for _, name, _ in candidates]}." + ) + joint_id, joint_name, joint_info = candidates[0] + child_link = str(getattr(joint_info, "child_link_name", "")) + parent_link = str(getattr(joint_info, "parent_link_name", "")) + if ( + child_link not in articulation.link_names + or parent_link not in articulation.link_names + ): + raise ValueError( + "Revolute joint metadata must identify live parent and child links." + ) + vertices, _ = articulation.get_link_vert_face(child_link) + vertices = torch.as_tensor( + vertices, dtype=torch.float32, device=self.env.device + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): + raise ValueError("TurnKnob child link has no valid grasp geometry.") + + child_pose = _batched_pose( + articulation.get_link_pose(child_link, to_matrix=True), self.env + ) + parent_pose = _batched_pose( + articulation.get_link_pose(parent_link, to_matrix=True), self.env + ) + raw_axis = joint_info.axis + raw_origin = joint_info.origin_pose + axis = torch.tensor( + raw_axis.tolist() if hasattr(raw_axis, "tolist") else raw_axis, + dtype=torch.float32, + device=self.env.device, + ).reshape(3) + origin = torch.tensor( + raw_origin.tolist() if hasattr(raw_origin, "tolist") else raw_origin, + dtype=torch.float32, + device=self.env.device, + ).reshape(4, 4) + joint_pose = torch.matmul(parent_pose, origin) + world_axis = torch.matmul(joint_pose[:, :3, :3], axis) + local_axis = torch.bmm( + child_pose[:, :3, :3].transpose(1, 2), + world_axis.unsqueeze(2), + ).squeeze(2) + local_axis = torch.nn.functional.normalize(local_axis, dim=1) + joint_origin_local = torch.bmm(torch.linalg.inv(child_pose), joint_pose)[ + :, :3, 3 + ] + if not torch.allclose( + local_axis, + local_axis[:1].expand_as(local_axis), + atol=1.0e-4, + rtol=1.0e-4, + ) or not torch.allclose( + joint_origin_local, + joint_origin_local[:1].expand_as(joint_origin_local), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched TurnKnob environments require shared local joint geometry." + ) + + agent_config = getattr(self.env, "agent_config", {}) + articulation_settings = ( + agent_config.get("articulation_settings", {}) + if isinstance(agent_config, Mapping) + else {} + ) + per_articulation = ( + articulation_settings.get(step.object_uid, {}) + if isinstance(articulation_settings, Mapping) + else {} + ) + setting_values = ( + per_articulation.get(joint_name, ()) + if isinstance(per_articulation, Mapping) + else () + ) + if ( + not isinstance(setting_values, Sequence) + or isinstance(setting_values, (str, bytes, bytearray)) + or not setting_values + ): + raise ValueError( + "TurnKnob requires explicit setting_values; ordinal settings " + "cannot be guessed from joint limits." + ) + values = torch.as_tensor( + setting_values, dtype=torch.float32, device=self.env.device + ) + if values.ndim != 1 or not torch.isfinite(values).all(): + raise ValueError("TurnKnob setting_values must be a finite list.") + setting = int(step.goal.get("target_setting", -1)) + if setting < 0 or setting >= values.numel(): + raise ValueError("TurnKnob target_setting is outside setting_values.") + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + target_qpos = values[setting] + if torch.any(target_qpos < limits[:, 0]) or torch.any( + target_qpos > limits[:, 1] + ): + raise ValueError("TurnKnob target setting violates revolute joint limits.") + qpos = articulation.get_qpos()[:, joint_id] + twist_angles = target_qpos - qpos + if not torch.allclose( + twist_angles, + twist_angles[:1].expand_as(twist_angles), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched TurnKnob environments require one shared twist angle." + ) + + axis_local = local_axis[0] + origin_local = joint_origin_local[0] + geometry_center = ( + vertices.min(dim=0).values + vertices.max(dim=0).values + ) * 0.5 + projections = torch.matmul(vertices - origin_local, axis_local) + axial_value = ( + projections.max() + if torch.abs(projections.max()) >= torch.abs(projections.min()) + else projections.min() + ) + center_projection = torch.dot(geometry_center - origin_local, axis_local) + grasp_position = geometry_center + ( + 0.8 * (axial_value - center_projection) * axis_local + ) + affordance = TwistAffordance( + object_label=f"{step.object_uid}:{child_link}", + grasp_position=tuple(float(value) for value in grasp_position), + axis_origin=tuple(float(value) for value in origin_local), + twist_axis=axis_local, + joint_name=joint_name, + joint_limits=(float(limits[0, 0]), float(limits[0, 1])), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={"mesh_vertices": vertices}, + label=f"{step.object_uid}:{child_link}", + ) + scoped_policy = dict(policy) + scoped_policy.update( + { + "twist_angle": float(twist_angles[0]), + "articulation_joint_name": joint_name, + "articulation_joint_id": joint_id, + "articulation_initial_qpos": qpos, + "articulation_target_qpos": target_qpos, + } + ) + log_info( + f"Twist grounding {step.id}/{arm}: joint={joint_name!r}, " + f"current={float(qpos[0]):.4f}, target={float(target_qpos):.4f}." + ) + return TwistGoal(semantics=semantics, target_pose=child_pose), scoped_policy + + def _relative_object_spacing( + self, + moved_uid: str, + reference_uid: str, + *, + axis: int | torch.Tensor, + nominal: float, + clearance: float, + ) -> float: + """Return deterministic center spacing from live object extents.""" + moved = _object(self.env, moved_uid) + reference = _object(self.env, reference_uid) + required = float(nominal) + for env_id in range(int(self.env.num_envs)): + moved_vertices = _world_vertices(moved, self.env, env_id) + reference_vertices = _world_vertices(reference, self.env, env_id) + if isinstance(axis, torch.Tensor): + direction = axis[env_id].to( + dtype=moved_vertices.dtype, + device=moved_vertices.device, + ) + moved_axis = moved_vertices[:, :2] @ direction + reference_axis = reference_vertices[:, :2] @ direction + else: + moved_axis = moved_vertices[:, axis] + reference_axis = reference_vertices[:, axis] + moved_half = (moved_axis.max() - moved_axis.min()) * 0.5 + reference_half = (reference_axis.max() - reference_axis.min()) * 0.5 + required = max( + required, + float(moved_half + reference_half) + float(clearance), + ) + return required + + def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: + axis = self._upright_local_axis(step) + entity = _object(self.env, step.object_uid) + vertices = _local_vertices(entity, self.env, 0) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + if axis == "long_axis": + axis_index = int(torch.argmax(extents).item()) + else: + axis_index = {"x": 0, "y": 1, "z": 2}[axis] + direction = torch.zeros(3, dtype=torch.float32, device=self.env.device) + direction[axis_index] = 1.0 + return direction + + def _uses_upright_yaw_search( + self, + step: SemanticStep, + constraint: OrientationConstraint, + ) -> bool: + """Preserve a live upright state as a planning preference. + + Explicit full-frame matching cannot admit yaw search. With no hard + orientation terms, yaw search is enabled only when the live object's + long axis is already upright, so a preceding upright operation remains + stable without turning that state into a sticky acceptance constraint. + """ + if constraint.allows_upright_yaw_search: + return True + if ( + constraint.terms + or constraint.planning_preference != "minimize_rotation_from_current" + ): + return False + entity = _object(self.env, step.object_uid) + vertices = _local_vertices(entity, self.env, 0) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + axis_index = int(torch.argmax(extents).item()) + pose = _live_pose(self.env, step.object_uid) + cosine = pose[:, 2, axis_index].abs().clamp(0.0, 1.0) + tolerance = float(self.runtime_policy.predicate_fallbacks["upright_max_tilt"]) + return bool(torch.all(torch.arccos(cosine) <= tolerance).item()) + + @staticmethod + def _upright_local_axis(step: SemanticStep) -> str: + align_terms = tuple( + term + for term in compile_orientation_constraint(step.goal).terms + if isinstance(term, AlignAxisConstraint) + ) + if align_terms: + return align_terms[0].local_axis + axis = str(step.goal.get("upright_local_axis", "auto")) + return "long_axis" if axis == "auto" else axis + + def _target_rotation( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + orientation_reference_pose: torch.Tensor | None = None, + ) -> torch.Tensor: + constraint = compile_orientation_constraint(step.goal) + if not constraint.terms: + return object_pose[:, :3, :3].clone() + if ( + len(constraint.terms) == 1 + and isinstance(constraint.terms[0], MatchRotationConstraint) + and constraint.terms[0].reference == "step_start" + ): + if orientation_reference_pose is not None: + reference = _batched_pose(orientation_reference_pose, self.env) + return reference[:, :3, :3].clone() + return object_pose[:, :3, :3].clone() + goal = str(step.goal.get("orientation_goal", "none")) + align_term = next( + ( + term + for term in constraint.terms + if isinstance(term, AlignAxisConstraint) + ), + None, + ) + if align_term is not None: + goal = "upright" + if goal not in {"upright", "lay_flat", "axis_align"}: + raise ValueError(f"Unsupported orientation_goal {goal!r}.") + + entity = _object(self.env, step.object_uid) + rotations = [] + for env_id in range(int(self.env.num_envs)): + vertices = _local_vertices(entity, self.env, env_id) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + longest_to_shortest = torch.argsort( + extents, + descending=True, + ).tolist() + if goal == "upright": + upright_axis = ( + align_term.local_axis + if align_term is not None + else self._upright_local_axis(step) + ) + vertical_axis = ( + int(longest_to_shortest[0]) + if upright_axis == "long_axis" + else {"x": 0, "y": 1, "z": 2}[upright_axis] + ) + horizontal_axis = next( + int(axis) + for axis in longest_to_shortest + if int(axis) != vertical_axis + ) + elif goal == "lay_flat": + vertical_axis = int(longest_to_shortest[-1]) + horizontal_axis = int(longest_to_shortest[0]) + else: + horizontal_axis = self._aligned_local_axis( + step, + longest_to_shortest, + ) + vertical_axis = next( + int(axis) + for axis in reversed(longest_to_shortest) + if int(axis) != horizontal_axis + ) + direction = self._horizontal_orientation( + step, + object_pose, + env_id, + horizontal_axis, + ) + rotations.append( + self._world_aligned_rotation( + direction, + horizontal_axis=horizontal_axis, + vertical_axis=vertical_axis, + ) + ) + return torch.stack(rotations) + + @staticmethod + def _aligned_local_axis( + step: SemanticStep, + longest_to_shortest: Sequence[int], + ) -> int: + axis = str(step.goal.get("orientation_axis", "long_axis")) + if axis == "x": + return 0 + if axis == "y": + return 1 + if axis == "long_axis": + return int(longest_to_shortest[0]) + if axis == "short_axis": + return int(longest_to_shortest[-1]) + raise ValueError(f"Unsupported axis_align orientation_axis {axis!r}.") + + def _horizontal_orientation( + self, + step: SemanticStep, + object_pose: torch.Tensor, + env_id: int, + local_axis: int, + ) -> torch.Tensor: + align_to = step.goal.get("orientation_reference_object") + if isinstance(align_to, str) and align_to: + reference = _object(self.env, align_to) + vertices = _local_vertices(reference, self.env, env_id) + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + ordered = torch.argsort(extents, descending=True) + requested = str(step.goal.get("orientation_axis", "long_axis")) + reference_axis = int( + ordered[-1] if requested == "short_axis" else ordered[0] + ) + reference_pose = _live_pose(self.env, align_to) + direction = reference_pose[env_id, :3, reference_axis].clone() + elif step.operator in {"arrange_line", "place_in_line"}: + arrangement = self.arrangements.get(step.id) + axis_index = 0 if arrangement is None else arrangement.axis_index + direction = torch.zeros( + 3, + dtype=object_pose.dtype, + device=object_pose.device, + ) + direction[axis_index] = 1.0 + elif str(step.goal.get("orientation_axis", "")) in {"y", "world_y"}: + direction = object_pose.new_tensor([0.0, 1.0, 0.0]) + elif str(step.goal.get("orientation_axis", "")) in {"x", "world_x"}: + direction = object_pose.new_tensor([1.0, 0.0, 0.0]) + else: + direction = object_pose[env_id, :3, local_axis].clone() + direction[2] = 0.0 + norm = torch.linalg.vector_norm(direction) + if float(norm) < 1.0e-6: + return object_pose.new_tensor([1.0, 0.0, 0.0]) + return direction / norm + + @staticmethod + def _world_aligned_rotation( + horizontal_direction: torch.Tensor, + *, + horizontal_axis: int, + vertical_axis: int, + ) -> torch.Tensor: + world_up = horizontal_direction.new_tensor([0.0, 0.0, 1.0]) + remaining_axis = ({0, 1, 2} - {horizontal_axis, vertical_axis}).pop() + columns = [torch.zeros_like(world_up) for _ in range(3)] + columns[horizontal_axis] = horizontal_direction + columns[vertical_axis] = world_up + columns[remaining_axis] = torch.linalg.cross( + world_up, + horizontal_direction, + ) + rotation = torch.stack(columns, dim=1) + if float(torch.linalg.det(rotation)) < 0.0: + rotation[:, remaining_axis] *= -1.0 + return rotation + + def _rotated_local_z_min( + self, + entity: Any, + rotation: torch.Tensor, + env_id: int, + ) -> torch.Tensor: + vertices = _local_vertices(entity, self.env, env_id) + return (vertices @ rotation.transpose(0, 1))[:, 2].min() + + def _current_eef_pose(self, arm: str) -> torch.Tensor: + """Return the live TCP pose for one logical Action Engine arm.""" + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a physical arm, got {arm!r}.") + if hasattr(self.env, "get_current_xpos_agent"): + left, right = self.env.get_current_xpos_agent() + value = left if arm == "left_arm" else right + if value is not None: + return _batched_pose(value, self.env) + + is_left = arm == "left_arm" + if not hasattr(self.env, "get_agent_arm_control_part"): + raise ValueError("Coordinated placement requires live TCP poses.") + part = self.env.get_agent_arm_control_part(is_left) + qpos = self._arm_qpos(arm) + return _batched_pose( + self.env.robot.compute_fk(qpos=qpos, name=part, to_matrix=True), + self.env, + ) + + def _press_goal( + self, + uid: str, + object_pose: torch.Tensor, + *, + semantics: ObjectSemantics | None = None, + terminal_state: str = "activated", + ) -> tuple[PressGoal, dict[str, Any]]: + """Ground the live top surface into the typed press-affordance contract.""" + if semantics is None: + semantics = self.semantics_factory(uid) + articulation = getattr( + self.env.sim, + "get_articulation", + lambda _uid: None, + )(uid) + if articulation is not None: + return self._articulation_press_goal( + uid, + articulation, + semantics, + terminal_state=terminal_state, + ) + entity = _object(self.env, uid) + reference_pose = object_pose[0] + world_position = reference_pose[:3, 3].clone() + world_position[2] = _world_vertices(entity, self.env, 0)[:, 2].max() + rotation = reference_pose[:3, :3] + local_position = rotation.transpose(0, 1) @ ( + world_position - reference_pose[:3, 3] + ) + local_axis = rotation.transpose(0, 1) @ torch.tensor( + [0.0, 0.0, -1.0], + dtype=rotation.dtype, + device=rotation.device, + ) + press_semantics = replace( + semantics, + affordance=PressAffordance( + press_axis=local_axis, + press_position=tuple(float(value) for value in local_position), + ), + ) + return ( + PressGoal( + semantics=press_semantics, + target_pose=object_pose.clone(), + ), + {}, + ) + + def _articulation_press_goal( + self, + uid: str, + articulation: Any, + semantics: ObjectSemantics, + *, + terminal_state: str, + ) -> tuple[PressGoal, dict[str, Any]]: + """Ground a calibrated prismatic button from live joint metadata.""" + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Button articulation exposes no joint metadata.") + backend = backend_entities[0] + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == "prismatic": + candidates.append((int(joint_id), joint_name, info)) + if len(candidates) != 1: + raise ValueError("Press requires exactly one active prismatic joint.") + joint_id, joint_name, joint_info = candidates[0] + child_link = str(getattr(joint_info, "child_link_name", "")) + parent_link = str(getattr(joint_info, "parent_link_name", "")) + if ( + child_link not in articulation.link_names + or parent_link not in articulation.link_names + ): + raise ValueError("Button joint must identify live parent and child links.") + + settings = getattr(self.env, "agent_config", {}).get( + "articulation_settings", {} + ) + per_articulation = ( + settings.get(uid, {}) if isinstance(settings, Mapping) else {} + ) + values = ( + per_articulation.get(joint_name, ()) + if isinstance(per_articulation, Mapping) + else () + ) + if ( + not isinstance(values, Sequence) + or isinstance(values, (str, bytes, bytearray)) + or len(values) < 2 + ): + raise ValueError( + "Press requires explicit inactive/activated joint settings." + ) + values_tensor = torch.as_tensor( + values, dtype=torch.float32, device=self.env.device + ) + if not torch.isfinite(values_tensor).all(): + raise ValueError("Press joint settings must be finite.") + if terminal_state == "activated": + target_qpos = values_tensor[-1] + elif terminal_state == "inactive": + target_qpos = values_tensor[0] + else: + raise ValueError(f"Unsupported button terminal state {terminal_state!r}.") + + qpos = articulation.get_qpos()[:, joint_id] + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + if torch.any(target_qpos < limits[:, 0]) or torch.any( + target_qpos > limits[:, 1] + ): + raise ValueError("Button target setting violates prismatic limits.") + distances = torch.abs(target_qpos - qpos) + if torch.any(distances <= 1.0e-5): + raise ValueError("Button is already at the requested terminal state.") + if not torch.allclose( + distances, + distances[:1].expand_as(distances), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError("Batched Press requires one shared press distance.") + + child_pose = _batched_pose( + articulation.get_link_pose(child_link, to_matrix=True), self.env + ) + parent_pose = _batched_pose( + articulation.get_link_pose(parent_link, to_matrix=True), self.env + ) + raw_axis = joint_info.axis + raw_origin = joint_info.origin_pose + axis = torch.tensor( + raw_axis.tolist() if hasattr(raw_axis, "tolist") else raw_axis, + dtype=torch.float32, + device=self.env.device, + ).reshape(3) + origin = torch.tensor( + raw_origin.tolist() if hasattr(raw_origin, "tolist") else raw_origin, + dtype=torch.float32, + device=self.env.device, + ).reshape(4, 4) + joint_pose = torch.matmul(parent_pose, origin) + world_axis = torch.matmul(joint_pose[:, :3, :3], axis) + direction_sign = torch.sign(target_qpos - qpos) + movement_world = torch.nn.functional.normalize( + world_axis * direction_sign[:, None], dim=1 + ) + movement_local = torch.bmm( + child_pose[:, :3, :3].transpose(1, 2), + movement_world.unsqueeze(2), + ).squeeze(2) + if not torch.allclose( + movement_local, + movement_local[:1].expand_as(movement_local), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError("Batched Press requires one shared movement axis.") + + vertices, _ = articulation.get_link_vert_face(child_link) + vertices = torch.as_tensor( + vertices, dtype=torch.float32, device=self.env.device + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): + raise ValueError("Button child link has no contact geometry.") + press_axis = movement_local[0] + geometry_center = ( + vertices.min(dim=0).values + vertices.max(dim=0).values + ) * 0.5 + projections = torch.matmul(vertices, press_axis) + contact_projection = projections.min() + center_projection = torch.dot(geometry_center, press_axis) + press_position = ( + geometry_center + (contact_projection - center_projection) * press_axis + ) + press_semantics = replace( + semantics, + affordance=PressAffordance( + press_axis=press_axis, + press_position=tuple(float(value) for value in press_position), + ), + ) + return ( + PressGoal(semantics=press_semantics, target_pose=child_pose), + { + "press_distance": float(distances[0]), + "articulation_joint_name": joint_name, + "articulation_joint_id": joint_id, + "articulation_initial_qpos": qpos, + "articulation_target_qpos": target_qpos, + }, + ) + + def _retreat_pose( + self, + arm: str, + policy: Mapping[str, Any], + reference: torch.Tensor | None, + *, + clear_exchange: bool = False, + ) -> torch.Tensor: + target = self._retreat_reference_pose(arm, reference).clone() + desired = float(self._policy_value(policy, "retreat_height")) + if clear_exchange: + _, lateral = robot_frame_axes(self.env) + direction = lateral if arm == "left_arm" else -lateral + target[:, :2, 3] += direction.to( + dtype=target.dtype, + device=target.device, + ) * float(policy.get("retreat_distance", 0.10)) + desired = max( + desired, + float(self._policy_value(policy, "minimum_retreat_height")), + ) + ceiling = float(self._policy_value(policy, "maximum_eef_height")) + height = torch.clamp(ceiling - target[:, 2, 3], min=0.0, max=desired) + target[:, 2, 3] += height + return target + + def _retreat_reference_pose( + self, + arm: str, + reference: torch.Tensor | None, + ) -> torch.Tensor: + """Resolve the live or speculative TCP pose from which retreat starts.""" + pose = reference + if pose is None and hasattr(self.env, "get_current_xpos_agent"): + left, right = self.env.get_current_xpos_agent() + pose = left if arm == "left_arm" else right + if pose is None: + raise ValueError("Retreat grounding requires a live end-effector pose.") + return _batched_pose(pose, self.env) + + def _joint_target( + self, + arm: str, + control: str, + source: str, + binding: Mapping[str, Any], + ) -> torch.Tensor: + if source in {"gripper_closed", "gripper_open"}: + value = ( + getattr(self.env, "close_state") + if source == "gripper_closed" + else getattr(self.env, "open_state") + ) + return torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if source == "joint_delta": + current = self._arm_qpos(arm).clone() + index = int(binding["joint_index"]) + current[:, index] += torch.deg2rad( + torch.tensor( + float(binding.get("delta_degrees", 0.0)), + device=current.device, + ) + ) + return current + initial = getattr(self.env, "init_qpos", self.env.robot.get_qpos()) + joint_ids = self._joint_ids(arm, control) + return torch.as_tensor(initial, device=self.env.device)[:, joint_ids] + + def _arm_qpos(self, arm: str) -> torch.Tensor: + if hasattr(self.env, "get_current_qpos_agent"): + left, right = self.env.get_current_qpos_agent() + return torch.as_tensor( + left if arm == "left_arm" else right, + dtype=torch.float32, + device=self.env.device, + ) + return self.env.robot.get_qpos()[:, self._joint_ids(arm, "arm")] + + def _joint_ids(self, arm: str, control: str) -> list[int]: + side = "left" if arm == "left_arm" else "right" + key = f"{side}_{'eef' if control == 'hand' else 'arm'}_joints" + return list(getattr(self.env, key, ())) + + def _explicit_pose( + self, + binding: Mapping[str, Any], + object_pose: torch.Tensor, + ) -> torch.Tensor: + reference = str(binding.get("reference", "absolute")) + target = object_pose.clone() + if reference == "absolute": + values = binding.get("position_by_env", binding.get("position")) + position = torch.as_tensor( + values, + dtype=target.dtype, + device=target.device, + ) + if position.ndim == 1: + position = position.unsqueeze(0).repeat(int(self.env.num_envs), 1) + target[:, :3, 3] = position + return target + offset = torch.as_tensor( + binding.get("offset", (0.0, 0.0, 0.0)), + dtype=target.dtype, + device=target.device, + ) + target[:, :3, 3] += offset + return target + + def _coordinated_grasps( + self, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build a deterministic opposing pair along the object's longest XY axis.""" + vertices = semantics.geometry.get("mesh_vertices") + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=self.env.device, + ) + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + axis = int(torch.argmax(upper[:2] - lower[:2]).item()) + center = (lower + upper) * 0.5 + grasp_policy = self.runtime_policy.grounding["coordinated_grasp"] + inset = max( + float(grasp_policy["minimum_inset"]), + float((upper[axis] - lower[axis]) * grasp_policy["inset_fraction"]), + ) + left = torch.eye(4, dtype=torch.float32, device=self.env.device) + right = left.clone() + left[:3, 3] = center + right[:3, 3] = center + left[axis, 3] = lower[axis] + inset + right[axis, 3] = upper[axis] - inset + # Keep TCP z horizontal and facing the object from opposite sides. + if axis == 0: + left[:3, :3] = torch.tensor( + [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=self.env.device, + ) + right[:3, :3] = torch.tensor( + [[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=self.env.device, + ) + else: + left[:3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]], + device=self.env.device, + ) + right[:3, :3] = torch.tensor( + [[-1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, -1.0, 0.0]], + device=self.env.device, + ) + batch = int(self.env.num_envs) + return left.unsqueeze(0).repeat(batch, 1, 1), right.unsqueeze(0).repeat( + batch, 1, 1 + ) diff --git a/embodichain/gen_sim/action_engine/runtime/loader.py b/embodichain/gen_sim/action_engine/runtime/loader.py new file mode 100644 index 000000000..a157e8dee --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/loader.py @@ -0,0 +1,296 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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) + for content in step.goal.get("contents", []): + value = content.get("object") if isinstance(content, Mapping) else content + if isinstance(value, str): + references.add(value) + unknown = references - known - {"self", "table", "table_center"} + if unknown: + raise ValueError( + "Execution Program references objects not present in the scene: " + f"{sorted(unknown)}. Regenerate the configuration bundle." + ) diff --git a/embodichain/gen_sim/action_engine/runtime/models.py b/embodichain/gen_sim/action_engine/runtime/models.py new file mode 100644 index 000000000..b4596a241 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/models.py @@ -0,0 +1,314 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Small typed runtime views over the serialized execution program.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.sim.atomic_actions import StateDelta + +from .state import ExecutionState + +__all__ = [ + "ActionOutcome", + "ExecutionEdge", + "ExecutionProgram", + "ExecutionReport", + "ExecutionResult", + "GroundedAction", + "SemanticStep", +] + + +@dataclass(frozen=True) +class ExecutionEdge: + """One executable DAG edge containing symbolic atomic actions.""" + + id: str + source: str + target: str + actions: tuple[dict[str, Any], ...] + depends_on: tuple[str, ...] = () + resources: tuple[str, ...] = () + + +@dataclass(frozen=True) +class SemanticStep: + """One closed-loop intent expanded into one or more execution edges.""" + + id: str + parent_step_id: str + operator: str + object_uid: str + actor: dict[str, Any] + goal: dict[str, Any] + depends_on: tuple[str, ...] + postcondition: dict[str, Any] + edge_ids: tuple[str, ...] + + +@dataclass(frozen=True) +class ExecutionProgram: + """Validated in-memory form of ``action_engine_execution_program_v1``.""" + + raw: dict[str, Any] + task: str + start: str + goal: str + nodes: tuple[dict[str, Any], ...] + edges: tuple[ExecutionEdge, ...] + semantic_steps: tuple[SemanticStep, ...] + allocation_groups: tuple[dict[str, Any], ...] + seed_graph: dict[str, Any] | None = None + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "ExecutionProgram": + """Construct an immutable runtime view from a validated mapping.""" + raw = deepcopy(dict(value)) + edges = tuple( + ExecutionEdge( + id=str(edge["id"]), + source=str(edge["source"]), + target=str(edge["target"]), + actions=tuple( + deepcopy(dict(action)) + for action in edge.get("actions", edge.get("symbolic_actions", ())) + ), + depends_on=tuple(str(item) for item in edge.get("depends_on", ())), + resources=tuple(str(item) for item in edge.get("resources", ())), + ) + for edge in raw["edges"] + ) + steps = tuple( + SemanticStep( + id=str(step["id"]), + parent_step_id=str(step["parent_step_id"]), + operator=str(step["operator"]), + object_uid=str(step.get("object", step.get("object_uid", ""))), + actor=deepcopy(dict(step["actor"])), + goal=deepcopy(dict(step.get("goal", {}))), + depends_on=tuple(str(item) for item in step.get("depends_on", ())), + postcondition=deepcopy(dict(step.get("postcondition", {}))), + edge_ids=tuple(str(item) for item in step["edge_ids"]), + ) + for step in raw["semantic_steps"] + ) + return cls( + raw=raw, + task=str(raw.get("task", raw.get("task_name", "task"))), + start=str(raw["start"]), + goal=str(raw["goal"]), + nodes=tuple(deepcopy(raw["nodes"])), + edges=edges, + semantic_steps=steps, + allocation_groups=tuple( + deepcopy(dict(group)) for group in raw.get("allocation_groups", ()) + ), + seed_graph=None, + ) + + +@dataclass(frozen=True) +class GroundedAction: + """A public atomic-action target resolved from the current simulator state.""" + + action_class: str + arm: str + control: str + target: Any + cfg: dict[str, Any] + object_pose: torch.Tensor | None = None + reference_pose: torch.Tensor | None = None + target_object_pose: torch.Tensor | None = None + motion_policy: dict[str, Any] = field(default_factory=dict) + object_uid: str | None = None + """Scene UID of the object whose semantic step produced this action.""" + + +@dataclass +class ActionOutcome: + """Planning output kept in full-robot coordinates.""" + + trajectory: torch.Tensor + success: torch.Tensor + next_state: ExecutionState + grounded: GroundedAction + prior_state: ExecutionState | None = None + expected_effects: StateDelta | None = None + planner_trace: dict[str, Any] = field(default_factory=dict) + + def state_after(self, verified: torch.Tensor) -> ExecutionState: + """Commit expected effects only for physically verified rows.""" + if self.prior_state is None or self.expected_effects is None: + return self.next_state + mask = torch.as_tensor( + verified, + dtype=torch.bool, + device=self.trajectory.device, + ).reshape(-1) + if mask.numel() != self.trajectory.shape[0]: + raise ValueError("Verified mask must match the ActionOutcome batch.") + terminal_qpos = ( + self.trajectory[:, -1] + if self.trajectory.shape[1] + else self.prior_state.last_qpos + ) + qpos = torch.where( + mask[:, None], + terminal_qpos, + self.prior_state.last_qpos, + ) + task = self.expected_effects.apply( + self.prior_state.to_task_state(), + mask, + ) + return ExecutionState.from_task_state(task, last_qpos=qpos) + + @property + def cost(self) -> torch.Tensor: + """Return joint-path length for each vectorized environment.""" + if self.trajectory.shape[1] < 2: + return torch.zeros( + self.trajectory.shape[0], + dtype=torch.float32, + device=self.trajectory.device, + ) + return torch.linalg.vector_norm( + torch.diff(self.trajectory, dim=1), + dim=-1, + ).sum(dim=1) + + +@dataclass +class ExecutionResult(Sequence[torch.Tensor]): + """Result marker used by the existing demonstration-runner contract.""" + + actions: list[torch.Tensor] + success: torch.Tensor + semantic_success: dict[str, torch.Tensor] + record_dir: str | None = None + already_executed: bool = True + retry_count: int = 0 + recovery_count: int = 0 + revision_count: int = 0 + failure_events: list[dict[str, Any]] = field(default_factory=list) + runtime_revisions: list[dict[str, Any]] = field(default_factory=list) + retry_counts: list[int] = field(default_factory=list) + + @property + def runtime_success(self) -> torch.Tensor: + return self.success + + @property + def runtime_graph_output_dir(self) -> str | None: + return self.record_dir + + def __len__(self) -> int: + return len(self.actions) + + def __iter__(self): + return iter(self.actions) + + def __getitem__(self, index): + return self.actions[index] + + +@dataclass(frozen=True) +class ExecutionReport: + """JSON-safe Task Engine result built from an ``ExecutionResult``. + + The runtime result deliberately keeps tensors because the legacy demo + runner consumes them. The Task Engine boundary instead exposes only a + compact, serializable audit view and never retains the action tensors. + """ + + task_id: str + plan_hash: str + action_graph_hash: str + status: str + run_id: str + episode_id: str + provenance: dict[str, Any] + environments: tuple[dict[str, Any], ...] = () + action_count: int = 0 + retry_count: int = 0 + recovery_count: int = 0 + revision_count: int = 0 + failure_events: tuple[dict[str, Any], ...] = () + graph_revisions: tuple[dict[str, Any], ...] = () + record_dir: str | None = None + error: str | None = None + schema_version: str = "action_engine_execution_report_v2" + + def as_mapping(self) -> dict[str, Any]: + """Return a detached mapping suitable for strict JSON serialization.""" + return { + "schema_version": self.schema_version, + "task_id": self.task_id, + "plan_hash": self.plan_hash, + "action_graph_hash": self.action_graph_hash, + "status": self.status, + "run_id": self.run_id, + "episode_id": self.episode_id, + "provenance": deepcopy(self.provenance), + "environments": deepcopy(list(self.environments)), + "action_count": self.action_count, + "retry_count": self.retry_count, + "recovery_count": self.recovery_count, + "revision_count": self.revision_count, + "failure_events": deepcopy(list(self.failure_events)), + "graph_revisions": deepcopy(list(self.graph_revisions)), + "record_dir": self.record_dir, + "error": self.error, + } + + def to_dict(self) -> dict[str, Any]: + """Compatibility spelling for artifact and CLI publishers.""" + return self.as_mapping() + + +def success_mask(value: bool | torch.Tensor, count: int, device: Any) -> torch.Tensor: + """Normalize a primitive's scalar or batched success result.""" + mask = torch.as_tensor(value, dtype=torch.bool, device=device).reshape(-1) + if mask.numel() == 1: + return mask.repeat(count) + if mask.numel() != count: + raise ValueError( + f"Atomic action success has {mask.numel()} values; expected {count}." + ) + return mask + + +def trajectory_cost_numpy(value: torch.Tensor) -> np.ndarray: + """Expose trajectory costs to assignment solvers without retaining gradients.""" + if value.shape[1] < 2: + return np.zeros(value.shape[0], dtype=np.float64) + diffs = torch.diff(value.detach(), dim=1) + return ( + torch.linalg.vector_norm(diffs, dim=-1) + .sum(dim=1) + .cpu() + .numpy() + .astype(np.float64) + ) diff --git a/embodichain/gen_sim/action_engine/runtime/motion_policy.py b/embodichain/gen_sim/action_engine/runtime/motion_policy.py new file mode 100644 index 000000000..8de7cd4bf --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/motion_policy.py @@ -0,0 +1,106 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from typing import Any + +from embodichain.gen_sim.action_engine.config import default_runtime_policy +from embodichain.gen_sim.action_engine.domain.motion import validate_motion_policy + +__all__ = ["resolve_motion_policy", "with_motion_modifiers"] + +_PROFILE_ALIASES = dict( + franka="dual_franka", ur3="dual_ur3", ur5="dual_ur5", ur10="dual_ur10" +) + + +def resolve_motion_policy( + robot_profile: str, + atomic_action: str, + policy_spec: Mapping[str, Any], + *, + motion_defaults: Mapping[str, Mapping[str, Any]] | None = None, + motion_modifiers: ( + Mapping[str, Mapping[str, Mapping[str, Mapping[str, Any]]]] | None + ) = None, + inline_overrides: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Resolve an action base policy plus its composable typed modifiers.""" + profile = _PROFILE_ALIASES.get(str(robot_profile), str(robot_profile)) + runtime_policy = ( + default_runtime_policy(profile) + if motion_defaults is None or motion_modifiers is None + else None + ) + defaults = ( + runtime_policy.motion_defaults + if motion_defaults is None and runtime_policy is not None + else motion_defaults + ) + modifiers = ( + runtime_policy.motion_modifiers + if motion_modifiers is None and runtime_policy is not None + else motion_modifiers + ) + if defaults is None or modifiers is None: + raise ValueError("Motion defaults and modifiers must be provided together.") + action = str(atomic_action) + if action not in defaults: + raise ValueError(f"Unknown Action Engine motion base {action!r}.") + + spec = validate_motion_policy(policy_spec) + policy = deepcopy(dict(defaults[action])) + modifier_values: dict[str, Any] = {} + modifier_sources: dict[str, tuple[str, str]] = {} + for modifier in spec["modifiers"]: + modifier_type = modifier["type"] + mode = modifier["mode"] + patch = modifiers.get(modifier_type, {}).get(mode, {}).get(action) + if not isinstance(patch, Mapping): + raise ValueError( + f"Motion modifier {(modifier_type, mode)!r} is not supported " + f"by AtomicAction {action!r}." + ) + for key, value in patch.items(): + if key in modifier_values and modifier_values[key] != value: + raise ValueError( + f"Motion modifiers {modifier_sources[key]!r} and " + f"{(modifier_type, mode)!r} conflict on parameter {key!r}." + ) + modifier_values[key] = deepcopy(value) + modifier_sources[key] = (modifier_type, mode) + policy.update(modifier_values) + if inline_overrides is not None: + policy.update(deepcopy(dict(inline_overrides))) + return policy + + +def with_motion_modifiers( + policy_spec: Mapping[str, Any], + *modifiers: tuple[str, str], +) -> dict[str, Any]: + """Return a validated policy reference with missing modifiers appended.""" + policy = validate_motion_policy(policy_spec) + existing = { + (modifier["type"], modifier["mode"]) for modifier in policy["modifiers"] + } + for modifier_type, mode in modifiers: + if (modifier_type, mode) not in existing: + policy["modifiers"].append({"type": modifier_type, "mode": mode}) + return validate_motion_policy(policy) diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py new file mode 100644 index 000000000..57fe8a774 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -0,0 +1,844 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Evaluate canonical closed-loop predicates against live environment state.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.config import default_runtime_policy + +from .frames import relation_axes +from .robot_parts import arm_control_part + +__all__ = ["PREDICATE_TYPES", "evaluate_predicate"] + +PREDICATE_TYPES = frozenset( + { + "both_arms_at_initial_qpos", + "both_grippers_open", + "coordinated_placed", + "grippers_clear_of_object", + "held_by_both_grippers", + "object_axis_near", + "object_axis_offset_near", + "object_held", + "object_held_by_both_grippers", + "object_held_by_gripper", + "object_in_container", + "object_lifted", + "object_not_fallen", + "object_on_object", + "object_supported_by", + "object_position_near", + "object_relative_position", + "object_upright", + "object_xy_near", + "objects_collinear", + "objects_ordered", + "pressed", + "poured", + } +) +_DEFAULT_PREDICATE_FALLBACKS = default_runtime_policy("dual_ur10").predicate_fallbacks + + +def _predicate_fallbacks(env: Any) -> Mapping[str, Any]: + policy = getattr(env, "runtime_policy", None) + value = getattr(policy, "predicate_fallbacks", None) + return value if isinstance(value, Mapping) else _DEFAULT_PREDICATE_FALLBACKS + + +def _constant(env: Any, value: bool) -> torch.Tensor: + return torch.full( + (int(env.num_envs),), + value, + dtype=torch.bool, + device=env.device, + ) + + +def _pose(env: Any, uid: str) -> torch.Tensor: + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=env.device, + ) + if pose.ndim == 2: + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + return pose + + +def _position(env: Any, uid: str) -> torch.Tensor: + return _pose(env, uid)[:, :3, 3] + + +def _world_vertices(env: Any, uid: str, env_id: int) -> torch.Tensor: + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + value = entity.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (tuple, list)): + value = value[0] + vertices = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") + pose = _pose(env, uid)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def _projected_center_of_mass( + env: Any, + uid: str, + env_id: int, + world_vertices: torch.Tensor, +) -> torch.Tensor: + """Return the live COM projection, with a geometry-center fallback.""" + entity = env.sim.get_rigid_object(uid) + body_data = None if entity is None else getattr(entity, "body_data", None) + com_pose = None if body_data is None else getattr(body_data, "com_pose", None) + if callable(com_pose): + com_pose = com_pose() + if com_pose is not None: + local_com = torch.as_tensor( + com_pose, + dtype=torch.float32, + device=env.device, + ) + if local_com.ndim == 1: + local_com = local_com.unsqueeze(0).repeat(int(env.num_envs), 1) + if local_com.ndim == 2 and local_com.shape[0] == int(env.num_envs): + pose = _pose(env, uid)[env_id] + return (pose[:3, :3] @ local_com[env_id, :3] + pose[:3, 3])[:2] + return ( + world_vertices[:, :2].min(dim=0).values + + world_vertices[:, :2].max(dim=0).values + ) * 0.5 + + +def _object_supported_by( + env: Any, + spec: Mapping[str, Any], + defaults: Mapping[str, Any], +) -> torch.Tensor: + """Evaluate one-frame geometric support without advancing simulation.""" + object_uid = _object(spec) + support_uid = str( + spec.get( + "support", + spec.get("reference_object", spec.get("reference", "")), + ) + ) + if not support_uid: + raise ValueError("Support predicate requires a support object uid.") + margin = float(spec.get("com_margin", defaults["support_com_margin"])) + max_gap = float(spec.get("max_vertical_gap", defaults["support_max_vertical_gap"])) + max_penetration = float( + spec.get("max_penetration", defaults["support_max_penetration"]) + ) + min_overlap = float( + spec.get("min_overlap_ratio", defaults["support_min_overlap_ratio"]) + ) + result = _constant(env, False) + for env_id in range(int(env.num_envs)): + moved = _world_vertices(env, object_uid, env_id) + support = _world_vertices(env, support_uid, env_id) + moved_lower = moved[:, :2].min(dim=0).values + moved_upper = moved[:, :2].max(dim=0).values + support_lower = support[:, :2].min(dim=0).values + support_upper = support[:, :2].max(dim=0).values + overlap_extent = torch.clamp( + torch.minimum(moved_upper, support_upper) + - torch.maximum(moved_lower, support_lower), + min=0.0, + ) + moved_extent = torch.clamp(moved_upper - moved_lower, min=1e-6) + overlap_ratio = torch.prod(overlap_extent) / torch.prod(moved_extent) + projected_center = _projected_center_of_mass( + env, + object_uid, + env_id, + moved, + ) + center_supported = torch.all( + projected_center >= support_lower + margin + ) & torch.all(projected_center <= support_upper - margin) + local_mask = torch.all( + (support[:, :2] >= moved_lower - margin) + & (support[:, :2] <= moved_upper + margin), + dim=1, + ) + if bool(local_mask.any()): + local_support_height = support[local_mask, 2].max() + else: + # Sparse meshes may have no vertex exactly under a small payload. + # Nearest vertices are a local fallback; using the mesh-wide peak + # would confuse a remote protrusion with the candidate support pose. + distances = torch.linalg.vector_norm( + support[:, :2] - projected_center, + dim=1, + ) + count = min(8, int(support.shape[0])) + local_support_height = support[ + torch.topk(distances, count, largest=False).indices, 2 + ].max() + vertical_gap = moved[:, 2].min() - local_support_height + result[env_id] = bool( + center_supported + and overlap_ratio >= min_overlap + and vertical_gap >= -max_penetration + and vertical_gap <= max_gap + ) + return result + + +def _objects(spec: Mapping[str, Any]) -> list[str]: + values = spec.get("objects", spec.get("object_uids")) + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise ValueError("Predicate requires a non-empty objects list.") + return [str(value) for value in values] + + +def _object(spec: Mapping[str, Any]) -> str: + value = spec.get("object", spec.get("object_uid")) + if not isinstance(value, str) or not value: + raise ValueError("Predicate requires a non-empty object uid.") + return value + + +def _local_axis_index(env: Any, uid: str, axis: Any) -> int: + name = str(axis).lower() + if name in {"x", "y", "z"}: + return {"x": 0, "y": 1, "z": 2}[name] + if name not in {"long", "long_axis", "longest"}: + raise ValueError(f"Unsupported upright local axis {axis!r}.") + entity = env.sim.get_rigid_object(uid) + if entity is None: + raise ValueError(f"Unknown rigid object {uid!r}.") + vertices = entity.get_vertices(env_ids=[0], scale=True) + if isinstance(vertices, (tuple, list)): + vertices = vertices[0] + vertices = torch.as_tensor(vertices, dtype=torch.float32, device=env.device) + if vertices.ndim == 3: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") + extents = vertices.max(dim=0).values - vertices.min(dim=0).values + return int(torch.argmax(extents).item()) + + +def _arm_values( + env: Any, kind: str +) -> tuple[torch.Tensor | None, torch.Tensor | None] | None: + getter = getattr(env, f"get_current_{kind}_agent", None) + if callable(getter): + left, right = getter() + values = [] + for value in (left, right): + if value is None: + values.append(None) + continue + item = torch.as_tensor(value, device=env.device) + if kind == "xpos" and item.ndim == 2: + item = item.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + elif kind == "gripper_state" and item.ndim == 1: + item = item.unsqueeze(0) + values.append(item) + return values[0], values[1] + if kind != "gripper_state": + return None + qpos = env.robot.get_qpos() + values = [] + for side in ("left", "right"): + ids = list(getattr(env, f"{side}_eef_joints", ())) + if not ids: + return None + values.append(qpos[:, ids]) + return values[0], values[1] + + +def _gripper_has_closed( + env: Any, + gripper: torch.Tensor, + *, + tolerance: float, +) -> torch.Tensor: + """Check closure intent without requiring an impossible empty-gripper pose.""" + gripper = gripper.to(device=env.device, dtype=torch.float32) + open_state = getattr(env, "open_state", None) + close_state = getattr(env, "close_state", None) + reference = open_state if open_state is not None else close_state + if reference is None: + return _constant(env, False) + expected = torch.as_tensor( + reference, + dtype=torch.float32, + device=env.device, + ).flatten() + repeats = (gripper.shape[-1] + expected.numel() - 1) // expected.numel() + expected = expected.repeat(repeats)[: gripper.shape[-1]] + distance = torch.linalg.vector_norm(gripper - expected, dim=-1) + if open_state is not None: + return distance > tolerance + return distance <= tolerance + + +def _object_held( + env: Any, + uid: str, + *, + owners: Mapping[str, Sequence[str | None]] | None, + states: Mapping[tuple[str, str], Any] | None, + position_tolerance: float, + gripper_tolerance: float, + required_arm: str | None = None, +) -> torch.Tensor: + """Verify registry ownership against live object, TCP, and gripper state.""" + result = _constant(env, False) + if owners is None or states is None or uid not in owners: + return result + eef_values = _arm_values(env, "xpos") + gripper_values = _arm_values(env, "gripper_state") + if eef_values is None or gripper_values is None: + return result + + object_pose = _pose(env, uid) + for arm_index, arm in enumerate(("left_arm", "right_arm")): + if required_arm is not None and arm != required_arm: + continue + state = states.get((uid, arm)) + held = ( + None if state is None else state.get_held_object(arm_control_part(env, arm)) + ) + actual_eef = eef_values[arm_index] + gripper = gripper_values[arm_index] + if held is None or actual_eef is None or gripper is None: + continue + label = getattr(held.semantics, "label", None) + if not label and held.semantics.entity is not None: + label = getattr(held.semantics.entity, "uid", None) + if label != uid: + continue + actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) + expected_eef = torch.bmm( + object_pose, + held.object_to_eef.to(device=env.device, dtype=object_pose.dtype), + ) + position_ok = ( + torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], dim=-1 + ) + <= position_tolerance + ) + closed = _gripper_has_closed( + env, + gripper, + tolerance=gripper_tolerance, + ) + owned = torch.tensor( + [item == arm for item in owners[uid]], + dtype=torch.bool, + device=env.device, + ) + result |= owned & position_ok & closed + return result + + +def _coordinated_held( + env: Any, + uid: str, + state: Any, + *, + position_tolerance: float, + gripper_tolerance: float, +) -> torch.Tensor: + result = _constant(env, False) + if state is None: + return result + held_relations = tuple( + state.get_held_object(arm_control_part(env, arm)) + for arm in ("left_arm", "right_arm") + ) + if any(held is None for held in held_relations): + return result + for held in held_relations: + assert held is not None + label = getattr(held.semantics, "label", None) + if not label and getattr(held.semantics, "entity", None) is not None: + label = getattr(held.semantics.entity, "uid", None) + if label != uid: + return result + eef_values = _arm_values(env, "xpos") + gripper_values = _arm_values(env, "gripper_state") + if eef_values is None or gripper_values is None: + return result + + object_pose = _pose(env, uid) + result = _constant(env, True) + for arm_index, held in enumerate(held_relations): + assert held is not None + if held.env_mask is not None: + result &= held.env_mask.to(device=env.device) + actual_eef = eef_values[arm_index] + gripper = gripper_values[arm_index] + if actual_eef is None or gripper is None: + return _constant(env, False) + transform = held.object_to_eef.to( + device=env.device, + dtype=object_pose.dtype, + ) + expected_eef = torch.bmm(object_pose, transform) + actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) + position_ok = ( + torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], + dim=-1, + ) + <= position_tolerance + ) + closed = _gripper_has_closed( + env, + gripper, + tolerance=gripper_tolerance, + ) + result &= position_ok & closed + return result + + +def evaluate_predicate( + env: Any, + spec: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None, + *, + held_owners: Mapping[str, Sequence[str | None]] | None = None, + held_states: Mapping[tuple[str, str], Any] | None = None, + coordinated_state: Any | None = None, +) -> torch.Tensor: + """Evaluate one typed predicate or a boolean predicate tree.""" + runtime = { + "held_owners": held_owners, + "held_states": held_states, + "coordinated_state": coordinated_state, + } + defaults = _predicate_fallbacks(env) + if spec is None: + return _constant(env, True) + if isinstance(spec, Sequence) and not isinstance(spec, (str, bytes, Mapping)): + result = _constant(env, True) + for term in spec: + result &= evaluate_predicate(env, term, **runtime) + return result + if not isinstance(spec, Mapping): + raise TypeError("Predicate must be a mapping or a sequence of mappings.") + op = str(spec.get("op", "")).lower() + if not op and "terms" in spec: + op = "all" + if op in {"all", "and"}: + return evaluate_predicate(env, list(spec.get("terms", ())), **runtime) + if op in {"any", "or"}: + result = _constant(env, False) + for term in spec.get("terms", ()): + result |= evaluate_predicate(env, term, **runtime) + return result + if op == "not": + return ~evaluate_predicate(env, spec.get("term"), **runtime) + + kind = str(spec.get("type", spec.get("kind", ""))).lower() + if kind in {"semantic_goal", "line_member_placed", "stack_layer_supported"}: + raise ValueError( + f"Predicate {kind!r} is a compiler marker and requires the " + "executor's grounded target." + ) + if kind in {"object_held", "object_held_by_gripper"}: + required_arm = spec.get("arm") + if required_arm in {"left", "right"}: + required_arm = f"{required_arm}_arm" + return _object_held( + env, + _object(spec), + owners=held_owners, + states=held_states, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + required_arm=str(required_arm) if required_arm else None, + ) + if kind == "handover_complete": + required_arm = spec.get("arm", "right_arm") + return _object_held( + env, + _object(spec), + owners=held_owners, + states=held_states, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + required_arm=str(required_arm), + ) + if kind in {"held_by_both_grippers", "object_held_by_both_grippers"}: + return _coordinated_held( + env, + _object(spec), + coordinated_state, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + ) + if kind in {"object_position_near", "position_near"}: + position = _position(env, _object(spec)) + target = torch.as_tensor( + spec.get("target_position", spec.get("target")), + dtype=position.dtype, + device=position.device, + ) + if target.ndim == 1: + target = target.unsqueeze(0) + return torch.linalg.vector_norm(position - target, dim=-1) <= float( + spec.get("tolerance", defaults["position_tolerance"]) + ) + if kind in {"object_xy_near", "xy_near"}: + position = _position(env, _object(spec))[:, :2] + target = torch.as_tensor( + spec.get("target_xy", spec.get("target")), + dtype=position.dtype, + device=position.device, + ).reshape(-1, 2) + return torch.linalg.vector_norm(position - target, dim=-1) <= float( + spec.get("tolerance", defaults["xy_tolerance"]) + ) + if kind in {"object_relative_position", "relative_position"}: + reference_uid = spec.get("reference_object", spec.get("reference")) + if not isinstance(reference_uid, str) or not reference_uid: + raise ValueError("Relative-position predicate requires a reference object.") + relation = str(spec.get("relation", "")) + axes = relation_axes( + env, + relation, + frame=str(spec.get("relation_frame", "world")), + ) + if not axes: + raise ValueError(f"Unsupported directional relation {relation!r}.") + delta = ( + _position(env, _object(spec))[:, :2] - _position(env, reference_uid)[:, :2] + ) + minimum_distance = float(spec.get("minimum_distance", 0.0)) + result = _constant(env, True) + for axis in axes: + projection = torch.sum( + delta * axis.to(dtype=delta.dtype, device=delta.device), dim=1 + ) + result &= projection >= minimum_distance + return result + if kind in {"object_in_container", "inside"}: + position = _position(env, _object(spec)) + container = _position( + env, str(spec.get("container", spec.get("reference_object"))) + ) + xy = torch.linalg.vector_norm(position[:, :2] - container[:, :2], dim=-1) + z = position[:, 2] - container[:, 2] + return ( + (xy <= float(spec.get("xy_radius", defaults["container_xy_radius"]))) + & (z >= float(spec.get("min_z_offset", defaults["container_min_z_offset"]))) + & (z <= float(spec.get("max_z_offset", defaults["container_max_z_offset"]))) + ) + if kind in {"object_supported_by", "object_on_object", "on"}: + return _object_supported_by(env, spec, defaults) + if kind == "object_not_fallen": + axis = _pose(env, _object(spec))[:, :3, 2] + cosine = axis[:, 2].clamp(-1.0, 1.0) + return torch.arccos(cosine) <= float( + spec.get("max_tilt", defaults["not_fallen_max_tilt"]) + ) + if kind == "object_upright": + uid = _object(spec) + local_axis = spec.get("local_axis", "long_axis") + axis_index = _local_axis_index( + env, + uid, + local_axis, + ) + axis = _pose(env, uid)[:, :3, axis_index] + cosine = axis[:, 2].clamp(-1.0, 1.0) + directed = spec.get( + "directed", + str(local_axis).lower() not in {"long", "long_axis", "longest"}, + ) + if not isinstance(directed, bool): + raise ValueError("object_upright directed must be a boolean.") + if not directed: + cosine = cosine.abs() + return torch.arccos(cosine) <= float( + spec.get("max_tilt", defaults["upright_max_tilt"]) + ) + if kind in {"object_axis_offset_near", "object_axis_near"}: + object_position = _position(env, _object(spec)) + axis = _axis_index(spec.get("axis", "x")) + reference_uid = spec.get( + "reference_object", + spec.get("reference", spec.get("support")), + ) + if isinstance(reference_uid, str) and reference_uid: + values = object_position[:, axis] - _position(env, reference_uid)[:, axis] + else: + values = object_position[:, axis] + target = spec.get( + "target_offset", + spec.get("offset", spec.get("target", 0.0)), + ) + target_value = torch.as_tensor( + target, + dtype=values.dtype, + device=values.device, + ) + return torch.abs(values - target_value) <= float( + spec.get("tolerance", defaults["axis_tolerance"]) + ) + if kind in {"objects_collinear", "collinear"}: + positions = torch.stack( + [_position(env, uid) for uid in _objects(spec)], + dim=1, + ) + axis = 0 if str(spec.get("axis", "x")) in {"x", "world_x"} else 1 + values = positions[:, :, 1 - axis] + return values.max(dim=1).values - values.min(dim=1).values <= float( + spec.get("tolerance", defaults["collinearity_tolerance"]) + ) + if kind in {"objects_ordered", "ordered"}: + positions = torch.stack( + [_position(env, uid) for uid in _objects(spec)], + dim=1, + ) + axis = 0 if str(spec.get("axis", "x")) in {"x", "world_x"} else 1 + differences = torch.diff(positions[:, :, axis], dim=1) + tolerance = float(spec.get("tolerance", defaults["ordering_tolerance"])) + if str(spec.get("direction", "ascending")) == "descending": + return torch.all(differences <= tolerance, dim=1) + return torch.all(differences >= -tolerance, dim=1) + if kind == "object_lifted": + position = _position(env, _object(spec))[:, 2] + initial = spec.get("initial_height") + if initial is None: + initial_pose = getattr(env, "agent_initial_object_poses", {}).get( + _object(spec) + ) + if initial_pose is None: + raise ValueError("object_lifted requires an initial object pose.") + initial = initial_pose[:, 2, 3] + initial = torch.as_tensor(initial, device=position.device) + return position >= initial + float( + spec.get("min_height", defaults["minimum_lift_height"]) + ) + if kind in {"both_arms_at_initial_qpos", "arms_home"}: + current = env.robot.get_qpos() + initial = getattr(env, "init_qpos", current) + return torch.all( + torch.abs(current - initial) + <= float(spec.get("tolerance", defaults["arm_initial_qpos_tolerance"])), + dim=-1, + ) + if kind in {"both_grippers_open", "grippers_open"}: + if not hasattr(env, "get_current_gripper_state_agent"): + return _constant(env, False) + left, right = env.get_current_gripper_state_agent() + expected = torch.as_tensor( + env.open_state, + dtype=torch.float32, + device=env.device, + ) + results = [] + for value in (left, right): + value = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if value.ndim == 1: + value = value.unsqueeze(0).repeat(int(env.num_envs), 1) + results.append( + torch.linalg.vector_norm(value - expected, dim=-1) + <= float(spec.get("tolerance", defaults["gripper_state_tolerance"])) + ) + return results[0] & results[1] + if kind == "grippers_clear_of_object": + eef_values = _arm_values(env, "xpos") + if eef_values is None: + return _constant(env, False) + object_position = _position(env, _object(spec)) + clearance = float( + spec.get( + "min_distance", + spec.get("clearance", defaults["gripper_clear_min_distance"]), + ) + ) + result = _constant(env, True) + for eef in eef_values: + if eef is None: + return _constant(env, False) + result &= ( + torch.linalg.vector_norm( + eef[:, :3, 3] - object_position, + dim=-1, + ) + >= clearance + ) + return result + if kind == "pressed": + checker = getattr(env, "is_object_pressed", None) + if callable(checker): + value = checker(_object(spec), spec.get("terminal_state", "activated")) + result = torch.as_tensor(value, dtype=torch.bool, device=env.device) + return ( + result.repeat(int(env.num_envs)) + if result.ndim == 0 + else result.reshape(-1) + ) + return _constant(env, False) + if kind == "poured": + raw_contents = spec.get("contents", ()) + if not isinstance(raw_contents, Sequence) or isinstance( + raw_contents, (str, bytes, bytearray) + ): + raise ValueError("poured contents must be a list of observable objects.") + contents = [ + item.get("object") if isinstance(item, Mapping) else item + for item in raw_contents + ] + if not contents or any(not isinstance(uid, str) or not uid for uid in contents): + raise ValueError( + "poured requires at least one independently observable content object." + ) + if len(contents) != len(set(contents)): + raise ValueError("poured content objects must be unique.") + target = spec.get("reference_object", spec.get("container")) + if not isinstance(target, str) or not target: + raise ValueError("poured requires a target reference_object.") + transferred = _constant(env, True) + for uid in contents: + transferred &= evaluate_predicate( + env, + { + "type": "object_in_container", + "object": uid, + "container": target, + }, + **runtime, + ) + return transferred + if kind == "articulation_joint_near": + uid = _object(spec) + articulation = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if articulation is None: + raise ValueError(f"Unknown articulation {uid!r}.") + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + backend = backend_entities[0] + expected_joint_type = "revolute" if "target_setting" in spec else "prismatic" + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == expected_joint_type: + candidates.append((int(joint_id), joint_name)) + requested = spec.get("joint_name") + if requested is not None: + candidates = [item for item in candidates if item[1] == str(requested)] + if len(candidates) != 1: + raise ValueError( + "articulation_joint_near requires exactly one matching " + f"{expected_joint_type} joint." + ) + joint_id, _ = candidates[0] + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + qpos = articulation.get_qpos()[:, joint_id] + target_state = spec.get("target_state") + if target_state == "open": + target = limits[:, 1] + elif target_state == "closed": + target = limits[:, 0] + elif "target_qpos" in spec: + target = torch.as_tensor( + spec["target_qpos"], dtype=torch.float32, device=env.device + ).expand_as(qpos) + elif "target_setting" in spec: + raw_values = spec.get("setting_values", ()) + if ( + not isinstance(raw_values, Sequence) + or isinstance(raw_values, (str, bytes, bytearray)) + or not raw_values + ): + raise ValueError( + "articulation_joint_near target_setting requires setting_values." + ) + values = torch.as_tensor(raw_values, dtype=torch.float32, device=env.device) + setting = int(spec["target_setting"]) + if setting < 0 or setting >= values.numel(): + raise ValueError( + "articulation_joint_near target_setting is outside setting_values." + ) + target = values[setting].expand_as(qpos) + else: + raise ValueError( + "articulation_joint_near requires open/closed target_state or " + "target_setting with setting_values." + ) + tolerance = float(spec.get("tolerance", defaults["axis_tolerance"])) + return torch.isfinite(qpos) & (torch.abs(qpos - target) <= tolerance) + if kind == "coordinated_placed": + relation = str(spec.get("relation", "on")) + reference = spec.get("support_object", spec.get("reference_object")) + translated = { + "type": ( + "object_in_container" if relation == "inside" else "object_supported_by" + ), + "object": _object(spec), + ("container" if relation == "inside" else "support"): reference, + } + return evaluate_predicate(env, translated, **runtime) + raise ValueError(f"Unsupported execution predicate {kind!r}.") + + +def _axis_index(value: Any) -> int: + axis = str(value).lower().replace("world_", "") + if axis not in {"x", "y", "z"}: + raise ValueError(f"Unsupported predicate axis {value!r}.") + return {"x": 0, "y": 1, "z": 2}[axis] diff --git a/embodichain/gen_sim/action_engine/runtime/recording.py b/embodichain/gen_sim/action_engine/runtime/recording.py new file mode 100644 index 000000000..4c38a64a2 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/recording.py @@ -0,0 +1,391 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Append compact per-environment execution events and a final summary.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from datetime import datetime, timezone +import json +import os +from pathlib import Path +import re +import tempfile +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.domain import ( + execution_program_hash, + seed_graph_hash, +) +from embodichain.utils.logger import log_warning + +from .models import ExecutionProgram, GroundedAction, SemanticStep + +__all__ = ["RuntimeRecorder"] + +_SAFE_NAME = re.compile(r"[^0-9A-Za-z._-]+") + + +def _safe_name(value: str) -> str: + result = _SAFE_NAME.sub("_", value).strip("._") + if not result: + raise ValueError("Runtime record path component must not be empty.") + return result + + +def _default_output_root() -> Path: + for parent in Path(__file__).resolve().parents: + if (parent / "setup.py").is_file() and (parent / "embodichain").is_dir(): + return parent / "outputs" / "action_engine" + return Path.cwd() / "outputs" / "action_engine" + + +def _jsonable(value: Any, env_id: int | None = None) -> Any: + if isinstance(value, torch.Tensor): + item = value + if env_id is not None and item.ndim > 0 and item.shape[0] > env_id: + item = item[env_id] + return item.detach().cpu().tolist() + if isinstance(value, dict): + return {str(key): _jsonable(item, env_id) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item, env_id) for item in value] + if isinstance(value, Path): + return value.as_posix() + return value + + +class RuntimeRecorder: + """Record execution decisions without copying the whole program per step.""" + + def __init__( + self, + program: ExecutionProgram, + *, + num_envs: int, + run_id: str | None = None, + episode_index: int = 0, + output_root: str | Path | None = None, + enabled: bool = True, + runtime_policy: Mapping[str, Any] | None = None, + runtime_policy_hash: str | None = None, + ) -> None: + self.enabled = enabled + self.num_envs = int(num_envs) + self.run_id = _safe_name( + run_id or datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + ) + root = ( + Path(output_root).expanduser().resolve() + if output_root is not None + else _default_output_root() + ) + self.output_dir = ( + root + / _safe_name(program.task) + / self.run_id + / f"episode_{int(episode_index):04d}" + ) + # The validated source graph remains untouched. Runtime documents are + # built from a detached copy and extend it only with a runtime envelope. + self.seed_topology = deepcopy(program.seed_graph or program.raw) + self.program_hash = ( + seed_graph_hash(program.seed_graph) + if program.seed_graph is not None + else execution_program_hash(program.raw) + ) + self.step_specs = { + str(step["id"]): deepcopy(step) for step in program.raw["semantic_steps"] + } + self.step_ordinals = { + step.id: index for index, step in enumerate(program.semantic_steps, start=1) + } + self.events: list[list[dict[str, Any]]] = [[] for _ in range(self.num_envs)] + self.program_metadata = { + "schema_version": "action_engine_runtime_record_v2", + "task": program.task, + "run_id": self.run_id, + "episode_index": int(episode_index), + "program_schema_version": self.seed_topology.get("schema_version"), + "seed_graph_hash": self.program_hash, + } + if runtime_policy is not None: + if not isinstance(runtime_policy_hash, str) or not runtime_policy_hash: + raise ValueError("Recorded runtime policy requires a non-empty hash.") + self.program_metadata["runtime_policy"] = deepcopy(dict(runtime_policy)) + self.program_metadata["runtime_policy_hash"] = runtime_policy_hash + + def register_step( + self, + step: SemanticStep, + spec: Mapping[str, Any], + ) -> None: + """Register a semantic step inserted by a runtime graph revision.""" + if not self.enabled: + return + raw = deepcopy(dict(spec)) + if str(raw.get("id")) != step.id: + raise ValueError("Runtime step spec ID must match the semantic step ID.") + existing = self.step_specs.get(step.id) + if existing is not None: + if existing != raw: + raise ValueError( + f"Runtime step {step.id!r} was registered with a different spec." + ) + return + self.step_specs[step.id] = raw + self.step_ordinals[step.id] = max(self.step_ordinals.values(), default=0) + 1 + + def edge( + self, + edge_id: str, + step: SemanticStep, + *, + assignments: list[str | None], + grounded: list[GroundedAction], + active: torch.Tensor, + failed: torch.Tensor, + action_steps: int, + planner_traces: Sequence[Mapping[str, Any]] = (), + diagnostics: Sequence[str] = (), + phase: str = "primary", + ) -> None: + if not self.enabled: + return + if phase not in {"primary", "recovery", "replay", "final_revalidation"}: + raise ValueError(f"Unknown execution phase {phase!r}.") + for env_id in range(self.num_envs): + event = { + "event": "edge", + "phase": phase, + "edge_id": edge_id, + "semantic_step_id": step.id, + "operator": step.operator, + "object": step.object_uid, + "arm": assignments[env_id], + "status": ( + "skipped" + if not bool(active[env_id]) + else ("failed" if bool(failed[env_id]) else "executed") + ), + "actions": [ + { + "class": item.action_class, + "control": item.control, + "target_object_pose": _jsonable( + item.target_object_pose, env_id + ), + "motion_policy": _jsonable(item.motion_policy), + } + for item in grounded + ], + "trajectory_steps": (int(action_steps) if bool(active[env_id]) else 0), + "time_utc": datetime.now(timezone.utc).isoformat(), + } + if diagnostics: + event["diagnostics"] = [str(item) for item in diagnostics] + if planner_traces: + event["planner_attempts"] = _jsonable(planner_traces, env_id) + self.events[env_id].append(event) + + def step( + self, + step: SemanticStep, + success: torch.Tensor, + *, + observed: torch.Tensor | None, + target: torch.Tensor | None, + metadata: Sequence[Mapping[str, Any]] | None = None, + phase: str = "primary", + ) -> None: + if not self.enabled: + return + if phase not in {"primary", "recovery", "replay", "final_revalidation"}: + raise ValueError(f"Unknown execution phase {phase!r}.") + if metadata is not None and len(metadata) != self.num_envs: + raise ValueError("Runtime step metadata must match num_envs.") + for env_id in range(self.num_envs): + event = { + "event": "semantic_step", + "phase": phase, + "semantic_step_id": step.id, + "status": "success" if bool(success[env_id]) else "failed", + "observed_position": _jsonable(observed, env_id), + "target_position": _jsonable(target, env_id), + "time_utc": datetime.now(timezone.utc).isoformat(), + } + if metadata is not None: + event.update(_jsonable(dict(metadata[env_id]))) + self.events[env_id].append(event) + self._write_step_checkpoint(env_id, step, event) + + def recovery( + self, + *, + failure_type: str, + failed_node_id: str, + active: torch.Tensor, + status: str, + recovery_group_id: str | None = None, + error: str | None = None, + semantic_step_id: str | None = None, + ) -> None: + """Record one bounded local-recovery phase for the selected rows.""" + if not self.enabled: + return + if status not in {"started", "succeeded", "failed", "rejected"}: + raise ValueError(f"Unknown recovery status {status!r}.") + for env_id in range(self.num_envs): + if not bool(active[env_id]): + continue + event = { + "event": "local_recovery", + "phase": "recovery", + "failure_type": str(failure_type), + "failed_node_id": str(failed_node_id), + "recovery_group_id": recovery_group_id, + "status": status, + "error": error, + "time_utc": datetime.now(timezone.utc).isoformat(), + } + if semantic_step_id is not None: + event["semantic_step_id"] = str(semantic_step_id) + self.events[env_id].append(event) + + def _env_dir(self, env_id: int) -> Path: + return self.output_dir / f"env_{env_id:04d}" + + def _write_step_checkpoint( + self, + env_id: int, + step: SemanticStep, + event: dict[str, Any], + ) -> None: + """Atomically publish one closed-loop semantic-step checkpoint.""" + related_events = [ + deepcopy(item) + for item in self.events[env_id] + if item.get("semantic_step_id") == step.id + ] + checkpoint = { + "schema_version": "action_engine_semantic_checkpoint_v2", + "seed_graph_hash": self.program_hash, + "task": self.program_metadata["task"], + "run_id": self.run_id, + "episode_index": self.program_metadata["episode_index"], + "env_id": env_id, + "semantic_step": deepcopy(self.step_specs[step.id]), + "status": event["status"], + "events": related_events, + "checkpointed_at_utc": event["time_utc"], + } + ordinal = self.step_ordinals[step.id] + filename = f"step_{ordinal:04d}_{_safe_name(step.id)}.json" + _write_json_atomic( + self._env_dir(env_id) / "checkpoints" / filename, + checkpoint, + ) + + def finalize( + self, + success: torch.Tensor, + *, + error: str | None = None, + ) -> str | None: + if not self.enabled: + return None + from embodichain.gen_sim.action_engine.graph_visualization import ( + render_task_graph_png, + ) + + finished_at = datetime.now(timezone.utc).isoformat() + for env_id in range(self.num_envs): + runtime = { + **self.program_metadata, + "env_id": env_id, + "status": ( + "aborted" + if error is not None + else ("success" if bool(success[env_id]) else "failed") + ), + "error": error, + "events": self.events[env_id], + "finished_at_utc": finished_at, + } + document = deepcopy(self.seed_topology) + document["runtime"] = runtime + env_dir = self._env_dir(env_id) + _write_json_atomic( + env_dir / "task_graph.json", + document, + ) + try: + png = render_task_graph_png(document) + if not isinstance(png, bytes): + raise TypeError("render_task_graph_png must return bytes.") + _write_bytes_atomic(env_dir / "task_graph.png", png) + except Exception as exc: + runtime["visualization_error"] = f"{type(exc).__name__}: {exc}" + document["runtime"] = runtime + _write_json_atomic(env_dir / "task_graph.json", document) + log_warning( + "Unable to render Action Engine runtime graph for " + f"env {env_id}: {exc}" + ) + return self.output_dir.as_posix() + + +def _write_json_atomic(path: Path, value: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temp_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + text=True, + ) + temp_path = Path(temp_name) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + json.dump(value, stream, ensure_ascii=False, indent=2) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_path, path) + finally: + temp_path.unlink(missing_ok=True) + + +def _write_bytes_atomic(path: Path, value: bytes) -> None: + """Write one binary artifact without exposing a partial destination.""" + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temp_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + ) + temp_path = Path(temp_name) + try: + with os.fdopen(descriptor, "wb") as stream: + stream.write(value) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temp_path, path) + finally: + temp_path.unlink(missing_ok=True) diff --git a/embodichain/gen_sim/action_engine/runtime/recovery.py b/embodichain/gen_sim/action_engine/runtime/recovery.py new file mode 100644 index 000000000..cbe3caa7a --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/recovery.py @@ -0,0 +1,649 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Bounded retry decisions and auditable RuntimeGraph revisions.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, + capability_precondition, +) +from embodichain.gen_sim.action_engine.domain import motion_policy, validate_seed_graph + +__all__ = [ + "FAILURE_TYPES", + "GraphRevision", + "RetryDecision", + "RuntimeGraph", + "build_upright_recovery", + "classify_failure", +] + +FAILURE_TYPES = frozenset( + { + "plan_failed", + "search_exhausted", + "grasp_missed", + "object_fallen", + "object_dropped", + "postcondition_failed", + } +) + + +@dataclass(frozen=True) +class RetryDecision: + """Per-environment result of one failed full-AtomicAction attempt.""" + + retry: torch.Tensor + recover: torch.Tensor + exhausted: torch.Tensor + attempts: tuple[int, ...] + + +@dataclass(frozen=True) +class GraphRevision: + """One immutable patch record over the original SeedGraph.""" + + revision: int + kind: str + reason: str + failed_node_id: str | None + inserted_group_ids: tuple[str, ...] + replaced_group_ids: tuple[str, ...] + active_env_ids: tuple[int, ...] = () + + +class RuntimeGraph: + """Keep SeedGraph immutable while applying bounded, validated revisions.""" + + def __init__( + self, + seed_graph: Mapping[str, Any], + *, + num_envs: int, + max_retries: int = 2, + max_revisions: int = 8, + max_recovery_actions: int = 12, + registry: AtomicCapabilityRegistry | None = None, + ) -> None: + if num_envs < 1: + raise ValueError("RuntimeGraph num_envs must be positive.") + self.registry = registry or build_atomic_capability_registry() + self.seed_graph = validate_seed_graph( + seed_graph, + known_actions=self.registry.names(), + ) + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + validate_persisted_contracts(self.seed_graph, self.registry) + self._graph = deepcopy(self.seed_graph) + self.num_envs = int(num_envs) + self.max_retries = int(max_retries) + self.max_revisions = int(max_revisions) + self.max_recovery_actions = int(max_recovery_actions) + if min(self.max_retries, self.max_revisions, self.max_recovery_actions) < 0: + raise ValueError("RuntimeGraph budgets must be non-negative.") + self._attempts: dict[str, list[int]] = {} + self._recovery_action_count = 0 + self.revisions: list[GraphRevision] = [] + + @property + def graph(self) -> dict[str, Any]: + """Return the current detached RuntimeGraph snapshot.""" + return deepcopy(self._graph) + + def record_failure( + self, + node_id: str, + failed: torch.Tensor, + *, + precondition_holds: torch.Tensor, + ) -> RetryDecision: + """Consume attempt budgets and distinguish retry from recovery.""" + failed = _mask(failed, self.num_envs) + precondition_holds = _mask(precondition_holds, self.num_envs) + attempts = self._attempts.setdefault(node_id, [1] * self.num_envs) + retry = torch.zeros_like(failed) + recover = torch.zeros_like(failed) + exhausted = torch.zeros_like(failed) + node = _node(self._graph, node_id) + capability = self.registry.get(str(node["atomic_action"])) + for env_id in torch.nonzero(failed, as_tuple=False).flatten().tolist(): + attempts[env_id] += 1 + can_retry = ( + capability.retry_mode != "non_retryable" + and bool(precondition_holds[env_id]) + and attempts[env_id] <= self.max_retries + 1 + ) + if can_retry: + retry[env_id] = True + elif capability.retry_mode != "non_retryable": + recover[env_id] = True + else: + exhausted[env_id] = True + return RetryDecision(retry, recover, exhausted, tuple(attempts)) + + def insert_recovery_subgraph( + self, + *, + failed_node_id: str, + recovery_nodes: Sequence[Mapping[str, Any]], + recovery_group: Mapping[str, Any], + failure_type: str, + active_env_ids: Sequence[int] | None = None, + preserve_failed_group_suffix: bool = False, + ) -> dict[str, Any]: + """Insert a complete recovery TaskGroup and rewire the unfinished suffix.""" + if failure_type not in FAILURE_TYPES: + raise ValueError(f"Unknown failure type {failure_type!r}.") + if len(self.revisions) >= self.max_revisions: + raise RuntimeError("RuntimeGraph revision budget exhausted.") + if ( + self._recovery_action_count + len(recovery_nodes) + > self.max_recovery_actions + ): + raise RuntimeError("RuntimeGraph recovery-action budget exhausted.") + env_ids = tuple( + sorted( + set( + range(self.num_envs) + if active_env_ids is None + else (int(env_id) for env_id in active_env_ids) + ) + ) + ) + if not env_ids or env_ids[0] < 0 or env_ids[-1] >= self.num_envs: + raise ValueError( + "Recovery active_env_ids are outside the environment range." + ) + failed_node = _node(self._graph, failed_node_id) + failed_group_id = str(failed_node["task_instance_id"]) + group = deepcopy(dict(recovery_group)) + if group.get("role") != "recovery": + raise ValueError("Inserted recovery TaskGroup must use role='recovery'.") + group_id = str(group.get("id", "")) + if not group_id: + raise ValueError("Inserted recovery TaskGroup requires an ID.") + if any(item["id"] == group_id for item in self._graph["task_groups"]): + raise ValueError(f"RuntimeGraph already contains TaskGroup {group_id!r}.") + nodes = [deepcopy(dict(node)) for node in recovery_nodes] + if not nodes: + raise ValueError("Recovery subgraph must contain at least one node.") + for node in nodes: + node.pop("contract", None) + node.pop("resources", None) + node["role"] = "cleanup" if node.get("role") == "cleanup" else "recovery" + node["task_instance_id"] = group_id + node["task_type"] = group["task_type"] + recovery_ids = {str(node["id"]) for node in nodes} + if len(recovery_ids) != len(nodes): + raise ValueError("Recovery node IDs must be unique.") + node_by_id = {str(node["id"]): node for node in self._graph["nodes"]} + children_by_id = {node_id: [] for node_id in node_by_id} + for node_id, node in node_by_id.items(): + for dependency in node["depends_on"]: + children_by_id[str(dependency)].append(node_id) + descendants: set[str] = set() + pending = list(children_by_id[failed_node_id]) + while pending: + node_id = pending.pop() + if node_id in descendants: + continue + descendants.add(node_id) + pending.extend(children_by_id[node_id]) + same_group_descendants = { + node_id + for node_id in descendants + if str(node_by_id[node_id]["task_instance_id"]) == failed_group_id + } + cleanup_suffix_ids: set[str] = set() + if ( + str(failed_node["atomic_action"]) == "HandOver" + and not preserve_failed_group_suffix + ): + # A failed handover leaves ownership indeterminate. Its + # transfer-arm retreat/home tail must not execute from a stale + # handover pose; recovery owns the cleanup before replanning. + cleanup_suffix_ids = { + node_id + for node_id in same_group_descendants + if node_by_id[node_id]["role"] == "cleanup" + } + non_cleanup_dependents = [ + node_id + for node_id in same_group_descendants - cleanup_suffix_ids + if any( + dependency in cleanup_suffix_ids + for dependency in node_by_id[node_id]["depends_on"] + ) + ] + if non_cleanup_dependents: + raise ValueError( + "Cannot remove the HandOver cleanup suffix because it feeds " + f"same-group non-cleanup nodes: {sorted(non_cleanup_dependents)}." + ) + first = [ + node + for node in nodes + if not any(dep in recovery_ids for dep in node.get("depends_on", [])) + ] + if not first: + raise ValueError("Recovery subgraph has no entry node.") + for node in first: + node["depends_on"] = list( + dict.fromkeys([*node.get("depends_on", []), failed_node_id]) + ) + terminal_ids = _terminal_ids(nodes) + + patched = deepcopy(self._graph) + for node in patched["nodes"]: + if node["id"] in recovery_ids: + raise ValueError(f"RuntimeGraph already contains node {node['id']!r}.") + node_id = str(node["id"]) + if ( + preserve_failed_group_suffix + or node_id not in descendants + or node_id in same_group_descendants + ): + continue + node["depends_on"] = list( + dict.fromkeys( + [ + dependency + for dependency in node["depends_on"] + if dependency != failed_node_id + and dependency not in cleanup_suffix_ids + ] + + terminal_ids + ) + ) + if cleanup_suffix_ids: + patched["nodes"] = [ + node + for node in patched["nodes"] + if str(node["id"]) not in cleanup_suffix_ids + ] + failed_group = next( + item + for item in patched["task_groups"] + if str(item["id"]) == failed_group_id + ) + failed_group["node_ids"] = [ + node_id + for node_id in failed_group["node_ids"] + if node_id not in cleanup_suffix_ids + ] + group["depends_on"] = list( + dict.fromkeys([failed_group_id, *group.get("depends_on", [])]) + ) + group.pop("contract", None) + group["node_ids"] = [str(node["id"]) for node in nodes] + for downstream in patched["task_groups"]: + if ( + not preserve_failed_group_suffix + and failed_group_id in downstream["depends_on"] + ): + downstream["depends_on"] = [ + dependency + for dependency in downstream["depends_on"] + if dependency != failed_group_id + ] + [group_id] + patched["nodes"].extend(nodes) + patched["task_groups"].append(group) + patched["metadata"] = { + **patched.get("metadata", {}), + "runtime_revision": len(self.revisions) + 1, + } + from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph + + self._graph = link_seed_graph( + patched, + registry=self.registry, + ) + self._recovery_action_count += len(nodes) + self.revisions.append( + GraphRevision( + revision=len(self.revisions) + 1, + kind="insert_recovery", + reason=failure_type, + failed_node_id=failed_node_id, + inserted_group_ids=(group_id,), + replaced_group_ids=(), + active_env_ids=env_ids, + ) + ) + return self.graph + + def insert_default_recovery( + self, + *, + failed_node_id: str, + failure_type: str, + active_env_ids: Sequence[int] | None = None, + resume_failed_group: bool = False, + ) -> dict[str, Any]: + """Insert one of the deliberately small built-in recovery strategies.""" + if failure_type != "object_fallen": + raise ValueError( + f"No bounded default recovery is registered for {failure_type!r}." + ) + nodes, group = build_upright_recovery( + self._graph, + failed_node_id=failed_node_id, + revision=len(self.revisions) + 1, + resume_failed_group=resume_failed_group, + ) + return self.insert_recovery_subgraph( + failed_node_id=failed_node_id, + recovery_nodes=nodes, + recovery_group=group, + failure_type=failure_type, + active_env_ids=active_env_ids, + preserve_failed_group_suffix=resume_failed_group, + ) + + def replace_unfinished_suffix( + self, + replacement: Mapping[str, Any], + *, + completed_group_ids: Sequence[str], + reason: str, + ) -> dict[str, Any]: + """Install a fully replanned suffix while preserving completed groups.""" + if len(self.revisions) >= self.max_revisions: + raise RuntimeError("RuntimeGraph revision budget exhausted.") + candidate = validate_seed_graph( + replacement, + known_actions=self.registry.names(), + ) + from embodichain.gen_sim.action_engine.planning.linker import ( + validate_persisted_contracts, + ) + + validate_persisted_contracts(candidate, self.registry) + if candidate["task_id"] != self.seed_graph["task_id"]: + raise ValueError("Suffix replanning cannot change the task_id.") + if candidate["capability_catalog_hash"] != self.registry.catalog_hash(): + raise ValueError( + "Replanned suffix capability catalog does not match runtime." + ) + current_groups = {group["id"]: group for group in self._graph["task_groups"]} + replacement_groups = {group["id"]: group for group in candidate["task_groups"]} + current_nodes = {node["id"]: node for node in self._graph["nodes"]} + replacement_nodes = {node["id"]: node for node in candidate["nodes"]} + completed = set(completed_group_ids) + for group_id in completed: + current_group = current_groups.get(group_id) + replacement_group = replacement_groups.get(group_id) + if current_group is None or replacement_group != current_group: + raise ValueError( + f"Replanning changed completed TaskGroup {group_id!r}." + ) + if any( + replacement_nodes.get(node_id) != current_nodes[node_id] + for node_id in current_group["node_ids"] + ): + raise ValueError( + f"Replanning changed nodes of completed TaskGroup {group_id!r}." + ) + replaced = tuple(sorted(set(current_groups) - completed)) + self._graph = candidate + self.revisions.append( + GraphRevision( + revision=len(self.revisions) + 1, + kind="replan_suffix", + reason=str(reason), + failed_node_id=None, + inserted_group_ids=tuple( + sorted(set(replacement_groups) - set(current_groups)) + ), + replaced_group_ids=replaced, + active_env_ids=tuple(range(self.num_envs)), + ) + ) + return self.graph + + +def classify_failure( + action_name: str, + *, + planning_succeeded: bool, + postcondition_succeeded: bool | None = None, + object_fallen: bool = False, + held_before: bool = False, + held_after: bool = False, + registry: AtomicCapabilityRegistry | None = None, +) -> str: + """Classify only the bounded common recovery cases supported by v2.""" + capability = (registry or build_atomic_capability_registry()).get(action_name) + if capability.failure_classifier_hook is not None: + result = capability.failure_classifier_hook( + action_name=action_name, + planning_succeeded=planning_succeeded, + postcondition_succeeded=postcondition_succeeded, + object_fallen=object_fallen, + held_before=held_before, + held_after=held_after, + ) + if result not in FAILURE_TYPES: + raise ValueError( + f"AtomicAction {action_name!r} failure classifier returned {result!r}." + ) + return result + if not planning_succeeded: + return "search_exhausted" + if object_fallen: + return "object_fallen" + if held_before and not held_after: + return "object_dropped" + if capability.failure_classifier == "grasp" and not held_after: + return "grasp_missed" + if postcondition_succeeded is False: + return "postcondition_failed" + return "postcondition_failed" + + +def build_upright_recovery( + graph: Mapping[str, Any], + *, + failed_node_id: str, + revision: int, + resume_failed_group: bool = False, +) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Build a coordinate-free E2 recovery group for a fallen rigid object.""" + failed = _node(graph, failed_node_id) + object_uid = str(failed["object_uid"]) + group_id = f"recovery_e2_{int(revision):02d}_{failed_node_id}" + actor = _recovery_actor(graph, failed) + held_consumer_arm = None + if not resume_failed_group: + held_consumer_arm = _downstream_held_consumer_arm(graph, failed, object_uid) + if held_consumer_arm is not None and not ( + actor.get("mode") == "required" and actor.get("arm") == held_consumer_arm + ): + raise ValueError( + "Recovery cannot satisfy the downstream held-object contract without " + "changing the failed TaskGroup actor; resume and replay the failed " + "TaskGroup instead." + ) + hold_for_downstream = ( + held_consumer_arm is not None + and actor.get("mode") == "required" + and actor.get("arm") == held_consumer_arm + ) + upright = motion_policy(("orientation", "upright")) + full_specs = ( + ("PickUp", {"kind": "object", "object": object_uid}, upright), + ( + "MoveHeldObject", + {"kind": "semantic_goal", "semantic_step": group_id, "phase": "final"}, + upright, + ), + ("Place", {"kind": "current_held_pose"}, upright), + ( + "MoveEndEffector", + {"kind": "policy_pose", "source": "release", "operation": "retreat"}, + upright, + ), + ( + "MoveJoints", + {"kind": "joint_state", "source": "initial"}, + motion_policy(), + ), + ) + specs = full_specs[:2] if hold_for_downstream else full_specs + nodes = [] + registry = build_atomic_capability_registry() + dependencies: list[str] = [] + for index, (action, binding, policy_spec) in enumerate(specs, start=1): + node_id = f"{group_id}__a{index:02d}" + node = { + "id": node_id, + "atomic_action": action, + "object_uid": object_uid, + "actor": actor, + "control": "arm", + "target_binding": binding, + "depends_on": dependencies, + "task_instance_id": group_id, + "task_type": "E2", + "role": "recovery" if index <= 3 else "cleanup", + "precondition": {}, + "postcondition": {}, + "motion_policy": deepcopy(dict(policy_spec)), + } + node["precondition"] = capability_precondition( + registry.get(action), + object_uid=object_uid, + actor=actor, + target_binding=binding, + ) + nodes.append(node) + dependencies = [node_id] + group = { + "id": group_id, + "task_type": "E2", + "role": "recovery", + "operator": "orient_object", + "object_uid": object_uid, + "actor": actor, + "goal": { + "relation": "none", + "reference_state": "live", + "orientation_goal": "upright", + "orientation_axis": "none", + "position_anchor": "live_xy", + "support_object": "table", + "upright_local_axis": "long_axis", + "terminal_behavior": "hold" if hold_for_downstream else "place", + }, + "depends_on": [], + "parent_task_instance_id": str(failed["task_instance_id"]), + "node_ids": [node["id"] for node in nodes], + "success": {"type": "object_upright", "object": object_uid}, + } + return nodes, group + + +def _recovery_actor( + graph: Mapping[str, Any], + failed: Mapping[str, Any], +) -> dict[str, Any]: + """Preserve the failed TaskGroup's arm-selection contract.""" + group_id = str(failed["task_instance_id"]) + group = next( + (item for item in graph["task_groups"] if str(item["id"]) == group_id), + None, + ) + source = (group or failed).get("actor", {"mode": "auto"}) + if not isinstance(source, Mapping): + raise ValueError(f"Failed TaskGroup {group_id!r} has an invalid actor.") + actor = deepcopy(dict(source)) + mode = str(actor.get("mode", "auto")) + if mode == "required": + if actor.get("arm") not in {"left_arm", "right_arm"}: + raise ValueError( + f"Failed TaskGroup {group_id!r} has an invalid required arm." + ) + elif mode == "auto": + actor = {"mode": "auto"} + elif mode == "coordinated": + raise ValueError( + "The single-arm upright recovery cannot inherit a coordinated actor." + ) + else: + raise ValueError(f"The upright recovery cannot inherit actor mode {mode!r}.") + return actor + + +def _downstream_held_consumer_arm( + graph: Mapping[str, Any], + failed: Mapping[str, Any], + object_uid: str, +) -> str | None: + failed_group_id = str(failed["task_instance_id"]) + nodes = {str(node["id"]): node for node in graph["nodes"]} + for group in graph["task_groups"]: + if failed_group_id not in {str(item) for item in group.get("depends_on", ())}: + continue + if str(group.get("object_uid")) != object_uid: + continue + node_ids = {str(item) for item in group["node_ids"]} + for node_id in group["node_ids"]: + node = nodes[str(node_id)] + if any(str(parent) in node_ids for parent in node["depends_on"]): + continue + for requirement in node.get("contract", {}).get("requires", ()): + if ( + requirement.get("predicate") == "object_held" + and requirement.get("object_uid") == object_uid + and requirement.get("arm") in {"left_arm", "right_arm"} + ): + return str(requirement["arm"]) + return None + + +def _node(graph: Mapping[str, Any], node_id: str) -> Mapping[str, Any]: + try: + return next(node for node in graph["nodes"] if node["id"] == node_id) + except StopIteration as error: + raise ValueError(f"RuntimeGraph contains no node {node_id!r}.") from error + + +def _terminal_ids(nodes: Sequence[Mapping[str, Any]]) -> list[str]: + depended = { + dependency for node in nodes for dependency in node.get("depends_on", []) + } + return [str(node["id"]) for node in nodes if node["id"] not in depended] + + +def _mask(value: torch.Tensor, count: int) -> torch.Tensor: + result = torch.as_tensor(value, dtype=torch.bool).reshape(-1) + if result.numel() != count: + raise ValueError(f"Expected a mask with {count} values.") + return result diff --git a/embodichain/gen_sim/action_engine/runtime/reporting.py b/embodichain/gen_sim/action_engine/runtime/reporting.py new file mode 100644 index 000000000..ea17ef151 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/reporting.py @@ -0,0 +1,348 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict validation and atomic publication for execution reports.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +import os +from pathlib import Path +import platform +import subprocess +import tempfile +from typing import Any + +from embodichain import __version__ as embodichain_version + +from .models import ExecutionReport + +__all__ = [ + "EXECUTION_REPORT_FILENAME", + "EXECUTION_REPORT_SCHEMA", + "build_execution_provenance", + "validate_execution_report", + "write_execution_report", +] + +EXECUTION_REPORT_SCHEMA = "action_engine_execution_report_v2" +EXECUTION_REPORT_FILENAME = "execution_report.json" + + +def build_execution_provenance( + *, + episode_seed: int | None = None, + runtime_arguments: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Capture the minimum code and runtime context needed to reproduce a run.""" + git_commit, git_dirty = _git_code_state() + provenance = { + "episode_seed": episode_seed, + "embodichain_version": str(embodichain_version), + "python_version": platform.python_version(), + "git_commit": git_commit, + "git_dirty": git_dirty, + "runtime_arguments": deepcopy(dict(runtime_arguments or {})), + } + return _validate_execution_provenance(provenance) + + +def validate_execution_report(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the tensor-free, strict-JSON Action Agent result protocol.""" + result = _mapping(value, "ExecutionReport") + keys = { + "schema_version", + "task_id", + "plan_hash", + "action_graph_hash", + "status", + "run_id", + "episode_id", + "provenance", + "environments", + "action_count", + "retry_count", + "recovery_count", + "revision_count", + "failure_events", + "graph_revisions", + "record_dir", + "error", + } + _keys(result, keys, "ExecutionReport") + if result.get("schema_version") != EXECUTION_REPORT_SCHEMA: + raise ValueError( + "ExecutionReport.schema_version must be " f"{EXECUTION_REPORT_SCHEMA!r}." + ) + for key in ("task_id", "run_id", "episode_id"): + result[key] = _nonempty(result.get(key), f"ExecutionReport.{key}") + result["provenance"] = _validate_execution_provenance(result.get("provenance")) + for key in ("plan_hash", "action_graph_hash"): + result[key] = _digest(result.get(key), f"ExecutionReport.{key}") + result["status"] = _enum( + result.get("status"), + {"succeeded", "failed", "rejected", "aborted"}, + "ExecutionReport.status", + ) + + env_keys = { + "env_id", + "success", + "semantic_success", + "action_count", + "retry_count", + "recovery_count", + "revision_count", + "failures", + } + environments = [] + for index, raw in enumerate( + _sequence(result.get("environments"), "ExecutionReport.environments") + ): + context = f"ExecutionReport.environments[{index}]" + environment = _mapping(raw, context) + _keys(environment, env_keys, context) + environment["env_id"] = _string(environment.get("env_id"), f"{context}.env_id") + if not isinstance(environment.get("success"), bool): + raise ValueError(f"{context}.success must be a boolean.") + semantic_success = _mapping( + environment.get("semantic_success"), f"{context}.semantic_success" + ) + if any(not isinstance(item, bool) for item in semantic_success.values()): + raise ValueError(f"{context}.semantic_success values must be booleans.") + environment["semantic_success"] = semantic_success + for key in ("action_count", "retry_count", "recovery_count", "revision_count"): + environment[key] = _integer( + environment.get(key), f"{context}.{key}", minimum=0 + ) + environment["failures"] = _mapping_sequence( + environment.get("failures"), f"{context}.failures" + ) + environments.append(environment) + result["environments"] = environments + + for key in ("action_count", "retry_count", "recovery_count", "revision_count"): + result[key] = _integer(result.get(key), f"ExecutionReport.{key}", minimum=0) + result["failure_events"] = _mapping_sequence( + result.get("failure_events"), "ExecutionReport.failure_events" + ) + result["graph_revisions"] = _mapping_sequence( + result.get("graph_revisions"), "ExecutionReport.graph_revisions" + ) + for key in ("record_dir", "error"): + if result.get(key) is not None: + result[key] = _string(result.get(key), f"ExecutionReport.{key}") + + if result["status"] == "rejected" and result["action_count"] != 0: + raise ValueError("A rejected ExecutionReport must have action_count=0.") + successes = [environment["success"] for environment in environments] + if result["status"] == "succeeded" and ( + not successes or not all(successes) or result.get("error") is not None + ): + raise ValueError( + "A succeeded ExecutionReport requires successful environments and no error." + ) + if result["status"] == "failed" and ( + not successes or all(successes) or result.get("error") is not None + ): + raise ValueError( + "A failed ExecutionReport requires at least one failed environment and no error." + ) + if result["status"] in {"rejected", "aborted"} and not result.get("error"): + raise ValueError( + f"A {result['status']} ExecutionReport requires a non-empty error." + ) + _json_safe(result, "ExecutionReport") + return result + + +def write_execution_report(output_dir: str | Path, value: Any) -> Path: + """Atomically write a validated execution report into a record directory.""" + payload = value.as_mapping() if isinstance(value, ExecutionReport) else value + validated = validate_execution_report(payload) + root = Path(output_dir).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + path = root / EXECUTION_REPORT_FILENAME + encoded = ( + json.dumps(validated, ensure_ascii=False, indent=2, allow_nan=False) + "\n" + ).encode("utf-8") + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=root, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as stream: + temporary = Path(stream.name) + stream.write(encoded) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + temporary = None + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + return path + + +def _validate_execution_provenance(value: Any) -> dict[str, Any]: + context = "ExecutionReport.provenance" + result = _mapping(value, context) + _keys( + result, + { + "episode_seed", + "embodichain_version", + "python_version", + "git_commit", + "git_dirty", + "runtime_arguments", + }, + context, + ) + seed = result.get("episode_seed") + if seed is not None and (not isinstance(seed, int) or isinstance(seed, bool)): + raise ValueError(f"{context}.episode_seed must be an integer or null.") + result["embodichain_version"] = _nonempty( + result.get("embodichain_version"), f"{context}.embodichain_version" + ) + result["python_version"] = _nonempty( + result.get("python_version"), f"{context}.python_version" + ) + commit = result.get("git_commit") + if commit is not None: + commit = _string(commit, f"{context}.git_commit") + if len(commit) not in {40, 64} or any( + character not in "0123456789abcdef" for character in commit + ): + raise ValueError( + f"{context}.git_commit must be a lowercase Git object ID or null." + ) + result["git_commit"] = commit + dirty = result.get("git_dirty") + if dirty is not None and not isinstance(dirty, bool): + raise ValueError(f"{context}.git_dirty must be a boolean or null.") + arguments = _mapping( + result.get("runtime_arguments"), f"{context}.runtime_arguments" + ) + if any(not isinstance(key, str) or not key for key in arguments): + raise ValueError(f"{context}.runtime_arguments keys must be non-empty strings.") + _json_safe(arguments, f"{context}.runtime_arguments") + result["runtime_arguments"] = arguments + return result + + +def _git_code_state() -> tuple[str | None, bool | None]: + repository = Path(__file__).resolve().parents[4] + try: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + cwd=repository, + check=False, + capture_output=True, + text=True, + timeout=2.0, + ) + except (OSError, subprocess.SubprocessError): + return None, None + commit_id = commit.stdout.strip().lower() + if commit.returncode != 0 or len(commit_id) not in {40, 64}: + return None, None + try: + status = subprocess.run( + ["git", "status", "--porcelain"], + cwd=repository, + check=False, + capture_output=True, + text=True, + timeout=2.0, + ) + except (OSError, subprocess.SubprocessError): + return commit_id, None + dirty = bool(status.stdout.strip()) if status.returncode == 0 else None + return commit_id, dirty + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"{context} must be a list.") + return list(value) + + +def _keys(value: Mapping[str, Any], expected: set[str], context: str) -> None: + if set(value) != expected: + raise ValueError( + f"{context} requires exactly fields {sorted(expected)}; " + f"received {sorted(value)}." + ) + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{context} must be a string.") + return value.strip() + + +def _nonempty(value: Any, context: str) -> str: + result = _string(value, context) + if not result: + raise ValueError(f"{context} must not be empty.") + return result + + +def _enum(value: Any, choices: set[str], context: str) -> str: + result = _string(value, context) + if result not in choices: + raise ValueError(f"{context} must be one of {sorted(choices)}.") + return result + + +def _integer(value: Any, context: str, *, minimum: int) -> int: + if not isinstance(value, int) or isinstance(value, bool) or value < minimum: + raise ValueError(f"{context} must be an integer in the allowed range.") + return value + + +def _digest(value: Any, context: str) -> str: + result = _string(value, context) + if len(result) != 64 or any( + character not in "0123456789abcdef" for character in result + ): + raise ValueError(f"{context} must be a lowercase SHA-256 digest.") + return result + + +def _mapping_sequence(value: Any, context: str) -> list[dict[str, Any]]: + result = [_mapping(item, context) for item in _sequence(value, context)] + _json_safe(result, context) + return result + + +def _json_safe(value: Any, context: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError) as error: + raise ValueError(f"{context} must be finite and JSON serializable.") from error diff --git a/embodichain/gen_sim/action_engine/runtime/robot_parts.py b/embodichain/gen_sim/action_engine/runtime/robot_parts.py new file mode 100644 index 000000000..fe6f71c58 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/robot_parts.py @@ -0,0 +1,34 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Resolve semantic Action Engine arms to physical robot control parts.""" + +from __future__ import annotations + +from typing import Any + +__all__ = ["arm_control_part"] + + +def arm_control_part(env: Any, arm: str) -> str: + """Return the physical arm control part for a semantic arm name.""" + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a semantic arm, got {arm!r}.") + if hasattr(env, "get_agent_arm_control_part"): + part = env.get_agent_arm_control_part(arm == "left_arm") + if part: + return str(part) + return arm diff --git a/embodichain/gen_sim/action_engine/runtime/solver_compat.py b/embodichain/gen_sim/action_engine/runtime/solver_compat.py new file mode 100644 index 000000000..810bc8cf7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/solver_compat.py @@ -0,0 +1,234 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Install solver compatibility corrections scoped to Action Engine.""" + +from __future__ import annotations + +from collections.abc import Mapping +import functools +import threading +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.sim.solvers import PytorchSolver, URSolver, URSolverCfg + +__all__ = [ + "install_action_engine_solver_compat", + "install_pytorch_solver_tcp_compat", + "install_ur5_solver_frame_compat", + "repair_action_engine_ur5_solver_cfg", +] + +_PYTORCH_INSTALL_MARKER = "_action_engine_tcp_inverse_compat_installed" +_UR5_INSTALL_MARKER = "_action_engine_ur5_frame_compat_installed" +_UR5_ANALYTIC_TO_URDF_EE = np.eye(4, dtype=np.float32) +_UR5_ANALYTIC_TO_URDF_EE[0, 3] = -0.01 +_UR_DH_FIELDS = ("d1", "a2", "a3", "d4", "d5", "d6") + + +def repair_action_engine_ur5_solver_cfg(robot_cfg: Any) -> int: + """Repair stale UR10 DH defaults before Action Engine creates a UR5 robot. + + ``SolverCfg.from_dict`` constructs a UR10 config before assigning a + non-default ``ur_type``. Generated UR5 Action Engine configs therefore + reach the environment with UR10 DH values. Repair only that exact stale + signature so explicitly calibrated parameters remain untouched. + + Args: + robot_cfg: Robot configuration whose solver configs will be inspected. + + Returns: + Number of unique solver configs repaired by this call. + """ + configured = getattr(robot_cfg, "solver_cfg", None) + candidates = ( + configured.values() if isinstance(configured, Mapping) else (configured,) + ) + stale_defaults = URSolverCfg() + stale_dh = tuple(float(getattr(stale_defaults, name)) for name in _UR_DH_FIELDS) + + repaired = 0 + visited: set[int] = set() + for solver_cfg in candidates: + cfg_id = id(solver_cfg) + if cfg_id in visited: + continue + visited.add(cfg_id) + if not isinstance(solver_cfg, URSolverCfg): + continue + ur_type = str(getattr(solver_cfg, "ur_type", "")) + if ur_type != "ur5": + continue + current_dh = tuple(float(getattr(solver_cfg, name)) for name in _UR_DH_FIELDS) + if not np.allclose(current_dh, stale_dh, rtol=0.0, atol=1.0e-12): + continue + canonical = URSolverCfg(ur_type=ur_type) + for name in _UR_DH_FIELDS: + setattr(solver_cfg, name, getattr(canonical, name)) + repaired += 1 + return repaired + + +def install_action_engine_solver_compat(robot: Any) -> int: + """Install all solver corrections required by the Action Engine runtime.""" + return install_pytorch_solver_tcp_compat(robot) + install_ur5_solver_frame_compat( + robot + ) + + +def install_pytorch_solver_tcp_compat(robot: Any) -> int: + """Correct TCP inversion on every PytorchSolver owned by ``robot``. + + The shared solver currently transposes a rotation into an overlapping + tensor view. This wrapper transforms the requested TCP pose with a proper + matrix inverse, temporarily presents an identity TCP to the original + implementation, and otherwise preserves its sampling and ranking behavior. + + Args: + robot: Initialized robot containing its private solver registry. + + Returns: + Number of solver instances wrapped by this call. + """ + solvers = getattr(robot, "_solvers", None) + if not isinstance(solvers, Mapping): + return 0 + + installed = 0 + visited: set[int] = set() + for solver in solvers.values(): + solver_id = id(solver) + if solver_id in visited: + continue + visited.add(solver_id) + if not isinstance(solver, PytorchSolver) or bool( + getattr(solver, _PYTORCH_INSTALL_MARKER, False) + ): + continue + _wrap_solver(solver) + installed += 1 + return installed + + +def _wrap_solver(solver: PytorchSolver) -> None: + original_get_ik = solver.get_ik + call_lock = threading.RLock() + + @functools.wraps(original_get_ik) + def corrected_get_ik( + target_xpos: torch.Tensor | np.ndarray, + *args: Any, + **kwargs: Any, + ) -> Any: + target = torch.as_tensor( + target_xpos, + dtype=torch.float32, + device=solver.device, + ) + tcp = torch.as_tensor( + solver.tcp_xpos, + dtype=torch.float32, + device=solver.device, + ) + link_target = target @ torch.linalg.inv(tcp) + + # The solver instance is shared by vectorized environments. Protect the + # temporary TCP substitution in case a caller plans from another thread. + with call_lock: + active_tcp = solver.tcp_xpos + solver.tcp_xpos = np.eye(4, dtype=np.float32) + try: + return original_get_ik( + target_xpos=link_target, + *args, + **kwargs, + ) + finally: + solver.tcp_xpos = active_tcp + + solver.get_ik = corrected_get_ik + setattr(solver, _PYTORCH_INSTALL_MARKER, True) + + +def install_ur5_solver_frame_compat(robot: Any) -> int: + """Align UR5 analytic IK targets with the URDF ``ee_link`` frame. + + The UR5 asset carries a fixed ``-0.01 m`` local-x offset on ``ee_link`` + that is absent from the analytic DH model. The correction is installed + only for UR5 solvers owned by an Action Engine environment. + + Args: + robot: Initialized robot containing its private solver registry. + + Returns: + Number of solver instances wrapped by this call. + """ + solvers = getattr(robot, "_solvers", None) + if not isinstance(solvers, Mapping): + return 0 + + installed = 0 + visited: set[int] = set() + for solver in solvers.values(): + solver_id = id(solver) + if solver_id in visited: + continue + visited.add(solver_id) + if ( + not isinstance(solver, URSolver) + or str(getattr(getattr(solver, "cfg", None), "ur_type", "")) != "ur5" + or bool(getattr(solver, _UR5_INSTALL_MARKER, False)) + ): + continue + _wrap_ur5_solver(solver) + installed += 1 + return installed + + +def _wrap_ur5_solver(solver: URSolver) -> None: + original_get_ik = solver.get_ik + + @functools.wraps(original_get_ik) + def corrected_get_ik( + target_xpos: torch.Tensor | np.ndarray, + *args: Any, + **kwargs: Any, + ) -> Any: + target = torch.as_tensor( + target_xpos, + dtype=torch.float32, + device=solver.device, + ) + tcp = torch.as_tensor( + solver.tcp_xpos, + dtype=torch.float32, + device=solver.device, + ) + analytic_to_urdf = torch.as_tensor( + _UR5_ANALYTIC_TO_URDF_EE, + dtype=torch.float32, + device=solver.device, + ) + corrected_target = ( + target @ torch.linalg.inv(tcp) @ torch.linalg.inv(analytic_to_urdf) @ tcp + ) + return original_get_ik(corrected_target, *args, **kwargs) + + solver.get_ik = corrected_get_ik + setattr(solver, _UR5_INSTALL_MARKER, True) diff --git a/embodichain/gen_sim/action_engine/runtime/state.py b/embodichain/gen_sim/action_engine/runtime/state.py new file mode 100644 index 000000000..dbafc4748 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/state.py @@ -0,0 +1,84 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Action Engine execution state at the atomic-planning boundary.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Mapping + +import torch + +from embodichain.lab.sim.atomic_actions import ( + HeldObjectState, + TaskState, +) + +__all__ = ["ExecutionState"] + + +@dataclass(slots=True, eq=False) +class ExecutionState: + """Projected task state paired with the next full-robot planning seed. + + The simulation atomic-action package deliberately no longer exposes the + legacy ``WorldState`` compatibility object. Action Engine keeps this narrow + orchestration state locally and converts it to immutable ``TaskState`` and + ``PlanningContext`` values immediately before invoking the shared planner. + """ + + last_qpos: torch.Tensor + held_objects: dict[str, HeldObjectState] = field(default_factory=dict) + + def get_held_object(self, control_part: str) -> HeldObjectState | None: + """Return the held-object relation for one control part.""" + return self.held_objects.get(control_part) + + def with_updates( + self, + *, + last_qpos: torch.Tensor | None = None, + held_objects: Mapping[str, HeldObjectState] | None = None, + ) -> ExecutionState: + """Return a detached successor state.""" + return ExecutionState( + last_qpos=self.last_qpos if last_qpos is None else last_qpos, + held_objects=dict( + self.held_objects if held_objects is None else held_objects + ), + ) + + def to_task_state(self) -> TaskState: + """Convert this state to the shared immutable symbolic task contract.""" + return TaskState( + batch_size=int(self.last_qpos.shape[0]), + device=self.last_qpos.device, + held_objects=self.held_objects, + ) + + @classmethod + def from_task_state( + cls, + task: TaskState, + *, + last_qpos: torch.Tensor, + ) -> ExecutionState: + """Build an orchestration state from a committed or projected task state.""" + return cls( + last_qpos=last_qpos, + held_objects=dict(task.held_objects), + ) diff --git a/embodichain/gen_sim/action_engine/tasks/__init__.py b/embodichain/gen_sim/action_engine/tasks/__init__.py new file mode 100644 index 000000000..843a55efa --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/__init__.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Task-first generation and scene hand-off for Action Engine v2.""" + +from __future__ import annotations + +from .assembly import GroundedTaskSpec +from .interpretation import ( + GroundingCaller, + INSTRUCTION_INTENT_SCHEMA, + InstructionDraftResult, + InstructionCaller, + InstructionIntent, + ground_instruction_draft, + interpret_instruction_draft, + interpret_and_ground_task_spec, + validate_instruction_intent, +) +from .recipes import instantiate_seed_graph +from .scene import SceneHandoff, validate_scene_handoff + +__all__ = [ + "GroundedTaskSpec", + "GroundingCaller", + "INSTRUCTION_INTENT_SCHEMA", + "InstructionDraftResult", + "InstructionCaller", + "InstructionIntent", + "SceneHandoff", + "ground_instruction_draft", + "instantiate_seed_graph", + "interpret_instruction_draft", + "interpret_and_ground_task_spec", + "validate_instruction_intent", + "validate_scene_handoff", +] diff --git a/embodichain/gen_sim/action_engine/tasks/assembly.py b/embodichain/gen_sim/action_engine/tasks/assembly.py new file mode 100644 index 000000000..7fdfe7246 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/assembly.py @@ -0,0 +1,425 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Language-neutral scene inventory and grounded TaskSpec assembly.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +from typing import Any + +from embodichain.gen_sim.action_engine.domain import ( + TASK_CONTRACTS, + task_contract, + task_success_type, + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.generation.config_builder import ( + canonical_robot_profile, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) + +__all__ = [ + "GroundedTaskBuilder", + "GroundedTaskSpec", + "SceneEntity", + "SceneInventory", + "validate_source_compatibility", + "validate_target_compatibility", +] + + +@dataclass(frozen=True) +class GroundedTaskSpec: + """One explicit TaskSpec plus verified scene role bindings.""" + + task_spec: dict[str, Any] + scene_requirements: dict[str, Any] + role_bindings: dict[str, str] + + +@dataclass(frozen=True) +class SceneEntity: + """One scene entity with source semantics preserved verbatim.""" + + uid: str + role: str + name: str + description: str + category: str + color: str | None + position: tuple[float, float, float] + affordances: frozenset[str] = frozenset() + initial_state: Mapping[str, Any] = field(default_factory=dict) + attributes: Mapping[str, Any] = field(default_factory=dict) + source_uid: str = "" + + +class SceneInventory: + """Structural scene index without natural-language matching rules.""" + + _PASSIVE_ROLES = frozenset( + { + "background", + "camera", + "light", + "robot", + "sensor", + "support_surface", + "table", + } + ) + + def __init__( + self, + scene_objects: Sequence[Mapping[str, Any]], + *, + robot_profile: str, + ) -> None: + self.profile = canonical_robot_profile(robot_profile) + self.entities = tuple(_scene_entity(item) for item in scene_objects) + self.by_uid = {entity.uid: entity for entity in self.entities} + if len(self.by_uid) != len(self.entities): + raise ValueError("Scene inventory contains duplicate runtime UIDs.") + self.support = tuple( + entity + for entity in self.entities + if entity.uid == "table" or entity.role in {"table", "support_surface"} + ) + self.passive = tuple( + entity + for entity in self.entities + if entity in self.support or entity.role in self._PASSIVE_ROLES + ) + self.interactive = tuple( + entity for entity in self.entities if entity not in self.passive + ) + if not self.interactive: + raise ValueError("Task planning requires at least one interaction object.") + + @property + def movable(self) -> tuple[SceneEntity, ...]: + """Compatibility alias for callers that mean source candidates.""" + return self.interactive + + def left_score(self, entity: SceneEntity) -> float: + """Return robot-relative lateral score; positive values are left. + + Generated dual-arm profiles share one final world layout: the semantic + left arm is on world ``-Y`` after all robot-level transforms. + """ + return -entity.position[1] + + +class GroundedTaskBuilder: + """Assemble grounded E1-E9 instances without parsing instruction text.""" + + def __init__( + self, + task_id: str, + instruction: str, + inventory: SceneInventory, + *, + planner: str = "structured_llm_v2", + ) -> None: + self.task_id = task_id + self.instruction = instruction + self.inventory = inventory + self.planner = planner + self.instances: list[dict[str, Any]] = [] + self.role_by_uid: dict[str, str] = {} + self.requirements: dict[str, dict[str, Any]] = {} + self.previous_object_uid: str | None = None + self.previous_arm: str | None = None + self.last_task_by_object_uid: dict[str, tuple[str, str]] = {} + + def add( + self, + task_type: str, + object_entity: SceneEntity, + *, + target: SceneEntity | None = None, + params: Mapping[str, Any] | None = None, + depends_on: Sequence[str] | None = None, + ) -> str: + values = deepcopy(dict(params or {})) + relation = str(values.get("relation", "none")) + validate_source_compatibility(task_type, (object_entity,)) + validate_target_compatibility(task_type, target, relation=relation) + + instance_id = f"task_{len(self.instances) + 1:02d}" + object_role = self._role( + object_entity, + required_affordances=task_contract(task_type).required_affordances, + initial_state={"orientation": "fallen"} if task_type == "E2" else {}, + ) + values = {"object_role": object_role, **values} + if task_type == "E3": + values["source_role"] = values.pop("object_role") + if target is not None: + values["target_role"] = self._role( + target, + required_affordances=_target_affordances(task_type, relation), + ) + if depends_on is None: + dependencies = [self.instances[-1]["id"]] if self.instances else [] + else: + dependencies = list(depends_on) + previous_for_object = self.last_task_by_object_uid.get(object_entity.uid) + if ( + task_type == "E4" + and previous_for_object is not None + and previous_for_object[1] == "E2" + and previous_for_object[0] not in dependencies + ): + dependencies.append(previous_for_object[0]) + self.instances.append( + { + "id": instance_id, + "task_type": task_type, + "params": values, + "depends_on": dependencies, + "role": "primary", + } + ) + self.last_task_by_object_uid[object_entity.uid] = (instance_id, task_type) + self.previous_object_uid = object_entity.uid + if task_type == "E4": + receive_arm = str(values.get("receive_arm", "")) + self.previous_arm = ( + receive_arm if receive_arm in {"left_arm", "right_arm"} else None + ) + elif str(values.get("required_arm", "")) in {"left_arm", "right_arm"}: + self.previous_arm = str(values["required_arm"]) + return instance_id + + def build(self) -> GroundedTaskSpec: + types = {item["task_type"] for item in self.instances} + if len(self.instances) == 1: + level = "L1" + elif len(types) == 1: + level = "L2" + else: + level = "L3" + success_terms = [ + { + "type": task_success_type(item["task_type"], item.get("params")), + "task_instance_id": item["id"], + } + for item in self.instances + ] + task = validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": self.task_id, + "level": level, + "instruction": self.instruction, + "reasoning_type": "none", + "task_instances": self.instances, + "success": {"op": "all", "terms": success_terms}, + "oracle": { + "task_order": [item["id"] for item in self.instances], + "role_bindings": dict(sorted(self.role_bindings().items())), + }, + "metadata": {"planner": self.planner}, + } + ) + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": self.task_id, + "objects": list(self.requirements.values()), + "cameras": [], + "spatial_constraints": [{"type": "preserve_source_scene"}], + "distractor_count": max( + 0, + len(self.inventory.interactive) - len(self.role_by_uid), + ), + "metadata": {"source": "existing_gym_project"}, + } + ) + return GroundedTaskSpec(task, requirements, self.role_bindings()) + + def role_bindings(self) -> dict[str, str]: + return {role: uid for uid, role in self.role_by_uid.items()} + + def _role( + self, + entity: SceneEntity, + task_type: str | None = None, + *, + required_affordances: Sequence[str] = (), + initial_state: Mapping[str, Any] | None = None, + ) -> str: + if task_type in TASK_CONTRACTS: + required_affordances = tuple( + set(required_affordances) + | set(task_contract(str(task_type)).required_affordances) + ) + if task_type == "E2": + initial_state = {"orientation": "fallen", **dict(initial_state or {})} + existing = self.role_by_uid.get(entity.uid) + if existing is not None: + requirement = self.requirements[existing] + requirement["affordances"] = sorted( + set(requirement["affordances"]) | set(required_affordances) + ) + requirement["initial_state"].update(dict(initial_state or {})) + return existing + role = f"object_{len(self.role_by_uid) + 1:02d}" + self.role_by_uid[entity.uid] = role + attributes = deepcopy(dict(entity.attributes)) + if entity.color is not None: + attributes.setdefault("color", entity.color) + self.requirements[role] = { + "role_id": role, + "category": entity.category or entity.role, + "count": 1, + "affordances": sorted(set(required_affordances)), + "initial_state": dict(initial_state or {}), + "attributes": attributes, + } + return role + + +def validate_source_compatibility( + task_type: str, + objects: Sequence[SceneEntity], +) -> None: + """Apply structural/explicit-affordance checks without a category taxonomy.""" + contract = task_contract(task_type) + if contract.source_structure == "articulation": + invalid = [entity.uid for entity in objects if entity.role != "articulation"] + else: + invalid = [ + entity.uid + for entity in objects + if entity.role not in {"object", "rigid_object"} + ] + if invalid: + structure_label = ( + "articulation" + if contract.source_structure == "articulation" + else "movable rigid-object" + ) + raise ValueError( + f"{task_type} requires {structure_label} structure; " + f"incompatible scene objects are {invalid}." + ) + required = set(contract.required_affordances) + for entity in objects: + if entity.affordances: + missing = required - set(entity.affordances) + if missing: + raise ValueError( + f"{task_type} is incompatible with scene object {entity.uid!r}; " + f"missing affordances {sorted(missing)}." + ) + + +def validate_target_compatibility( + task_type: str, + target: SceneEntity | None, + *, + relation: str, +) -> None: + """Reject only structural or explicitly declared target contradictions.""" + if task_type == "E1" and relation == "on" and target is not None: + # Support is a relation between two concrete bodies at a candidate + # pose. A positive affordance list is not a closed-world inventory, so + # omission of ``support_surface`` cannot prove incompatibility here. + return + requires_container = task_type == "E3" or ( + task_type == "E1" and relation == "inside" + ) + if requires_container and target is None: + raise ValueError( + f"{task_type} {relation} relation requires a target container." + ) + if not requires_container or target is None: + return + if target.role in SceneInventory._PASSIVE_ROLES: + raise ValueError( + f"{task_type} target {target.uid!r} is structurally incompatible " + "with containment." + ) + if target.affordances: + compatible = {"container", "fillable", "liquid_container", "receptacle"} + if set(target.affordances).isdisjoint(compatible): + raise ValueError( + f"{task_type} target {target.uid!r} has explicit affordances but " + f"none support containment; expected one of {sorted(compatible)}." + ) + + +def _target_affordances(task_type: str, relation: str) -> tuple[str, ...]: + if task_type == "E3" or (task_type == "E1" and relation == "inside"): + return ("container",) + return () + + +def _scene_entity(raw: Mapping[str, Any]) -> SceneEntity: + uid = str(raw.get("runtime_uid", raw.get("uid", ""))).strip() + if not uid: + raise ValueError("Every scene object requires a runtime UID.") + role = str(raw.get("role", raw.get("source_role", "object"))).strip().lower() + raw_category = raw.get("category", raw.get("object_category", "")) + category = "" if raw_category is None else str(raw_category).strip() + raw_color = raw.get("color") + attributes = raw.get("attributes", {}) + if not isinstance(attributes, Mapping): + raise ValueError(f"Scene object {uid!r} attributes must be a mapping.") + if raw_color is None: + raw_color = attributes.get("color") + color = str(raw_color).strip() if raw_color not in (None, "") else None + position = raw.get("init_pos", raw.get("position", (0.0, 0.0, 0.0))) + if ( + not isinstance(position, Sequence) + or isinstance(position, (str, bytes)) + or len(position) != 3 + ): + raise ValueError(f"Scene object {uid!r} requires a three-value init_pos.") + raw_affordances = raw.get("affordances", raw.get("capabilities", ())) + affordances = ( + frozenset( + str(item).strip().lower() for item in raw_affordances if str(item).strip() + ) + if isinstance(raw_affordances, Sequence) + and not isinstance(raw_affordances, (str, bytes)) + else frozenset() + ) + initial_state = raw.get("initial_state", raw.get("state", {})) + if not isinstance(initial_state, Mapping): + raise ValueError(f"Scene object {uid!r} initial_state must be a mapping.") + return SceneEntity( + uid=uid, + role=role, + name=str(raw.get("name", "")).strip(), + description=str(raw.get("description", "")).strip(), + category=category, + color=color, + position=tuple(float(value) for value in position), + affordances=affordances, + initial_state=dict(initial_state), + attributes=dict(attributes), + source_uid=str(raw.get("source_uid", "")).strip(), + ) diff --git a/embodichain/gen_sim/action_engine/tasks/grounding.py b/embodichain/gen_sim/action_engine/tasks/grounding.py new file mode 100644 index 000000000..de412e987 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/grounding.py @@ -0,0 +1,513 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Task-conditioned scene-UID grounding for structured instruction intents.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +import math +from time import perf_counter +from typing import Any + +from .assembly import SceneInventory + +__all__ = ["GroundingCaller", "GroundingResult", "ground_scene_references"] + +GroundingCaller = Callable[..., Mapping[str, Any]] + +_BINDING_KEYS = frozenset({"reference_id", "status", "uids", "confidence"}) +_QUANTIFIERS = frozenset({"one", "all", "count"}) +_REDACTED_KEYS = frozenset( + { + "absolute_position", + "bbox", + "bboxes", + "bounding_box", + "camera_matrix", + "center", + "centroid", + "coordinates", + "depth", + "dimensions", + "extrinsics", + "grasp_pose", + "init_local_pose", + "init_pos", + "init_rot", + "intrinsics", + "joint_positions", + "joints", + "keypoint", + "keypoints", + "location", + "matrix", + "pose", + "position", + "position_xyz", + "qpos", + "quaternion", + "rotation", + "scale", + "target_pose", + "trajectory", + "transform", + "translation", + "waypoints", + "world_x", + "world_y", + "world_z", + "x", + "y", + "z", + } +) + +_GROUNDING_SCHEMA: dict[str, Any] = { + "title": "ActionEngineSceneGrounding", + "type": "object", + "additionalProperties": False, + "required": ["bindings"], + "properties": { + "bindings": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": False, + "required": sorted(_BINDING_KEYS), + "properties": { + "reference_id": {"type": "string"}, + "status": { + "type": "string", + "enum": ["resolved", "ambiguous", "not_found"], + }, + "uids": { + "type": "array", + "items": {"type": "string"}, + "uniqueItems": True, + }, + "confidence": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0, + }, + }, + }, + } + }, +} + + +@dataclass(frozen=True) +class GroundingResult: + """Validated scene bindings and aggregate call statistics. + + Attributes: + bindings: Mapping from ``.`` to scene UIDs. + attempts: Number of grounding-model calls, including one repair call. + latency_seconds: Total elapsed wall-clock time across the grounding stage. + """ + + bindings: dict[str, tuple[str, ...]] + attempts: int + latency_seconds: float + + +def ground_scene_references( + instruction: str, + intent: Mapping[str, Any], + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], + model: str | None, + caller: GroundingCaller, +) -> GroundingResult: + """Resolve every ``scene_ref`` selector in one task-conditioned batch. + + The grounding model can only select stable UIDs from a redacted inventory. + Its output does not add affordances, physical state, coordinates, or poses. + One failed local validation is repaired with one additional model call. + + Args: + instruction: Original user instruction for task-level context. + intent: Validated structured instruction intent. + inventory: Structural scene inventory defining authoritative candidates. + scene_objects: Original semantic inventory used to retain open labels. + model: Model name forwarded unchanged to the injected caller. + caller: Structured model transport accepting ``prompt``, ``schema``, and + ``model`` keyword arguments. + + Returns: + Validated UID bindings together with call-count and latency statistics. + + Raises: + TypeError: If the intent or response has an invalid container type. + ValueError: If requests are malformed or grounding remains invalid after + one repair attempt. + """ + if not isinstance(instruction, str) or not instruction.strip(): + raise ValueError("Grounding instruction must be a non-empty string.") + if not callable(caller): + raise TypeError("Grounding caller must be callable.") + + requests = _collect_requests(intent) + prompt_inventory = _grounding_inventory(inventory, scene_objects) + prompt = _grounding_prompt(instruction.strip(), requests, prompt_inventory) + started = perf_counter() + first_error: Exception | None = None + + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nREPAIR OVERRIDE: the previous grounding JSON failed local " + "validation. Return one corrected JSON object only. Preserve the " + "exact output fields bindings/reference_id/status/uids/confidence, " + "cover every requested reference exactly once, and select only " + "UIDs from the supplied candidate inventory. Validation error: " + f"{first_error}" + ) + try: + response = caller( + prompt=current_prompt, + schema=deepcopy(_GROUNDING_SCHEMA), + model=model, + ) + bindings = _validate_response( + response, + requests=requests, + inventory=inventory, + ) + return GroundingResult( + bindings=bindings, + attempts=attempt + 1, + latency_seconds=perf_counter() - started, + ) + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "Scene grounding failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def _collect_requests(intent: Mapping[str, Any]) -> list[dict[str, Any]]: + if not isinstance(intent, Mapping): + raise TypeError("Instruction intent must be a mapping.") + steps = intent.get("steps") + if not isinstance(steps, Sequence) or isinstance(steps, (str, bytes)): + raise ValueError("Instruction intent steps must be a list.") + + requests: list[dict[str, Any]] = [] + request_ids: set[str] = set() + for step_index, step in enumerate(steps): + context = f"InstructionIntent.steps[{step_index}]" + if not isinstance(step, Mapping): + raise ValueError(f"{context} must be a mapping.") + step_id = step.get("id") + if not isinstance(step_id, str) or not step_id.strip(): + raise ValueError(f"{context}.id must be a non-empty string.") + task_type = step.get("task_type") + if not isinstance(task_type, str) or not task_type.strip(): + raise ValueError(f"{context}.task_type must be a non-empty string.") + relation = step.get("relation", "none") + if not isinstance(relation, str): + raise ValueError(f"{context}.relation must be a string.") + + for slot in ("object", "target"): + selector = step.get(slot) + if not isinstance(selector, Mapping): + raise ValueError(f"{context}.{slot} must be a mapping.") + if selector.get("kind") != "scene_ref": + continue + reference = selector.get("reference") + if not isinstance(reference, str) or not reference.strip(): + raise ValueError( + f"{context}.{slot}.reference must be a non-empty string." + ) + quantifier = selector.get("quantifier") + if quantifier not in _QUANTIFIERS: + raise ValueError( + f"{context}.{slot}.quantifier must be one of " + f"{sorted(_QUANTIFIERS)}." + ) + count = selector.get("count") + if not isinstance(count, int) or isinstance(count, bool) or count < 0: + raise ValueError(f"{context}.{slot}.count must be an integer >= 0.") + if quantifier == "count" and count < 1: + raise ValueError( + f"{context}.{slot} quantifier=count requires count>=1." + ) + if quantifier != "count" and count != 0: + raise ValueError( + f"{context}.{slot} quantifier={quantifier} requires count=0." + ) + + request_id = f"{step_id}.{slot}" + if request_id in request_ids: + raise ValueError(f"Duplicate grounding request ID {request_id!r}.") + request_ids.add(request_id) + requests.append( + { + "reference_id": request_id, + "step_id": step_id, + "slot": slot, + "task_type": task_type, + "relation": relation, + "reference": reference.strip(), + "quantifier": quantifier, + "count": count, + } + ) + return requests + + +def _grounding_inventory( + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + raw_by_uid: dict[str, Mapping[str, Any]] = {} + for item_index, raw in enumerate(scene_objects): + if not isinstance(raw, Mapping): + raise ValueError(f"Scene inventory item {item_index} must be a mapping.") + uid = str(raw.get("runtime_uid", raw.get("uid", ""))).strip() + if uid: + raw_by_uid[uid] = raw + + ranked = sorted( + inventory.entities, + key=lambda entity: (-inventory.left_score(entity), entity.uid), + ) + rank_by_uid = {entity.uid: rank for rank, entity in enumerate(ranked, start=1)} + payload = [] + for entity in sorted(inventory.entities, key=lambda item: item.uid): + raw = raw_by_uid.get(entity.uid, {}) + score = inventory.left_score(entity) + side = "left" if score > 0.0 else "right" if score < 0.0 else "center" + raw_category = raw.get( + "category", + raw.get("object_category", entity.category), + ) + attributes = _redact_semantic_mapping(entity.attributes) + if entity.color is not None: + attributes.setdefault("color", entity.color) + payload.append( + { + "uid": entity.uid, + "role": entity.role, + "name": str(raw.get("name", entity.name)).strip(), + "category": str(raw_category).strip() or entity.category, + "description": entity.description, + "affordances": sorted(entity.affordances), + "attributes": attributes, + "initial_state": _redact_semantic_mapping(entity.initial_state), + "side": side, + "rank": rank_by_uid[entity.uid], + } + ) + return payload + + +def _grounding_prompt( + instruction: str, + requests: Sequence[Mapping[str, Any]], + inventory: Sequence[Mapping[str, Any]], +) -> str: + return ( + "Ground the requested natural-language scene references to the supplied " + "scene inventory. Resolve all requests together using the original task, " + "step type, relation, quantifier, and reference text as context. Select " + "only exact inventory UIDs. The inventory's affordances and states are " + "source evidence only: never infer, add, authorize, or return an " + "affordance, capability, physical state, coordinate, pose, orientation, " + "path, or action. The side and rank fields are discrete robot-relative " + "labels; rank 1 is leftmost. Object requests may select only movable " + "inventory entities. Target requests may also select support surfaces. " + "Use status=ambiguous or status=not_found instead of guessing when the " + "evidence is insufficient. Return exactly one binding per reference_id " + "with only reference_id, status, uids, and confidence.\n\n" + f"Instruction:\n{instruction}\n\n" + "Grounding requests:\n" + f"{json.dumps(list(requests), ensure_ascii=False, sort_keys=True)}\n\n" + "Redacted scene inventory:\n" + f"{json.dumps(list(inventory), ensure_ascii=False, sort_keys=True)}" + ) + + +def _validate_response( + value: Mapping[str, Any], + *, + requests: Sequence[Mapping[str, Any]], + inventory: SceneInventory, +) -> dict[str, tuple[str, ...]]: + if not isinstance(value, Mapping): + raise TypeError("Scene grounding output must be a mapping.") + if set(value) != {"bindings"}: + raise ValueError( + "Scene grounding output must contain exactly the 'bindings' field." + ) + raw_bindings = value["bindings"] + if not isinstance(raw_bindings, Sequence) or isinstance(raw_bindings, (str, bytes)): + raise ValueError("Scene grounding bindings must be a list.") + + request_by_id = {str(request["reference_id"]): request for request in requests} + bindings: dict[str, tuple[str, ...]] = {} + for binding_index, raw in enumerate(raw_bindings): + context = f"SceneGrounding.bindings[{binding_index}]" + if not isinstance(raw, Mapping): + raise ValueError(f"{context} must be a mapping.") + if set(raw) != _BINDING_KEYS: + missing = sorted(_BINDING_KEYS - set(raw)) + extra = sorted(set(raw) - _BINDING_KEYS) + raise ValueError( + f"{context} fields must be exactly {sorted(_BINDING_KEYS)}; " + f"missing={missing}, unsupported={extra}." + ) + reference_id = raw["reference_id"] + if not isinstance(reference_id, str) or not reference_id: + raise ValueError(f"{context}.reference_id must be a non-empty string.") + if reference_id not in request_by_id: + raise ValueError(f"{context} references unknown request {reference_id!r}.") + if reference_id in bindings: + raise ValueError(f"Duplicate grounding binding for {reference_id!r}.") + + status = raw["status"] + if status not in {"resolved", "ambiguous", "not_found"}: + raise ValueError( + f"{context}.status must be resolved, ambiguous, or not_found." + ) + if status != "resolved": + raise ValueError( + f"Grounding request {reference_id!r} was not resolved: {status}." + ) + confidence = raw["confidence"] + if ( + not isinstance(confidence, (int, float)) + or isinstance(confidence, bool) + or not math.isfinite(float(confidence)) + or not 0.0 <= float(confidence) <= 1.0 + ): + raise ValueError(f"{context}.confidence must be a number in [0, 1].") + if float(confidence) < 0.5: + raise ValueError( + f"Grounding request {reference_id!r} confidence is below 0.5." + ) + + raw_uids = raw["uids"] + if not isinstance(raw_uids, Sequence) or isinstance(raw_uids, (str, bytes)): + raise ValueError(f"{context}.uids must be a list.") + uids = tuple(raw_uids) + if any(not isinstance(uid, str) or not uid for uid in uids): + raise ValueError(f"{context}.uids must contain non-empty strings.") + if len(set(uids)) != len(uids): + raise ValueError( + f"Grounding request {reference_id!r} contains duplicate UIDs." + ) + unknown = sorted(set(uids) - set(inventory.by_uid)) + if unknown: + raise ValueError( + f"Grounding request {reference_id!r} selected unknown UIDs {unknown}." + ) + + request = request_by_id[reference_id] + allowed = ( + {entity.uid for entity in inventory.interactive} + if request["slot"] == "object" + else {entity.uid for entity in (*inventory.interactive, *inventory.support)} + ) + disallowed = sorted(set(uids) - allowed) + if disallowed: + raise ValueError( + f"Grounding request {reference_id!r} selected UIDs outside its " + f"{request['slot']} candidate range: {disallowed}." + ) + _validate_cardinality(request, uids) + bindings[reference_id] = uids + + missing = sorted(set(request_by_id) - set(bindings)) + if missing: + raise ValueError(f"Scene grounding omitted requests {missing}.") + _reject_self_references(requests, bindings) + return bindings + + +def _validate_cardinality( + request: Mapping[str, Any], + uids: Sequence[str], +) -> None: + request_id = str(request["reference_id"]) + quantifier = str(request["quantifier"]) + if quantifier == "one" and len(uids) != 1: + raise ValueError( + f"Grounding request {request_id!r} quantifier=one requires exactly one UID." + ) + if quantifier == "count" and len(uids) != int(request["count"]): + raise ValueError( + f"Grounding request {request_id!r} requires exactly " + f"{request['count']} UIDs." + ) + if quantifier == "all" and not uids: + raise ValueError( + f"Grounding request {request_id!r} quantifier=all requires at " + "least one UID." + ) + + +def _reject_self_references( + requests: Sequence[Mapping[str, Any]], + bindings: Mapping[str, tuple[str, ...]], +) -> None: + slots_by_step: dict[str, dict[str, str]] = {} + for request in requests: + slots_by_step.setdefault(str(request["step_id"]), {})[str(request["slot"])] = ( + str(request["reference_id"]) + ) + for step_id, slots in slots_by_step.items(): + object_id = slots.get("object") + target_id = slots.get("target") + if object_id is None or target_id is None: + continue + overlap = sorted(set(bindings[object_id]) & set(bindings[target_id])) + if overlap: + raise ValueError( + f"Grounding step {step_id!r} uses the same UID as object and " + f"target: {overlap}." + ) + + +def _redact_semantic_mapping(value: Mapping[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + name = str(key) + normalized = name.strip().lower().replace("-", "_") + if normalized in _REDACTED_KEYS: + continue + if isinstance(child, Mapping): + nested = _redact_semantic_mapping(child) + if nested: + result[name] = nested + elif isinstance(child, (str, int, float, bool)) and not isinstance( + child, complex + ): + result[name] = child + elif isinstance(child, Sequence) and not isinstance(child, (str, bytes)): + semantic_values = [item for item in child if isinstance(item, (str, bool))] + if semantic_values and len(semantic_values) == len(child): + result[name] = semantic_values + return result diff --git a/embodichain/gen_sim/action_engine/tasks/interpretation.py b/embodichain/gen_sim/action_engine/tasks/interpretation.py new file mode 100644 index 000000000..f84f6d2bd --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/interpretation.py @@ -0,0 +1,404 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Compatibility bridge from Task Engine drafts to Action Engine TaskSpec v2.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from embodichain.gen_sim.task_engine.interpretation import ( + INSTRUCTION_INTENT_SCHEMA, + InstructionCaller, + InstructionDraftResult, + InstructionIntent, + _default_instruction_caller, + _instruction_prompt, + _instruction_selector_rules, + interpret_instruction_draft, + validate_instruction_intent, +) + +from .assembly import ( + GroundedTaskBuilder, + GroundedTaskSpec, + SceneEntity, + SceneInventory, + validate_source_compatibility, + validate_target_compatibility, +) +from .grounding import GroundingCaller, ground_scene_references + +__all__ = [ + "GroundingCaller", + "INSTRUCTION_INTENT_SCHEMA", + "InstructionCaller", + "InstructionDraftResult", + "InstructionIntent", + "ground_instruction_draft", + "interpret_and_ground_task_spec", + "interpret_instruction_draft", + "validate_instruction_intent", +] + + +def interpret_and_ground_task_spec( + task_name: str, + task_description: str, + scene_objects: Sequence[Mapping[str, Any]], + *, + robot_profile: str, + model: str | None = None, + caller: InstructionCaller | None = None, + grounding_caller: GroundingCaller | None = None, +) -> GroundedTaskSpec: + """Interpret through Task Engine, then ground through Action Engine.""" + task_id = str(task_name).strip() + instruction = str(task_description).strip() + if not task_id or not instruction: + raise ValueError("task_name and task_description must be non-empty.") + inventory = SceneInventory(scene_objects, robot_profile=robot_profile) + draft = interpret_instruction_draft(instruction, model=model, caller=caller) + invoke = caller or _default_instruction_caller + selected_model = None if draft.model == "injected_caller" else draft.model + grounding = ground_scene_references( + instruction=instruction, + intent=draft.intent, + inventory=inventory, + scene_objects=scene_objects, + model=selected_model, + caller=grounding_caller or invoke, + ) + grounded = _ground_intent( + task_id, + instruction, + draft.intent, + inventory, + grounding.bindings, + ) + grounded.task_spec["metadata"].update( + { + "instruction_interpreter": "structured_llm_v2", + "instruction_model": draft.model, + "instruction_call_count": draft.attempts, + "instruction_latency_seconds": draft.latency_seconds, + "scene_grounding_model": selected_model or "injected_caller", + "scene_grounding_call_count": grounding.attempts, + "scene_grounding_latency_seconds": grounding.latency_seconds, + } + ) + if draft.normalizations: + grounded.task_spec["metadata"]["instruction_intent_normalizations"] = list( + draft.normalizations + ) + return grounded + + +def ground_instruction_draft( + task_id: str, + instruction: str, + intent: Mapping[str, Any], + scene_objects: Sequence[Mapping[str, Any]], + *, + robot_profile: str, + reference_bindings: Mapping[str, Sequence[str]], +) -> GroundedTaskSpec: + """Lower a Task Engine draft using verified scene bindings.""" + normalized_task_id = str(task_id).strip() + normalized_instruction = str(instruction).strip() + if not normalized_task_id or not normalized_instruction: + raise ValueError("task_id and instruction must be non-empty.") + inventory = SceneInventory(scene_objects, robot_profile=robot_profile) + return _ground_intent( + normalized_task_id, + normalized_instruction, + validate_instruction_intent(intent), + inventory, + reference_bindings, + ) + + +def _ground_intent( + task_id: str, + instruction: str, + intent: Mapping[str, Any], + inventory: SceneInventory, + scene_bindings: Mapping[str, Sequence[str]], +) -> GroundedTaskSpec: + builder = GroundedTaskBuilder( + task_id, + instruction, + inventory, + planner="structured_llm_v2", + ) + objects_by_step: dict[str, list[SceneEntity]] = {} + task_ids_by_step: dict[str, list[str]] = {} + for step in _topological_steps(intent["steps"]): + step_id = str(step["id"]) + objects = _resolve_reference( + step["object"], + inventory, + objects_by_step, + context=f"instruction step {step_id!r} object", + reference_id=f"{step_id}.object", + scene_bindings=scene_bindings, + ) + validate_source_compatibility(str(step["task_type"]), objects) + target_objects = _resolve_reference( + step["target"], + inventory, + objects_by_step, + context=f"instruction step {step_id!r} target", + reference_id=f"{step_id}.target", + scene_bindings=scene_bindings, + allow_none=True, + exclude={item.uid for item in objects}, + allow_support=True, + ) + if len(target_objects) > 1: + raise ValueError(f"Instruction step {step_id!r} target is ambiguous.") + validate_target_compatibility( + str(step["task_type"]), + target_objects[0] if target_objects else None, + relation=str(step["relation"]), + ) + dependencies_by_step = list(step["depends_on"]) + for selector in (step["object"], step["target"]): + if selector["kind"] == "step_result": + reference = str(selector["step_id"]) + if reference not in dependencies_by_step: + dependencies_by_step.append(reference) + dependencies = [ + emitted_id + for dependency in dependencies_by_step + for emitted_id in task_ids_by_step[str(dependency)] + ] + emitted = _emit_step( + builder, + step, + objects, + target_objects[0] if target_objects else None, + dependencies, + ) + objects_by_step[step_id] = objects + task_ids_by_step[step_id] = emitted + return builder.build() + + +def _emit_step( + builder: GroundedTaskBuilder, + step: Mapping[str, Any], + objects: Sequence[SceneEntity], + target: SceneEntity | None, + dependencies: Sequence[str], +) -> list[str]: + task_type = str(step["task_type"]) + if step["layout"] == "line": + roles = [builder._role(entity, "E1") for entity in objects] + parent = str(step["id"]) + return [ + builder.add( + "E1", + entity, + params={ + "target_role": "table", + "relation": "on", + "layout": "line", + "objects_roles": roles, + "axis": "world_y" if step["axis"] == "none" else step["axis"], + "order_by": "explicit", + "order_direction": "given", + "order_constraint": "free", + "orientation_goal": step["orientation_goal"], + "orientation_axis": "none", + "nominal_slot_index": slot, + "slot_constraint": "free_reassignable", + "parent_task_instance_id": parent, + }, + depends_on=dependencies, + ) + for slot, entity in enumerate(objects) + ] + + emitted = [] + for entity in objects: + params: dict[str, Any] = {} + required_arm = str(step["required_arm"]) + if required_arm in {"left_arm", "right_arm"}: + params["required_arm"] = required_arm + if task_type == "E1": + relation = str(step["relation"]) + if relation == "none": + if target is None or target not in builder.inventory.support: + raise ValueError( + "E1 omitted relation is only valid for a unique table " + "support target." + ) + relation = "on" + params.update( + { + "relation": relation, + "relation_frame": "robot", + "orientation_goal": step["orientation_goal"], + "orientation_axis": "none", + } + ) + elif task_type == "E2": + params.update( + { + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + ) + elif task_type == "E3": + params.update({"relation": "above", "relation_frame": "robot"}) + elif task_type == "E4": + params.update( + { + "transfer_arm": step["transfer_arm"], + "receive_arm": step["receive_arm"], + "orientation_goal": step["orientation_goal"], + } + ) + elif task_type == "E5": + params.update( + { + "direction": step["direction"], + "terminal_behavior": step["terminal_behavior"], + "relation": step["relation"], + "relation_frame": "robot", + } + ) + elif task_type in {"E6", "E7"}: + params["target_state"] = step["target_state"] + elif task_type == "E8": + params["target_setting"] = int(step["target_setting"]) + elif task_type == "E9": + params["terminal_state"] = step["target_state"] + emitted.append( + builder.add( + task_type, + entity, + target=target, + params=params, + depends_on=dependencies, + ) + ) + return emitted + + +def _resolve_reference( + selector: Mapping[str, Any], + inventory: SceneInventory, + objects_by_step: Mapping[str, Sequence[SceneEntity]], + *, + context: str, + reference_id: str, + scene_bindings: Mapping[str, Sequence[str]], + allow_none: bool = False, + exclude: set[str] | None = None, + allow_support: bool = False, +) -> list[SceneEntity]: + kind = str(selector["kind"]) + if kind == "none": + if allow_none: + return [] + raise ValueError(f"{context} is required.") + if kind == "step_result": + step_id = str(selector["step_id"]) + if step_id not in objects_by_step: + raise ValueError(f"{context} references unavailable step {step_id!r}.") + objects = list(objects_by_step[step_id]) + if len(objects) != 1: + raise ValueError( + f"{context} references step {step_id!r}, which has {len(objects)} objects." + ) + if exclude and objects[0].uid in exclude: + raise ValueError( + f"{context} references the same object as its source; " + "self-referential placement is not allowed." + ) + return objects + + if reference_id not in scene_bindings: + raise ValueError(f"{context} has no verified scene-grounding binding.") + excluded = exclude or set() + source_uids = ( + {entity.uid for entity in inventory.entities} + if allow_support + else {entity.uid for entity in inventory.interactive} + ) + resolved_uids = tuple(str(uid) for uid in scene_bindings[reference_id]) + pool = [ + inventory.by_uid[uid] + for uid in resolved_uids + if uid in source_uids and uid not in excluded + ] + pool = sorted(pool, key=lambda item: item.uid) + if not pool: + raise ValueError(f"{context} did not bind an eligible scene object.") + quantifier = str(selector["quantifier"]) + count = int(selector["count"]) + if quantifier == "one" and len(pool) != 1: + raise ValueError( + f"{context} is ambiguous; matched scene UIDs {[item.uid for item in pool]}." + ) + if quantifier == "count" and (count < 1 or len(pool) != count): + raise ValueError( + f"{context} requested exactly {count} objects but matched {len(pool)}." + ) + if quantifier == "all" and count not in {0, len(pool)}: + raise ValueError( + f"{context} quantifier=all cannot carry count={count}; use count for an exact quantity." + ) + return pool + + +def _topological_steps( + steps: Sequence[Mapping[str, Any]], +) -> list[dict[str, Any]]: + by_id = {str(step["id"]): dict(step) for step in steps} + effective_dependencies: dict[str, tuple[str, ...]] = {} + for step_id, step in by_id.items(): + deps = [str(dep) for dep in step["depends_on"]] + for selector in (step["object"], step["target"]): + if selector["kind"] == "step_result": + reference = str(selector["step_id"]) + if reference not in deps: + deps.append(reference) + effective_dependencies[step_id] = tuple(deps) + pending = set(by_id) + ordered: list[dict[str, Any]] = [] + original = [str(step["id"]) for step in steps] + while pending: + ready = [ + step_id + for step_id in original + if step_id in pending + and all(str(dep) not in pending for dep in effective_dependencies[step_id]) + ] + if not ready: + raise ValueError("Instruction intent dependencies contain a cycle.") + # Select one earliest-ready step at a time. Emitting the whole ready + # frontier lets a later independent step leapfrog an earlier step that + # becomes ready after its predecessor, changing the instruction's + # resource-order tie break without any causal reason. + step_id = ready[0] + ordered.append(by_id[step_id]) + pending.remove(step_id) + return ordered diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py new file mode 100644 index 000000000..9beda20b6 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -0,0 +1,1061 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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"], + } + if incoming_held_arm is None and terminal_behavior == "place": + return ( + [ + _node( + group_id, + 1, + "AxisAlign", + task_type, + object_uid, + actor, + "arm", + {"kind": "object", "object": object_uid}, + dependencies, + role, + success, + motion_policy(), + ) + ], + "orient_object", + goal, + success, + ) + return ( + _single_arm_manipulation( + group_id, + task_type, + object_uid, + actor, + dependencies, + role=role, + already_held=incoming_held_arm is not None, + leave_held=terminal_behavior == "hold", + ), + "orient_object", + goal, + success, + ) + if task_type == "E3": + target = str(params["target_role"]) + contents = [{"object": str(uid)} for uid in params.get("content_roles", [])] + goal = { + "reference_object": target, + "relation": "above", + "amount": "task_defined", + "contents": deepcopy(contents), + } + success = { + "type": "poured", + "object": object_uid, + "reference_object": target, + "contents": deepcopy(contents), + } + specs = [] + if incoming_held_arm is None: + specs.append( + ( + "PickUp", + { + "kind": "object", + "object": object_uid, + "payloads": deepcopy(contents), + }, + ) + ) + specs.extend( + ( + ( + "MoveHeldObject", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "final", + "payloads": deepcopy(contents), + }, + ), + ( + "Pour", + { + "kind": "pour_goal", + "object": object_uid, + "reference_object": target, + "payloads": deepcopy(contents), + }, + ), + ) + ) + nodes = [] + previous = list(dependencies) + for index, (action, binding) in enumerate(specs, start=1): + node = _node( + group_id, + index, + action, + task_type, + object_uid, + actor, + "arm", + binding, + previous, + role, + success if action == "Pour" else {}, + motion_policy(), + ) + nodes.append(node) + previous = [node["id"]] + return ( + nodes, + "pour", + goal, + success, + ) + 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} + if task_type == "E8": + success["target_setting"] = int(params["target_setting"]) + else: + success["target_state"] = params["target_state"] + 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", "E3", "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 == "E3": + arm = str(params.get("required_arm", "")) + if arm in {"left_arm", "right_arm"}: + return object_uid, arm + if task_type == "E2" and str(params.get("terminal_behavior", "place")) == "hold": + arm = str(params.get("required_arm", "")) + if arm in {"left_arm", "right_arm"}: + return object_uid, arm + if task_type == "E5" and str(params.get("terminal_behavior", "hold")) == "hold": + return object_uid, "coordinated" + return None + + +def _validate_bindings( + task: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> dict[str, str]: + bindings = dict(role_bindings) + for role, uid in bindings.items(): + if not isinstance(role, str) or not role or not isinstance(uid, str) or not uid: + raise ValueError("role_bindings must map non-empty role IDs to scene UIDs.") + required = set() + for instance in task["task_instances"]: + required.update(_role_references(instance["params"])) + required.discard("table") + missing = sorted(required - set(bindings)) + if missing: + raise ValueError(f"Scene hand-off is missing role bindings: {missing}.") + if len(bindings.values()) != len(set(bindings.values())): + raise ValueError("Scene role bindings must resolve to unique object UIDs.") + return bindings + + +def _role_references(value: Any, key: str = "") -> set[str]: + if isinstance(value, Mapping): + return { + role + for child_key, child in value.items() + for role in _role_references(child, str(child_key)) + } + if isinstance(value, list): + return {role for child in value for role in _role_references(child, key)} + if isinstance(value, str) and (key.endswith("_role") or key.endswith("_roles")): + return {value} + return set() + + +def _resolve_params(value: Any, bindings: Mapping[str, str], key: str = "") -> Any: + if isinstance(value, Mapping): + return { + child_key: _resolve_params(child, bindings, str(child_key)) + for child_key, child in value.items() + } + if isinstance(value, list): + return [_resolve_params(child, bindings, key) for child in value] + if isinstance(value, str) and (key.endswith("_role") or key.endswith("_roles")): + return bindings.get(value, value) + return deepcopy(value) diff --git a/embodichain/gen_sim/action_engine/tasks/scene.py b/embodichain/gen_sim/action_engine/tasks/scene.py new file mode 100644 index 000000000..242b15dd4 --- /dev/null +++ b/embodichain/gen_sim/action_engine/tasks/scene.py @@ -0,0 +1,177 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Validate a Scene Engine result against task-first requirements.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +from embodichain.gen_sim.action_engine.domain import validate_scene_requirements + +__all__ = ["SceneHandoff", "validate_scene_handoff"] + + +@dataclass(frozen=True) +class SceneHandoff: + """Validated role-to-UID resolution returned by an external Scene Engine.""" + + task_id: str + role_bindings: dict[str, Any] + object_uids: tuple[str, ...] + camera_uids: tuple[str, ...] + + +def validate_scene_handoff( + requirements: Mapping[str, Any], + scene: Mapping[str, Any], + role_bindings: Mapping[str, str], +) -> SceneHandoff: + """Reject scenes that do not satisfy roles, affordances, state, or cameras.""" + required = validate_scene_requirements(requirements) + objects = scene.get("objects") + if not isinstance(objects, Sequence) or isinstance(objects, (str, bytes)): + raise ValueError("Scene hand-off requires an objects list.") + object_by_uid: dict[str, Mapping[str, Any]] = {} + for index, item in enumerate(objects): + if not isinstance(item, Mapping): + raise ValueError(f"Scene objects[{index}] must be a mapping.") + uid = item.get("uid") + if not isinstance(uid, str) or not uid: + raise ValueError(f"Scene objects[{index}] requires a UID.") + if uid in object_by_uid: + raise ValueError(f"Scene contains duplicate object UID {uid!r}.") + object_by_uid[uid] = item + + bindings = dict(role_bindings) + required_roles = {item["role_id"] for item in required["objects"]} + if set(bindings) != required_roles: + missing = sorted(required_roles - set(bindings)) + extra = sorted(set(bindings) - required_roles) + raise ValueError( + f"Scene role bindings mismatch; missing={missing}, extra={extra}." + ) + normalized_bindings: dict[str, str | tuple[str, ...]] = {} + assigned_uids: list[str] = [] + for requirement in required["objects"]: + role = requirement["role_id"] + count = int(requirement["count"]) + binding = bindings[role] + if isinstance(binding, str): + uids = [binding] + elif isinstance(binding, Sequence) and not isinstance(binding, (str, bytes)): + uids = [str(uid) for uid in binding] + else: + raise ValueError(f"Scene role {role!r} has an invalid UID binding.") + if len(uids) != count or any(not uid for uid in uids): + raise ValueError( + f"Scene role {role!r} requires exactly {count} UID binding(s)." + ) + normalized_bindings[role] = uids[0] if count == 1 else tuple(uids) + assigned_uids.extend(uids) + for uid in uids: + _validate_bound_object(object_by_uid, uid, role, requirement) + if len(assigned_uids) != len(set(assigned_uids)): + raise ValueError("Each scene requirement role must resolve to unique UIDs.") + + cameras = scene.get("cameras", []) + if not isinstance(cameras, Sequence) or isinstance(cameras, (str, bytes)): + raise ValueError("Scene cameras must be a list.") + camera_uids = [] + normalized_cameras = [] + for camera in cameras: + if not isinstance(camera, Mapping) or not isinstance(camera.get("uid"), str): + raise ValueError("Every scene camera requires a UID.") + camera_uids.append(str(camera["uid"])) + normalized_cameras.append(camera) + for camera_requirement in required["cameras"]: + modalities = set(camera_requirement.get("modalities", ())) + coverage = camera_requirement.get("coverage") + if not any( + modalities <= set(camera.get("modalities", ())) + and (coverage is None or camera.get("coverage") == coverage) + for camera in normalized_cameras + ): + raise ValueError( + "Scene cameras do not satisfy requirement " + f"{dict(camera_requirement)!r}." + ) + reported_constraints = scene.get("satisfied_spatial_constraints", []) + if not isinstance(reported_constraints, Sequence) or isinstance( + reported_constraints, (str, bytes) + ): + raise ValueError("Scene satisfied_spatial_constraints must be a list.") + reported = {_canonical(item) for item in reported_constraints} + missing_constraints = [ + constraint + for constraint in required["spatial_constraints"] + if _canonical(constraint) not in reported + ] + if missing_constraints: + raise ValueError( + "Scene does not satisfy spatial constraints: " f"{missing_constraints}." + ) + return SceneHandoff( + task_id=required["task_id"], + role_bindings=normalized_bindings, + object_uids=tuple(sorted(object_by_uid)), + camera_uids=tuple(sorted(camera_uids)), + ) + + +def _validate_bound_object( + object_by_uid: Mapping[str, Mapping[str, Any]], + uid: str, + role: str, + requirement: Mapping[str, Any], +) -> None: + if uid not in object_by_uid: + raise ValueError(f"Scene role {role!r} references unknown UID {uid!r}.") + actual = object_by_uid[uid] + if actual.get("category") != requirement["category"]: + raise ValueError( + f"Scene object {uid!r} category does not satisfy role {role!r}." + ) + missing_affordances = set(requirement["affordances"]) - set( + actual.get("affordances", ()) + ) + if missing_affordances: + raise ValueError( + f"Scene object {uid!r} lacks affordances {sorted(missing_affordances)}." + ) + for field in ("initial_state", "attributes"): + actual_values = actual.get(field, {}) + if not isinstance(actual_values, Mapping): + raise ValueError(f"Scene object {uid!r} {field} must be a mapping.") + mismatched = { + key: expected + for key, expected in requirement[field].items() + if actual_values.get(key) != expected + } + if mismatched: + raise ValueError( + f"Scene object {uid!r} does not satisfy {field} {mismatched}." + ) + + +def _canonical(value: Any) -> str: + import json + + if not isinstance(value, Mapping): + raise ValueError("Every satisfied spatial constraint must be a mapping.") + return json.dumps(dict(value), sort_keys=True, separators=(",", ":")) diff --git a/embodichain/gen_sim/action_engine/unbound.py b/embodichain/gen_sim/action_engine/unbound.py new file mode 100644 index 000000000..1b0a4af77 --- /dev/null +++ b/embodichain/gen_sim/action_engine/unbound.py @@ -0,0 +1,205 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scene-independent Action Engine draft produced before final UID binding.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +from typing import Any, Final, TypeAlias + +from embodichain.gen_sim.action_engine.domain.task_contracts import TASK_CONTRACTS + +__all__ = [ + "ActionCapabilityError", + "UNBOUND_ACTION_PLAN_SCHEMA", + "UnboundActionPlan", + "build_unbound_action_plan", + "validate_unbound_action_plan", +] + +UNBOUND_ACTION_PLAN_SCHEMA: Final = "embodichain.unbound-action-plan/v1" +UnboundActionPlan: TypeAlias = dict[str, Any] + + +class ActionCapabilityError(ValueError): + """A required AtomicAction is missing or not executable.""" + + +_PLAN_KEYS = frozenset( + { + "schema_version", + "task_id", + "candidate_id", + "instruction", + "steps", + "required_actions", + } +) +_STEP_KEYS = frozenset( + {"step_id", "task_type", "object", "target", "depends_on", "actions"} +) + + +def build_unbound_action_plan(candidate: Mapping[str, Any]) -> UnboundActionPlan: + """Lower a TaskCandidate into an Action-owned plan without scene UIDs. + + Args: + candidate: Validated Task Engine candidate or an equivalent mapping. + + Returns: + A strict JSON plan whose selectors remain logical references. + + Raises: + TypeError: If the candidate or draft is not a mapping. + ValueError: If the draft references an unsupported task type. + """ + value = _mapping(candidate, "candidate") + draft = _mapping(value.get("draft"), "candidate.draft") + task_id = _nonempty(draft.get("task_id"), "candidate.draft.task_id") + instruction = _nonempty(draft.get("instruction"), "candidate.draft.instruction") + candidate_id = _nonempty(value.get("candidate_id"), "candidate.candidate_id") + steps = [] + required_actions: set[str] = set() + for index, raw in enumerate(_sequence(draft.get("steps"), "candidate.draft.steps")): + step = _mapping(raw, f"candidate.draft.steps[{index}]") + task_type = _nonempty( + step.get("task_type"), f"candidate.draft.steps[{index}].task_type" + ) + contract = TASK_CONTRACTS.get(task_type) + if contract is None: + raise ValueError(f"Action Engine does not support task type {task_type!r}.") + actions = [str(name) for name in contract.core_actions] + required_actions.update(actions) + steps.append( + { + "step_id": _nonempty( + step.get("id"), f"candidate.draft.steps[{index}].id" + ), + "task_type": task_type, + "object": deepcopy(step.get("object")), + "target": deepcopy(step.get("target")), + "depends_on": deepcopy(step.get("depends_on", [])), + "actions": actions, + } + ) + return validate_unbound_action_plan( + { + "schema_version": UNBOUND_ACTION_PLAN_SCHEMA, + "task_id": task_id, + "candidate_id": candidate_id, + "instruction": instruction, + "steps": steps, + "required_actions": sorted(required_actions), + } + ) + + +def validate_unbound_action_plan( + value: Mapping[str, Any], +) -> UnboundActionPlan: + """Validate and detach one scene-independent Action plan. + + Args: + value: Candidate plan mapping. + + Returns: + A strict JSON-safe detached plan. + + Raises: + TypeError: If a mapping or sequence field has the wrong type. + ValueError: If the schema, dependency graph, or actions are invalid. + """ + result = _mapping(value, "UnboundActionPlan") + if set(result) != _PLAN_KEYS: + raise ValueError("UnboundActionPlan fields are invalid.") + if result.get("schema_version") != UNBOUND_ACTION_PLAN_SCHEMA: + raise ValueError("UnboundActionPlan.schema_version is invalid.") + for key in ("task_id", "candidate_id", "instruction"): + result[key] = _nonempty(result.get(key), f"UnboundActionPlan.{key}") + + steps = [] + seen: set[str] = set() + actions_used: set[str] = set() + for index, raw in enumerate(_sequence(result.get("steps"), "steps")): + context = f"UnboundActionPlan.steps[{index}]" + step = _mapping(raw, context) + if set(step) != _STEP_KEYS: + raise ValueError(f"{context} fields are invalid.") + step_id = _nonempty(step.get("step_id"), f"{context}.step_id") + if step_id in seen: + raise ValueError("UnboundActionPlan step IDs must be unique.") + task_type = _nonempty(step.get("task_type"), f"{context}.task_type") + contract = TASK_CONTRACTS.get(task_type) + if contract is None: + raise ValueError(f"{context}.task_type is unsupported.") + dependencies = _strings(step.get("depends_on"), f"{context}.depends_on") + if any(dependency not in seen for dependency in dependencies): + raise ValueError( + f"{context}.depends_on must reference preceding unbound steps." + ) + actions = _strings(step.get("actions"), f"{context}.actions") + if actions != [str(name) for name in contract.core_actions]: + raise ValueError(f"{context}.actions do not match the task contract.") + for selector_name in ("object", "target"): + if not isinstance(step.get(selector_name), Mapping): + raise TypeError(f"{context}.{selector_name} must be a mapping.") + step[selector_name] = deepcopy(dict(step[selector_name])) + step["step_id"] = step_id + step["task_type"] = task_type + step["depends_on"] = dependencies + step["actions"] = actions + steps.append(step) + seen.add(step_id) + actions_used.update(actions) + if not steps: + raise ValueError("UnboundActionPlan.steps must not be empty.") + required = _strings(result.get("required_actions"), "required_actions") + if required != sorted(actions_used): + raise ValueError("UnboundActionPlan.required_actions is not canonical.") + result["steps"] = steps + result["required_actions"] = required + json.dumps(result, ensure_ascii=False, allow_nan=False) + return result + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"{context} must be a sequence.") + return list(value) + + +def _strings(value: Any, context: str) -> list[str]: + result = _sequence(value, context) + if any(not isinstance(item, str) or not item for item in result): + raise ValueError(f"{context} must contain non-empty strings.") + if len(set(result)) != len(result): + raise ValueError(f"{context} must not contain duplicates.") + return list(result) + + +def _nonempty(value: Any, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{context} must be a non-empty string.") + return value.strip() diff --git a/tests/gen_sim/__init__.py b/tests/gen_sim/__init__.py index 355d915ff..cdeead7b0 100644 --- a/tests/gen_sim/__init__.py +++ b/tests/gen_sim/__init__.py @@ -14,4 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +"""Generative simulation tests.""" + from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/gen_sim/action_engine/__init__.py b/tests/gen_sim/action_engine/__init__.py new file mode 100644 index 000000000..e2bb4c0aa --- /dev/null +++ b/tests/gen_sim/action_engine/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine tests.""" diff --git a/tests/gen_sim/action_engine/capabilities/__init__.py b/tests/gen_sim/action_engine/capabilities/__init__.py new file mode 100644 index 000000000..046cb429b --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine capability tests.""" diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py new file mode 100644 index 000000000..07245b274 --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -0,0 +1,233 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from types import SimpleNamespace + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapability, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.runtime.loader import load_execution_program +from embodichain.gen_sim.action_engine.runtime.models import GroundedAction +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_spec +from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionOptions, + ActionPlan, + EndEffectorPoseGoal, + PlannerDiagnostics, + RuntimeCommandFrame, + TimedTrajectory, + TimedCommandSequence, +) + + +@dataclass(frozen=True, slots=True) +class _TestOptions(ActionOptions): + marker: str = "test" + + +class _TestAction: + skill_id = "test_retreat" + end_effector_roles: tuple[str, ...] = () + + +class _TestEngine: + binding_owner_id = "test-engine" + + def bind_control_parts(self, _skill_id, _endpoints): + return ActionBinding(owner_id=self.binding_owner_id) + + def plan(self, invocation, context): + assert isinstance(invocation.skill_options, _TestOptions) + positions = context.robot.qpos[:, None, :] + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=context.env_ids, + step_dt=context.require_control_dt(), + ) + return ActionPlan( + skill_id=invocation.skill_id, + plan_success=torch.ones(context.batch_size, dtype=torch.bool), + commands=TimedCommandSequence( + frames=( + RuntimeCommandFrame( + commands=(), + active_mask=torch.ones( + context.batch_size, + dtype=torch.bool, + ), + env_ids=context.env_ids, + hold_duration=trajectory.dt[:, 0], + ), + ), + env_ids=context.env_ids, + ), + joint_trajectory=trajectory, + recovery_policy=invocation.recovery_policy, + planned_scene_version=context.scene.version, + planned_collision_world_revision=(0,) * context.batch_size, + diagnostics=PlannerDiagnostics(backend="test"), + ) + + +class _Robot: + dof = 2 + uid = "test_robot" + control_parts = {"left_arm": [0], "right_arm": [1]} + + def get_qpos(self): + return torch.zeros((1, 2)) + + def get_joint_ids(self, *, name: str): + return self.control_parts.get(name, []) + + +class _Entity: + def get_local_pose(self, *, to_matrix: bool): + assert to_matrix + return torch.eye(4).unsqueeze(0) + + +class _Sim: + def get_rigid_object(self, _uid: str): + return _Entity() + + +def test_new_descriptor_reuses_loader_and_adapter_without_dispatch_changes() -> None: + registry = build_atomic_capability_registry() + calls = [] + + def target_hook(**kwargs): + calls.append("target") + pose = kwargs["object_pose"].clone() + return GroundedAction( + action_class="TestRetreat", + arm=kwargs["arm"], + control="arm", + target=EndEffectorPoseGoal(xpos=pose), + cfg=kwargs["policy"], + object_pose=pose, + target_object_pose=pose, + motion_policy=kwargs["policy"], + ) + + def config_hook(**_kwargs): + calls.append("config") + return _TestOptions() + + registry.register( + AtomicCapability( + "TestRetreat", + _TestAction, + _TestOptions, + frozenset({"policy_pose"}), + frozenset({"arm"}), + "single_arm", + "preserve", + "eef_pose", + motion_base="MoveEndEffector", + target_materializer_hook=target_hook, + config_materializer_hook=config_hook, + contract_resolver_hook=registry.get( + "MoveEndEffector" + ).contract_resolver_hook, + ) + ) + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + graph = instantiate_seed_graph(task, bindings) + graph = deepcopy(graph) + graph["capability_catalog_hash"] = registry.catalog_hash() + cleanup = next( + node for node in graph["nodes"] if node["atomic_action"] == "MoveEndEffector" + ) + cleanup["atomic_action"] = "TestRetreat" + cleanup.pop("contract") + for group in graph["task_groups"]: + group.pop("contract") + graph["metadata"].pop("action_contract_linker") + graph = link_seed_graph(graph, registry=registry) + + program = load_execution_program(graph, registry=registry) + assert any( + action["atomic_action_class"] == "TestRetreat" + for edge in program.edges + for action in edge.actions + ) + + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + robot=_Robot(), + sim=_Sim(), + agent_robot_profile="dual_ur10", + get_agent_arm_control_part=lambda is_left: ( + "left_arm" if is_left else "right_arm" + ), + get_agent_eef_control_part=lambda _is_left: None, + ) + adapter = AtomicActionAdapter( + env, + grasp_policy={}, + capability_registry=registry, + ) + adapter._atomic_engine = _TestEngine() + step = next( + step + for step in program.semantic_steps + if any( + action["atomic_action_class"] == "TestRetreat" + for edge_id in step.edge_ids + for action in next( + edge for edge in program.edges if edge.id == edge_id + ).actions + ) + ) + action = next( + action + for edge in program.edges + for action in edge.actions + if action["atomic_action_class"] == "TestRetreat" + ) + grounder = ActionGrounder( + program, + env, + lambda _uid: None, + capability_registry=registry, + ) + state = ExecutionState(last_qpos=torch.zeros((1, 2))) + grounded = grounder.ground(action, step, arm="left_arm", state=state) + outcome = adapter.plan( + grounded, + state, + ) + assert outcome.success.tolist() == [True] + assert calls == ["target", "config"] diff --git a/tests/gen_sim/action_engine/cli/test_run_agent.py b/tests/gen_sim/action_engine/cli/test_run_agent.py new file mode 100644 index 000000000..53b2dc99b --- /dev/null +++ b/tests/gen_sim/action_engine/cli/test_run_agent.py @@ -0,0 +1,219 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.action_engine.cli.run_agent import ( + _ABWorkerConfig, + _SerializedABBranch, + _capture_ab_initial_frame, + _prepare_ab_branches, + _publish_task_engine_report, + _task_engine_exit_code, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + build_execution_provenance, +) + + +class record_camera_data: + def __init__(self) -> None: + self.calls = [] + + def __call__(self, *args, **kwargs) -> None: + self.calls.append((args, kwargs)) + + +class _FakeEnv: + def __init__(self, recorder=None) -> None: + self.unwrapped = self + self.event_manager = SimpleNamespace( + _mode_functor_cfgs={ + "interval": ( + [ + SimpleNamespace( + func=recorder, + params={"name": "record_cam_audience_view"}, + ) + ] + if recorder is not None + else [] + ) + } + ) + + +def test_capture_ab_initial_frame_invokes_only_audience_recorder() -> None: + recorder = record_camera_data() + env = _FakeEnv(recorder) + + _capture_ab_initial_frame(env) + + assert len(recorder.calls) == 1 + args, kwargs = recorder.calls[0] + assert args == (env, None) + assert kwargs == {"name": "record_cam_audience_view"} + + +def test_capture_ab_initial_frame_requires_audience_recorder() -> None: + with pytest.raises(RuntimeError, match="audience recorder"): + _capture_ab_initial_frame(_FakeEnv()) + + +def _worker_config(route: str) -> _ABWorkerConfig: + return _ABWorkerConfig( + route=route, + gym_config={}, + env_options={}, + gym_id="ActionEngine-v1", + agent_config={}, + agent_config_path="agent_config.json", + task_name="smoke", + runtime_backend="independent", + seed=7, + camera_uids=("vlm_front",), + staging_dir=f"/tmp/ab/{route}/video", + ) + + +class _MemoryAwareFakeWorker: + instances = [] + + def __init__(self, config: _ABWorkerConfig) -> None: + self.config = config + self.closed = False + self.startup_snapshot = { + "robot_qpos": [0.0, 1.0], + "object_poses": {"object": [0.0, 0.0, 0.0]}, + } + self.startup_observation = {"route": config.route} + self.events = [] + self.instances.append(self) + if ( + config.route == "online" + and Path(config.staging_dir).parent.name == config.route + ): + raise RuntimeError("CUDA out of memory") + + def preflight(self, graph): + self.events.append(("preflight", graph)) + return True + + def run(self, graph, **kwargs): + self.events.append(("run", graph, kwargs)) + return SimpleNamespace(success=True) + + def finalize(self, branch_dir: Path, *, episode_index: int): + self.events.append(("finalize", branch_dir, episode_index)) + return [(branch_dir / "video.mp4").as_posix()] + + def close(self): + self.closed = True + + +def test_ab_serializes_workers_after_startup_oom() -> None: + _MemoryAwareFakeWorker.instances = [] + branches, snapshots = _prepare_ab_branches( + {"offline": _worker_config("offline"), "online": _worker_config("online")}, + worker_factory=_MemoryAwareFakeWorker, + prefer_serial=False, + ) + + assert set(branches) == {"offline", "online"} + assert all(isinstance(branch, _SerializedABBranch) for branch in branches.values()) + assert snapshots["offline"] == snapshots["online"] + for route, branch in branches.items(): + assert branch.preflight({"route": route}) is True + branch.run( + {"route": route}, + run_id=f"run-{route}", + episode_index=0, + record_root=Path("/tmp/ab/runtime"), + ) + assert branch.finalize(Path(f"/tmp/ab/{route}"), episode_index=0) == [ + f"/tmp/ab/{route}/video.mp4" + ] + branch.close() + + phases = [ + Path(worker.config.staging_dir).parent.name + for worker in _MemoryAwareFakeWorker.instances + if worker.config.route == "offline" + ] + assert phases == ["offline", "probe", "preflight", "execute"] + + +@pytest.mark.parametrize( + ("status", "success"), + [("succeeded", True), ("failed", False)], +) +def test_task_engine_report_is_mirrored_into_bundle_only_when_enabled( + tmp_path: Path, + status: str, + success: bool, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + agent_config = bundle / "agent_config.json" + report = ExecutionReport( + task_id="task", + plan_hash="0" * 64, + action_graph_hash="1" * 64, + status=status, + run_id="run", + episode_id="0", + provenance=build_execution_provenance(episode_seed=7), + environments=( + { + "env_id": "0", + "success": success, + "semantic_success": {"task_01": success}, + "action_count": 3, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + }, + ), + action_count=3, + record_dir=(tmp_path / "runtime-records").as_posix(), + ) + + assert _publish_task_engine_report(agent_config, report, enabled=False) is None + assert not (bundle / "execution_report.json").exists() + + path = _publish_task_engine_report(agent_config, report, enabled=True) + + assert path == bundle / "execution_report.json" + payload = json.loads(path.read_text(encoding="utf-8")) + assert payload["status"] == status + assert payload["record_dir"] == report.record_dir + + +def test_task_engine_exit_code_uses_report_status() -> None: + success = SimpleNamespace(status="succeeded") + failure = SimpleNamespace(status="failed") + + assert _task_engine_exit_code(False, [success]) == 0 + assert _task_engine_exit_code(False, [success, failure]) == 1 + assert _task_engine_exit_code(True, []) == 1 diff --git a/tests/gen_sim/action_engine/compiler/__init__.py b/tests/gen_sim/action_engine/compiler/__init__.py new file mode 100644 index 000000000..e7977347e --- /dev/null +++ b/tests/gen_sim/action_engine/compiler/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine compiler tests.""" diff --git a/tests/gen_sim/action_engine/compiler/test_compiler.py b/tests/gen_sim/action_engine/compiler/test_compiler.py new file mode 100644 index 000000000..6144fcc3f --- /dev/null +++ b/tests/gen_sim/action_engine/compiler/test_compiler.py @@ -0,0 +1,572 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import pytest + +from embodichain.gen_sim.action_engine.compiler import compile_task_agent +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + TASK_AGENT_SCHEMA, +) + + +def _program(step: Mapping[str, Any]) -> dict[str, Any]: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "operator_demo", + "goal": "Exercise one semantic operator.", + "semantic_steps": [dict(step)], + } + + +def test_place_relative_carries_payloads_through_single_arm_action_bindings() -> None: + execution = compile_task_agent( + _program( + { + "id": "s01_place_carrier", + "operator": "place_relative", + "object": "paper_cup", + "goal": { + "reference_object": "popcorn_bucket", + "relation": "on", + "payloads": [{"object": "glue_stick", "slot": "center"}], + }, + } + ) + ) + + step = execution["semantic_steps"][0] + assert step["goal"]["payloads"] == [{"object": "glue_stick", "slot": "center"}] + carrying_actions = [ + action + for edge in execution["edges"] + for action in edge["actions"] + if action["atomic_action_class"] in {"PickUp", "MoveHeldObject", "Place"} + ] + assert carrying_actions + assert all( + action["target_binding"]["payloads"] == step["goal"]["payloads"] + for action in carrying_actions + ) + + +@pytest.mark.parametrize( + ("step", "expected_action"), + [ + ( + { + "id": "s01_line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "goal": {"axis": "world_y", "anchor": "table_center"}, + }, + "MoveHeldObject", + ), + ( + { + "id": "s01_stack", + "operator": "build_stack", + "objects": ["block_a", "block_b"], + "goal": {"stack_mode": "on_top", "anchor": "table_center"}, + }, + "PickUp", + ), + ( + { + "id": "s01_place", + "operator": "place_relative", + "object": "cup", + "goal": {"reference_object": "tray", "relation": "on"}, + }, + "Place", + ), + ( + { + "id": "s01_hover", + "operator": "hold_hover", + "object": "cup", + "goal": {}, + }, + "MoveJoints", + ), + ( + { + "id": "s01_transport", + "operator": "coordinated_transport", + "object": "tray", + "goal": {"direction": "front", "terminal_behavior": "place"}, + }, + "CoordinatedPickment", + ), + ( + { + "id": "s01_orient", + "operator": "orient_object", + "object": "cup", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + }, + "MoveHeldObject", + ), + ( + { + "id": "s01_press", + "operator": "press", + "object": "button", + "goal": {"terminal_state": "activated"}, + }, + "Press", + ), + ( + { + "id": "s01_coordinated_place", + "operator": "coordinated_place", + "object": "cup", + "goal": {"support_object": "tray", "relation": "on"}, + }, + "CoordinatedPlacement", + ), + ], +) +def test_every_builtin_operator_compiles( + step: Mapping[str, Any], + expected_action: str, +) -> None: + execution = compile_task_agent(_program(step)) + action_classes = { + action["atomic_action_class"] + for edge in execution["edges"] + for action in edge["actions"] + } + + assert execution["schema_version"] == EXECUTION_PROGRAM_SCHEMA + assert expected_action in action_classes + assert execution["nodes"][0]["id"] == execution["start"] + assert execution["goal"] in {node["id"] for node in execution["nodes"]} + + +def test_collective_operator_expands_and_composes_with_press() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "arrange_then_press", + "goal": "Arrange both cans, then press the button.", + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "goal": {"axis": "world_y", "anchor": "table_center"}, + "depends_on": [], + }, + { + "id": "s02_press", + "operator": "press", + "object": "button", + "goal": {}, + "depends_on": ["s01_line"], + }, + ], + } + + execution = compile_task_agent(program) + steps = {step["id"]: step for step in execution["semantic_steps"]} + + assert list(steps) == ["s01_line__01", "s01_line__02", "s02_press"] + assert steps["s02_press"]["depends_on"] == [ + "s01_line__01", + "s01_line__02", + ] + assert execution["edges"][-1]["depends_on"] == [ + steps["s01_line__01"]["edge_ids"][-1], + steps["s01_line__02"]["edge_ids"][-1], + ] + assert "route" not in repr(execution) + + +def test_coordinated_place_picks_both_objects_before_placement() -> None: + execution = compile_task_agent( + _program( + { + "id": "s01_coordinated_place", + "operator": "coordinated_place", + "object": "cup", + "goal": {"support_object": "tray", "relation": "on"}, + } + ) + ) + step = execution["semantic_steps"][0] + first_edge, placement_edge = [ + next(edge for edge in execution["edges"] if edge["id"] == edge_id) + for edge_id in step["edge_ids"] + ] + + assert [action["atomic_action_class"] for action in first_edge["actions"]] == [ + "PickUp", + "PickUp", + ] + assert [action["actor"] for action in first_edge["actions"]] == [ + {"mode": "required", "arm": "left_arm"}, + {"mode": "required", "arm": "right_arm"}, + ] + assert [action["target_binding"]["object"] for action in first_edge["actions"]] == [ + "cup", + "tray", + ] + assert [action["motion_policy"] for action in first_edge["actions"]] == [ + {"modifiers": []}, + {"modifiers": []}, + ] + assert placement_edge["actions"][0]["atomic_action_class"] == ( + "CoordinatedPlacement" + ) + assert placement_edge["depends_on"] == [first_edge["id"]] + + +def test_independent_required_arms_create_allocation_group() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "parallel_place", + "goal": "Place two objects with opposite arms.", + "semantic_steps": [ + { + "id": "s01_left", + "operator": "place_relative", + "object": "left_object", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"reference_object": "left_tray", "relation": "on"}, + "depends_on": [], + }, + { + "id": "s02_right", + "operator": "place_relative", + "object": "right_object", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": {"reference_object": "right_tray", "relation": "on"}, + "depends_on": [], + }, + ], + } + + execution = compile_task_agent(program) + + assert execution["allocation_groups"] == [ + { + "id": "g01_distinct_arms", + "semantic_step_ids": ["s01_left", "s02_right"], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ] + + +def test_orient_object_composes_upright_motion_modifier() -> None: + execution = compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "upright", + "goal": "Stand the can upright.", + "semantic_steps": [ + { + "id": "s01_orient", + "operator": "orient_object", + "object": "can", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + ) + + assert [edge["actions"][0]["motion_policy"] for edge in execution["edges"]] == [ + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": [{"type": "orientation", "mode": "upright"}]}, + {"modifiers": []}, + ] + move_phases = [ + edge["actions"][0]["target_binding"]["phase"] + for edge in execution["edges"] + if edge["actions"][0]["atomic_action_class"] == "MoveHeldObject" + ] + assert move_phases == ["staging", "final"] + assert execution["semantic_steps"][0]["goal"] == { + "relation": "none", + "reference_state": "live", + "orientation_goal": "upright", + "orientation_axis": "none", + "position_anchor": "initial_xy", + "support_object": "table", + "upright_local_axis": "auto", + } + + +def test_auto_pickups_require_shared_explicit_allocation_group() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "dual_arm_basket", + "goal": "Use both arms to place the cube and cup in the basket.", + "semantic_steps": [ + { + "id": "s01_cube", + "operator": "place_relative", + "object": "cube", + "actor": { + "mode": "auto", + "allocation_group": "dual_arms_1", + }, + "goal": {"reference_object": "basket", "relation": "inside"}, + "depends_on": [], + }, + { + "id": "s02_cup", + "operator": "place_relative", + "object": "cup", + "actor": { + "mode": "auto", + "allocation_group": "dual_arms_1", + }, + "goal": {"reference_object": "basket", "relation": "inside"}, + "depends_on": [], + }, + ], + } + + execution = compile_task_agent(program) + pickup_edges = [ + next( + edge + for edge in execution["edges"] + if edge["semantic_step_id"] == step_id + and edge["actions"][0]["atomic_action_class"] == "PickUp" + ) + for step_id in ("s01_cube", "s02_cup") + ] + transport_edges = [ + next( + edge + for edge in execution["edges"] + if edge["semantic_step_id"] == step_id + and edge["actions"][0]["atomic_action_class"] == "MoveHeldObject" + ) + for step_id in ("s01_cube", "s02_cup") + ] + + assert execution["allocation_groups"] == [ + { + "id": "g01_distinct_arms", + "semantic_step_ids": ["s01_cube", "s02_cup"], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ] + assert all("workspace:basket" not in edge["resources"] for edge in pickup_edges) + assert all("workspace:basket" in edge["resources"] for edge in transport_edges) + + for step in program["semantic_steps"]: + step["actor"].pop("allocation_group") + assert compile_task_agent(program)["allocation_groups"] == [] + + +def test_unrelated_dependent_is_allowed_while_hold_reserves_arm() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "hold_then_press", + "goal": "Hold the cube and then press the button.", + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_press", + "operator": "press", + "object": "button", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": ["s01_hold"], + }, + ], + } + + execution = compile_task_agent(program) + steps = {step["id"]: step for step in execution["semantic_steps"]} + + assert steps["s01_hold"]["postcondition"] == { + "type": "object_held", + "object": "cube", + } + assert steps["s02_press"]["depends_on"] == ["s01_hold"] + + +def test_hold_may_follow_an_ancestor_that_previously_used_the_same_object() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "place_then_hold", + "goal": "Place the cube, then pick it up and keep holding it.", + "semantic_steps": [ + { + "id": "s01_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {"reference_object": "tray", "relation": "on"}, + "depends_on": [], + }, + { + "id": "s02_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": ["s01_place"], + }, + ], + } + + execution = compile_task_agent(program) + + assert execution["semantic_steps"][-1]["postcondition"]["type"] == "object_held" + + +@pytest.mark.parametrize( + ("operator", "actor", "goal"), + [ + ("press", {"mode": "required", "arm": "left_arm"}, {}), + ( + "coordinated_transport", + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + {"direction": "none", "terminal_behavior": "hold"}, + ), + ], +) +def test_required_hold_rejects_later_steps_that_need_its_arm( + operator: str, + actor: dict, + goal: dict, +) -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "occupied_arm", + "goal": "Keep holding the cube, then operate the button.", + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_other", + "operator": operator, + "object": "button", + "actor": actor, + "goal": goal, + "depends_on": ["s01_hold"], + }, + ], + } + + with pytest.raises(ValueError, match="reserves arm 'left_arm'"): + compile_task_agent(program) + + +def test_held_object_cannot_be_reused_by_an_independent_step() -> None: + program = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "conflicting_object_ownership", + "goal": "Hold and place the same cube.", + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {"reference_object": "tray", "relation": "on"}, + "depends_on": [], + }, + ], + } + + with pytest.raises(ValueError, match="reserves object 'cube'"): + compile_task_agent(program) + + +def test_unknown_operator_is_rejected_before_graph_construction() -> None: + with pytest.raises(ValueError, match="Unknown semantic operator"): + compile_task_agent( + _program( + { + "id": "s01_unknown", + "operator": "teleport", + "object": "cube", + "goal": {}, + } + ) + ) + + +def test_coordinated_transport_rejects_unknown_direction() -> None: + with pytest.raises(ValueError, match="direction"): + compile_task_agent( + _program( + { + "id": "s01_transport", + "operator": "coordinated_transport", + "object": "shared_box", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "direction": "somewhere_vague", + "terminal_behavior": "hold", + }, + "depends_on": [], + } + ) + ) diff --git a/tests/gen_sim/action_engine/compiler/test_v2.py b/tests/gen_sim/action_engine/compiler/test_v2.py new file mode 100644 index 000000000..6409a6a7a --- /dev/null +++ b/tests/gen_sim/action_engine/compiler/test_v2.py @@ -0,0 +1,172 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.compiler import ( + compile_task_agent, + compile_task_agent_v2, + seed_graph_to_execution_program, +) +from embodichain.gen_sim.action_engine.domain import TASK_AGENT_SCHEMA +from embodichain.gen_sim.action_engine.protocol import SEED_GRAPH_SCHEMA + + +def _task_agent() -> dict: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "place-cup", + "goal": "Place the cup in the tray.", + "semantic_steps": [ + { + "id": "s01_place_cup", + "operator": "place_relative", + "object": "cup", + "actor": {"mode": "auto"}, + "goal": { + "relation": "inside", + "reference_object": "tray", + "reference_state": "live", + "slot": "auto", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + +def test_v2_compiler_preserves_mature_atomic_action_topology() -> None: + known = {"cup", "tray"} + legacy = compile_task_agent(_task_agent(), known_objects=known) + seed = compile_task_agent_v2(_task_agent(), known_objects=known) + materialized = seed_graph_to_execution_program(seed, known_objects=known) + + legacy_actions = [ + action["atomic_action_class"] + for edge in legacy["edges"] + for action in edge["actions"] + ] + seed_actions = [node["atomic_action"] for node in seed["nodes"]] + materialized_actions = [ + action["atomic_action_class"] + for edge in materialized["edges"] + for action in edge["actions"] + ] + assert seed["schema_version"] == SEED_GRAPH_SCHEMA + assert seed_actions == legacy_actions + assert materialized_actions == legacy_actions + assert seed["task_groups"][0]["task_type"] == "E1" + + +@pytest.mark.parametrize( + ("operator", "objects", "goal", "actor"), + [ + ( + "orient_object", + ["can"], + { + "orientation_goal": "upright", + "orientation_axis": "long_axis", + "position_anchor": "initial_xy", + "support_object": "table", + "upright_local_axis": "long_axis", + }, + {"mode": "auto"}, + ), + ( + "coordinated_transport", + ["tray"], + { + "direction": "up", + "terminal_behavior": "hold", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + {"mode": "coordinated", "arms": ["left_arm", "right_arm"]}, + ), + ( + "build_stack", + ["cube_a", "cube_b"], + { + "anchor": "table_center", + "stack_mode": "on_top", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + {"mode": "auto"}, + ), + ( + "arrange_line", + ["cube_a", "cube_b"], + { + "anchor": "table_center", + "axis": "world_y", + "order_by": "explicit", + "order_constraint": "ordered", + "order_direction": "given", + "orientation_goal": "preserve", + "orientation_axis": "none", + "participation": "auto", + }, + {"mode": "auto"}, + ), + ], +) +def test_v2_preserves_all_current_task_recipe_topologies( + operator: str, + objects: list[str], + goal: dict, + actor: dict, +) -> None: + step = { + "id": "task_01", + "operator": operator, + "actor": actor, + "goal": goal, + "depends_on": [], + } + if operator in {"build_stack", "arrange_line"}: + step["objects"] = objects + else: + step["object"] = objects[0] + task = { + "schema_version": TASK_AGENT_SCHEMA, + "task": f"regression-{operator}", + "goal": f"Regression task for {operator}.", + "semantic_steps": [step], + "allocation_groups": [], + } + known = {*objects, "table"} + legacy = compile_task_agent(task, known_objects=known) + seed = compile_task_agent_v2(task, known_objects=known) + rematerialized = seed_graph_to_execution_program(seed, known_objects=known) + + def signature(program: dict) -> dict[str, list[list[str]]]: + edges = {edge["id"]: edge for edge in program["edges"]} + return { + step["id"]: [ + [action["atomic_action_class"] for action in edges[edge_id]["actions"]] + for edge_id in step["edge_ids"] + ] + for step in program["semantic_steps"] + } + + assert signature(rematerialized) == signature(legacy) diff --git a/tests/gen_sim/action_engine/config/__init__.py b/tests/gen_sim/action_engine/config/__init__.py new file mode 100644 index 000000000..355d915ff --- /dev/null +++ b/tests/gen_sim/action_engine/config/__init__.py @@ -0,0 +1,17 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py new file mode 100644 index 000000000..3e0b603d5 --- /dev/null +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -0,0 +1,308 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +import hashlib +import json + +import pytest + +from embodichain.gen_sim.action_engine.config import ( + ArmSelectionPolicyCfg, + RuntimePolicyCfg, + default_runtime_policy, + generation_defaults, + resolve_agent_runtime_policy, + runtime_policy_hash, +) +from embodichain.lab.sim.atomic_actions.primitives.place import PlaceOptions + + +def test_default_runtime_policy_preserves_current_arm_selection_behavior() -> None: + policy = default_runtime_policy("dual_ur10") + + assert policy.arm_selection.as_mapping() == { + "crossing_deadband_ratio": 0.08, + "allow_cross_side_fallback": False, + "pickup_crossing_weight": 1.0, + "placement_crossing_weight": 1.5, + "motion_cost_scale": pytest.approx(3.141592653589793), + "fallback_workspace_half_width": 0.5, + "orient_object_preferred_arm_deadband": 0.02, + } + + +def test_defaults_cover_current_execution_and_generation_policy() -> None: + runtime = default_runtime_policy("dual_ur10") + generation = generation_defaults() + + assert runtime.execution == { + "max_transitions": 1000, + "semantic_step_settle_steps": 10, + "max_retries_per_action": 2, + "max_graph_revisions": 8, + "max_recovery_actions": 12, + "support_stability_samples": 3, + "support_stability_interval_steps": 5, + "support_linear_velocity_tolerance": pytest.approx(0.02), + "support_angular_velocity_tolerance": pytest.approx(0.2), + } + assert runtime.planner == { + "backend": "curobo", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "fallback_strategy": "ik_interp", + "allow_fallback": True, + "dynamic_collision": False, + "static_obstacle_uids": [], + "dynamic_obstacle_uids": [], + "curobo": { + "log_level": "error", + "obstacle_representation": "cuboid", + "multi_env": False, + "use_cuda_graph": True, + "preserve_plan_samples": False, + "max_attempts": 5, + "collision_activation_distance": pytest.approx(0.01), + }, + } + assert runtime.grounding["arrangement"]["row_search_radius"] == 0.25 + assert runtime.grasp["antipodal_n_sample"] == 10000 + assert runtime.motion_modifiers["orientation"]["upright"]["MoveHeldObject"][ + "surface_clearance" + ] == pytest.approx(0.05) + assert runtime.motion_modifiers["handover_role"]["transfer"]["PickUp"] == { + "sample_interval": 80, + "hand_interp_steps": 5, + "pick_object_part": "top", + } + assert runtime.motion_defaults["HandOver"]["receive_pick_object_part"] == "bottom" + assert runtime.motion_defaults["CoordinatedPickment"][ + "middle_empty_ratio" + ] == pytest.approx(0.4) + assert ( + runtime.motion_defaults["CoordinatedPickment"]["is_filter_ground_collision"] + is False + ) + assert runtime.predicate_fallbacks["upright_max_tilt"] == pytest.approx( + 0.2617993877991494 + ) + assert generation["physics"]["rigid_object"]["mass"] == pytest.approx(0.1) + assert generation["environment"]["arm_aim_yaw_offset"] == { + "left": pytest.approx(0.0), + "right": pytest.approx(0.0), + } + assert generation["environment"]["recording"] == { + "enabled": True, + "resolution": [640, 360], + "interval_step": 1, + } + assert generation["scene"]["object_length_sample_points"] == 5000 + assert generation["dataset"]["control_frequency"] == 25 + assert generation["randomization"]["table_height_delta_range"] == [ + [-0.05], + [0.05], + ] + + +def test_place_defaults_fit_the_mainline_motion_sample_budget() -> None: + place = default_runtime_policy("dual_ur10").motion_defaults["Place"] + sample_count = int(place["sample_interval"]) + hand_steps = PlaceOptions().hand_interp_steps + motion_steps = sample_count - hand_steps + down_steps = int(round(motion_steps) * 0.6) + back_steps = motion_steps - down_steps + cartesian_count = int(place["cartesian_waypoint_count"]) + + assert 1 + 2 * cartesian_count <= down_steps + assert 1 + cartesian_count <= back_steps + + +def test_default_runtime_policy_returns_detached_profile_snapshots() -> None: + first = default_runtime_policy("dual_ur10") + second = default_runtime_policy("dual_ur10") + franka = default_runtime_policy("dual_franka") + + first.arm_selection.pickup_crossing_weight = 9.0 + first.motion_defaults["PickUp"]["lift_height"] = 9.0 + + assert second.arm_selection.pickup_crossing_weight == 1.0 + assert second.motion_defaults["PickUp"]["lift_height"] == 0.30 + assert franka.arm_selection.pickup_crossing_weight == 1.0 + assert franka.motion_defaults["MoveEndEffector"]["retreat_height"] == 0.10 + + +def test_generation_defaults_return_detached_values() -> None: + first = generation_defaults() + second = generation_defaults() + + first["physics"]["rigid_object"]["mass"] = 9.0 + + assert second["physics"]["rigid_object"]["mass"] == pytest.approx(0.1) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("crossing_deadband_ratio", 1.0, "crossing_deadband_ratio"), + ("pickup_crossing_weight", -0.1, "pickup_crossing_weight"), + ("placement_crossing_weight", -0.1, "placement_crossing_weight"), + ("motion_cost_scale", 0.0, "motion_cost_scale"), + ("fallback_workspace_half_width", 0.0, "fallback_workspace_half_width"), + ], +) +def test_arm_selection_policy_rejects_invalid_values( + field: str, + value: float, + message: str, +) -> None: + values = default_runtime_policy("dual_ur10").arm_selection.as_mapping() + values[field] = value + + with pytest.raises(ValueError, match=message): + ArmSelectionPolicyCfg.from_mapping(values) + + +def test_arm_selection_policy_requires_boolean_cross_side_fallback() -> None: + values = default_runtime_policy("dual_ur10").arm_selection.as_mapping() + values["allow_cross_side_fallback"] = "false" + + with pytest.raises(TypeError, match="allow_cross_side_fallback"): + ArmSelectionPolicyCfg.from_mapping(values) + + +def test_arm_selection_policy_loads_old_snapshot_without_fallback_field() -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["arm_selection"].pop("allow_cross_side_fallback") + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + policy = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert policy.arm_selection.allow_cross_side_fallback is False + + +def test_agent_policy_snapshot_is_hash_verified_and_legacy_config_falls_back() -> None: + policy = default_runtime_policy("dual_ur5") + snapshot = policy.as_mapping() + config = { + "robot_profile": "dual_ur5", + "runtime_policy": snapshot, + "runtime_policy_hash": runtime_policy_hash(policy), + } + + resolved = resolve_agent_runtime_policy(config) + assert resolved.as_mapping() == snapshot + + tampered = deepcopy(config) + tampered["runtime_policy"]["motion_defaults"]["PickUp"]["lift_height"] = 8.0 + with pytest.raises(ValueError, match="hash does not match"): + resolve_agent_runtime_policy(tampered) + + legacy = resolve_agent_runtime_policy({"robot_profile": "dual_ur5"}) + assert legacy.as_mapping() == snapshot + + +def test_narrow_v1_policy_snapshot_is_migrated_to_complete_runtime_policy() -> None: + snapshot = { + "schema_version": "action_engine_runtime_policy_v1", + "arm_selection": { + "crossing_deadband_ratio": 0.08, + "pickup_crossing_weight": 2.0, + "placement_crossing_weight": 1.5, + "motion_cost_scale": 3.141592653589793, + "fallback_workspace_half_width": 0.5, + }, + } + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + resolved = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert resolved.arm_selection.pickup_crossing_weight == 2.0 + assert resolved.motion_defaults["PickUp"]["lift_height"] == 0.30 + + +def test_v3_policy_snapshot_is_migrated_with_default_planner_policy() -> None: + expected = default_runtime_policy("dual_ur10") + snapshot = expected.as_mapping() + snapshot.pop("planner") + snapshot["schema_version"] = "action_engine_runtime_policy_v3" + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + resolved = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert resolved.schema_version == "action_engine_runtime_policy_v6" + assert resolved.planner == expected.planner + + +def test_curobo_policy_rejects_coordinated_motion_generation() -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["planner"]["coordinated_strategy"] = "motion_gen" + + with pytest.raises(ValueError, match="coordinated_strategy"): + RuntimePolicyCfg.from_mapping(snapshot) + + +@pytest.mark.parametrize( + ("patch", "message"), + [ + ({"fallback_strategy": "motion_gen"}, "fallback_strategy"), + ({"backend": "toppra", "dynamic_collision": True}, "dynamic_collision"), + ], +) +def test_planner_policy_rejects_unsupported_combinations( + patch: dict[str, object], + message: str, +) -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["planner"].update(patch) + + with pytest.raises(ValueError, match=message): + RuntimePolicyCfg.from_mapping(snapshot) diff --git a/tests/gen_sim/action_engine/domain/__init__.py b/tests/gen_sim/action_engine/domain/__init__.py new file mode 100644 index 000000000..c8e03f284 --- /dev/null +++ b/tests/gen_sim/action_engine/domain/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine domain tests.""" diff --git a/tests/gen_sim/action_engine/domain/test_programs.py b/tests/gen_sim/action_engine/domain/test_programs.py new file mode 100644 index 000000000..8fdf8074b --- /dev/null +++ b/tests/gen_sim/action_engine/domain/test_programs.py @@ -0,0 +1,150 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy + +import pytest + +from embodichain.gen_sim.action_engine.compiler import compile_task_agent +from embodichain.gen_sim.action_engine.domain import ( + TASK_AGENT_SCHEMA, + execution_program_hash, + validate_execution_program, + validate_task_agent, +) + + +def _task_agent() -> dict: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "place_demo", + "goal": "Place the cup on the tray.", + "semantic_steps": [ + { + "id": "s01_place", + "operator": "place_relative", + "object": "cup", + "actor": {"mode": "auto"}, + "goal": {"reference_object": "tray", "relation": "on"}, + "depends_on": [], + } + ], + } + + +def test_task_validation_is_detached_and_adds_unambiguous_defaults() -> None: + source = _task_agent() + del source["semantic_steps"][0]["actor"] + del source["semantic_steps"][0]["depends_on"] + + validated = validate_task_agent(source) + validated["semantic_steps"][0]["goal"]["relation"] = "inside" + + assert source["semantic_steps"][0]["goal"]["relation"] == "on" + assert validated["semantic_steps"][0]["actor"] == {"mode": "auto"} + assert validated["semantic_steps"][0]["depends_on"] == [] + + +@pytest.mark.parametrize( + "actor", + [ + {"mode": "auto", "allocation_group": "dual_arms_1"}, + { + "mode": "required", + "arm": "left_arm", + "allocation_group": "dual_arms_1", + }, + ], +) +def test_single_arm_allocation_group_is_validated_and_preserved(actor: dict) -> None: + source = _task_agent() + source["semantic_steps"][0]["actor"] = actor + + validated = validate_task_agent(source) + execution = compile_task_agent(validated) + + assert validated["semantic_steps"][0]["actor"] == actor + assert execution["semantic_steps"][0]["actor"] == actor + assert all( + action["actor"] == actor + for edge in execution["edges"] + for action in edge["actions"] + ) + + +def test_allocation_group_must_be_nonempty_and_single_arm_only() -> None: + source = _task_agent() + source["semantic_steps"][0]["actor"]["allocation_group"] = " " + with pytest.raises(ValueError, match="allocation_group"): + validate_task_agent(source) + + source["semantic_steps"][0]["actor"] = { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + "allocation_group": "dual_arms_1", + } + with pytest.raises(ValueError, match="unknown fields"): + validate_task_agent(source) + + +def test_task_validation_rejects_cycles_and_grounded_values() -> None: + cyclic = _task_agent() + cyclic["semantic_steps"].extend( + [ + { + "id": "s02", + "operator": "press", + "object": "button", + "depends_on": ["s03"], + }, + { + "id": "s03", + "operator": "press", + "object": "button", + "depends_on": ["s02"], + }, + ] + ) + with pytest.raises(ValueError, match="cycle"): + validate_task_agent(cyclic) + + grounded = _task_agent() + grounded["semantic_steps"][0]["goal"]["target_pose"] = [0.0, 0.0, 0.0] + with pytest.raises(ValueError, match="grounded runtime data"): + validate_task_agent(grounded) + + +def test_execution_hash_is_stable_and_validation_is_strict() -> None: + execution = compile_task_agent(_task_agent()) + reordered = {key: execution[key] for key in reversed(list(execution))} + + assert execution_program_hash(execution) == execution_program_hash(reordered) + assert len(execution_program_hash(execution)) == 64 + + broken = deepcopy(execution) + broken["edges"][0]["target_binding"] = {} + with pytest.raises(ValueError, match="unknown fields"): + validate_execution_program(broken) + + +def test_execution_validation_rejects_unowned_edges() -> None: + execution = compile_task_agent(_task_agent()) + execution["semantic_steps"][0]["edge_ids"].pop() + + with pytest.raises(ValueError, match="unowned edges"): + validate_execution_program(execution) diff --git a/tests/gen_sim/action_engine/domain/test_task_contracts.py b/tests/gen_sim/action_engine/domain/test_task_contracts.py new file mode 100644 index 000000000..765e05aa8 --- /dev/null +++ b/tests/gen_sim/action_engine/domain/test_task_contracts.py @@ -0,0 +1,71 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.domain import ( + RELATIONS, + TASK_CONTRACTS, + TASK_TYPES, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + normalize_placement_relation, + task_contract, + task_success_type, +) + + +def test_task_contract_catalog_covers_the_canonical_protocol() -> None: + assert set(TASK_CONTRACTS) == set(TASK_TYPES) + assert all(contract.core_actions for contract in TASK_CONTRACTS.values()) + assert {contract.source_structure for contract in TASK_CONTRACTS.values()} == { + "articulation", + "rigid_object", + } + assert task_contract("E2").success_type == "object_upright" + assert task_contract("E5").scene_affordances == { + "dual_graspable", + "rigid", + } + + +def test_e5_success_depends_only_on_terminal_behavior() -> None: + assert task_success_type("E5", {"terminal_behavior": "hold"}) == ( + "held_by_both_grippers" + ) + assert task_success_type("E5", {"terminal_behavior": "place"}) == "semantic_goal" + with pytest.raises(ValueError, match="terminal_behavior"): + task_success_type("E5", {"terminal_behavior": "none"}) + + +def test_symbolic_transport_values_are_language_neutral_protocol_enums() -> None: + assert {"on", "inside", "behind", "left_of"} <= RELATIONS + assert {"none", "up", "left", "world_y"} <= TRANSPORT_DIRECTIONS + assert TERMINAL_BEHAVIORS == {"none", "hold", "place"} + + +@pytest.mark.parametrize("relation", ["above", "on_top", "on_top_of"]) +def test_released_hover_and_legacy_support_relations_normalize_to_on( + relation: str, +) -> None: + assert normalize_placement_relation(relation) == "on" + + +def test_placement_relation_normalization_rejects_non_spatial_semantics() -> None: + with pytest.raises(ValueError, match="Unsupported placement relation"): + normalize_placement_relation("visual_slot") diff --git a/tests/gen_sim/action_engine/domain/test_v2.py b/tests/gen_sim/action_engine/domain/test_v2.py new file mode 100644 index 000000000..84cb02b7f --- /dev/null +++ b/tests/gen_sim/action_engine/domain/test_v2.py @@ -0,0 +1,280 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy + +import pytest + +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import ( + motion_policy, + public_task_spec, + seed_graph_hash, + validate_public_task_spec, + validate_scene_requirements, + validate_seed_graph, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph + + +def _seed_graph() -> dict: + registry = build_atomic_capability_registry() + draft = { + "schema_version": SEED_GRAPH_SCHEMA, + "task_id": "place-cup", + "instruction": "Place the cup in the tray.", + "level": "L1", + "reasoning_type": "none", + "planner_route": "offline", + "nodes": [ + { + "id": "pick_cup", + "atomic_action": "PickUp", + "object_uid": "cup", + "actor": {"mode": "auto"}, + "control": "arm", + "target_binding": {"kind": "object", "object": "cup"}, + "depends_on": [], + "task_instance_id": "e1_001", + "task_type": "E1", + "role": "primary", + "precondition": {"type": "object_not_fallen", "object": "cup"}, + "postcondition": {"type": "object_held", "object": "cup"}, + "motion_policy": motion_policy(), + }, + { + "id": "place_cup", + "atomic_action": "Place", + "object_uid": "cup", + "actor": {"mode": "auto"}, + "control": "arm", + "target_binding": {"kind": "current_held_pose"}, + "depends_on": ["pick_cup"], + "task_instance_id": "e1_001", + "task_type": "E1", + "role": "primary", + "precondition": {"type": "object_held", "object": "cup"}, + "postcondition": { + "type": "object_in_container", + "object": "cup", + "container": "tray", + }, + "motion_policy": motion_policy(), + }, + ], + "task_groups": [ + { + "id": "e1_001", + "task_type": "E1", + "role": "primary", + "operator": "place_relative", + "object_uid": "cup", + "actor": {"mode": "auto"}, + "goal": {"relation": "inside", "reference_object": "tray"}, + "depends_on": [], + "node_ids": ["pick_cup", "place_cup"], + "success": { + "type": "object_in_container", + "object": "cup", + "container": "tray", + }, + } + ], + "success": { + "type": "object_in_container", + "object": "cup", + "container": "tray", + }, + "capability_catalog_hash": registry.catalog_hash(), + "metadata": {}, + } + return link_seed_graph( + draft, + registry=registry, + task_order=["e1_001"], + known_objects={"cup", "tray"}, + ) + + +def test_seed_graph_validates_direct_atomic_action_nodes() -> None: + registry = build_atomic_capability_registry() + graph = validate_seed_graph( + _seed_graph(), + known_objects={"cup", "tray"}, + known_actions=registry.names(), + executable_actions=registry.executable_names(), + require_executable=True, + ) + assert [node["atomic_action"] for node in graph["nodes"]] == ["PickUp", "Place"] + assert graph["task_groups"][0]["node_ids"] == ["pick_cup", "place_cup"] + + +def test_seed_graph_rejects_grounded_motion_and_cycles() -> None: + grounded = _seed_graph() + grounded["nodes"][0]["target_binding"]["target_pose"] = [0.0, 0.0, 0.0] + with pytest.raises(ValueError, match="grounded motion data"): + validate_seed_graph(grounded) + + cyclic = _seed_graph() + cyclic["nodes"][0]["depends_on"] = ["place_cup"] + with pytest.raises(ValueError, match="dependency cycle"): + validate_seed_graph(cyclic) + + +def test_seed_graph_rejects_unknown_uids_illegal_groups_and_resource_conflicts() -> ( + None +): + with pytest.raises(ValueError, match="unknown object"): + validate_seed_graph(_seed_graph(), known_objects={"tray"}) + + illegal_group = _seed_graph() + illegal_group["task_groups"][0]["task_type"] = "E9" + for node in illegal_group["nodes"]: + node["task_type"] = "E9" + with pytest.raises(ValueError, match="core actions"): + validate_seed_graph(illegal_group) + + conflicting = _seed_graph() + conflicting["nodes"][1]["depends_on"] = [] + conflicting["task_groups"][0]["contract"]["entry_node_ids"] = [ + "pick_cup", + "place_cup", + ] + conflicting["task_groups"][0]["contract"]["terminal_node_ids"] = [ + "pick_cup", + "place_cup", + ] + with pytest.raises(ValueError, match="resource conflicts"): + validate_seed_graph(conflicting) + + +def test_seed_graph_hash_is_order_stable_and_detached() -> None: + graph = _seed_graph() + original = deepcopy(graph) + first = seed_graph_hash(graph) + second = seed_graph_hash({key: graph[key] for key in reversed(graph)}) + assert first == second + assert graph == original + + +def test_task_spec_enforces_reasoning_level_and_repetition_shape() -> None: + spec = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "upright-cans", + "level": "L2", + "instruction": "Stand both cans upright.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "e2_1", + "task_type": "E2", + "params": {"object_role": "can_1"}, + "depends_on": [], + }, + { + "id": "e2_2", + "task_type": "E2", + "params": {"object_role": "can_2"}, + "depends_on": [], + }, + ], + "success": {"type": "all_upright"}, + "oracle": {"object_roles": ["can_1", "can_2"]}, + "metadata": {}, + } + assert validate_task_spec(spec)["level"] == "L2" + spec["level"] = "L4" + with pytest.raises(ValueError, match="non-'none'"): + validate_task_spec(spec) + + +def test_public_l4_task_hides_oracle_and_reference_instances() -> None: + spec = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "complete-mouth", + "level": "L4", + "instruction": "Complete the missing mouth.", + "reasoning_type": "visual_semantics", + "task_instances": [ + { + "id": "hidden_e1", + "task_type": "E1", + "params": {"object_role": "mouth", "target_role": "face"}, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "visual_part_complete"}, + "oracle": {"missing_part": "mouth"}, + "metadata": {}, + } + + public = public_task_spec(spec) + + assert "oracle" not in public + assert "task_instances" not in public + assert validate_public_task_spec(public) == public + + +def test_scene_requirements_validate_task_first_handoff() -> None: + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": "upright-cans", + "objects": [ + { + "role_id": "can", + "category": "can", + "count": 2, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + "cameras": [{"role": "overview", "requires_rgb": True}], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {}, + } + ) + assert requirements["objects"][0]["count"] == 2 + + +def test_planning_only_capability_is_rejected_before_execution() -> None: + registry = build_atomic_capability_registry() + graph = _seed_graph() + graph["nodes"][0]["atomic_action"] = "TurnKnob" + graph["nodes"][0]["target_binding"] = { + "kind": "articulation_goal", + "object": "cup", + } + with pytest.raises(ValueError, match="planning-only"): + validate_seed_graph( + graph, + known_actions=registry.names(), + executable_actions=(set(registry.executable_names()) - {"TurnKnob"}), + require_executable=True, + ) diff --git a/tests/gen_sim/action_engine/evaluation/__init__.py b/tests/gen_sim/action_engine/evaluation/__init__.py new file mode 100644 index 000000000..adfe7f7b4 --- /dev/null +++ b/tests/gen_sim/action_engine/evaluation/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine evaluation tests.""" diff --git a/tests/gen_sim/action_engine/evaluation/test_ab.py b/tests/gen_sim/action_engine/evaluation/test_ab.py new file mode 100644 index 000000000..72baa5e47 --- /dev/null +++ b/tests/gen_sim/action_engine/evaluation/test_ab.py @@ -0,0 +1,445 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +import json +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.evaluation import run_strict_ab, state_digest +from embodichain.gen_sim.action_engine.evaluation.ab import _graph_difference +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_level, make_task_spec + + +class _Env: + def __init__(self, route: str, seed: int, config: dict) -> None: + self.route = route + self.seed = seed + self.config = config + self.closed = False + + def reset(self, *, seed: int) -> None: + self.seed = seed + + def close(self) -> None: + self.closed = True + + +class _Executor: + def __init__(self, graph: dict, env: _Env) -> None: + self.graph = graph + self.env = env + + def run(self, **_kwargs): + return SimpleNamespace( + success=torch.tensor([True]), + actions=[torch.tensor([[0.0]]), torch.tensor([[1.0]])], + retry_count=0, + recovery_count=0, + revision_count=0, + runtime_revisions=[], + record_dir=f"records/{self.env.route}", + ) + + +def _inputs(): + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + return task, offline, online + + +def test_state_digest_is_mapping_order_stable() -> None: + assert state_digest( + {"qpos": torch.tensor([1.0]), "objects": {"a": [2.0]}} + ) == state_digest({"objects": {"a": [2.0]}, "qpos": torch.tensor([1.0])}) + + +def test_strict_ab_writes_isolated_branches_and_comparison(tmp_path) -> None: + task, offline, online = _inputs() + created = [] + + def env_factory(**kwargs): + env = _Env(**kwargs) + created.append(env) + return env + + result = run_strict_ab( + task, + offline, + online, + env_factory=env_factory, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=123, + shared_config={"robot": "same"}, + planning_metrics={ + "offline": {"planning_seconds": 0.1, "vlm_call_count": 0}, + "online": {"planning_seconds": 0.2, "vlm_call_count": 2}, + }, + ) + + assert result.comparison_path.is_file() + assert (result.offline_dir / "seed_task_graph.json").is_file() + assert (result.online_dir / "seed_task_graph.json").is_file() + assert result.comparison["initial_state_digest"] == result.initial_state_digest + assert result.comparison["branches"]["offline"]["planning_seconds"] == 0.1 + assert result.comparison["branches"]["online"]["vlm_call_count"] == 2 + assert all(env.closed for env in created) + + +def test_graph_difference_reports_changed_task_group_fields() -> None: + _task_spec, offline, online = _inputs() + online["task_groups"][0]["goal"] = deepcopy(online["task_groups"][0]["goal"]) + online["task_groups"][0]["goal"]["relation"] = "right_of" + + difference = _graph_difference(offline, online) + + assert difference["changed_task_groups"] == [ + {"id": offline["task_groups"][0]["id"], "changed_fields": ["goal"]} + ] + assert ( + difference["task_group_difference"]["changed_groups"] + == difference["changed_task_groups"] + ) + + +def test_strict_ab_finalizes_two_branch_videos_and_revision_files(tmp_path) -> None: + task, offline, online = _inputs() + finalized = [] + + def finalizer(**kwargs): + route = kwargs["route"] + path = kwargs["branch_dir"] / "video.mp4" + path.write_bytes(route.encode("ascii")) + finalized.append(route) + return [path.as_posix()] + + result = run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + branch_finalizer=finalizer, + output_dir=tmp_path, + seed=11, + ) + + assert finalized == ["offline", "online"] + for route, branch_dir in ( + ("offline", result.offline_dir), + ("online", result.online_dir), + ): + assert (branch_dir / "video.mp4").read_bytes() == route.encode("ascii") + assert (branch_dir / "runtime_revisions.json").is_file() + assert result.comparison["branches"][route]["video_paths"] == [ + (branch_dir / "video.mp4").as_posix() + ] + + +def test_strict_ab_aborts_before_execution_on_state_mismatch(tmp_path) -> None: + task, offline, online = _inputs() + executions = [] + + with pytest.raises(RuntimeError, match="initial state mismatch"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: executions.append((graph, env)), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {env.route: torch.eye(4)}, + }, + output_dir=tmp_path, + seed=5, + ) + assert executions == [] + + +def test_strict_ab_rejects_incomplete_state_snapshot(tmp_path) -> None: + task, offline, online = _inputs() + + with pytest.raises(ValueError, match="missing required state"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: {"robot_qpos": torch.tensor([0.0])}, + output_dir=tmp_path, + seed=5, + ) + + +def test_strict_ab_requires_articulation_and_camera_digest_components(tmp_path) -> None: + task, offline, online = _inputs() + + with pytest.raises(ValueError, match="articulation_state.*camera_calibration"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=5, + strict_state_digest=True, + ) + + +def test_strict_ab_reuses_prepared_identical_resets(tmp_path) -> None: + task, offline, online = _inputs() + environments = { + route: _Env(route=route, seed=17, config={}) for route in ("offline", "online") + } + snapshots = { + route: { + "robot_qpos": torch.tensor([17.0]), + "object_poses": {"object": torch.eye(4)}, + } + for route in environments + } + + result = run_strict_ab( + task, + offline, + online, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda _env: pytest.fail("prepared snapshots must be reused"), + output_dir=tmp_path, + seed=17, + prepared_environments=environments, + prepared_snapshots=snapshots, + ) + + assert result.initial_state_digest == state_digest(snapshots["offline"]) + assert all(env.closed for env in environments.values()) + + +def test_strict_ab_stops_both_branches_when_global_preflight_fails(tmp_path) -> None: + task, offline, online = _inputs() + runs: list[str] = [] + + class PreflightExecutor(_Executor): + def preflight(self) -> bool: + if self.env.route == "online": + raise ValueError("online capability unavailable") + return True + + def run(self, **kwargs): + runs.append(self.env.route) + return super().run(**kwargs) + + with pytest.raises(RuntimeError, match="no branch was allowed to move"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: PreflightExecutor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=17, + ) + assert runs == [] + + +def test_strict_ab_keeps_other_branch_running_after_execution_failure(tmp_path) -> None: + task, offline, online = _inputs() + runs: list[str] = [] + + class IsolatedExecutor(_Executor): + def preflight(self) -> bool: + return True + + def run(self, **kwargs): + runs.append(self.env.route) + if self.env.route == "offline": + raise RuntimeError("offline execution failed") + return super().run(**kwargs) + + result = run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: IsolatedExecutor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=17, + ) + + assert runs == ["offline", "online"] + assert result.comparison["branches"]["offline"]["success_rate"] == 0.0 + assert result.comparison["branches"]["online"]["success_rate"] == 1.0 + + +def test_strict_ab_surfaces_video_finalizer_failure_and_closes(tmp_path) -> None: + task, offline, online = _inputs() + environments = [] + + def factory(**kwargs): + environment = _Env(**kwargs) + environments.append(environment) + return environment + + def finalizer(**kwargs): + if kwargs["route"] == "offline": + raise OSError("recorder did not produce a file") + return [] + + with pytest.raises(RuntimeError, match="branch video finalization failed"): + run_strict_ab( + task, + offline, + online, + env_factory=factory, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + branch_finalizer=finalizer, + output_dir=tmp_path, + seed=19, + ) + + assert len(environments) == 2 + assert all(environment.closed for environment in environments) + comparison = json.loads((tmp_path / "comparison.json").read_text()) + assert set(comparison["video_finalization_errors"]) == {"offline"} + + +def test_strict_ab_require_branch_videos_checks_normalized_artifact(tmp_path) -> None: + task, offline, online = _inputs() + + with pytest.raises(RuntimeError, match="video.mp4"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + branch_finalizer=lambda **_kwargs: [], + output_dir=tmp_path, + seed=19, + require_branch_videos=True, + ) + + +def test_strict_ab_closes_prepared_environments_on_graph_validation_error( + tmp_path, +) -> None: + task, offline, online = _inputs() + environments = { + route: _Env(route=route, seed=23, config={}) for route in ("offline", "online") + } + invalid_online = deepcopy(online) + invalid_online["planner_route"] = "offline" + + with pytest.raises(ValueError, match="explicit offline and online"): + run_strict_ab( + task, + offline, + invalid_online, + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda _env: pytest.fail("validation must happen first"), + output_dir=tmp_path, + seed=23, + prepared_environments=environments, + ) + + assert all(environment.closed for environment in environments.values()) + + +def test_strict_l4_ab_requires_and_records_private_oracle(tmp_path) -> None: + task, requirements = make_task_level("L4") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + + with pytest.raises(ValueError, match="private-oracle"): + run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + output_dir=tmp_path, + seed=7, + ) + + result = run_strict_ab( + task, + offline, + online, + env_factory=lambda **kwargs: _Env(**kwargs), + executor_factory=lambda graph, env: _Executor(graph, env), + snapshot_reader=lambda env: { + "robot_qpos": torch.tensor([float(env.seed)]), + "object_poses": {"object": torch.eye(4)}, + }, + success_evaluator=lambda **_kwargs: torch.tensor([True]), + output_dir=tmp_path, + seed=7, + ) + assert all( + branch["success_source"] == "private_oracle" + for branch in result.comparison["branches"].values() + ) diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py new file mode 100644 index 000000000..13f7d7f7d --- /dev/null +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -0,0 +1,1590 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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_explicit_named_legacy_scene_config_is_supported(gym_export: Path) -> None: + config_path = gym_export / "official_task_config.json" + config_path.write_text( + (gym_export / "gym_config.json").read_text(encoding="utf-8"), + encoding="utf-8", + ) + + resolved = resolve_source_scene(config_path) + scene = prepare_scene(config_path) + + assert resolved.path == config_path + assert resolved.source_format == "legacy_gym_config" + assert resolved.is_prompt2scene is False + assert scene.source_config_path == config_path + + +def test_explicit_robot_scene_is_centered_on_its_table_anchor( + gym_export: Path, +) -> None: + source = json.loads((gym_export / "gym_config.json").read_text(encoding="utf-8")) + source["robot"] = {"uid": "source_robot"} + source["background"][0]["init_pos"] = [1.0, 2.0, 0.0] + source["rigid_object"][0]["init_pos"] = [1.2, 2.3, 0.7] + config_path = gym_export / "official_task_config.json" + config_path.write_text(json.dumps(source), encoding="utf-8") + + scene = prepare_scene(config_path) + table = scene.background[0] + moved = scene.rigid_objects[0] + + assert scene.source_scene_xy_translation == pytest.approx((-1.0, -2.0)) + assert table["init_pos"][:2] == pytest.approx([0.0, 0.0]) + assert moved["init_pos"][:2] == pytest.approx([0.2, 0.3]) + + +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["entity_uids"] == ["interact_can"] + 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_agent_config_owns_articulation_setting_calibration( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + settings = {"microwave": {"timer_joint": [-1.0, 0.0, 1.0]}} + + agent = build_agent_config( + task_name="turn_knob", + robot_profile="franka", + execution_program_hash="f" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + articulation_settings=settings, + ) + settings["microwave"]["timer_joint"][0] = 99.0 + + assert agent["articulation_settings"] == { + "microwave": {"timer_joint": [-1.0, 0.0, 1.0]} + } + + +def test_ab_scene_requirements_declare_four_vlm_views() -> None: + requirements = { + "schema_version": "action_engine_scene_requirements_v2", + "task_id": "ab", + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable"], + "initial_state": {}, + "attributes": {}, + } + ], + "cameras": [], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {}, + } + output = _add_ab_camera_requirements(requirements) + assert [item["uid"] for item in output["cameras"]] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + assert all(item["modalities"] == ["rgb", "depth"] for item in output["cameras"]) + + +def test_ab_builder_rejects_noncanonical_vlm_camera_ids(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + with pytest.raises(ValueError, match="canonical"): + build_agent_config( + task_name="ab_invalid_cameras", + robot_profile="ur10", + execution_program_hash="f" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + planning_mode="ab", + vlm_camera_uids=["front", "left", "rear", "right"], + ) + + +@pytest.mark.parametrize( + ("profile", "robot_uid", "solver_type"), + [ + ("dual_ur3", "DualUR3", "ur3"), + ("dual_ur5", "DualUR5", "ur5"), + ("dual_ur10", "DualUR10", "ur10"), + ("dual_franka", "DualFrankaPanda", None), + ], +) +def test_fast_gym_config_supports_all_robot_profiles( + gym_export: Path, + profile: str, + robot_uid: str, + solver_type: str | None, +) -> None: + expected_tcp = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0], + ] + expected_hand_mount = [ + [0.0, 1.0, 0.0, 0.0], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="profile_task", + task_description="Profile smoke test.", + robot_profile=profile, + execution_program_hash="b" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + assert config["robot"]["uid"] == robot_uid + assert config["env"]["extensions"]["agent_robot_profile"] == profile + for arm in ("left_arm", "right_arm"): + assert config["robot"]["solver_cfg"][arm]["tcp"] == expected_tcp + components = { + component["component_type"]: component + for component in config["robot"]["urdf_cfg"]["components"] + } + for hand in ("left_hand", "right_hand"): + assert components[hand]["transform"] == expected_hand_mount + if solver_type is not None: + assert config["robot"]["solver_cfg"]["left_arm"]["ur_type"] == solver_type + + +@pytest.mark.parametrize( + ( + "profile", + "expected_position_xy", + "expected_rotation", + "expected_world_x", + ), + [ + ("ur10", [2.0, 0.0], [0.0, 0.0, 0.0], 0.9), + ("franka", [-0.7, 0.0], [0.0, 0.0, 180.0], 0.55), + ], +) +def test_dual_robot_profiles_use_identity_mounts_and_same_side_arm_names( + gym_export: Path, + profile: str, + expected_position_xy: list[float], + expected_rotation: list[float], + expected_world_x: float, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="dual_ur_frame_task", + task_description="Verify the Dual-UR world frame.", + robot_profile=profile, + execution_program_hash="c" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + robot = config["robot"] + robot_yaw = np.deg2rad(float(robot["init_rot"][2])) + robot_rotation = np.array( + [ + [np.cos(robot_yaw), -np.sin(robot_yaw), 0.0], + [np.sin(robot_yaw), np.cos(robot_yaw), 0.0], + [0.0, 0.0, 1.0], + ] + ) + robot_position = np.asarray(robot["init_pos"], dtype=np.float64) + components = { + component["component_type"]: np.asarray( + component["transform"], dtype=np.float64 + ) + for component in robot["urdf_cfg"]["components"] + if component["component_type"] in {"left_arm", "right_arm"} + } + world_transforms = {} + for side, component in components.items(): + world = np.eye(4) + world[:3, :3] = robot_rotation @ component[:3, :3] + world[:3, 3] = robot_position + robot_rotation @ component[:3, 3] + world_transforms[side] = world + + assert robot["init_pos"][:2] == pytest.approx(expected_position_xy) + assert robot["init_rot"] == pytest.approx(expected_rotation) + assert world_transforms["left_arm"][:3, 3] == pytest.approx( + [expected_world_x, -0.3, world_transforms["left_arm"][2, 3]] + ) + assert world_transforms["right_arm"][:3, 3] == pytest.approx( + [expected_world_x, 0.3, world_transforms["right_arm"][2, 3]] + ) + np.testing.assert_allclose(components["left_arm"][:3, :3], np.eye(3), atol=1.0e-12) + np.testing.assert_allclose(components["right_arm"][:3, :3], np.eye(3), atol=1.0e-12) + np.testing.assert_allclose( + world_transforms["left_arm"][:3, :3], robot_rotation, atol=1.0e-12 + ) + np.testing.assert_allclose( + world_transforms["right_arm"][:3, :3], robot_rotation, atol=1.0e-12 + ) + + +def test_fast_gym_config_keeps_scene_deterministic_by_default( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="deterministic_task", + task_description="Keep the source scene fixed.", + robot_profile="ur10", + execution_program_hash="d" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + events = config["env"]["events"] + assert "randomize_interact_can_pose" not in events + assert "randomize_table_height" not in events + + +@pytest.mark.parametrize( + ("alias", "canonical"), + [ + ("franka", "dual_franka"), + ("ur5", "dual_ur5"), + ("ur10", "dual_ur10"), + ], +) +def test_required_cli_robot_aliases_build_runnable_profiles( + gym_export: Path, + alias: str, + canonical: str, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="profile_alias_task", + task_description="Profile alias smoke test.", + robot_profile=alias, + execution_program_hash="c" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + assert config["env"]["extensions"]["agent_robot_profile"] == canonical + + +def test_source_scene_scale_policies_are_deterministic(gym_export: Path) -> None: + preserved = prepare_scene(gym_export) + multiplied = prepare_scene( + gym_export, + body_scale_policy="multiply", + body_scale=(2.0, 3.0, 4.0), + ) + absolute = prepare_scene( + gym_export, + body_scale_policy="absolute", + body_scale=(2.0, 3.0, 4.0), + ) + + assert preserved.body_scale_policy == "preserve" + assert multiplied.rigid_objects[0]["body_scale"] == [2.0, 3.0, 4.0] + assert absolute.rigid_objects[0]["body_scale"] == [2.0, 3.0, 4.0] + assert multiplied.asset_hashes == absolute.asset_hashes + + +def test_artifact_writer_refuses_implicit_overwrite(tmp_path: Path) -> None: + payload = {"value": 1} + paths = write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nold", + overwrite=False, + ) + assert json.loads(paths.gym_config.read_text(encoding="utf-8")) == payload + assert paths.seed_task_graph_png.read_bytes().startswith(b"\x89PNG") + + # A leftover PNG participates in the same preflight as every JSON artifact. + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.execution_program, + ): + path.unlink() + with pytest.raises(FileExistsError, match="--overwrite"): + write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nnew", + overwrite=False, + ) + + replaced = write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nnew", + overwrite=True, + ) + assert replaced.seed_task_graph_png.read_bytes().endswith(b"new") + + +def test_artifact_writer_creates_ab_branch_directory(tmp_path: Path) -> None: + payload = {"value": "ab"} + paths = write_generation_artifacts( + tmp_path, + gym_config=payload, + agent_config=payload, + task_spec=payload, + scene_requirements=payload, + seed_task_graph=payload, + seed_task_graph_png=b"\x89PNG\r\n\x1a\nab", + overwrite=False, + planning_mode="ab", + ) + + assert paths.seed_task_graph.parent == tmp_path / "offline" + assert json.loads(paths.seed_task_graph.read_text(encoding="utf-8")) == payload + + +def test_generation_calls_interpreter_recipe_and_renderer_once( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.gen_sim.action_engine import tasks + from embodichain.gen_sim.action_engine.generation import generator + + planner_call: dict[str, object] = {} + recipe_calls: list[tuple[object, object]] = [] + rendered: dict[str, object] = {} + published: dict[str, object] = {} + + def fake_interpret_and_ground(**kwargs): + planner_call.update(kwargs) + task_spec = _existing_v2_task_spec(str(kwargs["task_name"])) + task_spec["instruction"] = str(kwargs["task_description"]) + bindings = {"object_01": "interact_can"} + requirements = _scene_requirements_from_bindings( + str(kwargs["task_name"]), + kwargs["scene_objects"], + bindings, + ) + return GroundedTaskSpec(task_spec, requirements, bindings) + + monkeypatch.setattr( + tasks, + "interpret_and_ground_task_spec", + fake_interpret_and_ground, + ) + real_recipe = tasks.instantiate_seed_graph + + def capture_recipe(task_spec, role_bindings): + recipe_calls.append((task_spec, role_bindings)) + return real_recipe(task_spec, role_bindings) + + monkeypatch.setattr(tasks, "instantiate_seed_graph", capture_recipe) + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + + def fake_renderer(program): + rendered["program"] = program + return b"\x89PNG\r\n\x1a\nseed" + + renderer_module.render_seed_task_graph_png = fake_renderer + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + real_writer = generator.write_generation_artifacts + + def capture_writer(*args, **kwargs): + published["program"] = kwargs["seed_task_graph"] + return real_writer(*args, **kwargs) + + monkeypatch.setattr(generator, "write_generation_artifacts", capture_writer) + output_dir = tmp_path / "configs" + paths = generate_action_engine_config( + gym_export, + output_dir, + task_name="line_task", + task_description="test-instruction", + robot_profile="franka", + ) + + assert planner_call["task_name"] == "line_task" + assert planner_call["task_description"] == "test-instruction" + assert planner_call["robot_profile"] == "franka" + assert len(recipe_calls) == 1 + planner_objects = planner_call["scene_objects"] + assert isinstance(planner_objects, list) + assert {obj["uid"] for obj in planner_objects} == {"table", "interact_can"} + assert {path.name for path in output_dir.iterdir()} == { + "fast_gym_config.json", + "agent_config.json", + "task_spec.json", + "scene_requirements.json", + "seed_task_graph.json", + "seed_task_graph.png", + } + assert paths.seed_task_graph_png.read_bytes() == b"\x89PNG\r\n\x1a\nseed" + assert rendered["program"] is published["program"] + + agent_config = json.loads(paths.agent_config.read_text(encoding="utf-8")) + assert agent_config["schema_version"] == "action_engine_config_v2" + assert agent_config["task_spec"] == "task_spec.json" + assert agent_config["scene_requirements"] == "scene_requirements.json" + assert agent_config["seed_task_graph"] == "seed_task_graph.json" + assert len(agent_config["seed_task_graph_hash"]) == 64 + assert agent_config["runtime_policy"]["schema_version"] == ( + "action_engine_runtime_policy_v6" + ) + assert agent_config["runtime_policy"]["planner"]["dynamic_collision"] is True + assert agent_config["runtime_policy"]["planner"]["static_obstacle_uids"] == [ + "table" + ] + assert agent_config["runtime_policy"]["planner"]["dynamic_obstacle_uids"] == [ + "interact_can" + ] + assert len(agent_config["runtime_policy_hash"]) == 64 + assert "png" not in json.dumps(agent_config).lower() + + from embodichain.gen_sim.action_engine.runtime import ( + load_agent_execution_program, + ) + + regenerated = load_agent_execution_program( + agent_config, + agent_config_path=paths.agent_config, + regenerate=True, + ) + assert regenerated.task == "line_task" + assert regenerated.seed_graph is not None + + +def test_existing_v2_task_spec_bypasses_text_planner_and_derives_scene_requirements( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.gen_sim.action_engine import tasks + + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"direct-task-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + + def unexpected_text_planner(**_kwargs): + raise AssertionError("an existing TaskSpec must not invoke text planning") + + monkeypatch.setattr( + tasks, "interpret_and_ground_task_spec", unexpected_text_planner + ) + input_path = tmp_path / "task_spec.json" + input_path.write_text( + json.dumps(_existing_v2_task_spec()), + encoding="utf-8", + ) + + paths = generate_action_engine_config( + gym_export, + tmp_path / "generated", + task_name="direct_task", + task_spec=input_path, + robot_profile="ur10", + ) + + persisted_task = json.loads(paths.task_spec.read_text(encoding="utf-8")) + persisted_requirements = json.loads( + paths.scene_requirements.read_text(encoding="utf-8") + ) + assert persisted_task["metadata"]["role_bindings"] == {"object_01": "interact_can"} + assert [item["role_id"] for item in persisted_requirements["objects"]] == [ + "object_01" + ] + assert persisted_requirements["metadata"]["source"] == ("task_spec_role_bindings") + gym_config = json.loads(paths.gym_config.read_text(encoding="utf-8")) + assert ( + gym_config["env"]["dataset"]["lerobot"]["params"]["instruction"]["lang"] + == "test-instruction" + ) + + +def test_existing_v2_task_spec_uses_validated_scene_requirements_sidecar( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"sidecar-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + + input_dir = tmp_path / "task-first" + input_dir.mkdir() + task = _existing_v2_task_spec("sidecar_task") + requirements = { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": "sidecar_task", + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": "red"}, + } + ], + "cameras": [], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {"task_first": True}, + } + (input_dir / "task_spec.json").write_text( + json.dumps(task), + encoding="utf-8", + ) + (input_dir / "scene_requirements.json").write_text( + json.dumps(requirements), + encoding="utf-8", + ) + + paths = generate_action_engine_config( + gym_export, + tmp_path / "generated-sidecar", + task_name="sidecar_task", + task_spec=input_dir / "task_spec.json", + robot_profile="ur10", + ) + + assert json.loads(paths.scene_requirements.read_text(encoding="utf-8")) == ( + requirements + ) + + +def test_task_factory_style_sidecar_binds_roles_without_text_llm( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"task-first-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + source_path = gym_export / "gym_config.json" + source = json.loads(source_path.read_text(encoding="utf-8")) + source["rigid_object"][0]["category"] = "can" + source["rigid_object"][0]["attributes"] = {"color": "red"} + source["rigid_object"][0]["affordances"] = ["graspable", "orientable"] + source["rigid_object"][0]["initial_state"] = {"orientation": "fallen"} + source_path.write_text(json.dumps(source), encoding="utf-8") + + input_dir = tmp_path / "task-first-unbound" + input_dir.mkdir() + task = _existing_v2_task_spec("task_first_unbound") + task["metadata"] = {"fixture": "abstract-task"} + requirements = { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": "task_first_unbound", + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": "red"}, + } + ], + "cameras": [], + "spatial_constraints": [], + "distractor_count": 0, + "metadata": {"task_first": True}, + } + (input_dir / "task_spec.json").write_text(json.dumps(task), encoding="utf-8") + (input_dir / "scene_requirements.json").write_text( + json.dumps(requirements), encoding="utf-8" + ) + + paths = generate_action_engine_config( + gym_export, + tmp_path / "generated-unbound-sidecar", + task_name="task_first_unbound", + task_spec=input_dir / "task_spec.json", + robot_profile="ur10", + ) + + task_artifact = json.loads(paths.task_spec.read_text(encoding="utf-8")) + assert task_artifact["metadata"]["role_bindings"] == {"object_01": "interact_can"} + + +def test_task_spec_input_rejects_natural_language_conflict( + gym_export: Path, + tmp_path: Path, +) -> None: + task = _existing_v2_task_spec() + with pytest.raises(ValueError, match="task_spec cannot be combined"): + generate_action_engine_config( + gym_export, + tmp_path / "conflict-description", + task_name="direct_task", + task_description="do something", + task_spec=task, + robot_profile="ur10", + ) + + +def test_task_spec_role_binding_accepts_legacy_oracle_and_rejects_conflicts() -> None: + task = _existing_v2_task_spec() + task["metadata"] = {} + task["oracle"] = {"role_bindings": {"object_01": "interact_can"}} + assert _task_spec_role_bindings(task, ["table", "interact_can"]) == { + "object_01": "interact_can" + } + + task["metadata"] = {"role_bindings": {"object_01": "table"}} + with pytest.raises(ValueError, match="Conflicting role_bindings"): + _task_spec_role_bindings(task, ["table", "interact_can"]) + + +def test_task_spec_role_binding_merges_non_overlapping_handoffs() -> None: + task = _existing_v2_task_spec() + task["task_instances"][0]["params"]["target_role"] = "object_02" + task["metadata"] = {"role_bindings": {"object_01": "interact_can"}} + task["oracle"] = {"role_bindings": {"object_02": "interact_target"}} + + assert _task_spec_role_bindings( + task, + ["table", "interact_can", "interact_target"], + ) == {"object_01": "interact_can", "object_02": "interact_target"} + + +def test_task_factory_sidecar_requires_static_affordance_and_state_evidence() -> None: + task = _existing_v2_task_spec("missing-static-evidence") + task["metadata"] = {} + requirements = { + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": "red"}, + } + ] + } + scene = [ + { + "runtime_uid": "interact_can", + "description": "A red soda can.", + "init_pos": [0.0, 0.0, 0.7], + } + ] + + with pytest.raises(ValueError, match="requires one unambiguous scene match"): + _task_spec_role_bindings( + task, + ["interact_can"], + scene_requirements=requirements, + scene_objects=scene, + robot_profile="ur10", + ) + + +@pytest.mark.parametrize( + ("scene_metadata", "required_attributes"), + ( + ({}, {}), + ({"category": "can"}, {"color": "red"}), + ), +) +def test_task_factory_sidecar_does_not_infer_semantics_from_description( + scene_metadata: dict, + required_attributes: dict, +) -> None: + task = _existing_v2_task_spec("no-text-evidence") + task["metadata"] = {} + requirements = { + "objects": [ + { + "role_id": "object_01", + "category": "can", + "count": 1, + "affordances": [], + "initial_state": {}, + "attributes": required_attributes, + } + ] + } + scene = [ + { + "runtime_uid": "mystery_object", + "role": "rigid_object", + "description": "A red soda can.", + "init_pos": [0.0, 0.0, 0.7], + **scene_metadata, + } + ] + + with pytest.raises(ValueError, match="requires one unambiguous scene match"): + _task_spec_role_bindings( + task, + ["mystery_object"], + scene_requirements=requirements, + scene_objects=scene, + robot_profile="ur10", + ) + + +def test_ab_generation_writes_shared_and_offline_branch_artifacts( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + renderer_module = ModuleType( + "embodichain.gen_sim.action_engine.graph_visualization" + ) + renderer_module.render_seed_task_graph_png = lambda _program: b"ab-seed-png" + monkeypatch.setitem(sys.modules, renderer_module.__name__, renderer_module) + + output_dir = tmp_path / "ab-config" + paths = generate_action_engine_config( + gym_export, + output_dir, + task_name="ab_task", + task_spec=_existing_v2_task_spec("ab_task"), + robot_profile="ur10", + planning_mode="ab", + vlm_model="mimo-vlm", + ) + + assert paths.seed_task_graph == output_dir / "offline/seed_task_graph.json" + assert paths.seed_task_graph_png == output_dir / "offline/seed_task_graph.png" + assert not (output_dir / "seed_task_graph.json").exists() + agent_config = json.loads(paths.agent_config.read_text(encoding="utf-8")) + assert agent_config["planning_mode"] == "ab" + assert agent_config["offline_seed_task_graph"] == "offline/seed_task_graph.json" + assert agent_config["online_planning"]["vlm_model"] == "mimo-vlm" + scene_requirements = json.loads( + paths.scene_requirements.read_text(encoding="utf-8") + ) + assert [camera["uid"] for camera in scene_requirements["cameras"]] == [ + "vlm_front", + "vlm_left", + "vlm_rear", + "vlm_right", + ] + gym_config = json.loads(paths.gym_config.read_text(encoding="utf-8")) + assert [ + sensor["uid"] + for sensor in gym_config["sensor"] + if sensor["uid"].startswith("vlm_") + ] == ["vlm_front", "vlm_left", "vlm_rear", "vlm_right"] + + +def test_invalid_explicit_task_fails_before_output_asset_materialization( + gym_export: Path, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from embodichain.gen_sim.action_engine import tasks + from embodichain.gen_sim.action_engine.generation import generator + + normalized = False + recipe_called = False + writer_called = False + + def reject_task(**_kwargs): + raise ValueError("object selector is ambiguous") + + def record_normalization(*_args, **_kwargs): + nonlocal normalized + normalized = True + raise AssertionError("normalization must not run after planning failure") + + def unexpected_recipe(*_args, **_kwargs): + nonlocal recipe_called + recipe_called = True + raise AssertionError("recipe must not run after interpretation failure") + + def unexpected_writer(*_args, **_kwargs): + nonlocal writer_called + writer_called = True + raise AssertionError("writer must not run after interpretation failure") + + monkeypatch.setattr(tasks, "interpret_and_ground_task_spec", reject_task) + monkeypatch.setattr(tasks, "instantiate_seed_graph", unexpected_recipe) + monkeypatch.setattr(generator, "normalize_scene_assets", record_normalization) + monkeypatch.setattr(generator, "write_generation_artifacts", unexpected_writer) + output_dir = tmp_path / "invalid" + + with pytest.raises(ValueError, match="ambiguous"): + generate_action_engine_config( + gym_export, + output_dir, + task_name="invalid_task", + task_description="test-instruction", + robot_profile="franka", + ) + + assert normalized is False + assert recipe_called is False + assert writer_called is False + assert not output_dir.exists() + + +def test_agent_config_uses_relative_program_paths(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + config = build_agent_config( + task_name="line_task", + robot_profile="franka", + execution_program_hash="b" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + ) + assert config["task_spec"] == "task_spec.json" + assert config["scene_requirements"] == "scene_requirements.json" + assert config["seed_task_graph"] == "seed_task_graph.json" + assert config["runtime_policy"]["arm_selection"]["pickup_crossing_weight"] == 1.0 + assert config["runtime_policy"]["motion_defaults"]["PickUp"][ + "lift_height" + ] == pytest.approx(0.30) + assert config["runtime_policy"]["grasp"]["max_open_length"] == pytest.approx(0.15) + assert len(config["runtime_policy_hash"]) == 64 + + +def test_agent_config_anchors_absolute_motion_heights_to_tabletop( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + config = build_agent_config( + task_name="high_table_task", + robot_profile="franka", + execution_program_hash="c" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + table_top_z=1.05, + ) + + policy = config["runtime_policy"] + assert policy["motion_defaults"]["MoveEndEffector"][ + "maximum_eef_height" + ] == pytest.approx(1.45) + assert policy["grounding"]["handover"]["maximum_eef_height"] == pytest.approx(1.85) + + +def test_documented_cli_accepts_franka_profile() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/task4_2", + "--task_name", + "task4_2", + "--task_description", + "Arrange the cans in a line.", + "--robot-profile", + "franka", + "--overwrite", + ] + ) + assert args.robot_profile == "franka" + assert args.overwrite is True + + +def test_generation_cli_defaults_to_mature_robot_without_scene_randomization() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/task2_3", + "--task_name", + "task2_3", + "--task_description", + "Upright both objects.", + ] + ) + + assert args.robot_profile == "ur10" + assert args.randomize_scene is False + assert args.planning_mode == "offline" + assert not hasattr(args, "instruction_parser") + assert not hasattr(args, "task_agent") + + +def test_generation_cli_accepts_ab_models() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/ab", + "--task_name", + "ab", + "--task_description", + "test-instruction", + "--planning-mode", + "ab", + "--llm-model", + "text-model", + "--vlm-model", + "vision-model", + ] + ) + + assert args.planning_mode == "ab" + assert args.llm_model == "text-model" + assert args.vlm_model == "vision-model" + + +def test_generation_cli_accepts_existing_task_spec_without_description() -> None: + args = build_parser().parse_args( + [ + "--gym_project", + "gym_export", + "--output_dir", + "configs/direct", + "--task_name", + "direct_task", + "--task-spec", + "tasks/direct_task/task_spec.json", + ] + ) + + assert args.task_spec == "tasks/direct_task/task_spec.json" + assert cli_module._resolve_task_description(args) == "" + + +def test_generation_cli_reports_seed_png_path( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + paths = artifact_paths(tmp_path) + monkeypatch.setattr( + cli_module, + "generate_action_engine_config", + lambda *_args, **_kwargs: paths, + ) + monkeypatch.setattr( + sys, + "argv", + [ + "generate_action_agent_config", + "--gym_project", + "gym_export", + "--output_dir", + str(tmp_path), + "--task_name", + "task4_2", + "--task_description", + "Arrange cans.", + ], + ) + + cli_module.cli() + + assert ( + f"Generated Seed graph PNG: {paths.seed_task_graph_png}" + in capsys.readouterr().out + ) + + +@pytest.mark.parametrize( + "removed_args", + [ + ["--instruction-parser", "llm"], + ["--instruction_parser", "llm"], + ["--task-agent", "task-agent.json"], + ["--task_agent", "task-agent.json"], + ], +) +def test_generation_cli_rejects_removed_arguments(removed_args: list[str]) -> None: + base_args = [ + "--gym-project", + "gym_export", + "--output-dir", + "configs/task", + "--task-name", + "task", + "--task-description", + "Upright the can.", + ] + + with pytest.raises(SystemExit, match="2"): + build_parser().parse_args([*base_args, *removed_args]) + + +def test_removed_python_parameters_are_absent() -> None: + parameters = inspect.signature(generate_action_engine_config).parameters + + assert "instruction_parser" not in parameters + assert "task_agent" not in parameters diff --git a/tests/gen_sim/action_engine/planning/__init__.py b/tests/gen_sim/action_engine/planning/__init__.py new file mode 100644 index 000000000..de3758a5b --- /dev/null +++ b/tests/gen_sim/action_engine/planning/__init__.py @@ -0,0 +1,21 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Action Engine planning tests.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/tests/gen_sim/action_engine/planning/test_linker.py b/tests/gen_sim/action_engine/planning/test_linker.py new file mode 100644 index 000000000..2548a43a3 --- /dev/null +++ b/tests/gen_sim/action_engine/planning/test_linker.py @@ -0,0 +1,407 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy + +import pytest + +from embodichain.gen_sim.action_engine.domain import seed_graph_hash +from embodichain.gen_sim.action_engine.planning.linker import ( + link_seed_graph, + link_task_dependencies, +) +from embodichain.gen_sim.action_engine.protocol import ( + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.runtime.loader import load_execution_program +from embodichain.gen_sim.action_engine.tasks.recipes import instantiate_seed_graph + + +def _handover_task() -> dict: + return { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "handover_then_place", + "level": "L3", + "instruction": "Stand both cans, hand over the purple can, then place it.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "purple", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E2", + "params": { + "object_role": "orange", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_03", + "task_type": "E4", + "params": { + "object_role": "purple", + "transfer_arm": "right_arm", + "receive_arm": "left_arm", + }, + "depends_on": ["task_02"], + "role": "primary", + }, + { + "id": "task_04", + "task_type": "E1", + "params": { + "object_role": "purple", + "target_role": "orange", + "relation": "left_of", + "required_arm": "left_arm", + }, + "depends_on": ["task_03"], + "role": "primary", + }, + ], + "success": {"type": "all_complete"}, + "oracle": {}, + "metadata": {}, + } + + +def _handover_graph() -> dict: + return instantiate_seed_graph( + _handover_task(), + {"purple": "purple_can", "orange": "orange_can"}, + ) + + +def _unlink_for_rebuild(graph: dict) -> None: + graph["metadata"].pop("action_contract_linker", None) + for group in graph["task_groups"]: + group.pop("contract", None) + + +def test_task_linker_preserves_parallel_arms_and_waits_for_both_before_handover() -> ( + None +): + linked = link_task_dependencies( + _handover_task(), + {"purple": "purple_can", "orange": "orange_can"}, + ) + by_id = {item["id"]: item for item in linked["task_instances"]} + + assert by_id["task_01"]["depends_on"] == [] + assert by_id["task_02"]["depends_on"] == [] + assert by_id["task_03"]["depends_on"] == ["task_02", "task_01"] + + +def test_resource_dependency_provenance_is_persisted_in_seed_graph() -> None: + task = _handover_task() + task["task_instances"] = task["task_instances"][:3] + handover = task["task_instances"][2] + handover["params"] = { + "object_role": "orange", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + } + + graph = instantiate_seed_graph( + task, + {"purple": "purple_can", "orange": "orange_can"}, + ) + + provenance = graph["metadata"]["action_contract_task_linker"] + assert provenance["linked_dependencies"] == [ + { + "from": "task_01", + "to": "task_03", + "reason": "resource", + "detail": "arm:right_arm", + } + ] + + +def test_same_object_e2_handover_gets_direct_causal_edge_through_a_chain() -> None: + task = _handover_task() + task["task_instances"][1]["depends_on"] = ["task_01"] + linked = link_task_dependencies( + task, + {"purple": "purple_can", "orange": "orange_can"}, + ) + handover = next( + item for item in linked["task_instances"] if item["id"] == "task_03" + ) + + assert handover["depends_on"] == ["task_02", "task_01"] + + +def test_handover_ownership_flows_through_home_terminal_barrier() -> None: + graph = _handover_graph() + groups = {group["id"]: group for group in graph["task_groups"]} + nodes = {node["id"]: node for node in graph["nodes"]} + handover_group = groups["task_03"] + terminal_id = handover_group["contract"]["terminal_node_ids"][0] + terminal = nodes[terminal_id] + receiver_entry = nodes[groups["task_04"]["contract"]["entry_node_ids"][0]] + handover = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "HandOver" + ) + + assert terminal["atomic_action"] == "MoveJoints" + assert terminal["contract"]["completion"] == "terminal_barrier" + assert terminal["contract"]["failure_policy"] == "best_effort" + retreat = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveEndEffector" + ) + assert retreat["contract"]["failure_policy"] == "safety_required" + assert terminal_id in receiver_entry["depends_on"] + assert { + (effect["op"], effect["atom"]["predicate"], effect["atom"].get("arm")) + for effect in handover["contract"]["effects"] + } >= { + ("delete", "object_held", "right_arm"), + ("add", "object_held", "left_arm"), + } + + +def test_linker_is_idempotent_and_hash_stable() -> None: + graph = _handover_graph() + relinked = link_seed_graph( + graph, + task_order=["task_01", "task_02", "task_03", "task_04"], + known_objects={"purple_can", "orange_can", "table"}, + ) + + assert relinked == graph + assert seed_graph_hash(relinked) == seed_graph_hash(graph) + + +def test_linker_rejects_missing_cleanup_wrong_holder_and_duplicate_pickup() -> None: + missing_cleanup = deepcopy(_handover_graph()) + home = next( + node + for node in missing_cleanup["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveJoints" + ) + missing_cleanup["nodes"].remove(home) + next(group for group in missing_cleanup["task_groups"] if group["id"] == "task_03")[ + "node_ids" + ].remove(home["id"]) + for node in missing_cleanup["nodes"]: + node["depends_on"] = [ + dependency for dependency in node["depends_on"] if dependency != home["id"] + ] + _unlink_for_rebuild(missing_cleanup) + with pytest.raises(ValueError, match="terminal barrier"): + link_seed_graph(missing_cleanup) + + wrong_holder = deepcopy(_handover_graph()) + staging = next( + node + for node in wrong_holder["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveHeldObject" + ) + staging["actor"] = {"mode": "required", "arm": "left_arm"} + staging.pop("contract") + _unlink_for_rebuild(wrong_holder) + with pytest.raises(ValueError, match="no producer|unavailable state"): + link_seed_graph(wrong_holder) + + duplicate_pickup = deepcopy(_handover_graph()) + pickup = next( + node + for node in duplicate_pickup["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "PickUp" + ) + staging = next( + node + for node in duplicate_pickup["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveHeldObject" + ) + repeated = deepcopy(pickup) + repeated["id"] = "task_01__duplicate_pickup" + repeated["depends_on"] = [pickup["id"]] + repeated.pop("contract") + staging["depends_on"] = [repeated["id"]] + group = next( + group for group in duplicate_pickup["task_groups"] if group["id"] == "task_03" + ) + pickup_index = group["node_ids"].index(pickup["id"]) + group["node_ids"].insert(pickup_index + 1, repeated["id"]) + duplicate_pickup["nodes"].insert( + duplicate_pickup["nodes"].index(pickup) + 1, repeated + ) + _unlink_for_rebuild(duplicate_pickup) + with pytest.raises(ValueError, match="requires unavailable state"): + link_seed_graph(duplicate_pickup) + + +def test_unavailable_arm_reports_current_holder_and_requested_object() -> None: + task = _handover_task() + placement = task["task_instances"][3] + placement["params"].update( + { + "object_role": "orange", + "target_role": "purple", + "required_arm": "left_arm", + } + ) + + with pytest.raises( + ValueError, + match=( + "left_arm.*currently holds 'purple_can'.*" "primary object is 'orange_can'" + ), + ): + instantiate_seed_graph( + task, + {"purple": "purple_can", "orange": "orange_can"}, + ) + + +def test_readers_remain_parallel_and_writer_waits_for_both() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "read_write", + "level": "L3", + "instruction": "Inspect a shared target, then manipulate it.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "read_left", + "task_type": "E1", + "params": { + "object_role": "a", + "target_role": "target", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "read_right", + "task_type": "E1", + "params": { + "object_role": "b", + "target_role": "target", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "write_target", + "task_type": "E2", + "params": { + "object_role": "target", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + ], + "success": {}, + "oracle": {}, + "metadata": {}, + } + linked = link_task_dependencies( + task, + {"a": "object_a", "b": "object_b", "target": "shared_target"}, + ) + by_id = {item["id"]: item for item in linked["task_instances"]} + + assert by_id["read_left"]["depends_on"] == [] + assert by_id["read_right"]["depends_on"] == [] + assert by_id["write_target"]["depends_on"] == ["read_left", "read_right"] + + +def test_explicit_distinct_arm_allocation_keeps_auto_groups_parallel() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "allocated_auto", + "level": "L2", + "instruction": "Stand both objects upright in parallel.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "first", + "task_type": "E2", + "params": {"object_role": "first_object"}, + "depends_on": [], + "role": "primary", + }, + { + "id": "second", + "task_type": "E2", + "params": {"object_role": "second_object"}, + "depends_on": [], + "role": "primary", + }, + ], + "success": {}, + "oracle": {}, + "metadata": { + "allocation_groups": [ + { + "id": "distinct_pair", + "task_instance_ids": ["first", "second"], + "arm_constraint": "distinct_arms", + } + ] + }, + } + bindings = {"first_object": "first_uid", "second_object": "second_uid"} + linked = link_task_dependencies(task, bindings) + graph = instantiate_seed_graph(linked, bindings) + + assert all(not item["depends_on"] for item in linked["task_instances"]) + assert all(not group["depends_on"] for group in graph["task_groups"]) + + +def test_v2_and_resolver_mismatch_require_regeneration() -> None: + with pytest.raises( + ValueError, match="lacks persisted Action Contracts.*regenerate" + ): + load_execution_program({"schema_version": "action_engine_seed_graph_v2"}) + + graph = _handover_graph() + graph["nodes"][0]["contract"]["claims"][0]["access"] = "shared_read" + with pytest.raises( + ValueError, match="does not match the current capability resolver" + ): + load_execution_program(graph) + + +def test_seed_graph_schema_is_v3() -> None: + assert _handover_graph()["schema_version"] == SEED_GRAPH_SCHEMA diff --git a/tests/gen_sim/action_engine/planning/test_online_v2.py b/tests/gen_sim/action_engine/planning/test_online_v2.py new file mode 100644 index 000000000..290f667bd --- /dev/null +++ b/tests/gen_sim/action_engine/planning/test_online_v2.py @@ -0,0 +1,480 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from threading import Barrier + +import pytest +import torch + +import embodichain.gen_sim.action_engine.planning.online as online_module +import embodichain.gen_sim.action_engine.planning.planner as planner_module +import embodichain.gen_sim.action_engine.planning.vision as vision_module +from embodichain.gen_sim.action_engine.domain import public_task_spec +from embodichain.gen_sim.action_engine.planning import ( + CameraObservation, + SceneObservation, + analyze_visual_scene, + fuse_seed_graphs, + plan_candidates_parallel, + plan_online_seed_graph, + select_seed_graph, + validate_visual_facts, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_level + + +def _task(level: str, *, reasoning: str | None = None): + task, requirements = make_task_level(level, reasoning=reasoning) + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + return task, requirements, bindings + + +def test_online_planner_sees_public_task_and_returns_complete_seed_graph() -> None: + task, _, bindings = _task("L4", reasoning="visual_semantics") + offline = instantiate_seed_graph(task, bindings) + body = {key: deepcopy(offline[key]) for key in ("nodes", "task_groups", "success")} + for node in body["nodes"]: + node.pop("contract") + for group in body["task_groups"]: + group.pop("contract") + visual_move = next( + node for node in body["nodes"] if node["atomic_action"] == "MoveHeldObject" + ) + visual_move["target_binding"] = { + "kind": "visual_constraint", + "camera_uid": "front", + "normalized_keypoint": [0.2, 0.3], + } + camera = CameraObservation( + "front", + torch.zeros((8, 8, 3), dtype=torch.uint8), + None, + None, + None, + ) + observation = SceneObservation( + (camera,), + tuple({"uid": uid} for uid in bindings.values()), + ) + uid = next(iter(bindings.values())) + facts = { + "entities": [ + { + "uid": uid, + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "keypoints": {"center": [0.2, 0.3]}, + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + prompts = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return body + + graph, observed_facts = plan_online_seed_graph( + public_task_spec(task), + observation, + visual_facts=facts, + graph_caller=caller, + ) + + assert graph["planner_route"] == "online" + assert observed_facts == facts + assert "oracle" not in prompts[0] + assert '"task_instances"' not in prompts[0] + assert '"E4"' in prompts[0] + assert "Transfer one held object" in prompts[0] + assert graph["metadata"]["oracle_exposed"] is False + assert any( + node["target_binding"]["kind"] == "visual_constraint" for node in graph["nodes"] + ) + + +def test_offline_and_online_candidates_plan_concurrently_with_isolated_views() -> None: + task, _, bindings = _task("L1") + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + barrier = Barrier(2) + views = {} + + def offline_planner(*, task_spec): + views["offline"] = task_spec + barrier.wait(timeout=2.0) + return offline + + def online_planner(*, task_spec): + views["online"] = task_spec + barrier.wait(timeout=2.0) + return online + + pair = plan_candidates_parallel( + task, + offline_planner=offline_planner, + online_planner=online_planner, + ) + + assert "oracle" in views["offline"] + assert "oracle" not in views["online"] + assert pair.offline["planner_route"] == "offline" + assert pair.online["planner_route"] == "online" + + +def test_visual_facts_reject_unknown_uid_and_out_of_range_keypoint() -> None: + value = { + "entities": [ + { + "uid": "unknown", + "camera_uid": "front", + "bbox": [0.0, 0.0, 1.2, 1.0], + "keypoints": {}, + "confidence": 1.0, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 1.0, + } + with pytest.raises(ValueError, match="unknown UID"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_facts_reject_visible_entity_without_image_evidence() -> None: + value = { + "entities": [ + { + "uid": "known", + "camera_uid": "front", + "visible": True, + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="bbox or keypoint"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_facts_reject_non_numeric_image_coordinates() -> None: + value = { + "entities": [ + { + "uid": "known", + "camera_uid": "front", + "bbox": ["0.1", 0.2, 0.3, 0.4], + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="must be numeric"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_fact_caller_receives_rgb_depth_and_calibration_evidence() -> None: + task, _, bindings = _task("L1") + uid = next(iter(bindings.values())) + observation = SceneObservation( + ( + CameraObservation( + "front", + torch.zeros((4, 5, 3), dtype=torch.uint8), + torch.linspace(0.0, 1.0, 20, dtype=torch.float32).reshape(4, 5), + torch.eye(3), + torch.eye(4), + ), + ), + ({"uid": uid},), + ) + captured = {} + + def caller(**kwargs): + captured.update(kwargs) + return { + "entities": [ + { + "uid": uid, + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + facts = analyze_visual_scene(observation, task, caller=caller) + + assert facts["entities"][0]["uid"] == uid + assert len(captured["images"]) == 2 + assert '"depth_image_index": 1' in captured["prompt"] + assert '"intrinsics": [[1.0, 0.0, 0.0]' in captured["prompt"] + assert captured["schema"]["properties"]["task_predicates"]["maxItems"] == 0 + + +def test_visual_task_predicates_are_limited_to_the_current_task() -> None: + task, _, bindings = _task("L4", reasoning="visual_semantics") + uid = next(iter(bindings.values())) + observation = SceneObservation( + ( + CameraObservation( + "front", + torch.zeros((4, 5, 3), dtype=torch.uint8), + None, + None, + None, + ), + ), + ({"uid": uid},), + ) + captured = {} + + def caller(**kwargs): + captured.update(kwargs) + return { + "entities": [ + { + "uid": uid, + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [{"type": "mouth_completed", "confidence": 0.9}], + "confidence": 0.9, + } + + facts = analyze_visual_scene(observation, task, caller=caller) + + predicate_type = captured["schema"]["properties"]["task_predicates"]["items"][ + "properties" + ]["type"] + assert predicate_type["enum"] == ["mouth_completed"] + assert facts["task_predicates"][0]["type"] == "mouth_completed" + + +def test_visual_facts_reject_unrequested_task_predicate() -> None: + value = { + "entities": [], + "relations": [], + "task_predicates": [{"type": "mouth_completed", "confidence": 0.9}], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="task_predicates.*must be one of"): + validate_visual_facts( + value, + known_uids={"known"}, + camera_uids={"front"}, + ) + + +def test_production_online_graph_caller_receives_reset_time_multiview_evidence( + monkeypatch: pytest.MonkeyPatch, +) -> None: + observation = SceneObservation( + ( + CameraObservation( + "front", + torch.zeros((4, 5, 3), dtype=torch.uint8), + torch.zeros((4, 5), dtype=torch.float32), + torch.eye(3), + torch.eye(4), + ), + ), + ({"uid": "known"},), + ) + captured = {} + + def caller(**kwargs): + captured.update(kwargs) + return {"nodes": [], "task_groups": [], "success": {}} + + monkeypatch.setattr(vision_module, "_default_structured_caller", caller) + monkeypatch.setattr(vision_module, "_vlm_model", lambda model: f"resolved:{model}") + + result = online_module._default_graph_caller( + prompt="plan", + schema={"type": "object"}, + model="mimo", + observation=observation, + ) + + assert result == {"nodes": [], "task_groups": [], "success": {}} + assert captured["model"] == "resolved:mimo" + assert len(captured["images"]) == 2 + + +def test_default_vision_caller_disables_custom_socket_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import langchain_openai + + captured = {} + + class FakeRunnable: + def invoke(self, _messages): + return {"facts": []} + + class FakeChatOpenAI: + def __init__(self, **kwargs): + captured.update(kwargs) + + def with_structured_output(self, _schema, **_kwargs): + return FakeRunnable() + + monkeypatch.setattr(langchain_openai, "ChatOpenAI", FakeChatOpenAI) + monkeypatch.setattr( + planner_module, + "_load_llm_settings", + lambda *, model: { + "api_key": "test-key", + "model": model or "test-model", + "base_url": "https://example.test/v1", + "default_query": {}, + }, + ) + + vision_module._default_structured_caller( + prompt="inspect", + images=(), + schema={"type": "object"}, + model="test-model", + ) + + assert captured["http_socket_options"] == () + + +def test_visual_facts_reject_unstructured_entity_fields() -> None: + value = { + "entities": [ + { + "uid": "known", + "camera_uid": "front", + "bbox": [0.1, 0.2, 0.3, 0.4], + "semantic_label": "can", + "confidence": 0.9, + } + ], + "relations": [], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="unsupported fields"): + validate_visual_facts(value, known_uids={"known"}, camera_uids={"front"}) + + +def test_visual_facts_reject_noncanonical_relation_type() -> None: + value = { + "entities": [], + "relations": [ + {"type": "obstructs", "uids": ["box", "sign"], "confidence": 0.9} + ], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="relation type"): + validate_visual_facts( + value, + known_uids={"box", "sign"}, + camera_uids={"front"}, + ) + + +def test_visual_facts_require_ordered_relation_participants() -> None: + value = { + "entities": [], + "relations": [{"type": "occludes", "uids": ["box"], "confidence": 0.9}], + "task_predicates": [], + "confidence": 0.9, + } + + with pytest.raises(ValueError, match="exactly 2 UIDs"): + validate_visual_facts( + value, + known_uids={"box"}, + camera_uids={"front"}, + ) + + +def test_selection_prefers_exact_offline_and_l4_online() -> None: + task, _, bindings = _task("L1") + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + selected, evaluations = select_seed_graph( + offline, + online, + task, + known_objects=set(bindings.values()) | {"table"}, + exact_template_match=True, + ) + assert selected["metadata"]["selected_from"] == "offline" + assert evaluations["offline"].score > evaluations["online"].score + + l4, _, l4_bindings = _task("L4", reasoning="logic") + l4_offline = instantiate_seed_graph(l4, l4_bindings) + l4_online = deepcopy(l4_offline) + l4_online["planner_route"] = "online" + selected, _ = select_seed_graph( + l4_offline, + l4_online, + l4, + known_objects=set(l4_bindings.values()) | {"table"}, + visual_confidence=0.95, + ) + assert selected["metadata"]["selected_from"] == "online" + + +def test_fusion_keeps_whole_task_groups() -> None: + task, _, bindings = _task("L2") + offline = instantiate_seed_graph(task, bindings) + online = deepcopy(offline) + online["planner_route"] = "online" + routes = { + group["id"]: ("offline" if index % 2 == 0 else "online") + for index, group in enumerate(offline["task_groups"]) + } + fused = fuse_seed_graphs(offline, online, routes) + + assert fused["planner_route"] == "fused" + assert all( + all(node_id.startswith(routes[group["id"]]) for node_id in group["node_ids"]) + for group in fused["task_groups"] + ) diff --git a/tests/gen_sim/action_engine/planning/test_planner.py b/tests/gen_sim/action_engine/planning/test_planner.py new file mode 100644 index 000000000..c089ccef8 --- /dev/null +++ b/tests/gen_sim/action_engine/planning/test_planner.py @@ -0,0 +1,730 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import inspect +import json +from pathlib import Path +from typing import Any + +import pytest + +from embodichain.gen_sim.action_engine.domain import TASK_AGENT_SCHEMA +from embodichain.gen_sim.action_engine.planning import plan_task +from embodichain.gen_sim.action_engine.planning import planner as planner_module + + +def _scene() -> list[dict[str, Any]]: + return [ + { + "uid": "table", + "runtime_uid": "table", + "source_uid": "table", + "role": "background", + "description": "A table.", + }, + *[ + { + "uid": f"interact_soda_can_{index}_0", + "runtime_uid": f"interact_soda_can_{index}", + "source_uid": f"interact_soda_can_{index}_0", + "role": "rigid_object", + "description": "An aluminum soda can.", + } + for index in range(5) + ], + ] + + +def _dual_arm_scene() -> list[dict[str, Any]]: + return [ + { + "uid": uid, + "runtime_uid": uid, + "source_uid": uid, + "role": "rigid_object", + "description": description, + } + for uid, description in ( + ("cube", "A cube on the left side of the table."), + ("cup", "A paper cup on the right side of the table."), + ("basket", "A basket near the center of the table."), + ) + ] + + +def _stack_scene() -> list[dict[str, Any]]: + return [ + { + "uid": uid, + "runtime_uid": uid, + "source_uid": f"{uid}_0", + "role": role, + "description": description, + } + for uid, role, description in ( + ("table", "background", "A table."), + ("paper_cup", "rigid_object", "A paper cup."), + ("popcorn_bucket", "rigid_object", "A popcorn bucket."), + ("earbuds_case", "rigid_object", "A blue earbuds case."), + ) + ] + + +def test_injected_planner_returns_only_semantics_and_resolves_aliases() -> None: + observed: dict[str, Any] = {} + + def caller(*, prompt: str, model: str | None) -> dict[str, Any]: + observed.update(prompt=prompt, model=model) + return { + "semantic_steps": [ + { + "id": "s01_place", + "operator": "place_relative", + "object": "interact_soda_can_0_0", + "goal": {"reference_object": "table", "relation": "on"}, + }, + { + "id": "s02_orient", + "operator": "orient_object", + "object": "interact_soda_can_1_0", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + }, + ] + } + + program = plan_task( + task_name="injected", + task_description="Place one object and then orient another.", + scene_objects=_scene(), + model="test-model", + llm_caller=caller, + ) + + assert program["schema_version"] == TASK_AGENT_SCHEMA + assert program["semantic_steps"][0]["object"] == "interact_soda_can_0" + assert program["semantic_steps"][1]["depends_on"] == ["s01_place"] + assert "Do not select a task route" in observed["prompt"] + assert observed["model"] == "test-model" + + +def test_planner_repairs_a_non_visible_skill_once() -> None: + calls = 0 + + def caller(**_kwargs: Any) -> dict[str, Any]: + nonlocal calls + calls += 1 + if calls == 1: + return { + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "goal": {}, + } + ] + } + return { + "semantic_steps": [ + { + "id": "s01_place_cube", + "operator": "place_relative", + "object": "cube", + "goal": { + "reference_object": "basket", + "relation": "inside", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + "depends_on": [], + }, + { + "id": "s02_place_cup", + "operator": "place_relative", + "object": "cup", + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": [], + }, + ], + "allocation_groups": [ + { + "id": "dual_arms_1", + "semantic_step_ids": ["s01_place_cube", "s02_place_cup"], + "arm_constraint": "distinct_arms", + } + ], + } + + program = plan_task( + task_name="dual_arm_basket", + task_description="test-instruction", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + assert [ + (step["id"], step["operator"], step["object"], step["depends_on"]) + for step in program["semantic_steps"] + ] == [ + ("s01_place_cube", "place_relative", "cube", []), + ("s02_place_cup", "place_relative", "cup", []), + ] + assert calls == 2 + assert program["allocation_groups"][0]["arm_constraint"] == "distinct_arms" + + +def test_planner_repairs_build_stack_singular_object_contract() -> None: + prompts: list[str] = [] + + def caller(*, prompt: str, **_kwargs: Any) -> dict[str, Any]: + prompts.append(prompt) + if len(prompts) == 1: + return { + "semantic_steps": [ + { + "id": "s01_build_stack", + "operator": "build_stack", + "object": "paper_cup", + "goal": { + "anchor": "popcorn_bucket", + "stack_mode": "on_top", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + return { + "semantic_steps": [ + { + "id": "s01_build_stack", + "operator": "build_stack", + "objects": ["paper_cup", "earbuds_case"], + "goal": { + "anchor": "popcorn_bucket", + "stack_mode": "on_top", + "orientation_goal": "preserve", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="task3_2", + task_description="test-instruction", + scene_objects=_stack_scene(), + llm_caller=caller, + ) + + assert len(prompts) == 2 + assert "build_stack requires an 'objects' list" in prompts[1] + assert program["semantic_steps"][0]["objects"] == [ + "paper_cup", + "earbuds_case", + ] + assert program["semantic_steps"][0]["goal"]["anchor"] == "popcorn_bucket" + + +def test_planner_rejects_a_non_visible_skill_after_one_repair() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": ["s01_hold"], + }, + ] + } + + with pytest.raises(ValueError, match="after one repair"): + plan_task( + task_name="conflicting_arms", + task_description="Hold the cube with the left arm, then place it " + "with the right arm.", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + +def test_spatial_two_sided_phrase_does_not_invent_arm_constraint() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_left", + "operator": "orient_object", + "object": "cube", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + }, + { + "id": "s02_right", + "operator": "orient_object", + "object": "cup", + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + }, + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="two_sided_upright", + task_description="test-instruction", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + assert program["allocation_groups"] == [] + + +def test_planner_does_not_infer_arm_group_from_instruction_text() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": step_id, + "operator": "orient_object", + "object": object_uid, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for step_id, object_uid in (("s01", "cube"), ("s02", "cup")) + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="explicit_both_arms", + task_description="test-instruction", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + assert program["allocation_groups"] == [] + + +def test_planner_does_not_expose_internal_operator_contracts() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_hold", + "operator": "hold_hover", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s02_place", + "operator": "place_relative", + "object": "cube", + "actor": {"mode": "auto"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": ["s01_hold"], + }, + ] + } + + with pytest.raises(ValueError, match="after one repair"): + plan_task( + task_name="nondefault_hover", + task_description="Hold the cube in a special pose, then place it.", + scene_objects=_dual_arm_scene(), + llm_caller=caller, + ) + + +def test_plan_task_has_no_rule_fallback_parameter() -> None: + assert "deterministic_fallback" not in inspect.signature(plan_task).parameters + + +def test_arrange_line_preserves_structured_orientation_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "upright", + "orientation_axis": "long_axis", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="neutral_line", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + goal = program["semantic_steps"][0]["goal"] + assert goal["orientation_goal"] == "upright" + assert goal["orientation_axis"] == "long_axis" + + +def test_arrange_line_preserves_structured_axis_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "anchor": "table_center", + "axis": "world_x", + "order_constraint": "free", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="ambiguous_line_axis", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + assert program["semantic_steps"][0]["goal"]["axis"] == "world_x" + + +def test_instruction_text_does_not_override_structured_axis_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "anchor": "table_center", + "axis": "world_y", + "order_constraint": "free", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="front_to_back_line_axis", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + assert program["semantic_steps"][0]["goal"]["axis"] == "world_y" + + +def test_arrange_line_preserves_explicit_orientation_request() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return { + "semantic_steps": [ + { + "id": "s01_line", + "operator": "arrange_line", + "objects": [ + "interact_soda_can_0", + "interact_soda_can_1", + ], + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ], + "allocation_groups": [], + } + + program = plan_task( + task_name="upright_line", + task_description="test-instruction", + scene_objects=_scene(), + llm_caller=caller, + ) + + assert program["semantic_steps"][0]["goal"]["orientation_goal"] == "upright" + + +def test_planner_rejects_route_or_graph_output() -> None: + def caller(**_kwargs: Any) -> dict[str, Any]: + return {"route": "arrangement_line", "semantic_steps": []} + + with pytest.raises(ValueError, match="only 'semantic_steps'"): + plan_task( + task_name="bad", + task_description="Arrange objects.", + scene_objects=_scene(), + llm_caller=caller, + ) + + +def test_llm_settings_read_gen_sim_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "\n".join( + ( + "# Local Action Engine credentials", + 'export OPENAI_API_KEY="dotenv-key"', + "OPENAI_BASE_URL=https://dotenv.example/v1/", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + config_path = tmp_path / "gen_config.json" + config_path.write_text( + json.dumps( + { + "llm": { + "openai_compatible": { + "api_key": "json-key", + "base_url": "https://json.example/v1", + "model": "json-model", + "default_query": {"api-version": "test"}, + } + } + } + ), + encoding="utf-8", + ) + monkeypatch.setattr(planner_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr(planner_module, "_GEN_CONFIG_PATH", config_path) + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_MODEL", + "LLM_MODEL", + "LLM_URL", + ): + monkeypatch.delenv(name, raising=False) + + settings = planner_module._load_llm_settings(model=None) + + assert settings == { + "api_key": "dotenv-key", + "base_url": "https://dotenv.example/v1", + "model": "dotenv-model", + "default_query": {"api-version": "test"}, + } + + +def test_process_environment_and_model_argument_override_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "\n".join( + ( + "OPENAI_API_KEY=dotenv-key", + "OPENAI_BASE_URL=https://dotenv.example/v1", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + missing_config = tmp_path / "missing.json" + monkeypatch.setattr(planner_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr(planner_module, "_GEN_CONFIG_PATH", missing_config) + monkeypatch.setenv("OPENAI_API_KEY", "shell-key") + monkeypatch.setenv("OPENAI_API_BASE", "https://shell.example/v1/") + monkeypatch.setenv("LLM_MODEL", "shell-model") + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_MODEL", raising=False) + + settings = planner_module._load_llm_settings(model="argument-model") + + assert settings["api_key"] == "shell-key" + assert settings["base_url"] == "https://shell.example/v1" + assert settings["model"] == "argument-model" + + +def test_partial_process_transport_does_not_mix_with_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + env_path.write_text( + "\n".join( + ( + "OPENAI_API_KEY=dotenv-key", + "OPENAI_BASE_URL=https://dotenv.example/v1", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + monkeypatch.setattr(planner_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr( + planner_module, + "_GEN_CONFIG_PATH", + tmp_path / "missing.json", + ) + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "OPENAI_MODEL", + "LLM_MODEL", + "LLM_URL", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-process-key") + + settings = planner_module._load_llm_settings(model=None) + + assert settings["api_key"] == "dotenv-key" + assert settings["base_url"] == "https://dotenv.example/v1" + assert settings["model"] == "dotenv-model" + + +def test_default_llm_caller_disables_custom_socket_transport( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import langchain_openai + + captured: dict[str, Any] = {} + + class FakeRunnable: + def invoke(self, _messages: Any) -> dict[str, list[Any]]: + return {"semantic_steps": [], "allocation_groups": []} + + class FakeChatOpenAI: + def __init__(self, **kwargs: Any) -> None: + captured.update(kwargs) + + def with_structured_output( + self, + _schema: dict[str, Any], + **_kwargs: Any, + ) -> FakeRunnable: + return FakeRunnable() + + monkeypatch.setattr(langchain_openai, "ChatOpenAI", FakeChatOpenAI) + monkeypatch.setattr( + planner_module, + "_load_llm_settings", + lambda *, model: { + "api_key": "test-key", + "model": model or "test-model", + "base_url": "https://example.test/v1", + "default_query": {}, + }, + ) + + planner_module._default_llm_caller(prompt="plan", model="test-model") + + assert captured["http_socket_options"] == () + + +def test_structured_output_transport_selects_json_mode_only_for_mimo() -> None: + calls: list[dict[str, Any]] = [] + + class FakeClient: + def with_structured_output(self, schema: dict[str, Any], **kwargs: Any) -> str: + calls.append({"schema": schema, "kwargs": kwargs}) + return "structured" + + schema = {"type": "object"} + client = FakeClient() + mimo = planner_module._structured_output_runnable( + client, + schema, + settings={ + "model": "mimo-v2.5", + "base_url": "https://token-plan-cn.xiaomimimo.com/v1", + }, + ) + generic = planner_module._structured_output_runnable( + client, + schema, + settings={"model": "gpt-test", "base_url": "https://example.test/v1"}, + ) + + assert mimo == generic == "structured" + assert [call["kwargs"] for call in calls] == [ + {"method": "json_mode"}, + {"method": "json_schema"}, + ] diff --git a/tests/gen_sim/action_engine/runtime/__init__.py b/tests/gen_sim/action_engine/runtime/__init__.py new file mode 100644 index 000000000..66ef3ba1a --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Runtime contract tests for Action Engine.""" diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py new file mode 100644 index 000000000..ba2d35951 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -0,0 +1,835 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Focused contracts for the public atomic-action adapter.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime import actions +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.models import ( + ActionOutcome, + GroundedAction, +) +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.lab.sim.atomic_actions import ( + Affordance, + ActionBinding, + ActionPlan, + AntipodalAffordance, + CoordinatedPickGoal, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectState, + JointPositionGoal, + ObjectSemantics, + PlannerDiagnostics, + RecoveryPolicy, + RuntimeCommandFrame, + SceneSnapshot, + StateDelta, + TimedCommandSequence, + TimedTrajectory, +) +from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.toolkits.graspkit.pg_grasp import GraspGeneratorCfg + + +class _MeshEntity: + def get_vertices(self, *, env_ids: list[int], scale: bool) -> torch.Tensor: + assert env_ids == [0] + assert scale + return torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + ], + dtype=torch.float32, + ) + + def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: + assert env_ids == [0] + return torch.tensor([[0, 1, 2]], dtype=torch.int64) + + +class _PoseEntity: + def __init__(self, pose: torch.Tensor) -> None: + self.pose = pose + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return self.pose.clone() + + +class _PlannerRobot: + uid = "test_robot" + dof = 8 + + _ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": [2, 3], + "physical_right_arm": [4, 5], + "physical_right_eef": [6, 7], + } + control_parts = _ids + + def get_joint_ids(self, *, name: str) -> list[int]: + return list(self._ids[name]) + + def get_control_part_base_pose(self, *, name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).repeat(2, 1, 1) + pose[:, 1, 3] = 0.3 if name == "physical_left_arm" else -0.3 + return pose + + +def _commands_for(trajectory: TimedTrajectory) -> TimedCommandSequence: + """Build timing-only frames for retained test trajectories.""" + active = torch.ones( + trajectory.batch_size, + dtype=torch.bool, + device=trajectory.positions.device, + ) + frames = tuple( + RuntimeCommandFrame( + commands=(), + active_mask=active, + env_ids=trajectory.env_ids, + hold_duration=trajectory.dt[:, index], + ) + for index in range(trajectory.waypoint_count) + ) + return TimedCommandSequence(frames=frames, env_ids=trajectory.env_ids) + + +class _FakeEngine: + """Minimal endpoint-binding and planning surface for adapter unit tests.""" + + binding_owner_id = "action-engine-test" + + def __init__(self, plan=None) -> None: + self._plan = plan + + def bind_control_parts(self, _skill_id, _endpoints) -> ActionBinding: + return ActionBinding(owner_id=self.binding_owner_id) + + def plan(self, invocation, context) -> ActionPlan: + if self._plan is None: + raise AssertionError("This fake engine has no planning callback.") + return self._plan(invocation, context) + + +def _planner_env( + *, + table: Any | None = None, + rigid_objects: dict[str, Any] | None = None, +) -> SimpleNamespace: + entities = dict(rigid_objects or {}) + if table is not None: + entities["table"] = table + return SimpleNamespace( + num_envs=2, + device=torch.device("cpu"), + robot=_PlannerRobot(), + sim=SimpleNamespace(get_rigid_object=entities.get), + left_arm_joints=[0, 1], + left_eef_joints=[2, 3], + right_arm_joints=[4, 5], + right_eef_joints=[6, 7], + open_state=torch.zeros(2), + close_state=torch.ones(2), + get_agent_arm_control_part=lambda is_left: ( + "physical_left_arm" if is_left else "physical_right_arm" + ), + get_agent_eef_control_part=lambda is_left: ( + "physical_left_eef" if is_left else "physical_right_eef" + ), + ) + + +def test_semantics_prewarms_vhacd_cache_before_affordance( + monkeypatch: Any, +) -> None: + """The lazy shared checker must see V-HACD's pickle, never create CoACD.""" + events: list[str] = [] + observed: dict[str, Any] = {} + entity = _MeshEntity() + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + sim=SimpleNamespace( + get_rigid_object=lambda uid: entity if uid == "cube" else None + ), + agent_grasp_runtime_defaults={"max_decomposition_hulls": 8}, + ) + + def fake_prepare(**kwargs: Any) -> SimpleNamespace: + events.append("cache") + observed.update(kwargs) + return SimpleNamespace(status="hit") + + def fake_affordance(**kwargs: Any) -> Affordance: + events.append("affordance") + observed["generator_cfg"] = kwargs["generator_cfg"] + observed["gripper_collision_cfg"] = kwargs["gripper_collision_cfg"] + return Affordance() + + monkeypatch.setattr( + actions, + "ensure_vhacd_grasp_collision_cache", + fake_prepare, + ) + monkeypatch.setattr(actions, "AntipodalAffordance", fake_affordance) + + adapter = AtomicActionAdapter(env) + first = adapter.semantics("cube") + second = adapter.semantics("cube") + + assert first is second + assert events == ["cache", "affordance"] + assert observed["max_decomposition_hulls"] == 8 + assert observed["mesh_vertices"].dtype == torch.float32 + assert observed["mesh_triangles"].dtype == torch.int64 + assert observed["generator_cfg"].n_deviated_approach_directions == 4 + assert observed["gripper_collision_cfg"] is not None + + +def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None: + adapter = AtomicActionAdapter(_planner_env()) + adapter._atomic_engine = _FakeEngine() + goal = JointPositionGoal(target=torch.zeros(2, 2)) + + single = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "arm", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + coordinated_goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="tray", + geometry={}, + affordance=AntipodalAffordance(), + ), + object_target_pose=torch.eye(4), + object_initial_pose=torch.eye(4), + ) + coordinated = adapter._invocation( + GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + coordinated_goal, + {}, + ), + adapter.capabilities.get("CoordinatedPickment"), + ) + hand = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "hand", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + + assert adapter.planner_policy["backend"] == "curobo" + assert single.motion_policy.strategy == "motion_gen" + assert coordinated.motion_policy.strategy == "ik_interp" + assert torch.allclose( + coordinated.skill_options.left_to_right_arm_direction, + torch.tensor([0.0, -1.0, 0.0]), + ) + assert hand.motion_policy.strategy == "ik_interp" + + +def test_coordinated_pickment_scopes_ground_filter_to_gensim_goal_copy() -> None: + adapter = AtomicActionAdapter(_planner_env()) + adapter._atomic_engine = _FakeEngine() + original_cfg = GraspGeneratorCfg(is_filter_ground_collision=True) + affordance = AntipodalAffordance(generator_cfg=original_cfg) + goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="tray", + geometry={}, + affordance=affordance, + ), + object_target_pose=torch.eye(4), + object_initial_pose=torch.eye(4), + ) + grounded = GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + goal, + { + "middle_empty_ratio": 0.7, + "is_filter_ground_collision": False, + }, + ) + + invocation = adapter._invocation( + grounded, + adapter.capabilities.get("CoordinatedPickment"), + ) + + scoped_affordance = invocation.goal.semantics.affordance + assert isinstance(scoped_affordance, AntipodalAffordance) + assert scoped_affordance is not affordance + assert affordance.generator_cfg is original_cfg + assert original_cfg.is_filter_ground_collision is True + assert scoped_affordance.generator_cfg is not original_cfg + assert scoped_affordance.generator_cfg.is_filter_ground_collision is False + assert invocation.skill_options.middle_empty_ratio == pytest.approx(0.7) + + +def test_retreat_uses_row_local_motion_planner_reachability_search( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.05 + requested = reference.clone() + requested[:, 2, 3] = 1.35 + height_thresholds = torch.tensor([1.24, 1.00]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + height_reachable = target[:, 2, 3] <= height_thresholds + baseward_reachable = target[:, 1, 3] < -0.05 + success = height_reachable | baseward_reachable + terminal = target[:, 2, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "retreat_height": 0.30, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_reference_pose": reference, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert len(attempted_targets) > 1 + assert bool(outcome.success.all()) + selected_z = outcome.grounded.target.xpos[:, 2, 3] + assert selected_z.tolist() == pytest.approx([1.20, 1.35]) + assert outcome.grounded.target.xpos[:, 1, 3].tolist() == pytest.approx([0.0, -0.10]) + search = outcome.planner_trace["reachability_search"] + assert search["strategy"] == "bounded_motion_planner" + assert search["selected_target_z"].tolist() == pytest.approx([1.20, 1.35]) + assert len(search["attempts"]) == len(attempted_targets) + + +def test_curobo_generator_receives_generated_static_obstacles( + monkeypatch: Any, +) -> None: + table = object() + can = object() + captured: dict[str, Any] = {} + + def fake_motion_generator(*, cfg: Any) -> object: + captured["cfg"] = cfg + return object() + + monkeypatch.setattr(actions, "MotionGenerator", fake_motion_generator) + adapter = AtomicActionAdapter( + _planner_env(table=table, rigid_objects={"can": can}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["can"], + }, + ) + + generator = adapter._generator() + + assert generator is adapter._motion_generator + planner = captured["cfg"].planner_cfg + assert isinstance(planner, CuroboPlannerCfg) + assert planner.world.rigid_objects == {"table": table, "can": can} + assert planner.world.dynamic_obstacle_names == ["can"] + assert planner.world.obstacle_representation == "cuboid" + assert planner.world.collision_cache == {"cuboid": 8, "mesh": 2} + + +def test_curobo_generator_sizes_collision_cache_for_large_scene( + monkeypatch: Any, +) -> None: + rigid_objects = {f"object_{index:02d}": object() for index in range(13)} + captured: dict[str, Any] = {} + + def fake_motion_generator(*, cfg: Any) -> object: + captured["cfg"] = cfg + return object() + + monkeypatch.setattr(actions, "MotionGenerator", fake_motion_generator) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=rigid_objects), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(rigid_objects), + }, + ) + + adapter._generator() + + planner = captured["cfg"].planner_cfg + assert planner.world.collision_cache == {"cuboid": 13, "mesh": 2} + + +def test_dynamic_scene_parks_contact_target_and_held_rows() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + actual[:, 2, 3] = torch.tensor([0.7, 0.8]) + entities = {uid: _PoseEntity(actual.clone()) for uid in ("target", "held", "other")} + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=entities), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(entities), + }, + ) + held_semantics = ObjectSemantics( + label="held", + entity=entities["held"], + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=held_semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + env_mask=torch.tensor([True, False]), + ) + state = ExecutionState( + last_qpos=torch.zeros(2, 8), + held_objects={"physical_left_arm": held}, + ) + grounded = GroundedAction( + "PickUp", + "right_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="target", + ) + + scene = adapter._scene_snapshot(grounded, state) + + assert torch.equal( + scene.entities["target"].pose[:, 2, 3], + actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET, + ) + assert scene.entities["held"].pose[0, 2, 3] == ( + actual[0, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET + ) + assert scene.entities["held"].pose[1, 2, 3] == actual[1, 2, 3] + assert torch.equal(scene.entities["other"].pose, actual) + + +def test_released_object_returns_to_live_dynamic_collision_pose() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + actual[:, 0, 3] = torch.tensor([0.2, 0.4]) + entity = _PoseEntity(actual) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"released": entity}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["released"], + }, + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="released", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert torch.equal(scene.entities["released"].pose, actual) + + +def test_default_scene_provider_advances_only_after_material_change() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + entity = _PoseEntity(actual.clone()) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"can": entity}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["can"], + }, + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="can", + ) + state = ExecutionState(last_qpos=torch.zeros(2, 8)) + + first = adapter._scene_snapshot(grounded, state) + unchanged = adapter._scene_snapshot(grounded, state) + entity.pose[:, 0, 3] += 0.1 + changed = adapter._scene_snapshot(grounded, state) + + assert first.version == unchanged.version == 0 + assert changed.version == 1 + assert changed.collision_world_revisions(2) == (1, 1) + + +def test_external_scene_provider_is_used_by_planning_snapshot() -> None: + pose = torch.eye(4).repeat(2, 1, 1) + + class _Provider: + def snapshot(self, *, timestamp: float, env_ids: torch.Tensor) -> SceneSnapshot: + assert timestamp == 0.0 + assert torch.equal(env_ids, torch.tensor([0, 1])) + return SceneSnapshot( + timestamp=timestamp, + version=7, + entities={"can": actions.EntityState(pose)}, + ) + + adapter = AtomicActionAdapter( + _planner_env(), + scene_provider=_Provider(), + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="can", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert scene.version == 7 + assert torch.equal(scene.entities["can"].pose, pose) + + +def test_start_session_delegates_to_shared_atomic_engine(monkeypatch: Any) -> None: + adapter = AtomicActionAdapter(_planner_env()) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ) + state = ExecutionState(last_qpos=torch.zeros(2, 8)) + marker = object() + captured: dict[str, Any] = {} + + monkeypatch.setattr(adapter, "_planning_context", lambda *_args: "context") + monkeypatch.setattr(adapter, "_invocation", lambda *_args: "invocation") + + class _Engine: + def start(self, invocations: tuple[Any, ...], context: Any) -> object: + captured["invocations"] = invocations + captured["context"] = context + return marker + + monkeypatch.setattr(adapter, "_engine", lambda: _Engine()) + + result = adapter.start_session(grounded, state) + + assert result is marker + assert captured == {"invocations": ("invocation",), "context": "context"} + + +def test_retreat_parks_intentional_contact_objects() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + entities = { + uid: _PoseEntity(actual.clone()) for uid in ("released", "container", "other") + } + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=entities), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(entities), + }, + ) + grounded = GroundedAction( + "MoveEndEffector", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + motion_policy={ + "collision_exclusion_uids": ["released", "container"], + }, + object_uid="released", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + parked_z = actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET + assert torch.equal(scene.entities["released"].pose[:, 2, 3], parked_z) + assert torch.equal(scene.entities["container"].pose[:, 2, 3], parked_z) + assert torch.equal(scene.entities["other"].pose, actual) + + +def test_action_outcome_commits_state_delta_only_for_verified_rows() -> None: + semantics = ObjectSemantics( + label="cube", + entity=object(), + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + prior = ExecutionState(last_qpos=torch.zeros(2, 3)) + trajectory = torch.stack( + (torch.zeros(2, 3), torch.ones(2, 3)), + dim=1, + ) + delta = StateDelta(held_object_updates={"physical_left_arm": held}) + projected = ExecutionState.from_task_state( + delta.apply(prior.to_task_state(), torch.ones(2, dtype=torch.bool)), + last_qpos=trajectory[:, -1], + ) + grounded = GroundedAction( + "PickUp", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ) + outcome = ActionOutcome( + trajectory=trajectory, + success=torch.ones(2, dtype=torch.bool), + next_state=projected, + grounded=grounded, + prior_state=prior, + expected_effects=delta, + ) + + committed = outcome.state_after(torch.tensor([True, False])) + + assert torch.equal(committed.last_qpos[0], torch.ones(3)) + assert torch.equal(committed.last_qpos[1], torch.zeros(3)) + committed_held = committed.get_held_object("physical_left_arm") + assert committed_held is not None + assert torch.equal(committed_held.env_mask, torch.tensor([True, False])) + + +def test_fallback_rows_keep_the_fallback_plan_effects(monkeypatch: Any) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + semantics = ObjectSemantics( + label="cube", + entity=object(), + geometry={}, + affordance=Affordance(), + ) + + def held_at(x: float) -> HeldObjectState: + relation = torch.eye(4).repeat(2, 1, 1) + relation[:, 0, 3] = x + return HeldObjectState( + semantics=semantics, + object_to_eef=relation, + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + + def action_plan( + success: torch.Tensor, + terminal: float, + held: HeldObjectState, + ) -> ActionPlan: + positions = torch.full((2, 2, 8), terminal) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="pick_up", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta( + held_object_updates={"physical_left_arm": held} + ), + ) + + plans = iter( + ( + action_plan(torch.tensor([True, False]), 1.0, held_at(1.0)), + action_plan(torch.tensor([True, True]), 2.0, held_at(2.0)), + ) + ) + strategies: list[str] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + strategies.append(invocation.motion_policy.strategy) + return next(plans) + + monkeypatch.setattr( + adapter, + "_engine", + lambda: _FakeEngine(plan), + ) + grounded = GroundedAction( + "PickUp", + "left_arm", + "arm", + GraspGoal(semantics=semantics), + {}, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert strategies == ["motion_gen", "ik_interp"] + assert torch.equal(outcome.success, torch.tensor([True, True])) + assert torch.equal(outcome.next_state.last_qpos[0], torch.ones(8)) + assert torch.equal(outcome.next_state.last_qpos[1], torch.full((8,), 2.0)) + held = outcome.next_state.get_held_object("physical_left_arm") + assert held is not None + assert held.object_to_eef[0, 0, 3] == 1.0 + assert held.object_to_eef[1, 0, 3] == 2.0 + assert torch.equal( + outcome.planner_trace["primary_success"], torch.tensor([True, False]) + ) + assert torch.equal( + outcome.planner_trace["fallback_attempted"], torch.tensor([False, True]) + ) + assert torch.equal( + outcome.planner_trace["fallback_used"], torch.tensor([False, True]) + ) + + +def test_collision_required_cleanup_does_not_use_unsafe_fallback( + monkeypatch: Any, +) -> None: + pose = torch.eye(4).repeat(2, 1, 1) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"released": _PoseEntity(pose)}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["released"], + }, + ) + failed_trajectory = TimedTrajectory.from_uniform_step( + torch.zeros(2, 2, 8), + env_ids=torch.arange(2), + step_dt=0.01, + ) + failed_plan = ActionPlan( + skill_id="move_joints", + plan_success=torch.tensor([False, False]), + commands=_commands_for(failed_trajectory), + joint_trajectory=failed_trajectory, + recovery_policy=RecoveryPolicy(), + planned_scene_version=1, + planned_collision_world_revision=(1, 1), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + strategies: list[str] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + strategies.append(invocation.motion_policy.strategy) + assert invocation.motion_policy.dynamic_collision_mode.value == "required" + return failed_plan + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + motion_policy={"collision_safety": "required"}, + object_uid="released", + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert strategies == ["motion_gen"] + assert not bool(outcome.success.any()) + assert outcome.planner_trace["fallback_allowed"] is False + assert not bool(outcome.planner_trace["fallback_attempted"].any()) + assert not bool(outcome.planner_trace["fallback_used"].any()) + assert outcome.planner_trace["collision_obstacle_positions"]["released"].shape == ( + 2, + 3, + ) diff --git a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py new file mode 100644 index 000000000..f04c13b6e --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py @@ -0,0 +1,88 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ExactTargetMoveHeldObject, + ExactTargetMoveHeldObjectOptions, +) +from embodichain.lab.sim.atomic_actions import MoveHeldObject, MoveHeldObjectOptions + + +def test_exact_target_transport_only_disables_rotation_when_requested( + monkeypatch: pytest.MonkeyPatch, +) -> None: + applied = [] + result = object() + + def fake_apply(self, move_eef_xpos, end_arm_xpos) -> None: + del self, move_eef_xpos, end_arm_xpos + applied.append(True) + + def fake_plan(self, request, context): + del request, context + self._apply_automatic_transport_rotation(torch.eye(4), torch.eye(4)) + return result + + monkeypatch.setattr( + MoveHeldObject, + "_apply_automatic_transport_rotation", + fake_apply, + ) + monkeypatch.setattr(MoveHeldObject, "_plan", fake_plan) + action = ExactTargetMoveHeldObject() + assert type(action).__dict__["binding_contract"] is MoveHeldObject.binding_contract + disabled_request = SimpleNamespace( + skill_options=ExactTargetMoveHeldObjectOptions( + allow_automatic_transport_rotation=False, + ) + ) + enabled_request = SimpleNamespace( + skill_options=ExactTargetMoveHeldObjectOptions(), + ) + + assert action._plan(disabled_request, object()) is result + assert not applied + assert action._plan(enabled_request, object()) is result + assert applied == [True] + + +@pytest.mark.parametrize( + ("yaw_samples", "expected"), + [(1, True), (8, False)], +) +def test_semantic_transport_config_scopes_rotation_override( + yaw_samples: int, + expected: bool, +) -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + action = SimpleNamespace(cfg={"upright_yaw_samples": yaw_samples}) + capability = SimpleNamespace( + config_type=MoveHeldObjectOptions, + target_materializer="semantic_held_object", + ) + + options = adapter._build_single_arm_config(action, capability) + + assert isinstance(options, ExactTargetMoveHeldObjectOptions) + assert options.allow_automatic_transport_rotation is expected diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py b/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py new file mode 100644 index 000000000..a08771054 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py @@ -0,0 +1,354 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +import pickle +from typing import Callable + +import numpy as np +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime import grasp_collision_cache +from embodichain.gen_sim.action_engine.runtime.grasp_collision_cache import ( + GraspCollisionCacheError, + ensure_vhacd_grasp_collision_cache, + grasp_collision_cache_path, +) + + +def _tetrahedron() -> tuple[torch.Tensor, torch.Tensor]: + vertices = torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=torch.float32, + ) + triangles = torch.tensor( + [ + [0, 2, 1], + [0, 1, 3], + [0, 3, 2], + [1, 2, 3], + ], + dtype=torch.int64, + ) + return vertices, triangles + + +def _plane_equations() -> list[tuple[np.ndarray, np.ndarray]]: + return [ + ( + np.asarray( + [ + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ), + np.asarray([-1.0, -1.0, -1.0], dtype=np.float32), + ), + ( + np.asarray([[1.0, 1.0, 1.0]], dtype=np.float32), + np.asarray([-1.0], dtype=np.float32), + ), + ] + + +def _install_fake_decomposer( + monkeypatch: pytest.MonkeyPatch, +) -> list[tuple[tuple[int, ...], tuple[int, ...], int]]: + calls: list[tuple[tuple[int, ...], tuple[int, ...], int]] = [] + + def fake_decompose( + vertices: np.ndarray, + triangles: np.ndarray, + max_decomposition_hulls: int, + ) -> list[tuple[np.ndarray, np.ndarray]]: + calls.append( + ( + tuple(vertices.shape), + tuple(triangles.shape), + max_decomposition_hulls, + ) + ) + return _plane_equations() + + monkeypatch.setattr( + grasp_collision_cache, + "_compute_vhacd_plane_equations", + fake_decompose, + ) + return calls + + +def test_cache_key_and_payload_match_main_checker_contract( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + + result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + expected_hash = hashlib.md5( + vertices.numpy().tobytes() + triangles.numpy().tobytes() + ).hexdigest() + assert result.cache_path == tmp_path / f"{expected_hash}_16.pkl" + with result.cache_path.open("rb") as cache_file: + payload = pickle.load(cache_file) + assert set(payload) == {"plane_equations", "plane_equation_counts"} + assert payload["plane_equations"].shape == (2, 3, 4) + assert payload["plane_equations"].dtype == torch.float32 + assert payload["plane_equation_counts"].tolist() == [3, 1] + assert payload["plane_equation_counts"].dtype == torch.int32 + + +def test_main_checker_loads_prepared_cache_without_running_coacd( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import embodichain.lab.sim + from embodichain.toolkits.graspkit.pg_grasp import collision_checker + + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + def fail_coacd(*args: object, **kwargs: object) -> None: + raise AssertionError("The prepared V-HACD cache must bypass CoACD.") + + monkeypatch.setattr(embodichain.lab.sim, "CONVEX_DECOMP_DIR", tmp_path) + monkeypatch.setattr(collision_checker, "convex_decomposition_coacd", fail_coacd) + checker = collision_checker.ConvexCollisionChecker( + vertices, + triangles, + max_decomposition_hulls=16, + ) + + assert checker.cache_path == result.cache_path.as_posix() + assert checker.plane_equations["plane_equation_counts"].tolist() == [3, 1] + + +def test_matching_vhacd_metadata_returns_cache_hit( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + calls = _install_fake_decomposer(monkeypatch) + kwargs = { + "mesh_vertices": vertices, + "mesh_triangles": triangles, + "max_decomposition_hulls": 16, + "cache_dir": tmp_path, + } + + first = ensure_vhacd_grasp_collision_cache(**kwargs) + second = ensure_vhacd_grasp_collision_cache(**kwargs) + + assert first.status == "generated" + assert second.status == "hit" + assert len(calls) == 1 + + +def test_non_vhacd_metadata_forces_cache_replacement( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + calls = _install_fake_decomposer(monkeypatch) + kwargs = { + "mesh_vertices": vertices, + "mesh_triangles": triangles, + "max_decomposition_hulls": 16, + "cache_dir": tmp_path, + } + first = ensure_vhacd_grasp_collision_cache(**kwargs) + metadata = json.loads(first.metadata_path.read_text(encoding="utf-8")) + metadata["backend"] = "coacd" + first.metadata_path.write_text(json.dumps(metadata), encoding="utf-8") + + replaced = ensure_vhacd_grasp_collision_cache(**kwargs) + + assert replaced.status == "replaced" + assert len(calls) == 2 + repaired = json.loads(replaced.metadata_path.read_text(encoding="utf-8")) + assert repaired["backend"] == "vhacd" + + +def test_modified_cache_fails_checksum_and_is_rebuilt( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + calls = _install_fake_decomposer(monkeypatch) + kwargs = { + "mesh_vertices": vertices, + "mesh_triangles": triangles, + "max_decomposition_hulls": 16, + "cache_dir": tmp_path, + } + first = ensure_vhacd_grasp_collision_cache(**kwargs) + first.cache_path.write_bytes(b"not a valid collision cache") + + replaced = ensure_vhacd_grasp_collision_cache(**kwargs) + + assert replaced.status == "replaced" + assert len(calls) == 2 + with replaced.cache_path.open("rb") as cache_file: + assert "plane_equations" in pickle.load(cache_file) + + +def test_cache_and_metadata_are_published_by_atomic_replace( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + replacements: list[tuple[Path, Path]] = [] + real_replace: Callable[[os.PathLike[str], os.PathLike[str]], None] = os.replace + + def recording_replace( + source: os.PathLike[str], + destination: os.PathLike[str], + ) -> None: + replacements.append((Path(source), Path(destination))) + real_replace(source, destination) + + monkeypatch.setattr(grasp_collision_cache.os, "replace", recording_replace) + + result = ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + assert [destination for _, destination in replacements] == [ + result.cache_path, + result.metadata_path, + ] + assert all( + source.parent == destination.parent for source, destination in replacements + ) + assert all(not source.exists() for source, _ in replacements) + + +@pytest.mark.parametrize( + ("vertices", "triangles", "message"), + [ + ( + torch.empty((0, 3), dtype=torch.float32), + torch.tensor([[0, 1, 2]], dtype=torch.int64), + "mesh_vertices", + ), + ( + torch.zeros((3, 3), dtype=torch.float32), + torch.tensor([[0, 1]], dtype=torch.int64), + "mesh_triangles", + ), + ( + torch.tensor([[0.0, 0.0, 0.0], [1.0, float("nan"), 0.0], [0.0, 1.0, 0.0]]), + torch.tensor([[0, 1, 2]], dtype=torch.int64), + "finite", + ), + ( + torch.zeros((3, 3), dtype=torch.float32), + torch.tensor([[0, 1, 3]], dtype=torch.int64), + "indices", + ), + ], +) +def test_invalid_mesh_is_rejected_before_decomposition( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + vertices: torch.Tensor, + triangles: torch.Tensor, + message: str, +) -> None: + calls = _install_fake_decomposer(monkeypatch) + + with pytest.raises(ValueError, match=message): + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + assert calls == [] + + +@pytest.mark.parametrize("max_decomposition_hulls", [True, 0, -1, 1.5]) +def test_invalid_hull_limit_is_rejected( + tmp_path: Path, + max_decomposition_hulls: object, +) -> None: + vertices, triangles = _tetrahedron() + + with pytest.raises((TypeError, ValueError), match="max_decomposition_hulls"): + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=max_decomposition_hulls, # type: ignore[arg-type] + cache_dir=tmp_path, + ) + + +def test_symlinked_cache_path_is_refused( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, triangles = _tetrahedron() + _install_fake_decomposer(monkeypatch) + cache_path = grasp_collision_cache_path( + vertices, + triangles, + 16, + cache_dir=tmp_path, + ) + victim = tmp_path / "victim.pkl" + victim.write_bytes(b"do not overwrite") + cache_path.symlink_to(victim) + + with pytest.raises(GraspCollisionCacheError, match="symlink"): + ensure_vhacd_grasp_collision_cache( + mesh_vertices=vertices, + mesh_triangles=triangles, + max_decomposition_hulls=16, + cache_dir=tmp_path, + ) + + assert victim.read_bytes() == b"do not overwrite" diff --git a/tests/gen_sim/action_engine/runtime/test_recovery_v2.py b/tests/gen_sim/action_engine/runtime/test_recovery_v2.py new file mode 100644 index 000000000..5544aea45 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_recovery_v2.py @@ -0,0 +1,720 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +import embodichain.gen_sim.action_engine.runtime.executor as executor_module +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.runtime import ( + DynamicRecoveryController, + ProgramExecutor, + RuntimeGraph, + build_upright_recovery, + classify_failure, + load_execution_program, +) +from embodichain.gen_sim.action_engine.runtime.executor import _EdgeResult +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_spec + + +def _graph(task_type: str) -> dict: + task, requirements = make_task_spec(task_type) + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + return instantiate_seed_graph(task, bindings) + + +def _handover_then_place_graph() -> dict: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "handover_then_place_recovery", + "level": "L3", + "instruction": "Hand the yellow can from the left arm to the right arm.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E4", + "params": { + "object_role": "yellow_can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E1", + "params": { + "object_role": "yellow_can", + "target_role": "purple_can", + "relation": "right_of", + }, + "depends_on": ["task_01"], + "role": "primary", + }, + ], + "success": { + "op": "all", + "terms": [ + {"type": "handover_complete", "task_instance_id": "task_01"}, + {"type": "semantic_goal", "task_instance_id": "task_02"}, + ], + }, + "oracle": {"task_order": ["task_01", "task_02"]}, + "metadata": {}, + } + return instantiate_seed_graph( + task, + { + "yellow_can": "interact_yellow_can", + "purple_can": "interact_purple_can", + }, + ) + + +class _RecoveryRecorder: + def __init__(self) -> None: + self.recovery_events: list[dict[str, Any]] = [] + self.edge_events: list[dict[str, Any]] = [] + + def recovery(self, **event: Any) -> None: + self.recovery_events.append(event) + + def edge(self, edge_id: str, step: Any, **event: Any) -> None: + self.edge_events.append({"edge_id": edge_id, "step_id": step.id, **event}) + + def step(self, *_args: Any, **_kwargs: Any) -> None: + return None + + def register_step(self, *_args: Any, **_kwargs: Any) -> None: + return None + + +def _local_recovery_harness( + graph: dict[str, Any], + *, + num_envs: int, + max_transitions: int = 100, + max_revisions: int = 8, +) -> tuple[ProgramExecutor, Any, Any, list[tuple[str, str, list[bool]]]]: + program = load_execution_program(graph, require_executable=True) + step = next(item for item in program.semantic_steps if item.id == "task_01") + failed_node = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == step.id and node["atomic_action"] == "HandOver" + ) + edge = next( + item + for item in program.edges + if item.actions[0].get("seed_node_id") == failed_node["id"] + ) + executor = object.__new__(ProgramExecutor) + executor.runtime_graph = RuntimeGraph( + graph, + num_envs=num_envs, + max_revisions=max_revisions, + ) + executor.env = SimpleNamespace( + num_envs=num_envs, + device=torch.device("cpu"), + robot=SimpleNamespace(get_qpos=lambda: torch.zeros((num_envs, 4))), + ) + executor.capability_registry = None + executor.steps = {item.id: item for item in program.semantic_steps} + executor.edges = {item.id: item for item in program.edges} + executor.step_by_edge = { + edge_id: item for item in program.semantic_steps for edge_id in item.edge_ids + } + executor._assignments = {step.id: ["left_arm"] * num_envs} + executor._candidate_cache = {} + executor._candidate_failures = {} + executor._candidate_diagnostics = {} + executor._object_states = {} + executor._step_states = {} + executor._object_owners = {} + executor._arm_owners = { + "left_arm": [None] * num_envs, + "right_arm": [None] * num_envs, + } + executor._targets = {} + executor.record_runtime = False + executor.max_transitions = max_transitions + executor._transition_count = 0 + executor.retry_count = 0 + call_log: list[tuple[str, str, list[bool]]] = [] + + def execute_edge(current_edge: Any, current_step: Any, *, failed: torch.Tensor): + call_log.append((current_step.id, current_edge.id, failed.tolist())) + return _EdgeResult([], failed.clone(), []) + + def ensure_assignment(current_step: Any, failed: torch.Tensor) -> None: + actor = current_step.actor + assignment = ( + str(actor["arm"]) if actor.get("mode") == "required" else "right_arm" + ) + executor._assignments[current_step.id] = [ + None if bool(failed[index]) else assignment for index in range(num_envs) + ] + + executor._execute_edge_with_retries = execute_edge + executor._ensure_assignment = ensure_assignment + executor._clear_recovery_rows = lambda *_args, **_kwargs: None + executor._verify_step = lambda _step, failed: ( + failed.clone(), + ~failed, + torch.zeros((num_envs, 3)), + ) + return executor, step, edge, call_log + + +def test_runtime_graph_retries_twice_then_requests_recovery() -> None: + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=2, max_retries=2) + handover = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + failed = torch.tensor([True, False]) + holds = torch.tensor([True, True]) + + first = runtime.record_failure(handover["id"], failed, precondition_holds=holds) + second = runtime.record_failure(handover["id"], failed, precondition_holds=holds) + third = runtime.record_failure(handover["id"], failed, precondition_holds=holds) + + assert first.retry.tolist() == [True, False] + assert second.retry.tolist() == [True, False] + assert third.recover.tolist() == [True, False] + assert runtime.seed_graph == graph + + +def test_recovery_insertion_revises_runtime_graph_not_seed_graph() -> None: + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=1) + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + recovery_source = _graph("E2") + source_group = recovery_source["task_groups"][0] + recovery_group_id = "recovery_upright_01" + recovery_nodes = [] + id_map = { + node["id"]: f"recovery_{index:02d}" + for index, node in enumerate(recovery_source["nodes"], start=1) + } + for node in recovery_source["nodes"]: + item = deepcopy(node) + item["id"] = id_map[node["id"]] + item["object_uid"] = failed_node["object_uid"] + item["target_binding"] = deepcopy(item["target_binding"]) + if item["target_binding"].get("kind") == "object": + item["target_binding"]["object"] = failed_node["object_uid"] + item["depends_on"] = [id_map.get(dep, dep) for dep in node["depends_on"]] + recovery_nodes.append(item) + recovery_group = deepcopy(source_group) + recovery_group.update( + { + "id": recovery_group_id, + "role": "recovery", + "object_uid": failed_node["object_uid"], + "node_ids": [node["id"] for node in recovery_nodes], + "depends_on": [], + "parent_task_instance_id": failed_node["task_instance_id"], + } + ) + recovery_group["success"] = { + "type": "object_upright", + "object": failed_node["object_uid"], + } + + patched = runtime.insert_recovery_subgraph( + failed_node_id=failed_node["id"], + recovery_nodes=recovery_nodes, + recovery_group=recovery_group, + failure_type="object_fallen", + ) + + assert graph == runtime.seed_graph + assert any(group["id"] == recovery_group_id for group in patched["task_groups"]) + assert not any( + node["task_instance_id"] == failed_node["task_instance_id"] + and node["target_binding"].get("source") == "handover" + for node in patched["nodes"] + ) + assert runtime.revisions[0].kind == "insert_recovery" + assert ( + classify_failure("PickUp", planning_succeeded=True, held_after=False) + == "grasp_missed" + ) + + +def test_recovery_rejects_downstream_contract_that_requires_actor_switch() -> None: + graph = _handover_then_place_graph() + runtime = RuntimeGraph(graph, num_envs=1) + handover = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + cleanup_ids = { + node["id"] + for node in graph["nodes"] + if node["task_instance_id"] == handover["task_instance_id"] + and node["role"] == "cleanup" + } + assert cleanup_ids + + with pytest.raises(ValueError, match="without changing.*actor"): + runtime.insert_default_recovery( + failed_node_id=handover["id"], + failure_type="object_fallen", + ) + + assert runtime.graph == graph + assert runtime.revisions == [] + + +def test_failed_group_resume_recovery_places_before_prefix_replay() -> None: + graph = _handover_then_place_graph() + runtime = RuntimeGraph(graph, num_envs=1) + handover = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + + patched = runtime.insert_default_recovery( + failed_node_id=handover["id"], + failure_type="object_fallen", + resume_failed_group=True, + ) + + group_id = runtime.revisions[-1].inserted_group_ids[0] + group = next(item for item in patched["task_groups"] if item["id"] == group_id) + nodes = [node for node in patched["nodes"] if node["id"] in group["node_ids"]] + original_cleanup = { + node["id"] + for node in graph["nodes"] + if node["task_instance_id"] == handover["task_instance_id"] + and node["role"] == "cleanup" + } + assert group["goal"]["terminal_behavior"] == "place" + assert [node["atomic_action"] for node in nodes] == [ + "PickUp", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert original_cleanup <= {node["id"] for node in patched["nodes"]} + + +@pytest.mark.parametrize( + "actor", + ( + {"mode": "required", "arm": "left_arm"}, + {"mode": "required", "arm": "right_arm"}, + {"mode": "auto"}, + ), +) +def test_upright_recovery_inherits_failed_group_actor(actor: dict[str, Any]) -> None: + graph = _graph("E2") + failed_group = graph["task_groups"][0] + failed_group["actor"] = deepcopy(actor) + for node in graph["nodes"]: + if node["task_instance_id"] == failed_group["id"]: + node["actor"] = deepcopy(actor) + + nodes, recovery_group = build_upright_recovery( + graph, + failed_node_id=failed_group["node_ids"][0], + revision=1, + resume_failed_group=True, + ) + + assert recovery_group["actor"] == actor + assert all(node["actor"] == actor for node in nodes) + + +def test_local_recovery_replays_failed_group_prefix_and_preserves_seed_graph( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + original = deepcopy(graph) + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([True]), + recorder=recorder, + ) + + replayed = [edge_id for step_id, edge_id, _failed in calls if step_id == step.id] + expected_prefix = list(step.edge_ids[: step.edge_ids.index(edge.id) + 1]) + assert result.failed.tolist() == [False] + assert replayed == expected_prefix + assert executor.runtime_graph.seed_graph == original + assert graph == original + assert [event["status"] for event in recorder.recovery_events] == [ + "started", + "succeeded", + ] + recovery_edges = [ + event + for event in recorder.edge_events + if event["step_id"].startswith("recovery_e2_") + ] + replay_edges = [ + event for event in recorder.edge_events if event["step_id"] == step.id + ] + assert recovery_edges + assert all(event["phase"] == "recovery" for event in recovery_edges) + assert replay_edges + assert all(event["phase"] == "replay" for event in replay_edges) + + +def test_local_recovery_only_executes_and_rebinds_failed_vector_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=2) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([True, False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([False, True]), + [], + executed=torch.tensor([False, True]), + ), + inherited_failed=torch.tensor([False, False]), + fallen_transition=torch.tensor([False, True]), + recorder=recorder, + ) + + assert result.failed.tolist() == [False, False] + assert all(failed == [True, False] for _step_id, _edge_id, failed in calls) + assert executor._assignments[step.id] == ["left_arm", "left_arm"] + assert executor.runtime_graph.revisions[-1].active_env_ids == (1,) + assert all( + event["active"].tolist() == [False, True] for event in recorder.recovery_events + ) + + +def test_local_recovery_failure_does_not_replay_prefix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + executor._verify_step = lambda _step, failed: ( + torch.ones_like(failed), + torch.zeros_like(failed), + torch.zeros((1, 3)), + ) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([True]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert not any(step_id == step.id for step_id, _edge_id, _failed in calls) + assert recorder.recovery_events[-1]["status"] == "failed" + + +def test_local_recovery_budget_exhaustion_terminates_with_original_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness( + graph, + num_envs=1, + max_transitions=0, + ) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([False]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([True]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert calls == [] + assert recorder.recovery_events[-1]["status"] == "failed" + assert "max_transitions" in recorder.recovery_events[-1]["error"] + + +def test_non_fallen_failure_does_not_create_recovery_revision( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + recorder = _RecoveryRecorder() + monkeypatch.setattr( + executor_module, + "evaluate_predicate", + lambda *_args, **_kwargs: torch.tensor([True]), + ) + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([True]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert calls == [] + assert executor.runtime_graph.revisions == [] + assert recorder.recovery_events == [] + + +def test_initially_fallen_planning_failure_does_not_trigger_recovery() -> None: + graph = _handover_then_place_graph() + executor, step, edge, calls = _local_recovery_harness(graph, num_envs=1) + recorder = _RecoveryRecorder() + + result = executor._recover_object_fallen( + edge, + step, + _EdgeResult( + [], + torch.tensor([True]), + [], + executed=torch.tensor([False]), + ), + inherited_failed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + recorder=recorder, + ) + + assert result.failed.tolist() == [True] + assert calls == [] + assert executor.runtime_graph.revisions == [] + assert recorder.recovery_events == [] + + +def test_failure_provenance_distinguishes_planning_from_execution_caused_fall() -> None: + graph = _handover_then_place_graph() + executor, step, edge, _ = _local_recovery_harness(graph, num_envs=1) + executor.adapter = SimpleNamespace(capabilities=build_atomic_capability_registry()) + failed = torch.tensor([True]) + + planning = executor._failure_events( + edge, + step, + failed, + postcondition=False, + executed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + ) + execution = executor._failure_events( + edge, + step, + failed, + postcondition=False, + executed=torch.tensor([True]), + fallen_transition=torch.tensor([True]), + ) + + assert [event["failure_type"] for event in planning] == ["search_exhausted"] + assert [event["failure_type"] for event in execution] == ["object_fallen"] + + +def test_offline_and_online_dynamic_replanners_are_route_isolated() -> None: + for mode in ("offline_dynamic", "online_dynamic"): + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=1) + calls = [] + + def replanner(**kwargs): + calls.append((mode, kwargs["failure_type"])) + return kwargs["graph"] + + controller = DynamicRecoveryController( + runtime, + mode=mode, + offline_replanner=replanner if mode == "offline_dynamic" else None, + online_replanner=replanner if mode == "online_dynamic" else None, + ) + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + directive = controller.handle_failure( + failed_node_id=failed_node["id"], + failure_type="postcondition_failed", + ) + completed = [group["id"] for group in graph["task_groups"]] + controller.replan( + directive, + completed_group_ids=completed, + recovery_succeeded=False, + ) + + assert calls == [(mode, "postcondition_failed")] + assert runtime.revisions[-1].kind == "replan_suffix" + + +def test_dynamic_recovery_consumes_per_environment_failure_events() -> None: + graph = _graph("E4") + runtime = RuntimeGraph(graph, num_envs=2) + controller = DynamicRecoveryController( + runtime, + mode="offline_dynamic", + offline_replanner=lambda **kwargs: kwargs["graph"], + ) + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + result = SimpleNamespace( + failure_events=[ + { + "node_id": failed_node["id"], + "failure_type": "object_fallen", + "env_ids": [1], + } + ] + ) + + directive = controller.handle_execution_result(result) + + assert directive.active_env_ids == (1,) + assert runtime.revisions[-1].active_env_ids == (1,) + + +def test_runtime_graph_stops_at_revision_and_recovery_budgets() -> None: + graph = _graph("E4") + failed_node = next( + node for node in graph["nodes"] if node["atomic_action"] == "HandOver" + ) + + with pytest.raises(RuntimeError, match="revision budget"): + RuntimeGraph(graph, num_envs=1, max_revisions=0).insert_default_recovery( + failed_node_id=failed_node["id"], + failure_type="object_fallen", + ) + with pytest.raises(RuntimeError, match="recovery-action budget"): + RuntimeGraph(graph, num_envs=1, max_recovery_actions=0).insert_default_recovery( + failed_node_id=failed_node["id"], + failure_type="object_fallen", + ) + + +def test_visual_constraint_grounding_reads_fresh_camera_depth() -> None: + class Sensor: + def __init__(self) -> None: + self.depth = torch.ones((1, 4, 4, 1)) + + def get_data(self): + return {"depth": self.depth} + + def get_intrinsics(self): + return torch.tensor([[[2.0, 0.0, 1.5], [0.0, 2.0, 1.5], [0.0, 0.0, 1.0]]]) + + def get_arena_pose(self, *, to_matrix: bool): + assert to_matrix + return torch.eye(4).unsqueeze(0) + + sensor = Sensor() + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + sim=SimpleNamespace(get_sensor=lambda uid: sensor if uid == "front" else None), + get_current_xpos_agent=lambda: ( + torch.eye(4).unsqueeze(0), + torch.eye(4).unsqueeze(0), + ), + ) + grounder = object.__new__(ActionGrounder) + grounder.env = env + binding = {"camera_uid": "front", "normalized_keypoint": [0.5, 0.5]} + + first = grounder._visual_target(binding, "left_arm") + sensor.depth.fill_(2.0) + second = grounder._visual_target( + {"camera_uid": "front", "normalized_bbox": [0.4, 0.4, 0.6, 0.6]}, + "left_arm", + ) + + assert first[0, 2, 3].item() == 1.0 + assert second[0, 2, 3].item() == 2.0 diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py new file mode 100644 index 000000000..81db74238 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -0,0 +1,6253 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import replace +import hashlib +import json +from pathlib import Path +import sys +from types import ModuleType, SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + default_runtime_policy, + resolve_agent_runtime_policy, + runtime_policy_hash, +) +from embodichain.gen_sim.action_engine.compiler import ( + compile_task_agent, + compile_task_agent_v2, +) +from embodichain.gen_sim.action_engine.cli.run_agent import ( + build_parser as build_run_parser, +) +from embodichain.gen_sim.action_engine.domain import ( + TASK_AGENT_SCHEMA, + execution_program_hash, + motion_policy, + validate_execution_program, +) +from embodichain.gen_sim.action_engine.environment import agent_env as env_module +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.executor import ( + ProgramExecutor, + _EdgeResult, + _score_arm_candidate, +) +from embodichain.gen_sim.action_engine.runtime.frames import ( + relation_offset, + robot_frame_axes, +) +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.runtime.loader import ( + load_agent_execution_program, + load_execution_program as _load_execution_program, +) +from embodichain.gen_sim.action_engine.runtime.models import ( + ActionOutcome, + ExecutionEdge, + ExecutionProgram, + ExecutionResult, + GroundedAction, + SemanticStep, +) +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) +from embodichain.gen_sim.action_engine.runtime.predicates import evaluate_predicate +from embodichain.gen_sim.action_engine.runtime.recording import RuntimeRecorder +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.runtime import solver_compat +from embodichain.gen_sim.action_engine.protocol import ( + SEED_GRAPH_SCHEMA, + TASK_SPEC_SCHEMA, +) +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph +from embodichain.lab.gym.envs import EmbodiedEnv +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + Affordance, + AntipodalAffordance, + AxisAlignAffordance, + AxisAlignGoal, + CoordinatedPickGoal, + CoordinatedPlacementGoal, + CoordinatedPlacementOptions, + HandOverOptions, + HeldObjectPoseGoal, + HeldObjectState, + ObjectSemantics, + PickUpOptions, + PourGoal, + PressAffordance, + PressGoal, + PressOptions, + SlideAffordance, + SlideGoal, + TwistAffordance, + TwistGoal, +) + +from ..task_fixtures import make_task_spec +from embodichain.lab.sim.solvers import URSolverCfg + + +def _task_agent(*steps: dict[str, Any]) -> dict[str, Any]: + return { + "schema_version": TASK_AGENT_SCHEMA, + "task": "runtime_contract", + "goal": "Exercise the deterministic runtime contract.", + "semantic_steps": list(steps), + } + + +def load_execution_program(source: Any, **kwargs: Any) -> ExecutionProgram: + """Adapt legacy compiler fixtures without weakening the production loader.""" + if isinstance(source, dict) and source.get("schema_version") != SEED_GRAPH_SCHEMA: + return ExecutionProgram.from_mapping(validate_execution_program(source)) + return _load_execution_program(source, **kwargs) + + +def _hold_step(step_id: str, object_uid: str, arm: str) -> dict[str, Any]: + return { + "id": step_id, + "operator": "hold_hover", + "object": object_uid, + "actor": {"mode": "required", "arm": arm}, + "goal": {}, + "depends_on": [], + } + + +class _FakeEntity: + def __init__( + self, + uid: str, + pose: torch.Tensor, + vertices: torch.Tensor, + ) -> None: + self.uid = uid + self._pose = pose + self._vertices = vertices + self._triangles = torch.tensor( + [[0, 1, 2], [0, 2, 3]], + dtype=torch.int64, + ) + self.lin_vel = torch.zeros(pose.shape[0], 3) + self.ang_vel = torch.zeros(pose.shape[0], 3) + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return self._pose.clone() + + def get_vertices( + self, + *, + env_ids: list[int], + scale: bool, + ) -> torch.Tensor: + del env_ids, scale + return self._vertices.clone() + + def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: + del env_ids + return self._triangles.clone() + + +class _FakeArticulation: + def __init__(self, uid: str, qpos: float) -> None: + self.uid = uid + self.joint_names = ["slide_joint"] + self.active_joint_ids = [0] + self.link_names = ["base", "drawer_link", "handle"] + self.all_joint_names = ["slide_joint", "fixed_handle"] + self._qpos = torch.tensor([[qpos]], dtype=torch.float32) + self._limits = torch.tensor([[[0.0, 0.2]]], dtype=torch.float32) + self._pose = _pose(0.0, 0.0, 0.7) + self._vertices = _box_vertices(0.05) + self._triangles = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int64) + self._joint_info = SimpleNamespace( + joint_type=SimpleNamespace(name="PRISMATIC"), + child_link_name="drawer_link", + parent_link_name="base", + axis=torch.tensor([1.0, 0.0, 0.0]), + origin_pose=torch.eye(4), + ) + self._fixed_info = SimpleNamespace( + joint_type=SimpleNamespace(name="FIXED"), + child_link_name="handle", + parent_link_name="drawer_link", + axis=torch.tensor([1.0, 0.0, 0.0]), + origin_pose=torch.eye(4), + ) + self._entities = [ + SimpleNamespace( + get_joint_info=lambda name: ( + self._joint_info if name == "slide_joint" else self._fixed_info + ), + ) + ] + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return self._pose.clone() + + def get_link_pose(self, link_name: str, *, to_matrix: bool) -> torch.Tensor: + assert link_name in self.link_names + assert to_matrix + return self._pose.clone() + + def get_link_vert_face(self, link_name: str) -> tuple[torch.Tensor, torch.Tensor]: + assert link_name == "drawer_link" + return self._vertices.clone(), self._triangles.clone() + + def get_qpos(self) -> torch.Tensor: + return self._qpos.clone() + + def get_qpos_limits(self, *, joint_ids: list[int]) -> torch.Tensor: + return self._limits[:, joint_ids].clone() + + +class _FakeSim: + def __init__( + self, + entities: dict[str, _FakeEntity], + articulations: dict[str, _FakeArticulation] | None = None, + ) -> None: + self.entities = entities + self.articulations = articulations or {} + + def get_rigid_object(self, uid: str) -> _FakeEntity | None: + return self.entities.get(uid) + + def get_rigid_object_uid_list(self) -> list[str]: + return list(self.entities) + + def get_articulation(self, uid: str) -> _FakeArticulation | None: + return self.articulations.get(uid) + + def update(self, *, step: int) -> None: + del step + + +class _FakeRobot: + def __init__(self, num_envs: int = 1) -> None: + self.uid = "fake_robot" + self.dof = 8 + self._qpos = torch.zeros(num_envs, self.dof) + self.control_parts = { + "physical_left_arm": ["l0", "l1"], + "physical_left_eef": ["lh0", "lh1"], + "physical_right_arm": ["r0", "r1"], + "physical_right_eef": ["rh0", "rh1"], + "dual_arm": ["l0", "l1", "r0", "r1"], + } + self._ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": [2, 3], + "physical_right_arm": [4, 5], + "physical_right_eef": [6, 7], + "dual_arm": [0, 1, 4, 5], + } + + def get_qpos(self) -> torch.Tensor: + return self._qpos.clone() + + def get_joint_ids(self, *, name: str) -> list[int]: + return list(self._ids[name]) + + def get_control_part_base_pose(self, *, name: str, to_matrix: bool) -> torch.Tensor: + del name + assert to_matrix + return torch.eye(4).repeat(self._qpos.shape[0], 1, 1) + + def get_solver(self, *, name: str) -> SimpleNamespace: + return SimpleNamespace(root_link_name=name.replace("_arm", "_base")) + + def get_link_pose(self, *, link_name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).repeat(self._qpos.shape[0], 1, 1) + pose[:, 1, 3] = -0.3 if link_name == "physical_left_base" else 0.3 + return pose + + def compute_fk( + self, + qpos: torch.Tensor, + *, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + del name + assert to_matrix + return torch.eye(4).repeat(qpos.shape[0], 1, 1) + + +class _FakeEnv: + def __init__( + self, + entities: dict[str, _FakeEntity] | None = None, + articulations: dict[str, _FakeArticulation] | None = None, + ) -> None: + self.num_envs = 1 + self.device = torch.device("cpu") + self.robot = _FakeRobot(self.num_envs) + self.sim = _FakeSim(entities or {}, articulations) + self.left_arm_joints = [0, 1] + self.left_eef_joints = [2, 3] + self.right_arm_joints = [4, 5] + self.right_eef_joints = [6, 7] + self.open_state = torch.tensor([0.0, 0.0]) + self.close_state = torch.tensor([0.7, -0.7]) + + def get_agent_arm_control_part(self, is_left: bool) -> str: + return "physical_left_arm" if is_left else "physical_right_arm" + + def get_agent_eef_control_part(self, is_left: bool) -> str: + return "physical_left_eef" if is_left else "physical_right_eef" + + def get_current_xpos_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + left = torch.eye(4).repeat(self.num_envs, 1, 1) + right = left.clone() + left[:, 1, 3] = -0.2 + right[:, 1, 3] = 0.2 + return left, right + + def get_current_qpos_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return qpos[:, self.left_arm_joints], qpos[:, self.right_arm_joints] + + def get_current_gripper_state_agent(self) -> tuple[torch.Tensor, torch.Tensor]: + qpos = self.robot.get_qpos() + return qpos[:, self.left_eef_joints], qpos[:, self.right_eef_joints] + + +def _box_vertices(half_extent: float) -> torch.Tensor: + h = float(half_extent) + return torch.tensor( + [ + [-h, -h, -h], + [h, -h, -h], + [h, h, h], + [-h, h, h], + ], + dtype=torch.float32, + ) + + +def _rect_vertices(x: float, y: float, z: float) -> torch.Tensor: + return torch.tensor( + [ + [sx * x, sy * y, sz * z] + for sx in (-1.0, 1.0) + for sy in (-1.0, 1.0) + for sz in (-1.0, 1.0) + ], + dtype=torch.float32, + ) + + +def _pose(x: float, y: float, z: float) -> torch.Tensor: + result = torch.eye(4).unsqueeze(0) + result[:, :3, 3] = torch.tensor([x, y, z]) + return result + + +def test_press_grounding_adapts_top_surface_and_depth_to_mainline_contract() -> None: + entity = _FakeEntity("button", _pose(0.1, -0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"button": entity}) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "press", + "operator": "press", + "object": "button", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ) + ) + ) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label="button", + entity=entity, + ) + step = program.semantic_steps[0] + action = program.edges[0].actions[0] + + grounded = ActionGrounder(program, env, lambda _uid: semantics).ground( + action, + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, PressGoal) + assert isinstance(grounded.target.semantics.affordance, PressAffordance) + contact = grounded.target.semantics.affordance.get_press_pose( + grounded.target.target_pose + ) + assert torch.allclose(contact[0, :3, 3], torch.tensor([0.1, -0.2, 0.78])) + assert torch.allclose(contact[0, :3, 2], torch.tensor([0.0, 0.0, -1.0])) + + options = AtomicActionAdapter(env)._build_config(grounded, PressOptions) + + assert options.press_distance == pytest.approx(0.004) + + +def test_press_grounding_uses_calibrated_prismatic_button_state() -> None: + task, _ = make_task_spec("E9") + program = load_execution_program( + instantiate_seed_graph(task, {"object_01": "button"}) + ) + articulation = _FakeArticulation("button", 0.0) + articulation._joint_info.axis = torch.tensor([0.0, 0.0, -1.0]) + articulation._limits = torch.tensor([[[0.0, 0.02]]]) + env = _FakeEnv(articulations={"button": articulation}) + env.agent_config = { + "articulation_settings": {"button": {"slide_joint": [0.0, 0.02]}} + } + semantics = ObjectSemantics(affordance=Affordance(), geometry={}) + step = program.semantic_steps[0] + + grounded = ActionGrounder(program, env, lambda _uid: semantics).ground( + program.edges[0].actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, PressGoal) + affordance = grounded.target.semantics.affordance + assert isinstance(affordance, PressAffordance) + assert torch.allclose(affordance.press_axis, torch.tensor([0.0, 0.0, -1.0])) + assert grounded.cfg["press_distance"] == pytest.approx(0.02) + assert grounded.cfg["articulation_target_qpos"].item() == pytest.approx(0.02) + + +def test_pressed_provider_requires_live_calibrated_joint_state() -> None: + from embodichain.gen_sim.action_engine.environment.agent_env import ActionEngineEnv + + articulation = _FakeArticulation("button", 0.0) + articulation._limits = torch.tensor([[[0.0, 0.02]]]) + env = SimpleNamespace( + sim=_FakeSim({}, {"button": articulation}), + agent_config={ + "articulation_settings": {"button": {"slide_joint": [0.0, 0.02]}} + }, + runtime_policy=SimpleNamespace(predicate_fallbacks={"axis_tolerance": 0.03}), + num_envs=1, + device="cpu", + ) + + assert not bool(ActionEngineEnv.is_object_pressed(env, "button")[0]) + articulation._qpos[0, 0] = 0.019 + assert bool(ActionEngineEnv.is_object_pressed(env, "button")[0]) + + +def test_loader_regenerates_in_memory_without_execution_artifact( + tmp_path: Path, +) -> None: + task = _task_agent(_hold_step("hold", "can", "left_arm")) + graph = compile_task_agent_v2(task) + task_spec = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "runtime_contract", + "level": "L1", + "instruction": "Hold the can.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "hold", + "task_type": "E1", + "params": {"object_role": "can"}, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "object_held", "object": "can"}, + "oracle": {"reference_seed_graph": graph}, + "metadata": {"role_bindings": {"can": "can"}}, + } + task_path = tmp_path / "task_spec.json" + task_path.write_text(json.dumps(task_spec), encoding="utf-8") + agent_config = { + "schema_version": "action_engine_config_v2", + "task_spec": task_path.name, + "seed_task_graph": "not_written.json", + } + config_path = tmp_path / "agent_config.json" + config_path.write_text(json.dumps(agent_config), encoding="utf-8") + + program = load_agent_execution_program( + agent_config, + agent_config_path=config_path, + regenerate=True, + ) + + assert program.task == "runtime_contract" + assert program.semantic_steps[0].operator == "hold_hover" + assert not (tmp_path / "not_written.json").exists() + + +def test_production_loader_rejects_legacy_mapping() -> None: + legacy = compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + + with pytest.raises(ValueError, match="regenerate"): + _load_execution_program(legacy) + + +def test_documented_run_command_arguments_remain_compatible() -> None: + args = build_run_parser().parse_args( + [ + "--task_name", + "task4_2", + "--gym_config", + "/tmp/fast_gym_config.json", + "--agent_config", + "/tmp/agent_config.json", + "--regenerate", + "--headless", + "--seed", + "17", + ] + ) + + assert args.task_name == "task4_2" + assert args.regenerate is True + assert args.headless is True + assert args.seed == 17 + assert args.runtime_backend == "independent" + + +def test_dual_ur5_policy_uses_short_reach_upright_lifts() -> None: + upright = motion_policy(("orientation", "upright")) + ur5_pickup = resolve_motion_policy("dual_ur5", "PickUp", upright) + ur5_transport = resolve_motion_policy("dual_ur5", "MoveHeldObject", upright) + ur10_pickup = resolve_motion_policy("dual_ur10", "PickUp", upright) + ur10_transport = resolve_motion_policy("dual_ur10", "MoveHeldObject", upright) + + assert ur5_pickup["lift_height"] == pytest.approx(0.12) + assert ur5_transport["staging_lift_height"] == pytest.approx(0.12) + assert ur10_pickup["lift_height"] == pytest.approx(0.30) + assert ur10_transport["staging_lift_height"] == pytest.approx(0.25) + + +def test_joint_state_binding_selects_hand_timing_without_a_named_policy() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.agent_initial_object_poses = {"can": entity.get_local_pose(to_matrix=True)} + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + step = program.semantic_steps[0] + action = next( + action + for edge in program.edges + for action in edge.actions + if action["target_binding"].get("source") == "gripper_closed" + ) + grounder = ActionGrounder(program, env, lambda _uid: None) + + grounded = grounder.ground( + action, + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert grounded.cfg["sample_interval"] == 10 + + +def test_runtime_policy_discards_legacy_support_z_fallbacks() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["predicate_fallbacks"].update( + { + "support_min_z_offset": 0.02, + "support_max_z_offset": 0.35, + } + ) + + policy = RuntimePolicyCfg.from_mapping(snapshot) + + assert "support_min_z_offset" not in policy.predicate_fallbacks + assert "support_max_z_offset" not in policy.predicate_fallbacks + + +def test_runtime_policy_v4_migrates_grasp_direction_count() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["schema_version"] = "action_engine_runtime_policy_v4" + snapshot["grasp"].pop("n_deviated_approach_directions") + snapshot_hash = hashlib.sha256( + json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + ).hexdigest() + + policy = resolve_agent_runtime_policy( + { + "robot_profile": "dual_franka", + "runtime_policy": snapshot, + "runtime_policy_hash": snapshot_hash, + } + ) + + assert policy.schema_version == "action_engine_runtime_policy_v6" + assert policy.grasp["n_deviated_approach_directions"] == 4 + + +def test_runtime_policy_v5_migrates_support_geometry_thresholds() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["schema_version"] = "action_engine_runtime_policy_v5" + snapshot["grounding"]["placement"]["clearance"] = 0.019 + for key in ( + "candidate_count", + "candidate_offset_fraction", + "support_margin", + "recovery_attempts", + ): + snapshot["grounding"]["placement"].pop(key) + for key in ( + "support_stability_samples", + "support_stability_interval_steps", + "support_linear_velocity_tolerance", + "support_angular_velocity_tolerance", + ): + snapshot["execution"].pop(key) + for key in ( + "support_com_margin", + "support_max_vertical_gap", + "support_max_penetration", + "support_min_overlap_ratio", + ): + snapshot["predicate_fallbacks"].pop(key) + snapshot_hash = hashlib.sha256( + json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + ).hexdigest() + + policy = resolve_agent_runtime_policy( + { + "robot_profile": "dual_franka", + "runtime_policy": snapshot, + "runtime_policy_hash": snapshot_hash, + } + ) + + assert policy.schema_version == "action_engine_runtime_policy_v6" + assert policy.predicate_fallbacks["support_min_overlap_ratio"] == 0.25 + assert policy.grounding["placement"]["clearance"] == 0.019 + assert policy.grounding["placement"]["candidate_count"] == 5 + + +def test_runtime_recorder_writes_checkpoints_and_rendered_env_graphs( + tmp_path: Path, + monkeypatch: Any, +) -> None: + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + original_seed = deepcopy(program.raw) + runtime_policy = default_runtime_policy("dual_ur10") + recorder = RuntimeRecorder( + program, + num_envs=2, + run_id="run-1", + episode_index=3, + output_root=tmp_path, + runtime_policy=runtime_policy.as_mapping(), + runtime_policy_hash=runtime_policy_hash(runtime_policy), + ) + step = program.semantic_steps[0] + recorder.edge( + program.edges[0].id, + step, + assignments=["left_arm", None], + grounded=[ + GroundedAction( + action_class="PickUp", + arm="left_arm", + control="arm", + target=None, + cfg={}, + motion_policy={"obj_upright_direction": torch.tensor([0.0, 0.0, 1.0])}, + ) + ], + active=torch.tensor([True, False]), + failed=torch.tensor([False, True]), + action_steps=4, + planner_traces=[ + { + "primary_strategy": "motion_gen", + "primary_success": torch.tensor([True, False]), + "fallback_used": torch.tensor([False, True]), + "planned_trajectory": torch.arange(24, dtype=torch.float32).reshape( + 2, 3, 4 + ), + } + ], + ) + recorder.step( + step, + torch.tensor([True, False]), + observed=torch.tensor([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]), + target=torch.tensor([[0.1, 0.2, 0.3], [0.0, 0.0, 0.0]]), + metadata=[ + { + "assigned_arm": "left_arm", + "physical_control_part": "physical_right_arm", + }, + {"assigned_arm": None, "physical_control_part": None}, + ], + ) + + episode_dir = tmp_path / "runtime_contract" / "run-1" / "episode_0003" + checkpoint_paths = sorted(episode_dir.glob("env_*/checkpoints/*.json")) + assert len(checkpoint_paths) == 2 + checkpoint = json.loads(checkpoint_paths[0].read_text(encoding="utf-8")) + assert checkpoint["semantic_step"]["id"] == "hold" + assert checkpoint["status"] == "success" + assert [item["event"] for item in checkpoint["events"]] == [ + "edge", + "semantic_step", + ] + assert checkpoint["events"][0]["actions"][0]["motion_policy"][ + "obj_upright_direction" + ] == [0.0, 0.0, 1.0] + assert checkpoint["events"][0]["planner_attempts"] == [ + { + "primary_strategy": "motion_gen", + "primary_success": True, + "fallback_used": False, + "planned_trajectory": [ + [0.0, 1.0, 2.0, 3.0], + [4.0, 5.0, 6.0, 7.0], + [8.0, 9.0, 10.0, 11.0], + ], + } + ] + assert checkpoint["events"][1]["assigned_arm"] == "left_arm" + assert checkpoint["events"][1]["physical_control_part"] == "physical_right_arm" + + rendered_documents: list[dict[str, Any]] = [] + visualization = ModuleType("embodichain.gen_sim.action_engine.graph_visualization") + + def render_task_graph_png(document: dict[str, Any]) -> bytes: + rendered_documents.append(deepcopy(document)) + return b"\x89PNG\r\n\x1a\nruntime-graph" + + visualization.render_task_graph_png = render_task_graph_png + monkeypatch.setitem(sys.modules, visualization.__name__, visualization) + output_dir = recorder.finalize(torch.tensor([True, False])) + + assert output_dir == episode_dir.as_posix() + assert program.raw == original_seed + expected_hash = execution_program_hash(original_seed) + for env_id, expected_status in enumerate(("success", "failed")): + env_dir = episode_dir / f"env_{env_id:04d}" + document = json.loads((env_dir / "task_graph.json").read_text(encoding="utf-8")) + assert document["schema_version"] == original_seed["schema_version"] + assert document["nodes"] == original_seed["nodes"] + assert document["edges"] == original_seed["edges"] + assert document["runtime"]["status"] == expected_status + assert document["runtime"]["seed_graph_hash"] == expected_hash + assert document["runtime"]["runtime_policy"] == runtime_policy.as_mapping() + assert document["runtime"]["runtime_policy_hash"] == runtime_policy_hash( + runtime_policy + ) + assert (env_dir / "task_graph.png").read_bytes().startswith(b"\x89PNG") + assert len(rendered_documents) == 2 + assert not list(episode_dir.rglob("*.tmp")) + + +def test_runtime_recorder_separates_dynamic_recovery_and_replay_phases( + tmp_path: Path, +) -> None: + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + recorder = RuntimeRecorder( + program, + num_envs=1, + run_id="phased-recovery", + output_root=tmp_path, + ) + primary = program.semantic_steps[0] + recovery = replace( + primary, + id="recovery_e2_hold", + parent_step_id=primary.id, + ) + recovery_spec = deepcopy(program.raw["semantic_steps"][0]) + recovery_spec.update( + { + "id": recovery.id, + "parent_step_id": primary.id, + "role": "recovery", + } + ) + recorder.register_step(recovery, recovery_spec) + active = torch.tensor([True]) + recorder.edge( + "edge_recovery", + recovery, + assignments=["left_arm"], + grounded=[], + active=active, + failed=torch.tensor([False]), + action_steps=4, + phase="recovery", + ) + recorder.step( + recovery, + torch.tensor([True]), + observed=torch.zeros((1, 3)), + target=None, + phase="recovery", + ) + recorder.edge( + program.edges[0].id, + primary, + assignments=["left_arm"], + grounded=[], + active=active, + failed=torch.tensor([False]), + action_steps=3, + phase="replay", + ) + recorder.step( + primary, + torch.tensor([True]), + observed=torch.zeros((1, 3)), + target=None, + ) + + checkpoints = sorted( + (recorder.output_dir / "env_0000" / "checkpoints").glob("*.json") + ) + assert len(checkpoints) == 2 + recovery_checkpoint = next( + json.loads(path.read_text(encoding="utf-8")) + for path in checkpoints + if "recovery_e2_hold" in path.name + ) + primary_checkpoint = next( + json.loads(path.read_text(encoding="utf-8")) + for path in checkpoints + if path.name.endswith("_hold.json") and "recovery_e2" not in path.name + ) + assert {event["phase"] for event in recovery_checkpoint["events"]} == {"recovery"} + assert primary_checkpoint["events"][0]["phase"] == "replay" + assert primary_checkpoint["events"][-1]["phase"] == "primary" + + +def test_runtime_recorder_does_not_mask_execution_when_png_rendering_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + program = load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ) + recorder = RuntimeRecorder( + program, + num_envs=1, + run_id="render_failure", + output_root=tmp_path, + ) + visualization = ModuleType("embodichain.gen_sim.action_engine.graph_visualization") + + def fail_render(_document: dict[str, Any]) -> bytes: + raise ValueError("broken renderer") + + visualization.render_task_graph_png = fail_render + monkeypatch.setitem(sys.modules, visualization.__name__, visualization) + + output_dir = recorder.finalize(torch.tensor([False])) + + record = json.loads( + (Path(output_dir) / "env_0000" / "task_graph.json").read_text(encoding="utf-8") + ) + assert record["runtime"]["status"] == "failed" + assert record["runtime"]["visualization_error"] == ("ValueError: broken renderer") + + +def test_ready_scheduler_packs_only_declared_opposite_arm_pickups() -> None: + compiled = compile_task_agent( + _task_agent( + _hold_step("left", "can_a", "left_arm"), + _hold_step("right", "can_b", "right_arm"), + ) + ) + executor = ProgramExecutor( + load_execution_program(compiled), + _FakeEnv(), + record_runtime=False, + ) + ready = [edge for edge in executor.program.edges if not edge.depends_on] + + packed = executor._pack_ready_edges(ready) + + assert len(packed) == 2 + assert {executor.step_by_edge[edge.id].object_uid for edge in packed} == { + "can_a", + "can_b", + } + + +def test_ready_scheduler_serializes_contact_sensitive_orient_pickups() -> None: + steps = [ + { + "id": step_id, + "operator": "orient_object", + "object": object_uid, + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for step_id, object_uid in (("left", "can_a"), ("right", "can_b")) + ] + task_agent = _task_agent(*steps) + task_agent["allocation_groups"] = [ + { + "id": "dual_arms_1", + "semantic_step_ids": ["left", "right"], + "arm_constraint": "distinct_arms", + } + ] + executor = ProgramExecutor( + load_execution_program(compile_task_agent(task_agent)), + _FakeEnv(), + record_runtime=False, + ) + ready = [edge for edge in executor.program.edges if not edge.depends_on] + + assert len(executor._pack_ready_edges(ready)) == 1 + + +def test_ready_scheduler_defers_pickups_until_a_carried_payload_is_released() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + **{ + uid: _FakeEntity( + uid, + _pose(x, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ) + for uid, x in (("can_a", -0.2), ("can_b", 0.0), ("can_c", 0.2)) + }, + } + task = _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b", "can_c"], + "actor": {"mode": "auto"}, + "goal": {"axis": "world_x", "order_constraint": "free"}, + "depends_on": [], + } + ) + executor = ProgramExecutor( + load_execution_program(compile_task_agent(task)), + _FakeEnv(entities), + record_runtime=False, + ) + pickup_edges = [ + edge + for edge in executor.program.edges + if executor._parallel_pickup_candidate(edge) + ] + completed = {pickup_edges[0].id, pickup_edges[1].id} + ready = [ + edge + for edge in executor.program.edges + if edge.id not in completed and set(edge.depends_on) <= completed + ] + executor._arm_owners["left_arm"][0] = "can_a" + executor._arm_owners["right_arm"][0] = "can_b" + + packed = executor._pack_ready_edges(ready, completed=completed) + + assert not executor._parallel_pickup_candidate(packed[0]) + + executor._arm_owners["right_arm"][0] = None + packed = executor._pack_ready_edges(ready, completed=completed) + assert len(packed) == 1 + assert not executor._parallel_pickup_candidate(packed[0]) + + +def test_parallel_pickups_plan_each_arm_at_execution_time( + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = _hold_step("first", "can_a", "left_arm") + second = _hold_step("second", "can_b", "right_arm") + first["actor"] = {"mode": "auto"} + second["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(first, second))), + _FakeEnv( + { + "can_a": _FakeEntity( + "can_a", _pose(0.0, 0.2, 0.75), _box_vertices(0.03) + ), + "can_b": _FakeEntity( + "can_b", _pose(0.0, -0.2, 0.75), _box_vertices(0.03) + ), + } + ), + record_runtime=False, + ) + edges = tuple( + next( + edge + for edge in executor.program.edges + if edge.id in step.edge_ids + and edge.actions[0]["atomic_action_class"] == "PickUp" + ) + for step in executor.program.semantic_steps + ) + estimate = SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([0.0]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + live_calls: list[tuple[str, str]] = [] + + def plan_live(edge, step, arm): + live_calls.append((step.id, arm)) + grounded = GroundedAction( + action_class="PickUp", + arm=arm, + control="arm", + target=SimpleNamespace(), + cfg={}, + ) + return grounded, ActionOutcome( + trajectory=torch.zeros(1, 1, executor.env.robot.dof), + success=torch.tensor([True]), + next_state=ExecutionState(last_qpos=executor.env.robot.get_qpos()), + grounded=grounded, + ) + + monkeypatch.setattr(executor, "_plan_live_hold", plan_live) + monkeypatch.setattr(executor.adapter, "execute_trajectory", lambda *_a, **_k: []) + monkeypatch.setattr( + executor, "_physical_pickup", lambda _u, _a, _s, attempted: attempted + ) + monkeypatch.setattr( + executor, "_rebase_held_state", lambda _u, _a, state, *_args, **_kwargs: state + ) + monkeypatch.setattr(executor, "_update_ownership", lambda *_args, **_kwargs: None) + + _, failed = executor._execute_parallel_pickups( + edges, + failed=torch.tensor([False]), + ) + + assert set(live_calls) == { + ("first", "right_arm"), + ("second", "left_arm"), + } + assert not bool(failed[0]) + + +def test_required_arm_rejects_wrong_candidate_without_planning() -> None: + compiled = compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + executor = ProgramExecutor( + load_execution_program(compiled), + _FakeEnv(), + record_runtime=False, + ) + failed = torch.zeros(1, dtype=torch.bool) + + candidate = executor._candidate( + executor.program.semantic_steps[0], + "right_arm", + failed, + ) + + assert not bool(candidate.feasible.any()) + assert bool(torch.isinf(candidate.cost).all()) + + +def test_required_arm_speculative_failure_still_reaches_live_planning( + monkeypatch: pytest.MonkeyPatch, +) -> None: + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == ["left_arm"] + + +def test_auto_pickup_outside_deadband_requires_same_side_arm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, -0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def candidate(_step, arm, _failed): + cost = 10.0 if arm == "left_arm" else 1.0 + return SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([cost]), + ) + + monkeypatch.setattr(executor, "_candidate", candidate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._preferred_live_pickup_arm(step, 0) == "left_arm" + assert executor._assignments[step.id] == ["left_arm"] + + +def test_auto_pickup_inside_deadband_selects_lower_cost_arm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.01, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def candidate(_step, arm, _failed): + cost = 10.0 if arm == "left_arm" else 1.0 + return SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([cost]), + ) + + monkeypatch.setattr(executor, "_candidate", candidate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._preferred_live_pickup_arm(step, 0) is None + assert executor._assignments[step.id] == ["right_arm"] + + +def test_auto_pickup_retry_does_not_cross_sides_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + monkeypatch.setattr( + executor, + "_preferred_live_pickup_arm", + lambda *_args: "left_arm", + ) + + executor._ensure_assignment(step, torch.tensor([False])) + assert executor._assignments[step.id] == ["left_arm"] + + executor._pickup_retry_exclusions[(step.id, 0)] = {"left_arm"} + executor._assignments.pop(step.id) + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == [None] + + +def test_auto_pickup_retry_can_explicitly_cross_sides( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + executor.runtime_policy.arm_selection.allow_cross_side_fallback = True + step = executor.program.semantic_steps[0] + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + monkeypatch.setattr( + executor, + "_preferred_live_pickup_arm", + lambda *_args: "left_arm", + ) + + executor._pickup_retry_exclusions[(step.id, 0)] = {"left_arm"} + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == ["right_arm"] + + +def test_auto_pickup_runtime_retry_uses_the_other_arm( + monkeypatch: pytest.MonkeyPatch, +) -> None: + step_mapping = _hold_step("hold", "can", "left_arm") + step_mapping["actor"] = {"mode": "auto"} + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(step_mapping))), + _FakeEnv( + {"can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03))} + ), + record_runtime=False, + ) + executor.runtime_policy.arm_selection.allow_cross_side_fallback = True + step = executor.program.semantic_steps[0] + original_edge = executor.edges[step.edge_ids[0]] + action = {**original_edge.actions[0], "seed_node_id": "pickup_node"} + edge = replace(original_edge, actions=(action,)) + estimate = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: estimate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + monkeypatch.setattr( + executor, + "_preferred_live_pickup_arm", + lambda *_args: "left_arm", + ) + executor._ensure_assignment(step, torch.tensor([False])) + attempts: list[str | None] = [] + + def execute(_edge, _step, *, failed): + arm = executor._assignments[step.id][0] + attempts.append(arm) + return _EdgeResult( + actions=[], + failed=torch.tensor([arm == "left_arm"]) | failed, + grounded=[], + planner_traces=[], + executed=torch.tensor([False]), + ) + + decisions = 0 + + def record_failure(*_args, **_kwargs): + nonlocal decisions + decisions += 1 + return SimpleNamespace(retry=torch.tensor([decisions == 1])) + + executor.runtime_graph = SimpleNamespace( + graph={"nodes": [{"id": "pickup_node", "precondition": {}}]}, + record_failure=record_failure, + ) + monkeypatch.setattr(executor, "_execute_edge", execute) + + result = executor._execute_edge_with_retries( + edge, + step, + failed=torch.tensor([False]), + ) + + assert attempts == ["left_arm", "right_arm"] + assert executor.retry_count == 1 + assert not bool(result.failed[0]) + + +def _held_state( + env: _FakeEnv, + entity: _FakeEntity, + *, + arm: str = "left_arm", +) -> ExecutionState: + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=entity.uid, + entity=entity, + ) + left_eef, right_eef = env.get_current_xpos_agent() + eef = left_eef if arm == "left_arm" else right_eef + object_pose = entity.get_local_pose(to_matrix=True) + return ExecutionState( + last_qpos=env.robot.get_qpos(), + held_objects={ + f"physical_{arm}": HeldObjectState( + semantics=semantics, + object_to_eef=torch.bmm(torch.linalg.inv(object_pose), eef), + grasp_xpos=eef, + ) + }, + ) + + +def _coordinated_held_state( + env: _FakeEnv, + entity: _FakeEntity, +) -> ExecutionState: + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=entity.uid, + entity=entity, + ) + left_eef, right_eef = env.get_current_xpos_agent() + object_pose = entity.get_local_pose(to_matrix=True) + return ExecutionState( + last_qpos=env.robot.get_qpos(), + held_objects={ + "physical_left_arm": HeldObjectState( + semantics=semantics, + object_to_eef=torch.bmm(torch.linalg.inv(object_pose), left_eef), + grasp_xpos=left_eef, + env_mask=torch.ones(env.num_envs, dtype=torch.bool), + ), + "physical_right_arm": HeldObjectState( + semantics=semantics, + object_to_eef=torch.bmm(torch.linalg.inv(object_pose), right_eef), + grasp_xpos=right_eef, + env_mask=torch.ones(env.num_envs, dtype=torch.bool), + ), + }, + ) + + +@pytest.mark.parametrize( + ("opens", "expected_failed", "expect_held"), + ((True, False, False), (False, True, True)), +) +def test_explicit_dual_gripper_release_commits_only_after_both_hands_open( + opens: bool, + expected_failed: bool, + expect_held: bool, +) -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.2)) + env = _FakeEnv({"tray": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + env.robot._qpos[:, env.right_eef_joints] = env.close_state + state = _coordinated_held_state(env, entity) + executor = object.__new__(ProgramExecutor) + executor.env = env + executor._assignments = {"task_01": ["coordinated"]} + executor._step_states = {("task_01", "coordinated"): state} + executor._object_states = {} + executor._orientation_references = {} + + def ground( + action: dict[str, Any], + _step: Any, + *, + arm: str, + **_kwargs: Any, + ) -> GroundedAction: + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm=arm, + control="hand", + target=None, + cfg={}, + ) + + def plan(grounded: GroundedAction, current: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 2, env.robot.dof), + success=torch.ones(1, dtype=torch.bool), + next_state=current, + grounded=grounded, + ) + + def execute_trajectory( + _trajectory: torch.Tensor, + *, + active: torch.Tensor, + ) -> list[torch.Tensor]: + if opens and bool(active.any()): + env.robot._qpos[:, env.left_eef_joints] = env.open_state + env.robot._qpos[:, env.right_eef_joints] = env.open_state + elif bool(active.any()): + env.robot._qpos[:, env.left_eef_joints] = env.open_state + return [] + + executor.grounder = SimpleNamespace(ground=ground) + executor.adapter = SimpleNamespace( + plan=plan, + combine=lambda _outcomes, _masks: ( + torch.zeros(1, 2, env.robot.dof), + torch.ones(1, dtype=torch.bool), + ), + execute_trajectory=execute_trajectory, + ) + actions = [ + { + "atomic_action_class": "MoveJoints", + "actor": {"arm": arm}, + "control": "hand", + "target_binding": { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": role, + }, + } + for arm, role in ( + ("left_arm", "participant"), + ("right_arm", "commit"), + ) + ] + + result = executor._execute_explicit_dual( + SimpleNamespace(id="release", actions=actions), + SimpleNamespace(id="task_01"), + torch.zeros(1, dtype=torch.bool), + ) + + released_state = executor._step_states[("task_01", "coordinated")] + left_held = released_state.get_held_object("physical_left_arm") + right_held = released_state.get_held_object("physical_right_arm") + assert result.failed.tolist() == [expected_failed] + assert (left_held is not None and right_held is not None) is expect_held + + +def _handover_held_state( + env: _FakeEnv, + entity: _FakeEntity, + *, + arm: str = "left_arm", +) -> ExecutionState: + """Build a fixture grasp on the side assigned to the transfer arm.""" + state = _held_state(env, entity, arm=arm) + held = state.get_held_object(f"physical_{arm}") + assert held is not None + _, lateral = robot_frame_axes(env) + role_axis = lateral if arm == "left_arm" else -lateral + offset = torch.cat((role_axis, role_axis.new_zeros((int(env.num_envs), 1))), dim=1) + object_to_eef = held.object_to_eef.clone() + object_to_eef[:, :3, 3] = offset * 0.02 + replacement = HeldObjectState( + semantics=held.semantics, + object_to_eef=object_to_eef, + grasp_xpos=torch.bmm(entity.get_local_pose(to_matrix=True), object_to_eef), + env_mask=held.env_mask, + ) + held_objects = dict(state.held_objects) + held_objects[f"physical_{arm}"] = replacement + return state.with_updates(held_objects=held_objects) + + +def test_handover_grounding_uses_bottom_region_and_diagonal_receive() -> None: + entities = { + "can": _FakeEntity( + "can", + _pose(0.0, 0.2, 1.2), + _rect_vertices(0.03, 0.03, 0.10), + ), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "handover", + "level": "L1", + "instruction": "Hand over the can.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "handover_complete"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + step = program.semantic_steps[0] + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "HandOver" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + + grounded = grounder.ground( + edge.actions[0], + step, + arm="coordinated", + state=state, + ) + middle = grounded.cfg["middle_object_pose"] + final = grounded.cfg["final_object_pose"] + cfg = AtomicActionAdapter(env)._build_config(grounded, HandOverOptions) + + staging_edge = next( + edge + for edge in program.edges + if edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + staging = grounder.ground( + staging_edge.actions[0], + step, + arm="left_arm", + state=state, + ) + + assert middle[0, 1, 3] == pytest.approx(0.0) + torch.testing.assert_close(final, middle) + torch.testing.assert_close(cfg.middle_object_pose, cfg.final_object_pose) + assert cfg.receive_pick_object_part == "bottom" + assert cfg.receive_approach_direction[1] < 0.0 + assert cfg.receive_approach_direction[2] < 0.0 + assert staging.motion_policy["upright_yaw_samples"] == 8 + + +def test_handover_rejects_receiver_motion_during_internal_final_phase() -> None: + adapter = AtomicActionAdapter(_FakeEnv()) + grounded = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={"transfer_arm": "left_arm"}, + ) + options = HandOverOptions(retreat_steps=4) + trajectory = torch.zeros(1, 12, adapter.env.robot.dof) + + assert bool( + adapter._handover_receiver_hold_mask( + trajectory, + grounded, + options, + tolerance=1.0e-3, + )[0] + ) + + trajectory[0, -1, 4] = 0.02 + assert not bool( + adapter._handover_receiver_hold_mask( + trajectory, + grounded, + options, + tolerance=1.0e-3, + )[0] + ) + + +def _handover_then_place_task() -> dict[str, Any]: + return { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "handover_then_place", + "level": "L3", + "instruction": "Hand over the can and place it beside the target.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "handover", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "preserve", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "place", + "task_type": "E1", + "params": { + "object_role": "can", + "target_role": "target", + "relation": "right_of", + }, + "depends_on": ["handover"], + "role": "primary", + }, + ], + "success": { + "op": "all", + "terms": [ + {"type": "handover_complete", "task_instance_id": "handover"}, + {"type": "semantic_goal", "task_instance_id": "place"}, + ], + }, + "oracle": {}, + "metadata": {}, + } + + +def test_handover_continuation_uses_stable_upright_policies() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = deepcopy(_handover_then_place_task()) + task["task_instances"][1]["params"]["orientation_goal"] = "upright" + program = load_execution_program( + instantiate_seed_graph( + task, + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate + for candidate in program.semantic_steps + if candidate.operator == "place_relative" + ) + state = _held_state(env, entities["can"], arm="right_arm") + held = state.get_held_object("physical_right_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + edges = [edge for edge in program.edges if edge.id in step.edge_ids] + staging = next( + edge + for edge in edges + if edge.actions[0]["target_binding"].get("phase") == "staging" + ) + final = next( + edge + for edge in edges + if edge.actions[0]["target_binding"].get("phase") == "final" + ) + release = next( + edge for edge in edges if edge.actions[0]["atomic_action_class"] == "Place" + ) + retreat = next( + edge + for edge in edges + if edge.actions[0]["atomic_action_class"] == "MoveEndEffector" + ) + home = next( + edge for edge in edges if edge.actions[0]["atomic_action_class"] == "MoveJoints" + ) + + grounded_staging = grounder.ground( + staging.actions[0], step, arm="right_arm", state=state + ) + grounded_final = grounder.ground( + final.actions[0], step, arm="right_arm", state=state + ) + supported_reference = _pose(0.0, 0.0, 0.90) + grounded_final_with_reference = grounder.ground( + final.actions[0], + step, + arm="right_arm", + state=state, + orientation_reference_pose=supported_reference, + ) + grounded_release = grounder.ground( + release.actions[0], step, arm="right_arm", state=state + ) + grounded_retreat = grounder.ground( + retreat.actions[0], step, arm="right_arm", state=state + ) + grounded_home = grounder.ground(home.actions[0], step, arm="right_arm", state=state) + upright = motion_policy(("orientation", "upright")) + release_defaults = resolve_motion_policy("dual_ur10", "Place", upright) + retreat_defaults = resolve_motion_policy("dual_ur10", "MoveEndEffector", upright) + + assert grounded_staging.cfg["upright_yaw_samples"] == 8 + assert grounded_final.cfg["upright_yaw_samples"] == 8 + assert grounded_final_with_reference.target_object_pose is not None + assert grounded_final_with_reference.target_object_pose[0, 2, 3] == pytest.approx( + 0.90 + ) + assert ( + grounded_release.cfg["sample_interval"] == release_defaults["sample_interval"] + ) + assert ( + grounded_release.cfg["post_hold_steps"] == release_defaults["post_hold_steps"] + ) + assert ( + grounded_retreat.cfg["sample_interval"] == retreat_defaults["sample_interval"] + ) + assert grounded_retreat.cfg["retreat_height"] == pytest.approx( + retreat_defaults["retreat_height"] + ) + assert grounded_retreat.cfg["retreat_height"] == pytest.approx(0.30) + assert grounded_retreat.motion_policy["clearance_object_uid"] == "can" + assert grounded_retreat.motion_policy["collision_exclusion_uids"] == [ + "can", + "target", + ] + assert grounded_retreat.motion_policy["collision_safety"] == "required" + assert grounded_home.motion_policy["collision_safety"] == "required" + + +def test_preserve_handover_continuation_does_not_enable_yaw_search() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = deepcopy(_handover_then_place_task()) + task["task_instances"][1]["params"]["orientation_goal"] = "preserve" + program = load_execution_program( + instantiate_seed_graph(task, {"can": "can", "target": "target"}) + ) + step = next( + candidate + for candidate in program.semantic_steps + if candidate.operator == "place_relative" + ) + state = _held_state(env, entities["can"], arm="right_arm") + held = state.get_held_object("physical_right_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + final = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("phase") == "final" + ) + reference = _pose(0.0, 0.0, 0.90) + + grounded = grounder.ground( + final.actions[0], + step, + arm="right_arm", + state=state, + orientation_reference_pose=reference, + ) + + assert "upright_yaw_samples" not in grounded.cfg + assert grounded.target_object_pose is not None + torch.testing.assert_close( + grounded.target_object_pose[:, :3, :3], + reference[:, :3, :3], + ) + + +def test_dual_franka_handover_uses_explicit_exchange_clearance() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"], arm="left_arm") + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + handover_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_goal" + ) + staging = grounder.ground( + staging_edge.actions[0], step, arm="left_arm", state=state + ) + handover = grounder.ground( + handover_edge.actions[0], step, arm="coordinated", state=state + ) + + assert staging.motion_policy["exchange_clearance"] > 0.0 + assert handover.motion_policy["exchange_clearance"] > 0.0 + assert handover.motion_policy["lift_height"] > 0.0 + assert staging.target_object_pose is not None + live_object_pose = entities["can"].get_local_pose(to_matrix=True) + assert handover.motion_policy["middle_object_pose"][0, 2, 3] == pytest.approx( + live_object_pose[0, 2, 3] + ) + + +def test_handover_candidates_avoid_occupied_table_center_and_lift_payload() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.03), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.35, 0.0, 1.03), _box_vertices(0.03)), + "notebook": _FakeEntity("notebook", _pose(0.0, 0.0, 1.04), _box_vertices(0.05)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"], arm="left_arm") + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + handover_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_goal" + ) + + staging = grounder.ground( + staging_edge.actions[0], step, arm="left_arm", state=state + ) + candidates = grounder.ground_candidates( + handover_edge.actions[0], step, arm="coordinated", state=state + ) + + assert staging.target_object_pose is not None + assert torch.linalg.vector_norm(staging.target_object_pose[0, :2, 3]) > 0.10 + assert float(staging.target_object_pose[0, 2, 3]) >= 1.15 + assert len(candidates) == 4 + assert all( + torch.linalg.vector_norm(candidate.cfg["middle_object_pose"][0, :2, 3]) > 0.10 + for candidate in candidates[:2] + ) + + +def test_on_placement_grounding_samples_bounded_live_support_poses() -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.15, 0.01), + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edge = next( + item + for item in program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + grounder = ActionGrounder(program, env, lambda _uid: None) + + candidates = grounder.ground_candidates( + edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert len(candidates) == 5 + assert [item.motion_policy["placement_candidate_index"] for item in candidates] == [ + 0, + 1, + 2, + 3, + 4, + ] + offsets = [item.motion_policy["placement_xy_offset"][0] for item in candidates] + assert len({tuple(float(value) for value in offset) for offset in offsets}) == 5 + support_lower = torch.tensor([-0.20, -0.15]) + support_upper = torch.tensor([0.20, 0.15]) + for item in candidates: + center = item.target_object_pose[0, :2, 3] + assert torch.all(center >= support_lower) + assert torch.all(center <= support_upper) + + +def test_on_placement_candidates_respect_support_geometry_origin() -> None: + support_vertices = _rect_vertices(0.10, 0.08, 0.01) + support_vertices[:, 0] += 0.25 + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.10, 0.0, 0.75), + support_vertices, + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edge = next( + item + for item in program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + candidates = ActionGrounder(program, env, lambda _uid: None).ground_candidates( + edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + support_world = support_vertices[:, :2] + torch.tensor([0.10, 0.0]) + lower = support_world.min(dim=0).values + 0.002 + upper = support_world.max(dim=0).values - 0.002 + payload_local = entities["payload"]._vertices[:, :2] + for candidate in candidates: + origin = candidate.target_object_pose[0, :2, 3] + assert torch.all(origin + payload_local.min(dim=0).values >= lower) + assert torch.all(origin + payload_local.max(dim=0).values <= upper) + + +def test_build_stack_root_compiles_to_generic_table_support() -> None: + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "stack", + "operator": "build_stack", + "objects": ["base", "nested"], + "actor": {"mode": "auto"}, + "goal": { + "anchor": "table_center", + "stack_mode": "nested", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ) + + root, child = program.semantic_steps + assert root.goal["relation"] == "on" + assert root.goal["reference_object"] == "table" + assert root.postcondition["reference_object"] == "table" + assert child.goal["relation"] == "inside" + assert child.goal["reference_object"] == "base" + + +def test_executor_tries_next_placement_pose_after_planning_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.15, 0.01), + ), + } + env = _FakeEnv(entities) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + edge = next( + item + for item in executor.program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + state = ExecutionState(last_qpos=env.robot.get_qpos()) + + def plan(grounded: GroundedAction, _state: ExecutionState) -> ActionOutcome: + index = int(grounded.motion_policy["placement_candidate_index"]) + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([index == 1]), + next_state=state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.adapter, "plan", plan) + grounded, outcome = executor._ground_and_plan_candidates( + edge.actions[0], + step, + arm="left_arm", + state=state, + active=torch.tensor([True]), + ) + + assert bool(outcome.success[0]) + assert grounded.motion_policy["placement_candidate_index"] == 1 + assert outcome.planner_trace["selected_grounding_candidate"] == 1 + assert len(outcome.planner_trace["grounding_candidates"]) == 2 + + +def test_post_release_candidate_search_skips_the_released_pose( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, -0.20, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.15, 0.01), + ), + } + env = _FakeEnv(entities) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + edge = next( + item + for item in executor.program.edges + if item.id in step.edge_ids + and item.actions[0]["target_binding"].get("kind") == "semantic_goal" + and item.actions[0]["target_binding"].get("phase", "final") != "staging" + ) + state = ExecutionState(last_qpos=env.robot.get_qpos()) + executor._placement_candidate_history[(step.id, "left_arm")] = {0} + + def plan(grounded: GroundedAction, _state: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.adapter, "plan", plan) + grounded, outcome = executor._ground_and_plan_candidates( + edge.actions[0], + step, + arm="left_arm", + state=state, + active=torch.tensor([True]), + ) + + assert grounded.motion_policy["placement_candidate_index"] == 1 + assert outcome.planner_trace["grounding_candidates"][0] == { + "candidate_index": 0, + "status": "previously_released", + } + + +def test_unstable_placement_recovery_replays_pick_before_another_release( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + env = _FakeEnv(entities) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + replayed_actions: list[str] = [] + verification_count = 0 + + def ensure_assignment(_step: SemanticStep, failed: torch.Tensor) -> None: + executor._assignments[_step.id] = [ + None if bool(failed[env_id]) else "left_arm" + for env_id in range(len(failed)) + ] + + def execute( + edge: ExecutionEdge, + _step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + replayed_actions.append(str(edge.actions[0]["atomic_action_class"])) + return _EdgeResult([], failed.clone(), [], executed=~failed) + + def verify( + _step: SemanticStep, + failed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + nonlocal verification_count + verification_count += 1 + success = torch.tensor([verification_count == 2]) & ~failed + return failed | ~success, success, executor._entity_pose("payload")[:, :3, 3] + + monkeypatch.setattr(executor, "_ensure_assignment", ensure_assignment) + monkeypatch.setattr(executor, "_execute_edge_with_retries", execute) + monkeypatch.setattr(executor, "_verify_step", verify) + recorder = RuntimeRecorder( + executor.program, + num_envs=1, + enabled=False, + ) + + recovery = executor._recover_unstable_placement( + step, + torch.tensor([True]), + recorder=recorder, + ) + + first_action = str( + executor.edges[step.edge_ids[0]].actions[0]["atomic_action_class"] + ) + assert replayed_actions.count(first_action) == 2 + assert verification_count == 2 + assert bool(recovery.succeeded[0]) + assert not bool(recovery.failed[0]) + assert recovery.failure_events == [] + + +def test_unstable_placement_recovery_reports_its_own_planning_blocker( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def fail_assignment(*_args: Any, **_kwargs: Any) -> None: + raise RuntimeError("no plan") + + monkeypatch.setattr(executor, "_ensure_assignment", fail_assignment) + + recovery = executor._recover_unstable_placement( + step, + torch.tensor([True]), + recorder=RuntimeRecorder(executor.program, num_envs=1, enabled=False), + ) + + assert bool(recovery.failed[0]) + assert bool(recovery.covered_failures[0]) + assert len(recovery.failure_events) == 1 + event = recovery.failure_events[0] + assert event["failure_type"] == "search_exhausted" + assert event["phase"] == "recovery" + assert event["origin_edge_id"] == step.edge_ids[-1] + assert event["blocking_edge_id"] == step.edge_ids[0] + + +def test_handover_height_accounts_for_obstacle_and_tool_envelope() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.03), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.35, 0.0, 1.03), _box_vertices(0.03)), + "shelf": _FakeEntity( + "shelf", + _pose(0.0, 0.0, 1.05), + _rect_vertices(0.40, 0.35, 0.05), + ), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging = next( + edge.actions[0] + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + + grounded = grounder.ground(staging, step, arm="left_arm", state=state) + + assert grounded.target_object_pose is not None + obstacle_top = 1.10 + object_bottom = -0.03 + object_clearance = 0.06 + tool_vertical_envelope = 0.025 + 0.04 + expected_height = ( + obstacle_top + object_clearance + tool_vertical_envelope - object_bottom + ) + assert grounded.target_object_pose[0, 2, 3] == pytest.approx(expected_height) + + +def test_handover_workspace_rejects_points_outside_shared_reach() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + staging = next( + edge.actions[0] + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_staging" + ) + staging = { + **staging, + "motion_policy_config": {"exchange_maximum_reach": 0.20}, + } + + with pytest.raises(ValueError, match="reachable intersection"): + grounder.ground(staging, step, arm="left_arm", state=state) + + +def test_handover_grounding_preserves_the_original_object_affordance() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _handover_held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + held.object_to_eef[:, 1, 3] *= -1.0 + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + action = next( + edge.actions[0] + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("kind") == "handover_goal" + ) + + grounded = grounder.ground(action, step, arm="coordinated", state=state) + + assert grounded.target.semantics.affordance is held.semantics.affordance + + +def test_robot_relative_left_uses_live_right_to_left_arm_axis() -> None: + env = _FakeEnv() + + forward, lateral = robot_frame_axes(env) + offset = relation_offset( + env, + "left_of", + frame="robot", + forward_distance=0.10, + lateral_distance=0.12, + dtype=torch.float32, + device=env.device, + ) + + torch.testing.assert_close(forward, torch.tensor([[-1.0, 0.0]])) + torch.testing.assert_close(lateral, torch.tensor([[0.0, -1.0]])) + assert offset is not None + torch.testing.assert_close(offset, torch.tensor([[0.0, -0.12, 0.0]])) + + +def test_directional_verification_rejects_grounded_target_on_wrong_side() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.12, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + step = program.semantic_steps[0] + step.goal["relation_frame"] = "robot" + executor._targets[step.id] = entities["can"].get_local_pose(to_matrix=True)[ + :, :3, 3 + ] + executor._policies[step.id] = { + "postcondition_tolerance": 0.08, + "relation_clearance": 0.01, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(failed[0]) + assert not bool(success[0]) + + +def test_directional_verification_accepts_support_height_settling() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.12, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + step = program.semantic_steps[0] + step.goal["relation_frame"] = "robot" + executor._targets[step.id] = torch.tensor([[0.0, -0.12, 0.90]]) + executor._policies[step.id] = { + "postcondition_tolerance": 0.08, + "relation_clearance": 0.01, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + + +def test_legacy_released_above_relation_verifies_as_physical_support() -> None: + entities = { + "payload": _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ), + "support": _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + step = program.semantic_steps[0] + step.goal["relation"] = "above" + executor._targets[step.id] = torch.tensor([[0.0, 0.0, 1.0]]) + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + + +def test_handover_retreat_clears_exchange_toward_transfer_workspace() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.106), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_franka" + high_left = _pose(0.0, 0.2, 1.106) + right = _pose(0.0, -0.2, 0.8) + env.get_current_xpos_agent = lambda: (high_left, right) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + step = next( + candidate for candidate in program.semantic_steps if candidate.id == "handover" + ) + state = _held_state(env, entities["can"], arm="left_arm") + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + retreat_edge = next( + edge + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["target_binding"].get("source") == "handover" + ) + + grounded = grounder.ground( + retreat_edge.actions[0], step, arm="left_arm", state=state + ) + + torch.testing.assert_close( + grounded.target.xpos[0, :2, 3], + torch.tensor([0.0, 0.10]), + ) + assert grounded.target.xpos[0, 2, 3] == pytest.approx(1.206) + assert grounded.cfg["retreat_distance"] == pytest.approx(0.10) + assert grounded.cfg["maximum_eef_height"] == pytest.approx(1.50) + + +def test_handover_retreat_and_home_block_receiver_continuation() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + executor = ProgramExecutor(program, _FakeEnv(entities), record_runtime=False) + handover = next( + step for step in program.semantic_steps if step.operator == "handover" + ) + handover_edges = [edge for edge in program.edges if edge.id in handover.edge_ids] + retreat = next( + edge + for edge in handover_edges + if edge.actions[0]["target_binding"].get("source") == "handover" + ) + home = next( + edge + for edge in handover_edges + if edge.actions[0]["target_binding"].get("operation") == "handover_home" + ) + + assert executor._edge_failure_policy(retreat) == "safety_required" + assert executor._edge_failure_policy(home) == "best_effort" + + +def test_release_retreat_is_required_and_exact_home_is_best_effort() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + } + program = load_execution_program( + compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, _FakeEnv(entities), record_runtime=False) + step = program.semantic_steps[0] + edges = [edge for edge in program.edges if edge.id in step.edge_ids] + retreat = next( + edge + for edge in edges + if edge.actions[0]["atomic_action_class"] == "MoveEndEffector" + ) + home = next( + edge for edge in edges if edge.actions[0]["atomic_action_class"] == "MoveJoints" + ) + + assert executor._edge_failure_policy(retreat) == "safety_required" + assert executor._edge_failure_policy(home) == "best_effort" + + +def test_best_effort_home_does_not_veto_required_arm_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + graph = compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"reference_object": "target", "relation": "left_of"}, + "depends_on": [], + } + ) + ) + executor = ProgramExecutor( + load_execution_program(graph), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def ground(action: dict[str, Any], *_args: Any, **_kwargs: Any) -> GroundedAction: + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm="left_arm", + control=str(action["control"]), + target=SimpleNamespace(xpos=None), + cfg={}, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 1, executor.env.robot.dof), + success=torch.tensor([grounded.action_class != "MoveJoints"]), + next_state=state, + grounded=grounded, + planner_trace={"primary_strategy": "motion_gen"}, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor, "_with_downstream_targets", lambda *args: args[-1]) + monkeypatch.setattr(executor.adapter, "plan", plan) + + candidate = executor._candidate(step, "left_arm", torch.tensor([False])) + + assert bool(candidate.feasible[0]) + assert any("best-effort action degraded" in item for item in candidate.warnings) + + +def test_best_effort_home_exception_does_not_fail_semantic_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + graph = compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"reference_object": "target", "relation": "left_of"}, + "depends_on": [], + } + ) + ) + executor = ProgramExecutor( + load_execution_program(graph), + _FakeEnv(), + settle_steps=0, + record_runtime=False, + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, _failed: executor._assignments.setdefault(step.id, ["left_arm"]), + ) + + def execute(edge: ExecutionEdge, _step: SemanticStep, *, failed: torch.Tensor): + if executor._edge_failure_policy(edge) == "best_effort": + raise RuntimeError("home search failed") + return SimpleNamespace( + actions=[], + failed=failed.clone(), + grounded=[], + planner_traces=[], + executed=~failed, + ) + + monkeypatch.setattr(executor, "_execute_edge_with_retries", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda _step, failed: (failed, ~failed, torch.zeros(1, 3)), + ) + + result = executor.run() + + assert bool(result.success[0]) + assert len(result.failure_events) == 1 + assert result.failure_events[0]["failure_type"] == "search_exhausted" + assert result.failure_events[0]["failure_policy"] == "best_effort" + assert result.failure_events[0]["fatal"] is False + assert result.failure_events[0]["evidence"]["exception"].endswith( + "home search failed" + ) + + +def test_candidate_failure_reports_real_blocking_safety_edge( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "target", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + def ground(action: dict[str, Any], *_args: Any, **_kwargs: Any) -> GroundedAction: + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm="left_arm", + control=str(action["control"]), + target=SimpleNamespace(xpos=None), + cfg={}, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + return ActionOutcome( + trajectory=torch.zeros(1, 1, executor.env.robot.dof), + success=torch.tensor([grounded.action_class != "MoveEndEffector"]), + next_state=state, + grounded=grounded, + planner_trace={"primary_strategy": "motion_gen"}, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor, "_with_downstream_targets", lambda *args: args[-1]) + monkeypatch.setattr(executor.adapter, "plan", plan) + candidate = executor._candidate(step, "left_arm", torch.tensor([False])) + executor._assignments[step.id] = [None] + executor._report_candidates(step, (candidate,)) + first_edge = executor.edges[step.edge_ids[0]] + + events = executor._failure_events( + first_edge, + step, + torch.tensor([True]), + postcondition=False, + executed=torch.tensor([False]), + fallen_transition=torch.tensor([False]), + ) + + assert len(events) == 1 + event = events[0] + assert event["failure_type"] == "search_exhausted" + assert event["failure_policy"] == "safety_required" + assert event["atomic_action"] == "MoveEndEffector" + assert event["blocking_edge_id"] != first_edge.id + assert event["planning_stage"] == "candidate_suffix" + assert "not a geometric proof" in event["reason"] + + +def test_on_relation_rejects_preserve_orientation_drift() -> None: + rotated = _pose(0.0, 0.0, 0.82) + rotated[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "can": _FakeEntity("can", rotated, _rect_vertices(0.03, 0.03, 0.06)), + "notebook": _FakeEntity( + "notebook", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "notebook", + "relation": "on", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._orientation_references[step.id] = _pose(0.0, 0.0, 0.82) + executor._policies[step.id] = { + "preserve_orientation_tolerance": torch.pi / 12, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(failed[0]) + assert not bool(success[0]) + assert executor._orientation_errors[step.id][0] > torch.pi / 12 + + +def test_inside_relation_accepts_settling_orientation_drift() -> None: + rotated = _pose(0.02, -0.02, 0.72) + rotated[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "can": _FakeEntity("can", rotated, _rect_vertices(0.03, 0.03, 0.06)), + "basket": _FakeEntity( + "basket", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.10, 0.10, 0.08), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._orientation_references[step.id] = _pose(0.02, -0.02, 0.72) + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + assert step.id not in executor._orientation_errors + + +@pytest.mark.parametrize( + ("orientation_goal", "expected_success"), + (("upright", False), ("none", True)), +) +def test_on_relation_applies_only_the_requested_orientation_goal( + orientation_goal: str, + expected_success: bool, +) -> None: + fallen = _pose(0.0, 0.0, 0.79) + fallen[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "can": _FakeEntity("can", fallen, _rect_vertices(0.03, 0.03, 0.06)), + "notebook": _FakeEntity( + "notebook", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "notebook", + "relation": "on", + "orientation_goal": orientation_goal, + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(success[0]) is expected_success + assert bool(failed[0]) is not expected_success + + +def test_support_stability_window_rejects_motion_after_initial_contact() -> None: + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ) + env = _FakeEnv({"payload": payload, "support": support}) + update_count = 0 + + def update(*, step: int) -> None: + nonlocal update_count + del step + update_count += 1 + payload.lin_vel[:, 0] = 0.10 + + env.sim.update = update + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert update_count == executor.support_stability_samples - 1 + assert bool(failed[0]) + assert not bool(success[0]) + + +def test_support_stability_reads_real_rigid_object_body_state() -> None: + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ) + del payload.lin_vel + del payload.ang_vel + payload.body_state = torch.zeros(1, 13) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "payload", "left_arm"))) + ), + _FakeEnv({"payload": payload}), + settle_steps=0, + record_runtime=False, + ) + + assert bool(executor._entity_motion_stable("payload")[0]) + payload.body_state[:, 7] = 0.10 + assert not bool(executor._entity_motion_stable("payload")[0]) + + +def test_final_support_revalidation_detects_later_chain_damage() -> None: + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.79), + _rect_vertices(0.03, 0.03, 0.03), + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.08, 0.01), + ) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "payload", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "support", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv({"payload": payload, "support": support}), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + assert not bool(failed[0]) + assert bool(success[0]) + + payload._pose[:, 2, 3] += 0.20 + failures = executor._revalidate_support_relations() + + assert bool(failures[step.id][0]) + + +def test_support_relation_state_rejects_cycles() -> None: + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place_b", + "operator": "place_relative", + "object": "b", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "a", + "relation": "on", + "orientation_goal": "none", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv( + { + "a": _FakeEntity("a", _pose(0.0, 0.0, 0.75), _box_vertices(0.03)), + "b": _FakeEntity("b", _pose(0.0, 0.0, 0.81), _box_vertices(0.03)), + } + ), + settle_steps=0, + record_runtime=False, + ) + step_b = executor.program.semantic_steps[0] + step_a = replace(step_b, id="prior", object_uid="a") + executor._commit_support_relation(step_a, "b", torch.tensor([True])) + + cycle_free = executor._support_cycle_free("b", "a", torch.tensor([True])) + + assert not bool(cycle_free[0]) + + +def test_standalone_handover_assigns_its_pickup_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + executor = ProgramExecutor(program, _FakeEnv(entities), record_runtime=False) + step = program.semantic_steps[0] + calls: list[str] = [] + + def candidate(_step: SemanticStep, arm: str, _failed: torch.Tensor) -> Any: + calls.append(arm) + return SimpleNamespace(feasible=torch.tensor([True])) + + monkeypatch.setattr(executor, "_candidate", candidate) + monkeypatch.setattr(executor, "_report_candidates", lambda *_args: None) + + executor._ensure_assignment(step, torch.zeros(1, dtype=torch.bool)) + + assert calls == ["left_arm"] + assert executor._assignments[step.id] == ["left_arm"] + + +def test_standalone_handover_candidate_stops_before_coordinated_transfer( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + env = _FakeEnv(entities) + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + executor = ProgramExecutor(program, env, record_runtime=False) + step = program.semantic_steps[0] + planned_actions: list[str] = [] + + def ground( + action: dict[str, Any], + _step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> GroundedAction: + del state, reference_eef_pose, orientation_reference_pose + return GroundedAction( + action_class=str(action["atomic_action_class"]), + arm=arm, + control=str(action["control"]), + target=SimpleNamespace(), + cfg={}, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + planned_actions.append(grounded.action_class) + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor, "_with_downstream_targets", lambda *args: args[-1]) + monkeypatch.setattr(executor.adapter, "plan", plan) + + candidate = executor._candidate( + step, + "left_arm", + torch.zeros(1, dtype=torch.bool), + ) + + assert bool(candidate.feasible[0]) + assert planned_actions == ["PickUp", "MoveHeldObject"] + assert set(candidate.plans) == set(step.edge_ids[:2]) + + +def test_failed_handover_keeps_transfer_ownership( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + program = load_execution_program( + instantiate_seed_graph( + task, + {"can": "can"}, + ) + ) + step = program.semantic_steps[0] + edge = next( + candidate + for candidate in program.edges + if candidate.actions[0]["atomic_action_class"] == "HandOver" + ) + state = _held_state(env, entities["can"], arm="left_arm") + executor = ProgramExecutor(program, env, record_runtime=False) + executor._assignments[step.id] = ["coordinated"] + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = state + grounded = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={}, + ) + failed_outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([False]), + next_state=state, + grounded=grounded, + ) + monkeypatch.setattr( + executor.grounder, + "ground", + lambda *_args, **_kwargs: grounded, + ) + monkeypatch.setattr( + executor.adapter, "plan", lambda *_args, **_kwargs: failed_outcome + ) + monkeypatch.setattr( + executor.adapter, + "execute_trajectory", + lambda *_args, **_kwargs: [], + ) + + result = executor._execute_coordinated(edge, step, torch.tensor([False])) + + assert bool(result.failed[0]) + assert executor._object_owners["can"] == ["left_arm"] + assert executor._arm_owners["left_arm"] == ["can"] + assert executor._arm_owners["right_arm"] == [None] + assert ("can", "left_arm") in executor._object_states + assert ("can", "right_arm") not in executor._object_states + + +def test_handover_commits_receiver_ownership_only_after_physical_verification( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = _handover_then_place_task() + task["level"] = "L1" + task["task_instances"] = [task["task_instances"][0]] + task["success"] = {"type": "handover_complete"} + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + step = program.semantic_steps[0] + edge = next( + candidate + for candidate in program.edges + if candidate.actions[0]["atomic_action_class"] == "HandOver" + ) + transfer_state = _held_state(env, entities["can"], arm="left_arm") + receiver_state = _held_state(env, entities["can"], arm="right_arm") + executor = ProgramExecutor(program, env, record_runtime=False) + executor._assignments[step.id] = ["coordinated"] + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = transfer_state + grounded = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={}, + motion_policy={"held_position_tolerance": 0.03}, + ) + successful_outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=receiver_state, + grounded=grounded, + ) + observed_poses: list[torch.Tensor] = [] + + def ground_candidates(*_args, **_kwargs): + observed_poses.append(entities["can"].get_local_pose(to_matrix=True)) + return (grounded,) + + monkeypatch.setattr( + executor.grounder, + "ground_candidates", + ground_candidates, + ) + monkeypatch.setattr( + executor.adapter, "plan", lambda *_args, **_kwargs: successful_outcome + ) + monkeypatch.setattr( + executor.adapter, + "execute_trajectory", + lambda *_args, **_kwargs: [], + ) + + entities["can"]._pose[:, 0, 3] += 0.30 + result = executor._execute_coordinated(edge, step, torch.tensor([False])) + + assert observed_poses[0][0, 0, 3] == pytest.approx(0.30) + assert result.planner_traces[0]["execution_replanned_from_live_state"] is True + assert bool(result.failed[0]) + assert executor._object_owners["can"] == [None] + assert executor._arm_owners["left_arm"] == [None] + assert executor._arm_owners["right_arm"] == [None] + assert ("can", "right_arm") not in executor._object_states + + +def test_handover_defers_clearance_verification_to_retreat_action() -> None: + adapter = AtomicActionAdapter(_FakeEnv()) + + assert adapter.capabilities.get("HandOver").verifier_hook is None + assert adapter.capabilities.get("MoveEndEffector").verifier_hook is not None + + +def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv( + { + "can": entity, + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + ) + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "orient_then_handover", + "level": "L3", + "instruction": "Orient the can, then hand it over.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "can", + "required_arm": "right_arm", + "orientation_goal": "upright", + "support_role": "table", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "right_arm", + "receive_arm": "left_arm", + }, + "depends_on": ["task_01"], + "role": "primary", + }, + ], + "success": {"type": "handover_complete", "task_instance_id": "task_02"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + orient_step = next( + candidate for candidate in program.semantic_steps if candidate.id == "task_01" + ) + handover_step = next( + candidate for candidate in program.semantic_steps if candidate.id == "task_02" + ) + orient_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["atomic_action_class"] == "AxisAlign" + ) + handover_edge = next( + candidate + for candidate in program.edges + if candidate.id in handover_step.edge_ids + if candidate.actions[0]["atomic_action_class"] == "PickUp" + ) + vertices = _box_vertices(0.03) + semantics = ObjectSemantics( + affordance=AntipodalAffordance( + object_label="can", + mesh_vertices=vertices, + ), + geometry={"mesh_vertices": vertices}, + label="can", + entity=entity, + ) + grounder = ActionGrounder(program, env, lambda _uid: semantics) + + orient_alignment = grounder.ground( + orient_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + handover_pickup = grounder.ground( + handover_edge.actions[0], + handover_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert "approach_direction_mode" not in orient_alignment.cfg + assert "approach_direction_mode" not in handover_pickup.cfg + assert handover_pickup.cfg["pick_object_part"] == "top" + assert isinstance(orient_alignment.target, AxisAlignGoal) + assert orient_alignment.target.grasp_xpos is None + assert torch.equal( + orient_alignment.target.object_target_pose, + orient_alignment.target_object_pose, + ) + assert orient_alignment.target_object_pose is not None + assert isinstance( + orient_alignment.target.semantics.affordance, + AxisAlignAffordance, + ) + assert torch.equal( + orient_alignment.target.semantics.affordance.internal_axis, + torch.tensor([1.0, 0.0, 0.0]), + ) + assert handover_pickup.target.grasp_xpos is None + assert isinstance( + handover_pickup.target.semantics.affordance, + AntipodalAffordance, + ) + assert handover_pickup.target.semantics.affordance is semantics.affordance + + +def test_pour_grounding_targets_receiver_and_requires_physical_contents() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "pour_grounding", + "level": "L1", + "instruction": "Pour the ball from the cup into the bin.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E3", + "params": { + "source_role": "cup", + "target_role": "bin", + "content_roles": ["ball"], + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "poured"}, + "oracle": {}, + "metadata": {}, + } + bindings = {"cup": "source", "bin": "target", "ball": "content"} + program = load_execution_program(instantiate_seed_graph(task, bindings)) + step = program.semantic_steps[0] + edges = {edge.actions[0]["atomic_action_class"]: edge for edge in program.edges} + source = _FakeEntity("source", _pose(-0.2, 0.0, 0.7), _box_vertices(0.05)) + target = _FakeEntity("target", _pose(0.2, 0.0, 0.7), _box_vertices(0.10)) + content = _FakeEntity("content", _pose(-0.2, 0.0, 0.7), _box_vertices(0.01)) + env = _FakeEnv({"source": source, "target": target, "content": content}) + semantics = ObjectSemantics( + affordance=AntipodalAffordance( + object_label="source", + mesh_vertices=_box_vertices(0.05), + ), + geometry={}, + label="source", + entity=source, + ) + grounder = ActionGrounder(program, env, lambda _uid: semantics) + state = ExecutionState(last_qpos=env.robot.get_qpos()) + + pickup = grounder.ground( + edges["PickUp"].actions[0], step, arm="right_arm", state=state + ) + staging = grounder.ground( + edges["MoveHeldObject"].actions[0], + step, + arm="right_arm", + state=state, + ) + pouring = grounder.ground( + edges["Pour"].actions[0], step, arm="right_arm", state=state + ) + + assert isinstance(pickup.target.semantics.affordance, AxisAlignAffordance) + assert torch.allclose( + pickup.target.semantics.affordance.internal_axis, + torch.tensor([0.0, 1.0, 0.0]), + ) + assert staging.target_object_pose is not None + assert staging.target_object_pose[0, 0, 3] == pytest.approx(0.2) + assert staging.target_object_pose[0, 2, 3] > 0.8 + assert isinstance(pouring.target, PourGoal) + assert pouring.cfg["rotate_angle"] == pytest.approx(torch.pi / 2.0) + + task["task_instances"][0]["params"]["content_roles"] = [] + blocked = load_execution_program( + instantiate_seed_graph(task, {"cup": "source", "bin": "target"}) + ) + blocked_step = blocked.semantic_steps[0] + blocked_pour = next( + edge + for edge in blocked.edges + if edge.actions[0]["atomic_action_class"] == "Pour" + ) + blocked_grounder = ActionGrounder(blocked, env, lambda _uid: semantics) + with pytest.raises(ValueError, match="independently observable"): + blocked_grounder.ground( + blocked_pour.actions[0], + blocked_step, + arm="right_arm", + state=state, + ) + + +@pytest.mark.parametrize( + ("task_type", "initial_qpos", "direction", "target_state"), + (("E6", 0.0, "pull", "open"), ("E7", 0.2, "push", "closed")), +) +def test_articulation_grounding_reuses_slide_and_observes_joint_state( + task_type: str, + initial_qpos: float, + direction: str, + target_state: str, +) -> None: + task, _ = make_task_spec(task_type) + graph = instantiate_seed_graph(task, {"object_01": "drawer"}) + program = load_execution_program(graph) + articulation = _FakeArticulation("drawer", initial_qpos) + env = _FakeEnv(articulations={"drawer": articulation}) + grounder = ActionGrounder( + program, + env, + lambda _uid: ObjectSemantics(affordance=Affordance(), geometry={}), + ) + step = program.semantic_steps[0] + edge = program.edges[0] + + grounded = grounder.ground( + edge.actions[0], + step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, SlideGoal) + assert grounded.target.grasp_xpos is not None + assert isinstance(grounded.target.semantics.affordance, SlideAffordance) + assert grounded.cfg["direction"] == direction + assert grounded.cfg["translation_distance"] == pytest.approx(0.2) + assert grounded.cfg["articulation_joint_name"] == "slide_joint" + assert grounded.cfg["articulation_initial_qpos"].item() == pytest.approx( + initial_qpos + ) + assert torch.allclose( + grounded.target.semantics.affordance.translation_axis, + torch.tensor([-1.0, 0.0, 0.0]), + ) + assert edge.actions[0]["failure_policy"] == "task_required" + + predicate = { + "type": "articulation_joint_near", + "object": "drawer", + "target_state": target_state, + } + assert not bool(evaluate_predicate(env, predicate)[0]) + articulation._qpos[0, 0] = 0.19 if target_state == "open" else 0.01 + assert bool(evaluate_predicate(env, predicate)[0]) + + +def test_turn_knob_requires_setting_map_and_reuses_twist() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "turn_knob", + "level": "L1", + "instruction": "Turn the knob to setting two.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E8", + "params": { + "object_role": "knob", + "target_setting": 2, + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "articulation_joint_near"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program(instantiate_seed_graph(task, {"knob": "dial"})) + articulation = _FakeArticulation("dial", 0.0) + articulation._joint_info.joint_type = SimpleNamespace(name="REVOLUTE") + articulation._limits = torch.tensor([[[-1.0, 1.0]]]) + env = _FakeEnv(articulations={"dial": articulation}) + env.agent_config = { + "articulation_settings": {"dial": {"slide_joint": [-1.0, 0.0, 1.0]}} + } + grounder = ActionGrounder( + program, + env, + lambda _uid: ObjectSemantics(affordance=Affordance(), geometry={}), + ) + step = program.semantic_steps[0] + + grounded = grounder.ground( + program.edges[0].actions[0], + step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, TwistGoal) + assert isinstance(grounded.target.semantics.affordance, TwistAffordance) + assert grounded.cfg["twist_angle"] == pytest.approx(1.0) + assert grounded.cfg["articulation_joint_name"] == "slide_joint" + assert grounded.cfg["articulation_initial_qpos"].item() == pytest.approx(0.0) + assert grounded.target.semantics.affordance.grasp_position == pytest.approx( + (0.04, 0.0, 0.0) + ) + predicate = { + **step.postcondition, + "joint_name": "slide_joint", + "target_qpos": 1.0, + } + assert not bool(evaluate_predicate(env, predicate)[0]) + articulation._qpos[0, 0] = 0.99 + assert bool(evaluate_predicate(env, predicate)[0]) + + env.agent_config = {"articulation_settings": {}} + with pytest.raises(ValueError, match="explicit setting_values"): + ActionGrounder( + program, + env, + lambda _uid: ObjectSemantics(affordance=Affordance(), geometry={}), + ).ground( + program.edges[0].actions[0], + program.semantic_steps[0], + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + +def test_handover_clearance_verifier_checks_distance_and_transfer_side() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, -0.2, 1.0), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program( + instantiate_seed_graph( + _handover_then_place_task(), + {"can": "can", "target": "target"}, + ) + ) + executor = ProgramExecutor(program, env, record_runtime=False) + hook = executor.adapter.capabilities.get("MoveEndEffector").verifier_hook + assert hook is not None + _, lateral = robot_frame_axes(env) + grounded = GroundedAction( + action_class="MoveEndEffector", + arm="left_arm", + control="arm", + target=SimpleNamespace(), + cfg={}, + motion_policy={ + "clearance_object_uid": "can", + "transfer_arm": "left_arm", + "transfer_role_axis": torch.cat( + (lateral, lateral.new_zeros((1, 1))), dim=1 + ), + "minimum_transfer_clearance": 0.10, + "minimum_transfer_lateral_clearance": 0.06, + }, + ) + outcome = SimpleNamespace(grounded=grounded) + attempted = torch.tensor([True]) + + assert not bool( + hook( + executor=executor, + step=program.semantic_steps[0], + arm="left_arm", + outcome=outcome, + attempted=attempted, + )[0] + ) + + clear_left = _pose(0.0, -0.4, 1.0) + env.get_current_xpos_agent = lambda: (clear_left, _pose(0.0, 0.2, 1.0)) + assert bool( + hook( + executor=executor, + step=program.semantic_steps[0], + arm="left_arm", + outcome=outcome, + attempted=attempted, + )[0] + ) + + +@pytest.mark.parametrize("arm", ["left_arm", "right_arm"]) +def test_handover_source_policy_uses_pickup_default_top_down_approach( + arm: str, +) -> None: + action = GroundedAction( + action_class="PickUp", + arm=arm, + control="arm", + target=SimpleNamespace(), + cfg={"pick_object_part": "top"}, + ) + + cfg = AtomicActionAdapter(_FakeEnv())._build_config(action, PickUpOptions) + + assert cfg.pick_object_part == "top" + torch.testing.assert_close( + cfg.approach_direction, + torch.tensor([0.0, 0.0, -1.0]), + ) + + +@pytest.mark.parametrize( + ("transfer_arm", "expected_world_y"), + [("left_arm", -1.0), ("right_arm", 1.0)], +) +def test_handover_receiver_uses_the_mirrored_diagonal_approach( + transfer_arm: str, + expected_world_y: float, +) -> None: + env = _FakeEnv() + action = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={ + "transfer_arm": transfer_arm, + "middle_object_pose": torch.eye(4).unsqueeze(0), + "final_object_pose": torch.eye(4).unsqueeze(0), + }, + ) + + cfg = AtomicActionAdapter(env)._build_config(action, HandOverOptions) + + diagonal = 2.0**-0.5 + assert cfg.receive_approach_direction[0] == pytest.approx(0.0) + assert cfg.receive_approach_direction[1] == pytest.approx( + expected_world_y * diagonal + ) + assert cfg.receive_approach_direction[2] == pytest.approx(-diagonal) + _, lateral = robot_frame_axes(env) + receive_side = "right_arm" if transfer_arm == "left_arm" else "left_arm" + receiver_outward = lateral[0] if receive_side == "left_arm" else -lateral[0] + pre_grasp_offset = -cfg.receive_approach_direction[:2] * cfg.pre_grasp_distance + assert torch.dot(pre_grasp_offset, receiver_outward) > 0.0 + + +@pytest.mark.parametrize( + ("transfer_arm", "expected_x"), + [("left_arm", 1.0), ("right_arm", -1.0)], +) +def test_handover_receiver_approach_tracks_rotated_robot_lateral_axis( + monkeypatch: pytest.MonkeyPatch, + transfer_arm: str, + expected_x: float, +) -> None: + env = _FakeEnv() + + def get_link_pose(*, link_name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).unsqueeze(0) + pose[:, 0, 3] = 0.3 if link_name == "physical_left_base" else -0.3 + return pose + + monkeypatch.setattr(env.robot, "get_link_pose", get_link_pose) + action = GroundedAction( + action_class="HandOver", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={ + "transfer_arm": transfer_arm, + "middle_object_pose": torch.eye(4).unsqueeze(0), + "final_object_pose": torch.eye(4).unsqueeze(0), + }, + ) + + cfg = AtomicActionAdapter(env)._build_config(action, HandOverOptions) + + diagonal = 2.0**-0.5 + assert cfg.receive_approach_direction[0] == pytest.approx(expected_x * diagonal) + assert cfg.receive_approach_direction[1] == pytest.approx(0.0) + assert cfg.receive_approach_direction[2] == pytest.approx(-diagonal) + + +@pytest.mark.parametrize( + ("arm", "outward_x"), + [("left_arm", 1.0), ("right_arm", -1.0)], +) +def test_legacy_handover_transfer_mode_tracks_a_rotated_live_base_line( + monkeypatch: pytest.MonkeyPatch, + arm: str, + outward_x: float, +) -> None: + env = _FakeEnv() + + def get_link_pose(*, link_name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).unsqueeze(0) + pose[:, 0, 3] = 0.3 if link_name == "physical_left_base" else -0.3 + return pose + + monkeypatch.setattr(env.robot, "get_link_pose", get_link_pose) + action = GroundedAction( + action_class="PickUp", + arm=arm, + control="arm", + target=SimpleNamespace(), + cfg={"approach_direction_mode": "handover_transfer"}, + ) + + cfg = AtomicActionAdapter(env)._build_config(action, PickUpOptions) + + outward = torch.tensor([outward_x, 0.0]) + assert torch.dot(cfg.approach_direction[:2], outward) < 0.0 + assert cfg.approach_direction[2] < 0.0 + + +def test_pickup_is_replanned_from_live_pose_and_screens_downstream_targets( + monkeypatch: Any, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + plan_calls: list[GroundedAction] = [] + + def ground( + action, + step, + *, + arm, + state, + reference_eef_pose=None, + orientation_reference_pose=None, + ): + del reference_eef_pose, orientation_reference_pose + action_class = action["atomic_action_class"] + target_pose = ( + _pose(0.0, 0.2, 0.85) if action_class == "MoveHeldObject" else None + ) + return GroundedAction( + action_class=action_class, + arm=arm, + control=str(action.get("control", "arm")), + target=SimpleNamespace(xpos=None), + cfg={"planned_object_pose": entity.get_local_pose(to_matrix=True)}, + target_object_pose=target_pose, + ) + + def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: + plan_calls.append(grounded) + next_state = ( + _held_state(env, entity) if grounded.action_class == "PickUp" else state + ) + return ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=next_state, + grounded=grounded, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor.adapter, "plan", plan) + monkeypatch.setattr(executor.adapter, "execute_trajectory", lambda *a, **k: []) + step = executor.program.semantic_steps[0] + failed = torch.tensor([False]) + + executor._ensure_assignment(step, failed) + planned_call_count = len(plan_calls) + executor._candidate_cache.clear() + entity._pose[:, 0, 3] += 0.25 + edge_result = executor._execute_edge( + executor.edges[step.edge_ids[0]], step, failed=failed + ) + + assert len(plan_calls) == planned_call_count + 1 + assert planned_call_count == len(step.edge_ids) + assert len(plan_calls[0].cfg["downstream_object_target_poses"]) == 1 + assert plan_calls[-1].cfg["planned_object_pose"][0, 0, 3] == pytest.approx(0.25) + assert plan_calls[-1].cfg["downstream_object_target_poses"] + assert edge_result.planner_traces[0]["execution_replanned_from_live_state"] + assert not edge_result.planner_traces[0]["speculative_candidate_available"] + assert bool(edge_result.failed[0]) + assert executor._object_owners["can"] == [None] + + +def test_live_pickup_planning_exception_is_a_retryable_edge_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + _FakeEnv({"can": entity}), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + edge = executor.edges[step.edge_ids[0]] + executor._assignments[step.id] = ["left_arm"] + monkeypatch.setattr( + executor.grounder, + "ground", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("no IK")), + ) + + result = executor._execute_edge(edge, step, failed=torch.tensor([False])) + + assert bool(result.failed[0]) + assert result.actions == [] + assert result.planner_traces[0]["primary_strategy"] == "live_pickup_replan" + assert result.planner_traces[0]["exception"] == "RuntimeError: no IK" + + +def test_pickup_candidate_screens_handover_successor_target( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A later handover staging pose participates in pickup grasp screening.""" + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("pickup", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + pickup_step = executor.program.semantic_steps[0] + handover_step = SemanticStep( + id="handover", + parent_step_id="handover", + operator="handover", + object_uid="can", + actor={"mode": "required", "arm": "left_arm"}, + goal={"transfer_arm": "left_arm", "receive_arm": "right_arm"}, + depends_on=(pickup_step.id,), + postcondition={}, + edge_ids=("handover_staging",), + ) + handover_edge = ExecutionEdge( + id="handover_staging", + source="pickup_done", + target="handover_done", + actions=( + { + "atomic_action_class": "MoveHeldObject", + "actor": {"mode": "required", "arm": "left_arm"}, + "control": "arm", + "target_binding": { + "kind": "handover_staging", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "motion_policy": motion_policy(), + }, + ), + ) + executor.steps[handover_step.id] = handover_step + executor.edges[handover_edge.id] = handover_edge + + existing_target = _pose(0.0, 0.2, 0.85) + handover_target = _pose(0.0, 0.0, 1.15) + grounded = GroundedAction( + action_class="PickUp", + arm="left_arm", + control="arm", + target=SimpleNamespace(), + cfg={"downstream_object_target_poses": (existing_target,)}, + ) + + def ground( + _action: Any, + candidate: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> GroundedAction: + del arm, state, reference_eef_pose, orientation_reference_pose + target = handover_target if candidate.id == handover_step.id else None + return GroundedAction( + action_class="MoveHeldObject", + arm="left_arm", + control="arm", + target=SimpleNamespace(), + cfg={}, + target_object_pose=target, + ) + + monkeypatch.setattr(executor.grounder, "ground", ground) + result = executor._with_downstream_targets( + pickup_step, + pickup_step.edge_ids[0], + "left_arm", + ExecutionState(last_qpos=env.robot.get_qpos()), + grounded, + ) + + targets = result.cfg["downstream_object_target_poses"] + assert len(targets) == 2 + assert torch.equal(targets[0], existing_target) + assert torch.equal(targets[1], handover_target) + + +def test_object_held_predicate_checks_live_gripper_and_tcp_geometry() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + left_eef, _ = env.get_current_xpos_agent() + env.get_current_xpos_agent = lambda: (left_eef, None) + + held = evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + ) + + assert bool(held[0]) + env.robot._qpos[:, env.left_eef_joints] = env.open_state + assert not bool( + evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + )[0] + ) + env.robot._qpos[:, env.left_eef_joints] = (env.open_state + env.close_state) / 2 + assert bool( + evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + )[0] + ) + env.close_state = torch.tensor([0.0, 0.0]) + env.open_state = torch.tensor([0.04, 0.04]) + env.robot._qpos[:, env.left_eef_joints] = env.open_state + assert not bool( + evaluate_predicate( + env, + {"type": "object_held", "object": "can"}, + held_owners={"can": ["left_arm"]}, + held_states={("can", "left_arm"): state}, + )[0] + ) + + +def test_coordinated_held_predicate_uses_per_arm_held_relations() -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.2)) + env = _FakeEnv({"tray": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + env.robot._qpos[:, env.right_eef_joints] = env.close_state + state = _coordinated_held_state(env, entity) + predicate = {"type": "object_held_by_both_grippers", "object": "tray"} + + assert bool(evaluate_predicate(env, predicate, coordinated_state=state)[0]) + + env.robot._qpos[:, env.right_eef_joints] = env.open_state + assert not bool(evaluate_predicate(env, predicate, coordinated_state=state)[0]) + + +def test_object_supported_by_requires_overlap_and_vertical_contact() -> None: + support_z = 0.75 + payload_z = support_z + 0.05 + 0.02 + 0.005 + payload = _FakeEntity( + "payload", + _pose(0.002, -0.002, payload_z), + _box_vertices(0.02), + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, support_z), + _box_vertices(0.05), + ) + env = _FakeEnv({"payload": payload, "support": support}) + predicate = { + "type": "object_supported_by", + "object": "payload", + "support": "support", + } + + assert bool(evaluate_predicate(env, predicate)[0]) + + payload._pose = _pose(0.002, -0.002, payload_z - 1.0) + assert not bool(evaluate_predicate(env, predicate)[0]) + + payload._pose = _pose(0.002, -0.002, payload_z + 1.0) + assert not bool(evaluate_predicate(env, predicate)[0]) + + payload._pose = _pose(0.081, 0.0, payload_z) + assert not bool(evaluate_predicate(env, predicate)[0]) + + +def test_object_supported_by_uses_local_not_mesh_wide_support_height() -> None: + support_vertices = torch.tensor( + [ + [-0.10, -0.10, -0.05], + [-0.02, -0.02, 0.05], + [0.02, -0.02, 0.05], + [0.02, 0.02, 0.05], + [-0.02, 0.02, 0.05], + [0.40, 0.00, 0.40], + ], + dtype=torch.float32, + ) + payload = _FakeEntity( + "payload", + _pose(0.0, 0.0, 0.075), + _box_vertices(0.02), + ) + support = _FakeEntity("support", _pose(0.0, 0.0, 0.0), support_vertices) + env = _FakeEnv({"payload": payload, "support": support}) + + supported = evaluate_predicate( + env, + { + "type": "object_supported_by", + "object": "payload", + "support": "support", + }, + ) + + assert bool(supported[0]) + + +def test_poured_predicate_requires_observed_contents_inside_target() -> None: + source = _FakeEntity("source", _pose(-0.2, 0.0, 0.7), _box_vertices(0.08)) + target = _FakeEntity("target", _pose(0.2, 0.0, 0.7), _box_vertices(0.12)) + content = _FakeEntity("content", _pose(0.2, 0.0, 0.7), _box_vertices(0.01)) + env = _FakeEnv({"source": source, "target": target, "content": content}) + predicate = { + "type": "poured", + "object": "source", + "reference_object": "target", + "contents": [{"object": "content"}], + } + + assert bool(evaluate_predicate(env, predicate)[0]) + + content._pose = _pose(-0.2, 0.0, 0.7) + assert not bool(evaluate_predicate(env, predicate)[0]) + + with pytest.raises(ValueError, match="independently observable"): + evaluate_predicate(env, {**predicate, "contents": []}) + + +def test_object_supported_by_uses_live_center_of_mass_projection() -> None: + payload = _FakeEntity( + "payload", + _pose(0.04, 0.0, 0.125), + _rect_vertices(0.08, 0.02, 0.02), + ) + payload.body_data = SimpleNamespace( + com_pose=torch.tensor([[0.04, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]]) + ) + support = _FakeEntity( + "support", + _pose(0.0, 0.0, 0.05), + _rect_vertices(0.05, 0.05, 0.05), + ) + env = _FakeEnv({"payload": payload, "support": support}) + + supported = evaluate_predicate( + env, + { + "type": "object_supported_by", + "object": "payload", + "support": "support", + }, + ) + + assert not bool(supported[0]) + + +def test_physical_pickup_rebases_a_compliant_grasp_from_live_pose() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + entity._pose[:, 0, 3] += 0.055 + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + physical = executor._physical_pickup( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + state = executor._rebase_held_state( + "can", + "left_arm", + state, + physical, + from_planned_qpos=False, + ) + + assert bool(physical[0]) + left_eef, _ = env.get_current_xpos_agent() + rebased_eef = torch.bmm( + entity.get_local_pose(to_matrix=True), + state.get_held_object("physical_left_arm").object_to_eef, + ) + assert torch.allclose(rebased_eef, left_eef) + + +def test_physical_hold_accepts_configured_held_position_tolerance() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + executor._object_owners["can"] = ["left_arm"] + entity._pose[:, 0, 3] += 0.055 + + held = executor._physical_hold( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + assert bool(held[0]) + + +def test_physical_pickup_rejects_large_grasp_slip() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + entity._pose[:, 0, 3] += 0.08 + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + physical = executor._physical_pickup( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + assert not bool(physical[0]) + + +def test_physical_pickup_rejects_offset_even_when_object_was_lifted() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + entity._pose[:, 0, 3] += 0.08 + entity._pose[:, 2, 3] += 0.08 + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + physical = executor._physical_pickup( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + assert not bool(physical[0]) + + +def test_physical_hold_detects_loss_and_releases_runtime_ownership() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + env.robot._qpos[:, env.left_eef_joints] = env.close_state + state = _held_state(env, entity) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = state + entity._pose[:, 0, 3] += 0.08 + + held = executor._physical_hold( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + executor._release_ownership("can", "left_arm", ~held) + + assert not bool(held[0]) + assert executor._object_owners["can"] == [None] + assert executor._arm_owners["left_arm"] == [None] + assert ("can", "left_arm") not in executor._object_states + + +def test_rebase_held_state_uses_fk_qpos_not_stale_eef_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + state = _held_state(env, entity, arm="left_arm") + expected_eef = _pose(0.31, -0.17, 1.06) + + def compute_fk( + qpos: torch.Tensor, + *, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + del name + assert to_matrix + return expected_eef.repeat(qpos.shape[0], 1, 1) + + monkeypatch.setattr(env.robot, "compute_fk", compute_fk) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + + state = executor._rebase_held_state( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + held = state.get_held_object("physical_left_arm") + assert held is not None + expected_relation = torch.bmm( + torch.linalg.inv(entity.get_local_pose(to_matrix=True)), + expected_eef, + ) + assert torch.allclose(held.object_to_eef, expected_relation) + + +def test_upright_transport_state_tracks_selected_target_pose( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + state = _held_state(env, entity, arm="left_arm") + expected_eef = _pose(0.27, -0.11, 1.04) + target_pose = _pose(0.05, 0.01, 0.92) + + def compute_fk( + qpos: torch.Tensor, + *, + name: str, + to_matrix: bool, + ) -> torch.Tensor: + del name + assert to_matrix + return expected_eef.repeat(qpos.shape[0], 1, 1) + + monkeypatch.setattr(env.robot, "compute_fk", compute_fk) + entity._pose = target_pose.clone() + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + synchronized = executor._rebase_held_state( + "can", + "left_arm", + state, + torch.tensor([True]), + ) + + held = synchronized.get_held_object("physical_left_arm") + assert held is not None + assert torch.allclose( + held.object_to_eef, + torch.bmm(torch.linalg.inv(target_pose), expected_eef), + ) + assert torch.allclose(held.grasp_xpos, expected_eef) + + +def test_existing_object_owner_reserves_same_arm(monkeypatch: Any) -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "left_arm"))) + ), + env, + record_runtime=False, + ) + original = executor.program.semantic_steps[0] + continuation = replace( + original, + id="continuation", + actor={"mode": "auto"}, + ) + executor._object_owners["can"] = ["left_arm"] + executor._arm_owners["left_arm"] = ["can"] + executor._object_states[("can", "left_arm")] = _held_state(env, entity) + + monkeypatch.setattr( + executor, + "_candidate", + lambda step, arm, failed: SimpleNamespace( + feasible=torch.tensor([True]), + cost=torch.tensor([0.0 if arm == "right_arm" else 10.0]), + warnings=(), + ), + ) + executor._ensure_assignment(continuation, torch.tensor([False])) + + assert executor._assignments["continuation"] == ["left_arm"] + assert bool(executor._resource_conflicts(continuation, "right_arm")[0]) + + +def test_new_task_group_hydrates_predecessor_held_state() -> None: + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + env = _FakeEnv({"can": entity}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent(_task_agent(_hold_step("hold", "can", "right_arm"))) + ), + env, + record_runtime=False, + ) + held_state = _held_state(env, entity, arm="right_arm") + executor._object_states[("can", "right_arm")] = held_state + continuation = replace( + executor.program.semantic_steps[0], + id="place_after_handover", + actor={"mode": "required", "arm": "right_arm"}, + ) + + hydrated = executor._state_for(continuation, "right_arm") + + assert hydrated.get_held_object("physical_right_arm") is not None + assert torch.equal(hydrated.last_qpos, env.robot.get_qpos()) + + +def test_place_uses_preceding_or_live_eef_pose_not_original_grasp() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "basket": _FakeEntity("basket", _pose(0.0, 0.0, 0.70), _box_vertices(0.10)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": [], + } + ) + ) + ) + state = _held_state(env, entities["can"]) + held_object = state.get_held_object("physical_left_arm") + assert held_object is not None + replacement = HeldObjectState( + semantics=held_object.semantics, + object_to_eef=held_object.object_to_eef, + grasp_xpos=_pose(0.0, -0.3, 0.75), + env_mask=held_object.env_mask, + ) + held_objects = dict(state.held_objects) + held_objects["physical_left_arm"] = replacement + state = state.with_updates(held_objects=held_objects) + held_object = replacement + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "Place" + ) + grounder = ActionGrounder( + program, + env, + lambda uid: held_object.semantics, + ) + reference = _pose(0.0, 0.4, 0.85) + + planned = grounder.ground( + edge.actions[0], + program.semantic_steps[0], + arm="left_arm", + state=state, + reference_eef_pose=reference, + ) + live = grounder.ground( + edge.actions[0], + program.semantic_steps[0], + arm="left_arm", + state=state, + ) + + assert torch.equal(planned.target.xpos, reference) + assert torch.equal(live.target.xpos, env.get_current_xpos_agent()[0]) + assert not torch.equal(live.target.xpos, held_object.grasp_xpos) + + +def test_inside_target_preserves_pre_pick_supported_height() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 1.40), _box_vertices(0.03)), + "basket": _FakeEntity("basket", _pose(0.0, 0.0, 0.70), _box_vertices(0.10)), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "basket", + "relation": "inside", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + final = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + ) + state = _held_state(env, entities["can"]) + held = state.get_held_object("physical_left_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + supported_pose = _pose(0.0, 0.2, 0.75) + + grounded = grounder.ground( + final.actions[0], + step, + arm="left_arm", + state=state, + orientation_reference_pose=supported_pose, + ) + + assert grounded.target_object_pose is not None + assert grounded.target_object_pose[0, 2, 3] == pytest.approx(0.75) + + +def test_coordinated_step_rejects_an_arm_reserved_by_terminal_hold() -> None: + entity = _FakeEntity("shared_box", _pose(0.0, 0.0, 0.75), _box_vertices(0.05)) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "transport", + "operator": "coordinated_transport", + "object": "shared_box", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "direction": "front", + "terminal_behavior": "hold", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv({"shared_box": entity}), + record_runtime=False, + ) + executor._arm_owners["left_arm"] = ["held_can"] + step = executor.program.semantic_steps[0] + + executor._ensure_assignment(step, torch.tensor([False])) + + assert executor._assignments[step.id] == [None] + + +def test_failure_propagates_only_to_dependent_branch(monkeypatch: Any) -> None: + program = load_execution_program( + compile_task_agent( + _task_agent( + _hold_step("left", "can_a", "left_arm"), + _hold_step("right", "can_b", "right_arm"), + ) + ) + ) + executor = ProgramExecutor( + program, _FakeEnv(), settle_steps=0, record_runtime=False + ) + monkeypatch.setattr( + executor, + "_pack_ready_edges", + lambda ready, **_kwargs: (ready[0],), + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, failed: executor._assignments.setdefault( + step.id, [step.actor["arm"]] + ), + ) + active_by_step: dict[str, list[bool]] = {"left": [], "right": []} + + def execute(edge, step, *, failed): + active_by_step[step.id].append(not bool(failed[0])) + action_failed = failed.clone() + if step.id == "left" and edge.id == step.edge_ids[0]: + action_failed[:] = True + return SimpleNamespace(actions=[], failed=action_failed, grounded=[]) + + monkeypatch.setattr(executor, "_execute_edge", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda step, failed: ( + failed, + ~failed, + torch.zeros(1, 3), + ), + ) + + result = executor.run() + + assert any(active_by_step["right"]) + assert bool(result.semantic_success["right"][0]) + assert not bool(result.success[0]) + + +def test_resource_ordering_waits_without_propagating_semantic_failure() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "resource_ordering", + "level": "L3", + "instruction": "Stand both cans, then hand over the second can.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "first", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E2", + "params": { + "object_role": "second", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_03", + "task_type": "E4", + "params": { + "object_role": "second", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + }, + "depends_on": ["task_02"], + "role": "primary", + }, + ], + "success": {"type": "all_complete"}, + "oracle": {}, + "metadata": {}, + } + program = load_execution_program( + instantiate_seed_graph( + task, + {"first": "first_can", "second": "second_can"}, + ) + ) + executor = ProgramExecutor(program, _FakeEnv(), record_runtime=False) + handover_entry = next( + edge + for edge in program.edges + if executor.step_by_edge[edge.id].id == "task_03" + and all( + executor.step_by_edge[dependency].id != "task_03" + for dependency in edge.depends_on + ) + ) + dependencies = { + executor.step_by_edge[dependency].id: dependency + for dependency in handover_entry.depends_on + } + failures = { + dependency: torch.tensor([step_id == "task_01"]) + for step_id, dependency in dependencies.items() + } + + assert not bool(executor._dependency_failures(handover_entry, failures)[0]) + + failures[dependencies["task_02"]][:] = True + assert bool(executor._dependency_failures(handover_entry, failures)[0]) + + +def test_v2_executor_retries_one_complete_atomic_action_twice( + monkeypatch: Any, +) -> None: + task, requirements = make_task_spec("E9") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + program = load_execution_program(instantiate_seed_graph(task, bindings)) + executor = ProgramExecutor( + program, + _FakeEnv(), + settle_steps=0, + record_runtime=False, + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, failed: executor._assignments.setdefault(step.id, ["left_arm"]), + ) + attempts = 0 + + def execute(_edge, _step, *, failed): + nonlocal attempts + attempts += 1 + action_failed = failed.clone() + if attempts < 3: + action_failed[:] = True + return SimpleNamespace(actions=[], failed=action_failed, grounded=[]) + + monkeypatch.setattr(executor, "_execute_edge", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda step, failed: (failed, ~failed, torch.zeros(1, 3)), + ) + + result = executor.run() + + assert attempts == 3 + assert result.retry_count == 2 + assert bool(result.success[0]) + assert result.failure_events == [] + + +def test_v2_executor_stops_at_transition_budget() -> None: + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + program = load_execution_program(instantiate_seed_graph(task, bindings)) + executor = ProgramExecutor( + program, + _FakeEnv(), + max_transitions=0, + settle_steps=0, + record_runtime=False, + ) + + with pytest.raises(RuntimeError, match="max_transitions"): + executor.run() + + +def test_failed_arrangement_records_candidate_diagnostics_without_marker_error( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(-0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_x", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor( + program, + _FakeEnv(entities), + settle_steps=0, + record_root=tmp_path, + ) + infeasible = SimpleNamespace( + feasible=torch.tensor([False]), + cost=torch.tensor([torch.inf]), + plans={}, + warnings=("No IK solutions found for downstream target poses.",), + ) + monkeypatch.setattr(executor, "_candidate", lambda *_args, **_kwargs: infeasible) + + result = executor.run(run_id="no_candidate") + + assert not bool(result.success[0]) + record_path = Path(result.record_dir) / "env_0000" / "task_graph.json" + record = json.loads(record_path.read_text(encoding="utf-8")) + assert record["runtime"]["status"] == "failed" + first_event = record["runtime"]["events"][0] + assert first_event["status"] == "failed" + assert first_event["diagnostics"] == [ + "No IK solutions found for downstream target poses." + ] + + +def test_arrange_line_builds_live_slots_for_compiler_operator_name() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(-0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + compiled = compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_x", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + + executor = ProgramExecutor( + load_execution_program(compiled), + _FakeEnv(entities), + record_runtime=False, + ) + + assert executor.arrangement is not None + assert executor.arrangement.positions.shape == (1, 2, 3) + assert {step.operator for step in executor.program.semantic_steps} == { + "arrange_line" + } + + +def test_free_arrangement_matches_live_object_order_without_crossing() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(0.0, 0.20, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.0, 0.00, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_c": _FakeEntity( + "can_c", + _pose(0.0, -0.20, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b", "can_c"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + record_runtime=False, + ) + arrangement = executor.arrangement + + assert arrangement is not None + assert int(arrangement.assignments["line__01"][0]) == 2 + assert int(arrangement.assignments["line__02"][0]) == 1 + assert int(arrangement.assignments["line__03"][0]) == 0 + assert arrangement.spacing[0] == pytest.approx(0.1648528) + + +def test_arm_candidate_score_softly_penalizes_cross_zone_motion() -> None: + source = _pose(0.0, -0.30, 0.78) + target = _pose(0.0, -0.20, 0.78) + kwargs = { + "motion_cost": torch.tensor([torch.pi]), + "source_pose": source, + "target_pose": target, + "workspace_center_xy": torch.tensor([[0.0, 0.0]]), + "workspace_half_width": torch.tensor([0.40]), + "robot_lateral_axis": torch.tensor([[0.0, -1.0]]), + "policy": default_runtime_policy("dual_ur10").arm_selection, + } + + left = _score_arm_candidate(arm="left_arm", **kwargs) + right = _score_arm_candidate(arm="right_arm", **kwargs) + + assert left["normalized_motion_cost"][0] == pytest.approx(1.0) + assert left["pickup_crossing_penalty"][0] == pytest.approx(0.0) + assert left["placement_crossing_penalty"][0] == pytest.approx(0.0) + assert right["pickup_crossing_penalty"][0] > 0.0 + assert right["placement_crossing_penalty"][0] > 0.0 + assert right["total_cost"][0] > left["total_cost"][0] + + +def test_preserve_grounding_uses_pre_pickup_orientation_reference() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(0.0, -0.10, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.0, 0.10, 0.78), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + final_edge = next( + edge + for edge in executor.program.edges + if edge.id in step.edge_ids + and edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + and edge.actions[0]["target_binding"].get("phase") == "final" + ) + reference = entities[step.object_uid].get_local_pose(to_matrix=True) + disturbed = reference.clone() + disturbed[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities[step.object_uid]._pose = disturbed + + grounded = executor.grounder.ground( + final_edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=executor.env.robot.get_qpos()), + orientation_reference_pose=reference, + ) + + assert grounded.target_object_pose is not None + assert torch.allclose( + grounded.target_object_pose[:, :3, :3], + reference[:, :3, :3], + ) + + +def test_arrange_line_verifies_planar_slot_without_height_coupling() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + _pose(-0.055, -0.190, 0.755), + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._targets[step.id] = torch.tensor([[0.0, -0.216, 0.842]]) + executor._policies[step.id] = { + "line_axis_tolerance": 0.06, + "line_perpendicular_tolerance": 0.06, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert not bool(failed[0]) + assert bool(success[0]) + + +def test_arrange_line_rejects_preserve_orientation_drift() -> None: + rotated = _pose(0.0, -0.190, 0.755) + rotated[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "can_a": _FakeEntity( + "can_a", + rotated, + _rect_vertices(0.03, 0.03, 0.06), + ), + "can_b": _FakeEntity( + "can_b", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.06), + ), + } + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "line", + "operator": "arrange_line", + "objects": ["can_a", "can_b"], + "actor": {"mode": "auto"}, + "goal": { + "axis": "world_y", + "order_constraint": "free", + "orientation_goal": "preserve", + }, + "depends_on": [], + } + ) + ) + ), + _FakeEnv(entities), + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._targets[step.id] = rotated[:, :3, 3].clone() + executor._orientation_references[step.id] = _pose(0.0, -0.190, 0.755) + executor._policies[step.id] = { + "line_axis_tolerance": 0.06, + "line_perpendicular_tolerance": 0.06, + "preserve_orientation_tolerance": torch.pi / 12, + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(failed[0]) + assert not bool(success[0]) + + +def test_shared_container_placements_receive_non_overlapping_live_slots() -> None: + entities = { + "basket": _FakeEntity( + "basket", + _pose(0.0, 0.0, 0.72), + _rect_vertices(0.25, 0.18, 0.08), + ), + "cube": _FakeEntity( + "cube", + _pose(-0.15, 0.0, 0.76), + _rect_vertices(0.03, 0.03, 0.03), + ), + "cup": _FakeEntity( + "cup", + _pose(0.15, 0.0, 0.76), + _rect_vertices(0.035, 0.035, 0.06), + ), + } + steps = [ + { + "id": f"place_{uid}", + "operator": "place_relative", + "object": uid, + "actor": {"mode": "auto"}, + "goal": {"relation": "inside", "reference_object": "basket"}, + "depends_on": [], + } + for uid in ("cube", "cup") + ] + task = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "two_in_basket", + "goal": "Place both objects in the basket.", + "semantic_steps": steps, + } + executor = ProgramExecutor( + load_execution_program(compile_task_agent(task)), + _FakeEnv(entities), + record_runtime=False, + ) + targets = [ + executor.placements[step.id].positions[step.id][0, :2] + for step in executor.program.semantic_steps + ] + + assert set(executor.placements) == {"place_cube", "place_cup"} + assert torch.linalg.vector_norm(targets[0] - targets[1]) > 0.05 + + +@pytest.mark.parametrize( + ("direction", "expected_position"), + ( + ("front_left", (0.16, 0.16, 0.85)), + ("up", (0.0, 0.0, 0.91)), + ), +) +def test_coordinated_transport_direction_is_grounded_from_live_pose( + direction: str, + expected_position: tuple[float, float, float], +) -> None: + entities = { + "shared_box": _FakeEntity( + "shared_box", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.10, 0.06, 0.03), + ) + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "transport", + "operator": "coordinated_transport", + "object": "shared_box", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "direction": direction, + "terminal_behavior": "hold", + }, + "depends_on": [], + } + ) + ) + ) + + def semantics(uid: str) -> ObjectSemantics: + entity = entities[uid] + return ObjectSemantics( + affordance=Affordance(), + geometry={ + "mesh_vertices": entity.get_vertices(env_ids=[0], scale=True), + "mesh_triangles": entity.get_triangles(env_ids=[0]), + }, + label=uid, + entity=entity, + ) + + step = program.semantic_steps[0] + edge = program.edges[0] + grounded = ActionGrounder(program, env, semantics).ground( + edge.actions[0], + step, + arm="coordinated", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, CoordinatedPickGoal) + assert torch.allclose( + grounded.target.object_target_pose[0, :3, 3], + torch.tensor(expected_position), + ) + + +def test_coordinated_payload_monitor_rejects_drift_and_carrier_tilt() -> None: + class _BatchedVerticesEntity(_FakeEntity): + def get_vertices( + self, + *, + env_ids: list[int], + scale: bool, + ) -> torch.Tensor: + return super().get_vertices(env_ids=env_ids, scale=scale).unsqueeze(0) + + entities = { + "tray": _BatchedVerticesEntity( + "tray", + _pose(0.0, 0.0, 0.75), + _rect_vertices(0.20, 0.14, 0.02), + ), + "bottle": _FakeEntity( + "bottle", + _pose(0.0, 0.0, 0.80), + _rect_vertices(0.03, 0.03, 0.08), + ), + } + program = compile_task_agent( + _task_agent( + { + "id": "carry", + "operator": "coordinated_transport", + "object": "tray", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "terminal_behavior": "place", + "payloads": [{"object": "bottle", "slot": "center"}], + }, + "depends_on": [], + } + ) + ) + executor = ProgramExecutor( + load_execution_program(program), + _FakeEnv(entities), + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._capture_payloads(step) + + assert bool(executor._verify_payloads(step)[0]) + entities["bottle"]._pose[:, 0, 3] += 0.20 + assert not bool(executor._verify_payloads(step)[0]) + entities["bottle"]._pose = _pose(0.0, 0.0, 0.80) + entities["tray"]._pose[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + assert not bool(executor._verify_payloads(step)[0]) + + +def test_lay_flat_surface_height_uses_rotated_live_mesh() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "rod": _FakeEntity( + "rod", + _pose(0.2, 0.0, 0.80), + _rect_vertices(0.02, 0.03, 0.10), + ), + } + env = _FakeEnv(entities) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "rod", + "actor": {"mode": "auto"}, + "goal": { + "reference_object": "table", + "relation": "on", + "orientation_goal": "lay_flat", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + ) + grounder = ActionGrounder( + program, + env, + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ), + ) + grounded = grounder.ground( + edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, HeldObjectPoseGoal) + table_top = 0.70 + 0.02 + rotated_rod_half_height = 0.02 + surface_clearance = 0.005 + expected_surface_z = table_top + rotated_rod_half_height + surface_clearance + assert grounded.target.object_target_pose[0, 2, 3] == pytest.approx( + expected_surface_z, + abs=1.0e-5, + ) + + +def test_orient_object_anchors_final_pose_to_support_not_live_lift_height() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "bottle": _FakeEntity( + "bottle", + _pose(0.10, 0.20, 1.30), + _rect_vertices(0.02, 0.03, 0.10), + ), + } + env = _FakeEnv(entities) + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + edges = { + edge.actions[0]["target_binding"].get("phase"): edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + } + grounder = ActionGrounder( + program, + env, + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ), + ) + + staging = grounder.ground( + edges["staging"].actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + final = grounder.ground( + edges["final"].actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + expected_final_z = 0.70 + 0.02 + 0.10 + 0.05 + assert final.target_object_pose[0, 2, 3] == pytest.approx(expected_final_z) + assert final.target_object_pose[0, :2, 3].tolist() == pytest.approx([0.10, 0.20]) + assert staging.target_object_pose[0, 2, 3] > final.target_object_pose[0, 2, 3] + assert staging.target_object_pose[0, :2, 3].tolist() == pytest.approx([0.10, 0.20]) + + +def test_orient_grounding_uses_mature_robot_profile_policy() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(0.0, 0.0, 0.70), + _rect_vertices(0.60, 0.40, 0.02), + ), + "bottle": _FakeEntity( + "bottle", + _pose(0.10, 0.20, 0.78), + _rect_vertices(0.02, 0.03, 0.10), + ), + } + env = _FakeEnv(entities) + env.agent_robot_profile = "dual_ur10" + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + "support_object": "table", + }, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + pickup_edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "PickUp" + ) + final_edge = next( + edge + for edge in program.edges + if edge.actions[0]["target_binding"].get("phase") == "final" + ) + grounder = ActionGrounder( + program, + env, + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ), + ) + + pickup = grounder.ground( + pickup_edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + final = grounder.ground( + final_edge.actions[0], + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + table_top = 0.72 + bottle_half_height = 0.10 + expected_z = table_top + bottle_half_height + 0.05 + assert torch.equal( + pickup.motion_policy["obj_upright_direction"], + torch.tensor([0.0, 0.0, 1.0]), + ) + assert pickup.motion_policy["rotate_upright"] == pytest.approx(torch.pi / 4) + assert pickup.motion_policy["upright_yaw_samples"] == 8 + assert final.target_object_pose[0, 2, 3] == pytest.approx(expected_z) + assert final.motion_policy["upright_local_axis"] == "long_axis" + assert final.motion_policy["upright_yaw_samples"] == 8 + + +def test_orient_verification_requires_upright_pose_near_initial_xy() -> None: + bottle = _FakeEntity( + "bottle", + _pose(0.10, 0.20, 0.823), + _rect_vertices(0.02, 0.03, 0.10), + ) + env = _FakeEnv({"bottle": bottle}) + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._policies[step.id] = { + "upright_max_tilt": torch.pi / 12, + "upright_xy_tolerance": 0.05, + "upright_local_axis": "z", + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + assert bool(success[0]) + assert not bool(failed[0]) + + bottle._pose[:, :3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] + ) + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + assert not bool(success[0]) + assert bool(failed[0]) + + +def test_orient_verification_accepts_grounded_live_xy_anchor() -> None: + bottle = _FakeEntity( + "bottle", + _pose(0.15, -0.10, 0.823), + _rect_vertices(0.02, 0.03, 0.10), + ) + env = _FakeEnv({"bottle": bottle}) + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + "position_anchor": "live_xy", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._targets[step.id] = torch.tensor([[0.15, -0.10, 0.823]]) + executor._policies[step.id] = { + "upright_max_tilt": torch.pi / 12, + "upright_xy_tolerance": 0.05, + "upright_local_axis": "z", + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(success[0]) + assert not bool(failed[0]) + + +def test_long_axis_upright_is_undirected_but_explicit_axis_is_not() -> None: + pose = _pose(0.0, 0.0, 0.75) + pose[:, :3, 1] = torch.tensor([0.0, 0.0, -1.0]) + pose[:, :3, 2] = torch.tensor([0.0, 1.0, 0.0]) + entity = _FakeEntity("can", pose, _rect_vertices(0.03, 0.10, 0.03)) + env = _FakeEnv({"can": entity}) + + assert bool( + evaluate_predicate( + env, + { + "type": "object_upright", + "object": "can", + "local_axis": "long_axis", + }, + )[0] + ) + assert not bool( + evaluate_predicate( + env, + { + "type": "object_upright", + "object": "can", + "local_axis": "long_axis", + "directed": True, + }, + )[0] + ) + assert not bool( + evaluate_predicate( + env, + { + "type": "object_upright", + "object": "can", + "local_axis": "y", + }, + )[0] + ) + + +def test_orient_object_uses_solver_roots_when_control_groups_share_root() -> None: + entities = { + "left_object": _FakeEntity( + "left_object", + _pose(0.0, -0.20, 0.8), + _rect_vertices(0.02, 0.02, 0.08), + ), + "right_object": _FakeEntity( + "right_object", + _pose(0.0, 0.20, 0.8), + _rect_vertices(0.02, 0.02, 0.08), + ), + } + env = _FakeEnv(entities) + env.agent_initial_object_poses = { + uid: entity.get_local_pose(to_matrix=True) for uid, entity in entities.items() + } + execution = compile_task_agent( + _task_agent( + *[ + { + "id": uid, + "operator": "orient_object", + "object": uid, + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for uid in ("left_object", "right_object") + ] + ) + ) + executor = ProgramExecutor( + load_execution_program(execution), env, record_runtime=False + ) + + torch.testing.assert_close( + env.robot.get_control_part_base_pose(name="physical_left_arm", to_matrix=True), + env.robot.get_control_part_base_pose(name="physical_right_arm", to_matrix=True), + ) + assert executor._preferred_in_place_arm(executor.steps["left_object"], 0) == ( + "left_arm" + ) + assert executor._preferred_in_place_arm(executor.steps["right_object"], 0) == ( + "right_arm" + ) + + +def test_orient_object_arm_preference_follows_translated_robot_and_table() -> None: + entities = { + "table": _FakeEntity( + "table", + _pose(1.50, -0.70, 0.70), + _rect_vertices(0.50, 0.40, 0.02), + ), + "left_object": _FakeEntity( + "left_object", + _pose(1.70, -0.70, 0.80), + _rect_vertices(0.02, 0.02, 0.08), + ), + "right_object": _FakeEntity( + "right_object", + _pose(1.30, -0.70, 0.80), + _rect_vertices(0.02, 0.02, 0.08), + ), + } + env = _FakeEnv(entities) + env.robot.get_link_pose = lambda *, link_name, to_matrix: _pose( + 1.80 if link_name == "physical_left_base" else 1.20, + -0.70, + 0.0, + ) + env.agent_initial_object_poses = { + uid: entity.get_local_pose(to_matrix=True) for uid, entity in entities.items() + } + execution = compile_task_agent( + _task_agent( + *[ + { + "id": uid, + "operator": "orient_object", + "object": uid, + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + for uid in ("left_object", "right_object") + ] + ) + ) + executor = ProgramExecutor( + load_execution_program(execution), env, record_runtime=False + ) + + center, _, lateral = executor._arm_selection_workspace( + executor.steps["left_object"] + ) + torch.testing.assert_close(center, torch.tensor([[1.50, -0.70]])) + torch.testing.assert_close(lateral, torch.tensor([[1.0, 0.0]])) + assert executor._preferred_in_place_arm(executor.steps["left_object"], 0) == ( + "left_arm" + ) + assert executor._preferred_in_place_arm(executor.steps["right_object"], 0) == ( + "right_arm" + ) + + +def test_coordinated_placement_uses_live_typed_target_and_profile_parts() -> None: + entities = { + "placing": _FakeEntity( + "placing", + _pose(0.0, 0.1, 0.75), + _box_vertices(0.04), + ), + "support": _FakeEntity( + "support", + _pose(0.0, -0.1, 0.75), + _box_vertices(0.06), + ), + } + env = _FakeEnv(entities) + compiled = compile_task_agent( + _task_agent( + { + "id": "place", + "operator": "coordinated_place", + "object": "placing", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": { + "support_object": "support", + "relation": "on", + "release": True, + }, + "depends_on": [], + } + ) + ) + program = load_execution_program(compiled) + + def semantics(uid: str) -> ObjectSemantics: + return ObjectSemantics( + affordance=Affordance(), + geometry={}, + label=uid, + entity=entities[uid], + ) + + step = program.semantic_steps[0] + assert [action["atomic_action_class"] for action in program.edges[0].actions] == [ + "PickUp", + "PickUp", + ] + edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "CoordinatedPlacement" + ) + grounder = ActionGrounder(program, env, semantics) + grounded = grounder.ground( + edge.actions[0], + step, + arm="coordinated", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + + assert isinstance(grounded.target, CoordinatedPlacementGoal) + assert grounded.target.release is True + assert torch.allclose( + grounded.target.support_object_target_pose, + entities["support"].get_local_pose(to_matrix=True), + ) + + adapter = AtomicActionAdapter(env) + cfg = adapter._build_config(grounded, CoordinatedPlacementOptions) + bound_endpoints: dict[str, dict[str, str]] = {} + + class _BindingEngine: + binding_owner_id = "runtime-contract-test" + + def bind_control_parts( + self, + _skill_id: str, + endpoints: dict[str, dict[str, str]], + ) -> ActionBinding: + bound_endpoints.update(deepcopy(endpoints)) + return ActionBinding(owner_id=self.binding_owner_id) + + adapter._atomic_engine = _BindingEngine() + binding = adapter._binding( + grounded, + adapter.capabilities.get("CoordinatedPlacement"), + ) + assert binding.owner_id == "runtime-contract-test" + assert bound_endpoints == { + "placing": { + "motion": "physical_left_arm", + "grasp": "physical_left_eef", + }, + "support": { + "motion": "physical_right_arm", + "grasp": "physical_right_eef", + }, + } + assert cfg.release is True + + +def test_online_environment_preserves_result_and_disables_terminations( + monkeypatch: Any, +) -> None: + installed: list[Any] = [] + initialization_order: list[tuple[str, Any]] = [] + + def fake_super_init(self: Any, cfg: Any, **kwargs: Any) -> None: + del kwargs + initialization_order.append(("super", cfg.robot)) + self.cfg = cfg + self.robot = object() + self.ignore_terminations_during_agent = True + + def fake_repair(robot_cfg: Any) -> int: + initialization_order.append(("repair", robot_cfg)) + return 1 + + def fake_install(robot: Any) -> int: + initialization_order.append(("install", robot)) + installed.append(robot) + return 1 + + monkeypatch.setattr(EmbodiedEnv, "__init__", fake_super_init) + monkeypatch.setattr( + env_module, + "repair_action_engine_ur5_solver_cfg", + fake_repair, + ) + monkeypatch.setattr( + env_module, + "install_action_engine_solver_compat", + fake_install, + ) + monkeypatch.setattr( + env_module.ActionEngineEnv, + "_capture_runtime_state", + lambda self: None, + ) + robot_cfg = object() + cfg = SimpleNamespace(ignore_terminations=False, robot=robot_cfg) + env = env_module.ActionEngineEnv( + cfg, + agent_config={"schema_version": "action_engine_config_v2"}, + task_name="task", + agent_config_path="/tmp/agent_config.json", + ) + result = ExecutionResult( + actions=[], + success=torch.tensor([True]), + semantic_success={}, + ) + + assert cfg.ignore_terminations is True + assert installed == [env.robot] + assert [name for name, _ in initialization_order] == [ + "repair", + "super", + "install", + ] + assert initialization_order[0][1] is robot_cfg + assert env._normalize_demo_action_list(result) is result + + +def test_solver_compat_repairs_only_stale_action_engine_ur_dh_defaults() -> None: + stale_ur5 = URSolverCfg() + stale_ur5.ur_type = "ur5" + custom_ur5 = URSolverCfg(ur_type="ur5") + custom_ur5.d1 = 0.1 + ur10 = URSolverCfg() + robot_cfg = SimpleNamespace( + solver_cfg={ + "left": stale_ur5, + "left_alias": stale_ur5, + "custom": custom_ur5, + "right": ur10, + } + ) + expected = URSolverCfg(ur_type="ur5") + dh_fields = ("d1", "a2", "a3", "d4", "d5", "d6") + + assert solver_compat.repair_action_engine_ur5_solver_cfg(robot_cfg) == 1 + assert tuple(getattr(stale_ur5, name) for name in dh_fields) == pytest.approx( + tuple(getattr(expected, name) for name in dh_fields) + ) + assert custom_ur5.d1 == pytest.approx(0.1) + assert ur10.ur_type == "ur10" + assert solver_compat.repair_action_engine_ur5_solver_cfg(robot_cfg) == 0 + + +def test_solver_compat_uses_true_tcp_inverse_and_restores_solver( + monkeypatch: Any, +) -> None: + class FakeSolver: + def __init__(self) -> None: + self.device = torch.device("cpu") + self.tcp_xpos = np.array( + [ + [0.0, -1.0, 0.0, 0.1], + [1.0, 0.0, 0.0, 0.2], + [0.0, 0.0, 1.0, 0.3], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + self.received: torch.Tensor | None = None + self.saw_identity = False + + def get_ik(self, target_xpos: torch.Tensor, **kwargs: Any) -> str: + del kwargs + self.received = target_xpos + self.saw_identity = np.allclose(self.tcp_xpos, np.eye(4)) + return "ok" + + monkeypatch.setattr(solver_compat, "PytorchSolver", FakeSolver) + solver = FakeSolver() + original_tcp = solver.tcp_xpos.copy() + robot = SimpleNamespace(_solvers={"left": solver, "alias": solver}) + target = torch.eye(4).unsqueeze(0) + + assert solver_compat.install_pytorch_solver_tcp_compat(robot) == 1 + assert solver.get_ik(target_xpos=target) == "ok" + assert solver.saw_identity + assert torch.allclose( + solver.received, + target @ torch.linalg.inv(torch.as_tensor(original_tcp)), + ) + assert np.allclose(solver.tcp_xpos, original_tcp) + assert solver_compat.install_pytorch_solver_tcp_compat(robot) == 0 + + +def test_solver_compat_aligns_ur5_analytic_ik_with_urdf_ee_frame( + monkeypatch: Any, +) -> None: + class FakeSolver: + def __init__(self, ur_type: str) -> None: + self.cfg = SimpleNamespace(ur_type=ur_type) + self.device = torch.device("cpu") + self.tcp_xpos = np.array( + [ + [0.0, -1.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.2], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + self.received: torch.Tensor | None = None + + def get_ik( + self, + target_xpos: torch.Tensor, + qpos_seed: torch.Tensor | None = None, + **kwargs: Any, + ) -> str: + del kwargs, qpos_seed + self.received = target_xpos + return "ok" + + monkeypatch.setattr(solver_compat, "URSolver", FakeSolver) + ur5 = FakeSolver("ur5") + ur10 = FakeSolver("ur10") + robot = SimpleNamespace(_solvers={"left": ur5, "alias": ur5, "right": ur10}) + target = torch.eye(4).unsqueeze(0) + target[:, :3, 3] = torch.tensor([0.3, -0.2, 0.8]) + qpos_seed = torch.zeros((1, 6)) + + assert solver_compat.install_ur5_solver_frame_compat(robot) == 1 + assert ur5.get_ik(target, qpos_seed) == "ok" + + tcp = torch.as_tensor(ur5.tcp_xpos) + analytic_to_urdf = torch.eye(4) + analytic_to_urdf[0, 3] = -0.01 + expected = target @ torch.linalg.inv(tcp) @ torch.linalg.inv(analytic_to_urdf) @ tcp + assert torch.allclose(ur5.received, expected) + assert ur10.received is None + assert solver_compat.install_ur5_solver_frame_compat(robot) == 0 diff --git a/tests/gen_sim/action_engine/task_fixtures.py b/tests/gen_sim/action_engine/task_fixtures.py new file mode 100644 index 000000000..c845cdfab --- /dev/null +++ b/tests/gen_sim/action_engine/task_fixtures.py @@ -0,0 +1,229 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Language-neutral structured fixtures for Action Engine tests.""" + +from __future__ import annotations + +from typing import Any + +from embodichain.gen_sim.action_engine.domain import ( + task_success_type, + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.protocol import ( + SCENE_REQUIREMENTS_SCHEMA, + TASK_SPEC_SCHEMA, +) + +__all__ = ["make_task_level", "make_task_spec"] + +_OBJECT_FIXTURES = { + "E1": ("can", ["graspable", "placeable"], {}), + "E2": ("can", ["graspable", "orientable"], {"orientation": "fallen"}), + "E3": ("container", ["graspable", "pourable"], {"held_by": "left_arm"}), + "E4": ("cup", ["graspable", "handover"], {}), + "E5": ("tray", ["dual_graspable", "rigid"], {}), + "E6": ("drawer", ["articulated", "pullable"], {"joint_state": "closed"}), + "E7": ("drawer", ["articulated", "pushable"], {"joint_state": "open"}), + "E8": ("knob", ["turnable"], {}), + "E9": ("button", ["pressable"], {"activation": "inactive"}), +} + + +def make_task_spec( + task_type: str = "E1", + *, + task_id: str | None = None, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Build one validated L1 TaskSpec and matching scene requirements.""" + if task_type not in _OBJECT_FIXTURES: + raise ValueError(f"Unsupported fixture task type {task_type!r}.") + category, affordances, initial_state = _OBJECT_FIXTURES[task_type] + object_role = "object_01" + params: dict[str, Any] = {"object_role": object_role} + objects = [ + { + "role_id": object_role, + "category": category, + "count": 1, + "affordances": affordances, + "initial_state": initial_state, + "attributes": {}, + } + ] + if task_type in {"E1", "E3"}: + target_role = "target_01" + params.update({"target_role": target_role, "relation": "inside"}) + if task_type == "E3": + params["source_role"] = params.pop("object_role") + objects.append( + { + "role_id": target_role, + "category": "container", + "count": 1, + "affordances": ["container", "support_surface"], + "initial_state": {}, + "attributes": {}, + } + ) + elif task_type == "E2": + params.update( + { + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + ) + elif task_type == "E4": + params.update( + { + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "none", + } + ) + elif task_type == "E5": + params.update({"direction": "up", "terminal_behavior": "hold"}) + elif task_type == "E6": + params["target_state"] = "open" + elif task_type == "E7": + params["target_state"] = "closed" + elif task_type == "E8": + params["target_setting"] = 2 + elif task_type == "E9": + params["target_state"] = "activated" + + effective_id = task_id or f"fixture-{task_type.lower()}" + task = validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": effective_id, + "level": "L1", + "instruction": "test-instruction", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": task_type, + "params": params, + "depends_on": [], + "role": "primary", + } + ], + "success": { + "type": task_success_type(task_type, params), + "task_instance_id": "task_01", + }, + "oracle": {}, + "metadata": {"fixture": True}, + } + ) + requirements = validate_scene_requirements( + { + "schema_version": SCENE_REQUIREMENTS_SCHEMA, + "task_id": effective_id, + "objects": objects, + "cameras": [], + "spatial_constraints": [ + {"type": "reachable", "roles": "all_interaction_objects"} + ], + "distractor_count": 0, + "metadata": {"fixture": True}, + } + ) + return task, requirements + + +def make_task_level( + level: str, + *, + reasoning: str | None = None, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Build a validated fixture for one public TaskSpec level.""" + if level == "L1": + return make_task_spec("E1") + first, requirements = make_task_spec("E1", task_id=f"fixture-{level.lower()}") + if level == "L2": + second = { + "id": "task_02", + "task_type": "E1", + "params": { + "object_role": "object_02", + "target_role": "target_02", + "relation": "inside", + }, + "depends_on": ["task_01"], + "role": "primary", + } + first["level"] = "L2" + first["task_instances"].append(second) + first["success"] = { + "op": "all", + "terms": [ + {"type": "semantic_goal", "task_instance_id": "task_01"}, + {"type": "semantic_goal", "task_instance_id": "task_02"}, + ], + } + requirements["objects"].extend( + [ + { + "role_id": "object_02", + "category": "can", + "count": 1, + "affordances": ["graspable", "placeable"], + "initial_state": {}, + "attributes": {}, + }, + { + "role_id": "target_02", + "category": "container", + "count": 1, + "affordances": ["container", "support_surface"], + "initial_state": {}, + "attributes": {}, + }, + ] + ) + return validate_task_spec(first), validate_scene_requirements(requirements) + if level == "L4": + first["level"] = "L4" + first["reasoning_type"] = reasoning or "visual_semantics" + first["success"] = { + "visual_semantics": { + "type": "visual_relation", + "relation": "mouth_completed", + }, + "pattern": { + "type": "visual_relation", + "relation": "pattern_completed", + }, + "logic": {"type": "sum_equals", "value": 5}, + "memory": {"type": "original_order_restored"}, + "common_sense": {"type": "functional_place_setting"}, + "constraint": {"type": "stable_unobstructed"}, + }[first["reasoning_type"]] + first["oracle"] = {"fixture": True} + requirements["cameras"] = [ + { + "role": "reasoning_view", + "modalities": ["rgb", "depth"], + "coverage": "all_interaction_objects", + } + ] + return validate_task_spec(first), validate_scene_requirements(requirements) + raise ValueError(f"Unsupported fixture task level {level!r}.") diff --git a/tests/gen_sim/action_engine/tasks/__init__.py b/tests/gen_sim/action_engine/tasks/__init__.py new file mode 100644 index 000000000..d9480994f --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine task generation tests.""" diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py new file mode 100644 index 000000000..b9af243ee --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -0,0 +1,522 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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] == ["AxisAlign"] + assert orient_nodes[0]["motion_policy"] == {"modifiers": []} + 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[-1]["contract"]["failure_policy"] == "task_required" + 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_pour_recipe_establishes_hold_and_requires_observable_contents() -> None: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "pour_contents", + "level": "L1", + "instruction": "Pour the ball from the cup into the bin.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "pour", + "task_type": "E3", + "params": { + "source_role": "cup", + "target_role": "bin", + "content_roles": ["ball"], + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "poured"}, + "oracle": {}, + "metadata": {}, + } + + graph = instantiate_seed_graph( + task, + {"cup": "source_cup", "bin": "target_bin", "ball": "content_ball"}, + ) + nodes = graph["nodes"] + group = graph["task_groups"][0] + + assert [node["atomic_action"] for node in nodes] == [ + "PickUp", + "MoveHeldObject", + "Pour", + ] + assert nodes[1]["depends_on"] == [nodes[0]["id"]] + assert nodes[2]["depends_on"] == [nodes[1]["id"]] + assert nodes[2]["contract"]["completion"] == "terminal_barrier" + assert nodes[2]["contract"]["failure_policy"] == "task_required" + assert group["success"]["contents"] == [{"object": "content_ball"}] + assert nodes[2]["target_binding"]["payloads"] == [{"object": "content_ball"}] + + +def test_handover_to_place_uses_receiver_hold_without_repickup() -> None: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "handover_then_place", + "level": "L3", + "instruction": "test-instruction-handover-place", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E4", + "params": { + "object_role": "yellow_can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "preserve", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E1", + "params": { + "object_role": "yellow_can", + "target_role": "purple_can", + "relation": "right_of", + }, + "depends_on": ["task_01"], + "role": "primary", + }, + ], + "success": { + "op": "all", + "terms": [ + {"type": "handover_complete", "task_instance_id": "task_01"}, + {"type": "semantic_goal", "task_instance_id": "task_02"}, + ], + }, + "oracle": {"task_order": ["task_01", "task_02"]}, + "metadata": {}, + } + + graph = instantiate_seed_graph( + task, + { + "yellow_can": "interact_yellow_can", + "purple_can": "interact_purple_can", + }, + ) + + handover = next(group for group in graph["task_groups"] if group["id"] == "task_01") + placement = next( + group for group in graph["task_groups"] if group["id"] == "task_02" + ) + placement_nodes = [ + node for node in graph["nodes"] if node["task_instance_id"] == "task_02" + ] + + assert [ + node["atomic_action"] + for node in graph["nodes"] + if node["task_instance_id"] == "task_01" + ] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + ] + assert handover["actor"] == {"mode": "required", "arm": "left_arm"} + assert graph["nodes"][0]["motion_policy"] == { + "modifiers": [{"type": "handover_role", "mode": "transfer"}] + } + assert graph["nodes"][1]["target_binding"]["kind"] == "handover_staging" + assert graph["nodes"][2]["motion_policy"] == {"modifiers": []} + handover_retreat = graph["nodes"][3] + assert handover_retreat["actor"] == {"mode": "required", "arm": "left_arm"} + assert handover_retreat["target_binding"] == { + "kind": "policy_pose", + "source": "handover", + "operation": "retreat", + } + assert handover_retreat["motion_policy"] == {"modifiers": []} + assert handover_retreat["depends_on"] == [graph["nodes"][2]["id"]] + handover_home = graph["nodes"][4] + assert handover_home["actor"] == {"mode": "required", "arm": "left_arm"} + assert handover_home["target_binding"] == { + "kind": "joint_state", + "source": "initial", + "operation": "handover_home", + } + assert handover_home["motion_policy"] == {"modifiers": []} + assert handover_home["depends_on"] == [handover_retreat["id"]] + assert [node["atomic_action"] for node in placement_nodes] == [ + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert placement["actor"] == {"mode": "required", "arm": "right_arm"} + assert placement_nodes[0]["precondition"] == { + "type": "object_held", + "object": "interact_yellow_can", + "arm": "right_arm", + } + assert placement_nodes[0]["depends_on"] == [handover["node_ids"][-1]] + + +def test_structured_draft_grounds_handover_then_receiver_placement() -> None: + scene = [ + { + "runtime_uid": "table", + "role": "background", + "description": "A table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "interact_yellow_can", + "role": "rigid_object", + "description": "A yellow soda can.", + "init_pos": [0.0, -0.25, 0.75], + }, + { + "runtime_uid": "interact_purple_can", + "role": "rigid_object", + "description": "A purple soda can.", + "init_pos": [0.0, 0.25, 0.75], + }, + ] + + planned = _ground_draft( + "handover_then_place", + "test-instruction-handover-place", + scene, + [ + _intent_step( + "handover", + "E4", + _selector("scene_ref", reference="object-alpha"), + transfer_arm="left_arm", + receive_arm="right_arm", + ), + _intent_step( + "place", + "E1", + _selector("step_result", step_id="handover"), + target=_selector("scene_ref", reference="object-beta"), + relation="right_of", + required_arm="right_arm", + ), + ], + { + "handover.object": ["interact_yellow_can"], + "place.target": ["interact_purple_can"], + }, + ) + graph = instantiate_seed_graph(planned.task_spec, planned.role_bindings) + + assert planned.task_spec["level"] == "L3" + assert [item["task_type"] for item in planned.task_spec["task_instances"]] == [ + "E4", + "E1", + ] + assert planned.role_bindings == { + "object_01": "interact_yellow_can", + "object_02": "interact_purple_can", + } + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert graph["task_groups"][1]["goal"]["relation"] == "right_of" + assert graph["task_groups"][1]["goal"]["relation_frame"] == "robot" + + +def test_seed_graph_adds_missing_same_object_e2_handover_dependency() -> None: + scene = [ + { + "runtime_uid": "purple_can", + "role": "rigid_object", + "description": "A purple soda can.", + "init_pos": [0.0, -0.25, 0.7], + }, + { + "runtime_uid": "orange_can", + "role": "rigid_object", + "description": "An orange soda can.", + "init_pos": [0.0, 0.2, 0.7], + }, + ] + planned = _ground_draft( + "missing_same_object_edge", + "test-instruction-multi-step", + scene, + [ + _intent_step( + "orient_purple", + "E2", + _selector("scene_ref", reference="object-alpha"), + required_arm="right_arm", + ), + _intent_step( + "orient_orange", + "E2", + _selector("scene_ref", reference="object-beta"), + required_arm="left_arm", + depends_on=["orient_purple"], + ), + _intent_step( + "handover_purple", + "E4", + _selector("step_result", step_id="orient_purple"), + transfer_arm="right_arm", + receive_arm="left_arm", + depends_on=["orient_orange"], + ), + _intent_step( + "place_purple", + "E1", + _selector("step_result", step_id="handover_purple"), + target=_selector("scene_ref", reference="object-beta"), + relation="left_of", + required_arm="left_arm", + ), + ], + { + "orient_purple.object": ["purple_can"], + "orient_orange.object": ["orange_can"], + "place_purple.target": ["orange_can"], + }, + ) + underconstrained = deepcopy(planned.task_spec) + underconstrained["task_instances"][2]["depends_on"] = ["task_02"] + + graph = instantiate_seed_graph(underconstrained, planned.role_bindings) + handover = next(group for group in graph["task_groups"] if group["id"] == "task_03") + purple = next(group for group in graph["task_groups"] if group["id"] == "task_01") + staging = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "MoveHeldObject" + ) + assert handover["depends_on"] == ["task_02", "task_01"] + assert [ + node["atomic_action"] + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + ] == ["PickUp", "MoveHeldObject", "HandOver", "MoveEndEffector", "MoveJoints"] + pickup = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "PickUp" + ) + assert purple["node_ids"][-1] in pickup["depends_on"] + assert staging["depends_on"] == [pickup["id"]] + + +def test_structured_draft_treats_table_as_support_in_generic_line_task() -> None: + scene = [ + { + "runtime_uid": "table", + "role": "background", + "description": "A table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "interact_red_can", + "role": "rigid_object", + "description": "A red soda can.", + "init_pos": [0.0, -0.25, 0.75], + }, + { + "runtime_uid": "interact_blue_cup", + "role": "rigid_object", + "description": "A blue cup.", + "init_pos": [0.0, 0.25, 0.75], + }, + ] + + planned = _ground_draft( + "arrange_line", + "test-instruction-line", + scene, + [ + _intent_step( + "line", + "E1", + _selector( + "scene_ref", + reference="object-set", + quantifier="all", + ), + layout="line", + ) + ], + {"line.object": ["interact_red_can", "interact_blue_cup"]}, + ) + graph = instantiate_seed_graph(planned.task_spec, planned.role_bindings) + + assert planned.task_spec["level"] == "L2" + assert set(planned.role_bindings.values()) == { + "interact_red_can", + "interact_blue_cup", + } + assert all(group["operator"] == "arrange_line" for group in graph["task_groups"]) diff --git a/tests/gen_sim/action_engine/tasks/test_grounding.py b/tests/gen_sim/action_engine/tasks/test_grounding.py new file mode 100644 index 000000000..e927c5fb6 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_grounding.py @@ -0,0 +1,339 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +import json + +import pytest + +from embodichain.gen_sim.action_engine.tasks.grounding import ( + ground_scene_references, +) +from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + +def _selector( + reference: str, + *, + quantifier: str = "one", + count: int = 0, +) -> dict: + return { + "kind": "scene_ref", + "step_id": "", + "reference": reference, + "quantifier": quantifier, + "count": count, + } + + +def _scene() -> list[dict]: + return [ + { + "runtime_uid": "table", + "uid": "table", + "role": "background", + "category": "dining_table", + "name": "work table", + "description": "A rectangular work table.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "cutting_board", + "uid": "cutting_board", + "role": "rigid_object", + "category": "cutting_board", + "name": "wood board", + "description": "A large rectangular wooden cutting board.", + "attributes": { + "size": "large", + "geometry": {"position": [0.0, 0.2, 0.7], "note": "flat"}, + }, + "initial_state": {"orientation": "fallen"}, + "init_pos": [0.0, 0.2, 0.7], + }, + { + "runtime_uid": "salt_shaker", + "uid": "salt_shaker", + "role": "rigid_object", + "category": "salt_shaker", + "description": "A small glass salt shaker.", + "affordances": ["graspable"], + "init_pos": [0.0, -0.2, 0.7], + }, + ] + + +def _intent( + *, + object_selector: dict | None = None, + target_selector: dict | None = None, +) -> dict: + return { + "steps": [ + { + "id": "move", + "task_type": "E1", + "object": object_selector or _selector("object-alpha"), + "target": target_selector or _selector("target-alpha"), + "relation": "on", + } + ] + } + + +def _binding( + reference_id: str, + uids: list[str], + *, + status: str = "resolved", + confidence: float = 1.0, + **extra: object, +) -> dict: + return { + "reference_id": reference_id, + "status": status, + "uids": uids, + "confidence": confidence, + **extra, + } + + +def _run(intent: dict, caller) -> object: + scene = _scene() + return ground_scene_references( + instruction="test-instruction", + intent=intent, + inventory=SceneInventory(scene, robot_profile="franka"), + scene_objects=scene, + model="test-model", + caller=caller, + ) + + +def test_grounding_prompt_preserves_open_semantics_and_redacts_geometry() -> None: + captured: dict[str, object] = {} + + def caller(**kwargs): + captured.update(kwargs) + return { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.target", ["table"]), + ] + } + + result = _run(_intent(), caller) + + prompt = str(captured["prompt"]) + assert result.bindings == { + "move.object": ("cutting_board",), + "move.target": ("table",), + } + assert '"category": "cutting_board"' in prompt + assert '"category": "salt_shaker"' in prompt + assert '"name": "wood board"' in prompt + assert '"orientation": "fallen"' in prompt + assert '"size": "large"' in prompt + prompt_inventory = json.loads(prompt.split("Redacted scene inventory:\n", 1)[1]) + side_by_uid = {item["uid"]: item["side"] for item in prompt_inventory} + assert side_by_uid["cutting_board"] == "right" + assert side_by_uid["salt_shaker"] == "left" + assert '"position"' not in prompt + assert '"init_pos"' not in prompt + + +@pytest.mark.parametrize("robot_profile", ["ur5", "ur10", "franka"]) +def test_scene_inventory_uses_the_shared_final_world_lateral_axis( + robot_profile: str, +) -> None: + inventory = SceneInventory(_scene(), robot_profile=robot_profile) + + assert inventory.left_score(inventory.by_uid["salt_shaker"]) > 0.0 + assert inventory.left_score(inventory.by_uid["cutting_board"]) < 0.0 + + +def test_grounding_repairs_one_invalid_uid_in_the_same_batch() -> None: + responses = [ + { + "bindings": [ + _binding("move.object", ["invented"]), + _binding("move.target", ["table"]), + ] + }, + { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.target", ["table"]), + ] + }, + ] + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(responses[len(prompts) - 1]) + + result = _run(_intent(), caller) + + assert result.attempts == 2 + assert "previous grounding JSON failed" in prompts[1] + assert result.bindings["move.object"] == ("cutting_board",) + + +@pytest.mark.parametrize( + "response,error", + [ + ( + { + "bindings": [ + _binding( + "move.object", + ["cutting_board"], + status="ambiguous", + ), + _binding("move.target", ["table"]), + ] + }, + "was not resolved", + ), + ( + { + "bindings": [ + _binding( + "move.object", + [], + status="not_found", + confidence=0.0, + ), + _binding("move.target", ["table"]), + ] + }, + "was not resolved", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board"], confidence=0.49), + _binding("move.target", ["table"]), + ] + }, + "confidence is below", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board", "cutting_board"]), + _binding("move.target", ["table"]), + ] + }, + "duplicate UIDs", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.object", ["salt_shaker"]), + _binding("move.target", ["table"]), + ] + }, + "Duplicate grounding binding", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board", "salt_shaker"]), + _binding("move.target", ["table"]), + ] + }, + "quantifier=one requires exactly one UID", + ), + ( + {"bindings": [_binding("move.object", ["cutting_board"])]}, + "omitted requests", + ), + ( + { + "bindings": [ + _binding("move.object", ["table"]), + _binding("move.target", ["cutting_board"]), + ] + }, + "candidate range", + ), + ( + { + "bindings": [ + _binding("move.object", ["cutting_board"]), + _binding("move.target", ["cutting_board"]), + ] + }, + "same UID", + ), + ( + { + "bindings": [ + _binding( + "move.object", + ["cutting_board"], + affordances=["graspable"], + ), + _binding("move.target", ["table"]), + ] + }, + "unsupported", + ), + ], +) +def test_grounding_fails_closed_after_one_repair(response: dict, error: str) -> None: + with pytest.raises(ValueError, match=f"after one repair.*{error}"): + _run(_intent(), lambda **_kwargs: deepcopy(response)) + + +def test_grounding_enforces_count_and_accepts_an_open_world_set() -> None: + intent = _intent( + object_selector=_selector("object-set", quantifier="count", count=2) + ) + response = { + "bindings": [ + _binding("move.object", ["cutting_board", "salt_shaker"]), + _binding("move.target", ["table"]), + ] + } + + result = _run(intent, lambda **_kwargs: response) + assert result.bindings["move.object"] == ("cutting_board", "salt_shaker") + + invalid = deepcopy(response) + invalid["bindings"][0]["uids"] = ["cutting_board"] + with pytest.raises(ValueError, match="requires exactly 2 UIDs"): + _run(intent, lambda **_kwargs: invalid) + + +def test_grounding_accepts_a_nonempty_all_binding() -> None: + intent = _intent(object_selector=_selector("object-set", quantifier="all")) + response = { + "bindings": [ + _binding("move.object", ["cutting_board", "salt_shaker"]), + _binding("move.target", ["table"]), + ] + } + + result = _run(intent, lambda **_kwargs: response) + + assert result.bindings["move.object"] == ("cutting_board", "salt_shaker") diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py new file mode 100644 index 000000000..3cc64182b --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -0,0 +1,1939 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 == ["AxisAlign"] + 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_e3_target_repair_preserves_pour_semantics() -> None: + invalid = { + "steps": [ + _step( + "pour", + "E3", + _selector("scene_ref", reference="the source cup"), + relation="above", + ) + ] + } + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(invalid) + + with pytest.raises(ValueError, match="after one repair"): + task_interpretation_module.interpret_instruction_draft( + "Grab the cup and pour its contents into the bin.", + model="test-model", + caller=caller, + ) + + assert "exactly one E3 step" in prompts[0] + assert "Missing-target repair rule for E3" in prompts[1] + assert "must not be reclassified as E1" in prompts[1] + + +def test_missing_e6_object_reaches_targeted_repair() -> None: + invalid = { + "steps": [ + _step( + "open_drawer", + "E6", + _selector(), + target_state="open", + required_arm="right_arm", + ) + ] + } + repaired = deepcopy(invalid) + repaired["steps"][0]["object"] = _selector("scene_ref", reference="the drawer") + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(invalid if len(prompts) == 1 else repaired) + + result = task_interpretation_module.interpret_instruction_draft( + "Open the drawer with the right arm.", + model="test-model", + caller=caller, + ) + + assert result.attempts == 2 + assert result.intent == repaired + assert "Opening or pulling out a drawer is E6" in prompts[0] + assert "Missing-object repair rule" in prompts[1] + + +def test_missing_target_completion_rejects_other_semantic_disagreement() -> None: + invalid_intent = _two_object_handover_intent_with_missing_place_target() + invalid_intent["steps"][-1]["required_arm"] = "right_arm" + + with pytest.raises(ValueError, match="after one repair"): + interpret_and_ground_task_spec( + "unsafe_target_completion", + "test-instruction-missing-target", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(invalid_intent), + ) + + +def test_second_invalid_intent_fails_without_rule_fallback() -> None: + with pytest.raises(ValueError, match="after one repair"): + interpret_and_ground_task_spec( + "invalid", + "test-instruction-invalid", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: {"steps": []}, + ) + + +def test_intent_infers_pronoun_dependency_from_canonical_symbols() -> None: + intent = { + "steps": [ + _step( + "handover", + "E4", + _selector("scene_ref", reference="object-alpha"), + transfer_arm="right_arm", + receive_arm="left_arm", + ), + _step( + "place", + "E1", + _selector("step_result", step_id="handover"), + target=_selector("scene_ref", reference="object-beta"), + relation="left_of", + required_arm="left_arm", + depends_on=["handover"], + ), + ] + } + + grounded = interpret_and_ground_task_spec( + "implicit_dependency", + "test-instruction-pronoun-dependency", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "handover.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + + instances = grounded.task_spec["task_instances"] + assert [item["task_type"] for item in instances] == ["E4", "E1"] + assert instances[1]["depends_on"] == [instances[0]["id"]] + assert instances[1]["params"]["relation"] == "left_of" + assert instances[1]["params"]["required_arm"] == "left_arm" + + +def test_scene_grounding_rejects_unknown_uid() -> None: + intent = { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + } + with pytest.raises(ValueError, match="after one repair.*unknown UIDs"): + interpret_and_ground_task_spec( + "unknown_uid", + "test-instruction-unknown-uid", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller(**{"orient.object": "invented_uid"}), + ) + + +def test_instruction_intent_rejects_legacy_selector_protocol() -> None: + intent = _handover_intent() + intent["steps"][0]["object"] = { + "kind": "selector", + "step_id": "", + "uid": "purple_can", + "category": "can", + "color": "purple", + "side": "none", + "quantifier": "one", + "count": 0, + } + with pytest.raises(ValueError, match="requires exactly fields"): + validate_instruction_intent(intent) + + +def test_step_result_must_reference_a_preceding_step() -> None: + intent = { + "steps": [ + _step( + "handover", + "E4", + _selector("step_result", step_id="orient"), + transfer_arm="right_arm", + receive_arm="left_arm", + ), + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ), + ] + } + with pytest.raises(ValueError, match="preceding step"): + validate_instruction_intent(intent) + + +def test_step_result_selector_rejects_object_constraints() -> None: + intent = _handover_intent() + intent["steps"][1]["object"]["reference"] = "object-alpha" + + with pytest.raises(ValueError, match="may identify only a prior step_id"): + validate_instruction_intent(intent) + + +def test_instruction_intent_rejects_non_e_specific_parameters() -> None: + invalid_e9 = _step( + "press", + "E9", + _selector("selector", category="button"), + target_state="activated", + orientation_goal="upright", + ) + with pytest.raises(ValueError, match="orientation_goal"): + validate_instruction_intent({"steps": [invalid_e9]}) + + invalid_line = _step( + "line", + "E1", + _selector("selector", category="can", quantifier="all"), + layout="line", + relation="on", + ) + with pytest.raises(ValueError, match="line arrangement cannot carry a relation"): + validate_instruction_intent({"steps": [invalid_line]}) + + +def test_implicit_e1_relation_requires_an_unambiguous_support_target() -> None: + intent = { + "steps": [ + _step( + "place", + "E1", + _selector("selector", category="can", color="purple"), + target=_selector("selector", category="can", color="orange"), + ) + ] + } + + with pytest.raises(ValueError, match="omitted relation"): + interpret_and_ground_task_spec( + "ambiguous_implicit_place", + "test-instruction-implicit-relation", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "place.object": "purple_can", + "place.target": "orange_can", + } + ), + ) + + +def test_instruction_and_grounding_prompts_keep_their_boundaries() -> None: + captured: dict[str, dict] = {} + intent = { + "steps": [ + _step( + "place", + "E1", + _selector("selector", category="can", color="purple"), + target=_selector("selector", uid="table", category="table"), + relation="on", + ) + ] + } + + def caller(**kwargs): + captured["intent"] = kwargs + return intent + + def grounding_caller(**kwargs): + captured["grounding"] = kwargs + return _grounding(**{"place.object": "purple_can", "place.target": "table"}) + + grounded = interpret_and_ground_task_spec( + "onto_table", + "Put the purple can on the table.", + _scene_with_table(), + robot_profile="ur10", + caller=caller, + grounding_caller=grounding_caller, + ) + + assert '"uid": "table"' not in captured["intent"]["prompt"] + assert '"uid": "table"' in captured["grounding"]["prompt"] + assert '"core_actions"' not in captured["intent"]["prompt"] + assert captured["intent"]["schema"] == INSTRUCTION_INTENT_SCHEMA + assert grounded.role_bindings["object_02"] == "table" + + +def test_instruction_intent_schema_declares_every_required_selector_field() -> None: + selector_schema = INSTRUCTION_INTENT_SCHEMA["properties"]["steps"]["items"][ + "properties" + ]["object"] + + assert set(selector_schema["required"]) == set(selector_schema["properties"]) + assert "quantifier" in selector_schema["properties"] + + +def test_grounding_prompt_redacts_nested_scene_geometry() -> None: + scene = _scene() + scene[0]["attributes"] = { + "label": "purple", + "geometry": {"position": [0.0, 0.0, 0.7], "note": "can"}, + } + captured: dict[str, str] = {} + + def grounding_caller(**kwargs): + captured["prompt"] = kwargs["prompt"] + return _grounding(**{"orient.object": "purple_can"}) + + interpret_and_ground_task_spec( + "redacted_inventory", + "test-instruction-grounding-redaction", + scene, + robot_profile="ur10", + caller=lambda **_kwargs: { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + }, + grounding_caller=grounding_caller, + ) + assert '"position"' not in captured["prompt"] + assert '"label": "purple"' in captured["prompt"] + + +def test_default_llm_parser_requires_the_documented_model_configuration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("ACTION_ENGINE_LLM_MODEL", raising=False) + monkeypatch.delenv("OPENAI_MODEL", raising=False) + monkeypatch.setattr(task_interpretation_module, "_load_local_env", lambda: {}) + + with pytest.raises(ValueError, match="text LLM model is required"): + interpret_and_ground_task_spec( + "missing_model", + "test-instruction-model-config", + _scene(), + robot_profile="ur10", + ) + + +def test_injected_caller_skips_production_model_resolution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def unexpected_model_resolution(_explicit: str | None) -> str | None: + raise AssertionError( + "injected callers must not resolve production model config" + ) + + monkeypatch.setattr( + task_interpretation_module, + "_instruction_model", + unexpected_model_resolution, + ) + + grounded = interpret_and_ground_task_spec( + "injected_caller", + "test-instruction-injected-caller", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: { + "steps": [ + _step( + "orient", + "E2", + _selector("scene_ref", reference="object-alpha"), + ) + ] + }, + grounding_caller=_grounding_caller(**{"orient.object": "purple_can"}), + ) + + assert grounded.task_spec["metadata"]["instruction_model"] == "injected_caller" + + +def test_mimo_instruction_caller_uses_json_mode_and_disables_thinking( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """MiMo-compatible endpoints must not use the lossy JSON-schema route.""" + import langchain_openai + + calls: list[dict] = [] + responses = [ + { + "steps": [ + { + "id": "orient", + "task_type": "E2", + "object": _selector("scene_ref", reference="object-alpha"), + } + ] + }, + _handover_intent(), + _grounding( + **{ + "orient.object": "purple_can", + "place.target": "orange_can", + } + ), + ] + + class FakeRunnable: + def invoke(self, messages): + calls[-1]["messages"] = messages + return deepcopy(responses.pop(0)) + + class FakeChatOpenAI: + def __init__(self, **kwargs): + calls.append({"kwargs": kwargs}) + + def with_structured_output(self, schema, **kwargs): + calls[-1]["schema"] = schema + calls[-1]["structured_kwargs"] = kwargs + return FakeRunnable() + + monkeypatch.setattr(langchain_openai, "ChatOpenAI", FakeChatOpenAI) + monkeypatch.setattr( + task_interpretation_module, + "_load_llm_settings", + lambda *, model: { + "api_key": "test-key", + "model": model or "mimo-v2.5", + "base_url": "https://token-plan-cn.xiaomimimo.com/v1", + "default_query": {}, + }, + ) + + grounded = interpret_and_ground_task_spec( + "mimo_repair", + "test-instruction-json-mode", + _scene(), + robot_profile="ur10", + model="mimo-v2.5", + ) + + assert [item["task_type"] for item in grounded.task_spec["task_instances"]] == [ + "E2", + "E4", + "E1", + ] + assert len(calls) == 3 + for call in calls: + assert call["structured_kwargs"] == {"method": "json_mode"} + assert call["kwargs"]["http_socket_options"] == () + assert call["kwargs"]["max_completion_tokens"] == 4096 + assert call["kwargs"]["extra_body"] == {"thinking": {"type": "disabled"}} + repair_messages = calls[1]["messages"] + assert "previous JSON was invalid" in repair_messages[1].content + + +def test_instruction_prompt_contains_a_complete_shape_example() -> None: + prompt = interpretation_module._instruction_prompt("instruction-marker") + selector_rules = interpretation_module._instruction_selector_rules() + assert '"target_setting": 0' in prompt + assert '"depends_on": []' in prompt + assert "every step has all 16 step keys" in prompt + assert "step_result" in prompt + assert "open scene_ref.reference" in prompt + assert "Do not classify it or emit a scene UID" in prompt + assert "example object A" in prompt + assert "stale-object-reference" not in prompt + assert "step_result" in selector_rules + assert "step_id" in selector_rules + assert "reference" in selector_rules + + +def test_scene_export_spatial_descriptions_do_not_create_false_supports() -> None: + index = SceneInventory(_scene_export_style_scene(), robot_profile="franka") + + assert [entity.uid for entity in index.support] == ["table"] + assert {entity.uid for entity in index.movable} == { + "carrot_001", + "cutting_board_001", + "peeler_001", + } + + +def test_scene_export_exact_uids_ground_pick_and_place() -> None: + intent = { + "steps": [ + _step( + "step_1", + "E1", + _selector("selector", uid="carrot_001"), + target=_selector("selector", uid="cutting_board_001"), + relation="on", + required_arm="left_arm", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "scene_export_pick_place", + "test-instruction-exact-uids", + _scene_export_style_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: intent, + grounding_caller=_grounding_caller( + **{ + "step_1.object": "carrot_001", + "step_1.target": "cutting_board_001", + } + ), + ) + + assert set(grounded.role_bindings.values()) == { + "carrot_001", + "cutting_board_001", + } + assert grounded.task_spec["task_instances"][0]["params"]["required_arm"] == ( + "left_arm" + ) + assert {item["category"] for item in grounded.scene_requirements["objects"]} == { + "carrot", + "cutting_board", + } + + +def test_multi_object_handover_keeps_both_order_and_holder_dependencies() -> None: + intent = _two_object_handover_intent() + grounded = interpret_and_ground_task_spec( + "multi_object_handover", + "test-instruction-multi-object-handover", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "orient_purple.object": "purple_can", + "orient_orange.object": "orange_can", + "place.target": "orange_can", + } + ), + ) + + instances = grounded.task_spec["task_instances"] + assert instances[2]["depends_on"] == ["task_02", "task_01"] + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + handover = next(group for group in graph["task_groups"] if group["id"] == "task_03") + handover_nodes = [ + node for node in graph["nodes"] if node["task_instance_id"] == "task_03" + ] + assert handover["depends_on"] == ["task_02", "task_01"] + assert [node["atomic_action"] for node in handover_nodes] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + ] + purple = next(group for group in graph["task_groups"] if group["id"] == "task_01") + orange = next(group for group in graph["task_groups"] if group["id"] == "task_02") + assert handover_nodes[0]["depends_on"] == [ + orange["node_ids"][-1], + purple["node_ids"][-1], + ] + + +def test_single_arm_e1_propagates_direct_payload_into_goal_and_contracts() -> None: + intent = { + "steps": [ + _step( + "handover_glue", + "E4", + _selector("selector", uid="glue_stick"), + required_arm="left_arm", + transfer_arm="left_arm", + receive_arm="right_arm", + ), + _step( + "place_glue", + "E1", + _selector("step_result", step_id="handover_glue"), + target=_selector("selector", uid="paper_cup"), + relation="on", + required_arm="right_arm", + depends_on=["handover_glue"], + ), + _step( + "place_cup", + "E1", + _selector("selector", uid="paper_cup"), + target=_selector("selector", uid="popcorn_bucket"), + relation="on", + required_arm="right_arm", + depends_on=["place_glue"], + ), + ] + } + grounded = interpret_and_ground_task_spec( + "payload_chain", + "test-instruction-payload-propagation", + _payload_scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "handover_glue.object": "glue_stick", + "place_glue.target": "paper_cup", + "place_cup.object": "paper_cup", + "place_cup.target": "popcorn_bucket", + } + ), + ) + + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + carrier_group = next( + group for group in graph["task_groups"] if group["id"] == "task_03" + ) + assert carrier_group["goal"]["payloads"] == [ + {"object": "glue_stick", "slot": "center"} + ] + carrier_nodes = [ + node + for node in graph["nodes"] + if node["task_instance_id"] == carrier_group["id"] + and node["atomic_action"] in {"PickUp", "MoveHeldObject", "Place"} + ] + assert carrier_nodes + for node in carrier_nodes: + assert node["target_binding"]["payloads"] == carrier_group["goal"]["payloads"] + assert any( + claim["resource"] == "object:glue_stick" and claim["access"] == "exclusive" + for claim in node["contract"]["claims"] + ) + + +def test_seed_graph_repairs_missing_e2_handover_lifecycle_edge() -> None: + intent = _two_object_handover_intent() + grounded = interpret_and_ground_task_spec( + "missing_lifecycle_edge", + "test-instruction-lifecycle-repair", + _scene(), + robot_profile="ur10", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + **{ + "orient_purple.object": "purple_can", + "orient_orange.object": "orange_can", + "place.target": "orange_can", + } + ), + ) + underconstrained = deepcopy(grounded.task_spec) + underconstrained["task_instances"][2]["depends_on"] = ["task_02"] + + graph = instantiate_seed_graph(underconstrained, grounded.role_bindings) + handover = next(group for group in graph["task_groups"] if group["id"] == "task_03") + purple = next(group for group in graph["task_groups"] if group["id"] == "task_01") + assert handover["depends_on"] == ["task_02", "task_01"] + pickup = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" and node["atomic_action"] == "PickUp" + ) + assert purple["node_ids"][-1] in pickup["depends_on"] diff --git a/tests/gen_sim/action_engine/tasks/test_language_decoupling.py b/tests/gen_sim/action_engine/tasks/test_language_decoupling.py new file mode 100644 index 000000000..b9f0768d3 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_language_decoupling.py @@ -0,0 +1,420 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Acceptance tests for the structured-LLM language boundary.""" + +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path + +import pytest + +import embodichain.gen_sim.action_engine.tasks as action_engine_tasks +from embodichain.gen_sim.action_engine.tasks import ( + instantiate_seed_graph, + interpret_and_ground_task_spec, + validate_instruction_intent, +) +from embodichain.gen_sim.action_engine.tasks.assembly import SceneInventory + + +def _selector( + kind: str = "none", + *, + reference: str = "", + step_id: str = "", +) -> dict: + return { + "kind": kind, + "step_id": step_id, + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _step(step_id: str, task_type: str, reference: str, **updates: object) -> dict: + step = { + "id": step_id, + "task_type": task_type, + "object": _selector("scene_ref", reference=reference), + "target": _selector(), + "relation": "none", + "required_arm": "none", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright" if task_type == "E2" else "preserve", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + step.update(updates) + return step + + +def _binding(reference_id: str, *uids: str) -> dict: + return { + "reference_id": reference_id, + "status": "resolved", + "uids": list(uids), + "confidence": 1.0, + } + + +def _grounding_caller(*bindings: dict): + response = {"bindings": list(bindings)} + return lambda **_kwargs: deepcopy(response) + + +def _open_scene() -> list[dict]: + return [ + { + "runtime_uid": "work_surface", + "uid": "work_surface", + "role": "support_surface", + "category": "obsidian_dock", + "name": "the landing ledge", + "description": "A flat black ledge used as a work surface.", + "init_pos": [0.0, 0.0, 0.0], + }, + { + "runtime_uid": "aerogel_fixture_7", + "uid": "aerogel_fixture_7", + "role": "rigid_object", + "category": "aerogel_fixture", + "name": "translucent fixture", + "description": "A translucent rectangular fixture with a frosted edge.", + "init_pos": [0.0, 0.1, 0.7], + }, + { + "runtime_uid": "plantain_marker", + "uid": "plantain_marker", + "role": "rigid_object", + "category": "plantain_marker", + "description": "A curved yellow marker behind the fixture.", + "init_pos": [0.1, -0.2, 0.7], + }, + ] + + +def test_scene_inventory_preserves_open_category_labels() -> None: + scene = _open_scene() + scene[1]["category"] = "Prototype.Fixture/V2" + inventory = SceneInventory(scene, robot_profile="franka") + + assert inventory.by_uid["aerogel_fixture_7"].category == ("Prototype.Fixture/V2") + + +@pytest.mark.parametrize( + ("step", "invalid_field"), + [ + ( + _step( + "place", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="invalid-relation", + ), + "relation", + ), + ( + _step("orient", "E2", "object-alpha", required_arm="invalid-arm"), + "required_arm", + ), + ( + _step( + "orient", + "E2", + "object-alpha", + orientation_goal="invalid-orientation", + ), + "orientation_goal", + ), + ], +) +def test_llm_intent_rejects_natural_language_aliases( + step: dict, + invalid_field: str, +) -> None: + """Canonical protocol fields are not a second local language parser.""" + with pytest.raises(ValueError, match=invalid_field): + validate_instruction_intent({"steps": [step]}) + + +def test_noncanonical_llm_value_is_repaired_instead_of_locally_normalized() -> None: + invalid = { + "steps": [ + _step( + "place", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="invalid-relation", + ) + ] + } + valid = deepcopy(invalid) + valid["steps"][0]["relation"] = "left_of" + responses = [invalid, valid] + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(responses[len(prompts) - 1]) + + grounded = interpret_and_ground_task_spec( + "strict_canonical_repair", + "test-instruction-invalid-relation", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=caller, + grounding_caller=_grounding_caller( + _binding("place.object", "aerogel_fixture_7"), + _binding("place.target", "work_surface"), + ), + ) + + assert len(prompts) == 2 + assert "previous JSON was invalid" in prompts[1] + assert grounded.task_spec["metadata"]["instruction_call_count"] == 2 + assert grounded.task_spec["task_instances"][0]["params"]["relation"] == ("left_of") + assert "instruction_intent_normalizations" not in grounded.task_spec["metadata"] + + +def test_two_noncanonical_llm_responses_fail_without_grounding_or_rule_fallback() -> ( + None +): + invalid = { + "steps": [ + _step( + "place", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="invalid-relation", + ) + ] + } + grounding_called = False + + def unexpected_grounding(**_kwargs): + nonlocal grounding_called + grounding_called = True + raise AssertionError("invalid canonical intent must not reach grounding") + + with pytest.raises(ValueError, match="after one repair.*relation"): + interpret_and_ground_task_spec( + "strict_canonical_failure", + "test-instruction-invalid-relation", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: deepcopy(invalid), + grounding_caller=unexpected_grounding, + ) + + assert grounding_called is False + + +def test_legacy_instruction_parser_modules_and_api_are_absent() -> None: + tasks_dir = Path(action_engine_tasks.__file__).resolve().parent + + assert not (tasks_dir / "deterministic.py").exists() + assert not (tasks_dir / "planning.py").exists() + assert not hasattr(action_engine_tasks, "plan_grounded_task_spec") + + +def test_production_sources_do_not_reference_legacy_instruction_parser() -> None: + action_engine_dir = Path(action_engine_tasks.__file__).resolve().parent.parent + forbidden = ( + "tasks.deterministic", + "tasks.planning", + "plan_grounded_task_spec", + "instruction_parser", + "deterministic_fallback", + ) + offenders: dict[str, list[str]] = {} + for path in action_engine_dir.rglob("*.py"): + source = path.read_text(encoding="utf-8") + matches = [term for term in forbidden if term in source] + if matches: + offenders[str(path.relative_to(action_engine_dir))] = matches + + assert offenders == {} + + +def test_llm_caller_exception_propagates_without_scene_grounding() -> None: + expected = RuntimeError("model unavailable") + grounding_called = False + + def fail_model(**_kwargs): + raise expected + + def unexpected_grounding(**_kwargs): + nonlocal grounding_called + grounding_called = True + raise AssertionError("failed interpretation must not reach grounding") + + with pytest.raises(RuntimeError) as caught: + interpret_and_ground_task_spec( + "model_failure", + "test-instruction-caller-error", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=fail_model, + grounding_caller=unexpected_grounding, + ) + + assert caught.value is expected + assert grounding_called is False + + +def test_unfamiliar_wording_and_categories_flow_through_injected_llm_stages() -> None: + intent = { + "steps": [ + _step( + "relocate_fixture", + "E1", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="on", + required_arm="auto", + ) + ] + } + + grounded = interpret_and_ground_task_spec( + "open_world_fixture", + "test-instruction-open-reference", + _open_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: deepcopy(intent), + grounding_caller=_grounding_caller( + _binding("relocate_fixture.object", "aerogel_fixture_7"), + _binding("relocate_fixture.target", "work_surface"), + ), + ) + + assert set(grounded.role_bindings.values()) == { + "aerogel_fixture_7", + "work_surface", + } + assert {item["category"] for item in grounded.scene_requirements["objects"]} == { + "aerogel_fixture", + "obsidian_dock", + } + assert grounded.task_spec["metadata"]["instruction_call_count"] == 1 + assert grounded.task_spec["metadata"]["scene_grounding_call_count"] == 1 + + +@pytest.mark.parametrize( + ("name", "instruction", "step", "bindings", "actions", "success"), + [ + ( + "dual_lift", + "test-instruction-hold", + _step( + "lift_fixture", + "E5", + "object-alpha", + terminal_behavior="hold", + ), + [_binding("lift_fixture.object", "aerogel_fixture_7")], + ["CoordinatedPickment"], + "held_by_both_grippers", + ), + ( + "dual_move_place", + "test-instruction-directional-place", + _step( + "move_fixture", + "E5", + "object-alpha", + direction="left", + terminal_behavior="place", + ), + [_binding("move_fixture.object", "aerogel_fixture_7")], + ["CoordinatedPickment", "MoveJoints", "MoveJoints"], + "semantic_goal", + ), + ( + "dual_relative", + "test-instruction-relative-place", + _step( + "move_relative", + "E5", + "object-alpha", + target=_selector("scene_ref", reference="target-alpha"), + relation="behind", + terminal_behavior="hold", + ), + [ + _binding("move_relative.object", "aerogel_fixture_7"), + _binding("move_relative.target", "plantain_marker"), + ], + ["CoordinatedPickment"], + "held_by_both_grippers", + ), + ], +) +def test_e5_symbolic_intent_reaches_the_seed_graph( + name: str, + instruction: str, + step: dict, + bindings: list[dict], + actions: list[str], + success: str, +) -> None: + grounded = interpret_and_ground_task_spec( + name, + instruction, + _open_scene(), + robot_profile="franka", + model="test-model", + caller=lambda **_kwargs: {"steps": [deepcopy(step)]}, + grounding_caller=_grounding_caller(*bindings), + ) + instance = grounded.task_spec["task_instances"][0] + graph = instantiate_seed_graph(grounded.task_spec, grounded.role_bindings) + + assert [node["atomic_action"] for node in graph["nodes"]] == actions + assert grounded.task_spec["success"]["terms"] == [ + {"type": success, "task_instance_id": instance["id"]} + ] + assert instance["params"].get("direction") == ( + "up" if name == "dual_lift" else step["direction"] + ) + if name == "dual_relative": + assert graph["task_groups"][0]["goal"]["reference_object"] == ( + "plantain_marker" + ) + assert graph["task_groups"][0]["goal"]["relation"] == "behind" + if name == "dual_move_place": + release_nodes = graph["nodes"][1:] + assert {node["actor"]["arm"] for node in release_nodes} == { + "left_arm", + "right_arm", + } + assert len({node["sync_group"] for node in release_nodes}) == 1 diff --git a/tests/gen_sim/action_engine/test_agent.py b/tests/gen_sim/action_engine/test_agent.py new file mode 100644 index 000000000..d42256ffa --- /dev/null +++ b/tests/gen_sim/action_engine/test_agent.py @@ -0,0 +1,219 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Action Agent compilation, preflight, and report boundary tests.""" + +from __future__ import annotations + +from dataclasses import replace +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +import embodichain.gen_sim.action_engine.agent as module +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapabilityRegistry, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.domain import seed_graph_hash +from embodichain.gen_sim.action_engine.runtime import ExecutionResult +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from .task_fixtures import make_task_spec + + +def _bindings(requirements: dict) -> dict[str, str]: + return { + item["role_id"]: f"scene_{item['role_id']}" for item in requirements["objects"] + } + + +def _task_of_type(task_type: str) -> tuple[dict, dict]: + return make_task_spec(task_type) + + +def _registry_with_planning_only(action_name: str) -> AtomicCapabilityRegistry: + source = build_atomic_capability_registry() + registry = AtomicCapabilityRegistry() + for name in source.names(): + capability = source.get(name) + if name == action_name: + capability = replace( + capability, + runtime_available=False, + unavailable_reason="test-only unavailable runtime", + ) + registry.register(capability) + return registry + + +def test_plan_hash_matches_direct_seed_graph_instantiation(monkeypatch) -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + grounded_plan = { + "task_spec": task, + "role_bindings": {"role_bindings": bindings}, + } + monkeypatch.setattr( + module, + "_validate_grounded_plan", + lambda value: dict(value), + ) + + graph = ActionAgent().plan(grounded_plan) + direct = instantiate_seed_graph(task, bindings) + + assert seed_graph_hash(graph) == seed_graph_hash(direct) + + +def test_planning_only_graph_is_rejected_before_executor_construction() -> None: + task, requirements = _task_of_type("E6") + bindings = _bindings(requirements) + registry = _registry_with_planning_only("PullArticulatedPart") + graph = instantiate_seed_graph(task, bindings, registry=registry) + constructed = False + + def executor_factory(*args, **kwargs): + nonlocal constructed + constructed = True + raise AssertionError("preflight must reject before executor construction") + + report = ActionAgent( + registry=registry, + executor_factory=executor_factory, + ).execute( + graph, + SimpleNamespace(num_envs=2), + known_uids=set(bindings.values()), + run_id="preflight-test", + ) + + assert report.status == "rejected" + assert report.action_count == 0 + assert "planning-only" in (report.error or "") + assert not constructed + + +def test_execution_report_is_strictly_json_serializable(tmp_path: Path) -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + + class FakeExecutor: + def __init__(self, program, env, **kwargs) -> None: + self.program = program + self.env = env + + def run(self, **kwargs) -> ExecutionResult: + return ExecutionResult( + actions=[torch.ones((2, 3), dtype=torch.float32)], + success=torch.tensor([True, False]), + semantic_success={ + "task_01": torch.tensor([True, False]), + }, + record_dir=str(tmp_path), + retry_count=1, + retry_counts=[0, 1], + failure_events=[ + { + "failure_type": "plan_failed", + "env_ids": torch.tensor([1]), + } + ], + ) + + report = ActionAgent(executor_factory=FakeExecutor).execute( + graph, + SimpleNamespace(num_envs=2), + known_uids=set(bindings.values()), + run_id="json-test", + ) + payload = report.as_mapping() + + assert report.status == "failed" + assert payload["environments"][0]["semantic_success"] == {"task_01": True} + assert payload["environments"][1]["semantic_success"] == {"task_01": False} + assert [item["retry_count"] for item in payload["environments"]] == [0, 1] + assert "actions" not in payload + json.dumps(payload, allow_nan=False) + assert ( + json.loads((tmp_path / "execution_report.json").read_text(encoding="utf-8")) + == payload + ) + trajectory = torch.load(tmp_path / "executed_trajectory.pt", weights_only=True) + assert torch.equal(trajectory["actions"][0], torch.ones((2, 3))) + trajectory_manifest = json.loads( + (tmp_path / "executed_trajectory.json").read_text(encoding="utf-8") + ) + assert trajectory_manifest["actions"][0]["shape"] == [2, 3] + + +def test_existing_execution_result_can_be_reported_without_reexecution() -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + result = ExecutionResult( + actions=[torch.zeros((1, 2), dtype=torch.float32)], + success=torch.tensor([True]), + semantic_success={"task_01": torch.tensor([True])}, + ) + + report = ActionAgent().report_execution_result( + result, + action_graph=graph, + run_id="legacy-run", + episode_index=3, + ) + + assert report.status == "succeeded" + assert report.episode_id == "3" + assert report.action_count == 1 + + +def test_runtime_exception_is_reported_as_aborted() -> None: + task, requirements = make_task_spec("E1") + bindings = _bindings(requirements) + graph = instantiate_seed_graph(task, bindings) + + def fail_executor(*_args, **_kwargs): + raise RuntimeError("simulator stopped") + + report = ActionAgent(executor_factory=fail_executor).execute( + graph, + SimpleNamespace(num_envs=1), + known_uids=set(bindings.values()), + run_id="aborted-test", + ) + + assert report.status == "aborted" + assert report.action_count == 0 + assert report.error == "RuntimeError: simulator stopped" + + +def test_preflight_raises_for_planning_only_graph() -> None: + task, requirements = _task_of_type("E8") + bindings = _bindings(requirements) + registry = _registry_with_planning_only("TurnKnob") + + with pytest.raises(ValueError, match="planning-only"): + ActionAgent(registry=registry).preflight( + instantiate_seed_graph(task, bindings, registry=registry), + known_uids=set(bindings.values()), + ) diff --git a/tests/gen_sim/action_engine/test_architecture.py b/tests/gen_sim/action_engine/test_architecture.py new file mode 100644 index 000000000..689743318 --- /dev/null +++ b/tests/gen_sim/action_engine/test_architecture.py @@ -0,0 +1,155 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Guard the migration boundaries that make the rewrite meaningful.""" + +from __future__ import annotations + +import ast +from pathlib import Path + +import embodichain.gen_sim.action_engine as action_engine_package +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, + build_default_registry, +) +from embodichain.gen_sim.action_engine.protocol import ( + ACTION_ENGINE_CONFIG_SCHEMA, + ACTION_ENGINE_ENV_ID, + EXECUTION_PROGRAM_FILENAME, + SCENE_REQUIREMENTS_FILENAME, + SCENE_REQUIREMENTS_SCHEMA, + SEED_GRAPH_SCHEMA, + TASK_SPEC_FILENAME, + TASK_SPEC_SCHEMA, +) + +_PACKAGE_ROOT = Path(action_engine_package.__file__).resolve().parent +_LEGACY_PACKAGE = "embodichain.gen_sim.action_agent_pipeline" + + +def _production_python_files() -> list[Path]: + return sorted( + path + for path in _PACKAGE_ROOT.rglob("*.py") + if "tests" not in path.relative_to(_PACKAGE_ROOT).parts + ) + + +def test_production_code_has_no_legacy_pipeline_imports() -> None: + offenders: list[str] = [] + for path in _production_python_files(): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + names = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + names = [node.module or ""] + else: + continue + if any(name.startswith(_LEGACY_PACKAGE) for name in names): + offenders.append(path.relative_to(_PACKAGE_ROOT).as_posix()) + break + assert offenders == [] + + +def test_protocol_identifiers_are_new_and_stable() -> None: + assert ACTION_ENGINE_ENV_ID == "ActionEngine-v1" + assert ACTION_ENGINE_CONFIG_SCHEMA == "action_engine_config_v2" + assert SEED_GRAPH_SCHEMA == "action_engine_seed_graph_v3" + assert TASK_SPEC_SCHEMA == "action_engine_task_spec_v2" + assert SCENE_REQUIREMENTS_SCHEMA == "action_engine_scene_requirements_v2" + assert EXECUTION_PROGRAM_FILENAME == "seed_task_graph.json" + assert TASK_SPEC_FILENAME == "task_spec.json" + assert SCENE_REQUIREMENTS_FILENAME == "scene_requirements.json" + + +def test_planner_exposes_exactly_the_first_phase_skill_catalog() -> None: + assert set(build_default_registry().operator_names()) == { + "arrange_line", + "build_stack", + "coordinated_transport", + "orient_object", + "place_relative", + } + + +def test_atomic_actions_have_one_runtime_capability_catalog() -> None: + registry = build_atomic_capability_registry() + assert set(registry.names()) == { + "AxisAlign", + "CoordinatedPickment", + "CoordinatedPlacement", + "HandOver", + "MoveEndEffector", + "MoveHeldObject", + "MoveJoints", + "PickUp", + "Place", + "Pour", + "Press", + "PullArticulatedPart", + "PushArticulatedPart", + "TurnKnob", + } + assert set(registry.executable_names()) == { + "AxisAlign", + "CoordinatedPickment", + "CoordinatedPlacement", + "HandOver", + "MoveEndEffector", + "MoveHeldObject", + "MoveJoints", + "PickUp", + "Place", + "Pour", + "Press", + "PullArticulatedPart", + "PushArticulatedPart", + "TurnKnob", + } + + +def test_action_class_dispatch_is_not_duplicated_across_runtime_layers() -> None: + offenders = [] + for path in _production_python_files(): + if path.name == "atomic.py" and path.parent.name == "capabilities": + continue + source = path.read_text(encoding="utf-8") + if "_ACTION_TYPES" in source: + offenders.append(path.relative_to(_PACKAGE_ROOT).as_posix()) + assert offenders == [] + + +def test_runtime_core_has_no_action_name_dispatch_branches() -> None: + action_names = set(build_atomic_capability_registry().executable_names()) + offenders = {} + for relative in ( + "runtime/actions.py", + "runtime/executor.py", + "runtime/grounding.py", + ): + path = _PACKAGE_ROOT / relative + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + literals = { + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + duplicated = sorted(literals & action_names) + if duplicated: + offenders[relative] = duplicated + assert offenders == {} diff --git a/tests/gen_sim/action_engine/test_graph_visualization.py b/tests/gen_sim/action_engine/test_graph_visualization.py new file mode 100644 index 000000000..efb1474ac --- /dev/null +++ b/tests/gen_sim/action_engine/test_graph_visualization.py @@ -0,0 +1,362 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from io import BytesIO + +from PIL import Image, ImageStat +import pytest + +from embodichain.gen_sim.action_engine.compiler import ( + compile_task_agent, + compile_task_agent_v2, +) +from embodichain.gen_sim.action_engine.domain import ( + EXECUTION_PROGRAM_SCHEMA, + MOTION_POLICY_VERSION, + TASK_AGENT_SCHEMA, + validate_execution_program, +) +from embodichain.gen_sim.action_engine.graph_visualization import ( + _RuntimeOverlay, + _dag_levels, + _dag_positions, + _dependency_pairs, + _graph_data, + render_seed_task_graph_png, + render_task_graph_png, +) + +_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" + + +def _image(payload: bytes) -> Image.Image: + assert payload.startswith(_PNG_SIGNATURE) + image = Image.open(BytesIO(payload)).convert("RGB") + extrema = ImageStat.Stat(image).extrema + assert any(low != high for low, high in extrema) + return image + + +def _contains_color( + image: Image.Image, + color: str, + *, + minimum_pixels: int = 8, + tolerance: int = 4, +) -> bool: + target = tuple(bytes.fromhex(color.removeprefix("#"))) + matches = 0 + payload = image.tobytes() + for offset in range(0, len(payload), 3): + pixel = payload[offset : offset + 3] + if all( + abs(channel - expected) <= tolerance + for channel, expected in zip(pixel, target) + ): + matches += 1 + if matches >= minimum_pixels: + return True + return False + + +def _chain_program() -> dict[str, object]: + return compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "unicode-λ-task", + "goal": "Pick up the cup and keep it hovering.", + "semantic_steps": [ + { + "id": "s01_hover", + "operator": "hold_hover", + "object": "cup", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ], + } + ) + + +def _action( + action_class: str, + arm: str | None, + target: str, +) -> dict[str, object]: + actor = {"mode": "auto"} if arm is None else {"mode": "required", "arm": arm} + return { + "atomic_action_class": action_class, + "actor": actor, + "control": "arm", + "target_binding": {"kind": "object", "object": target}, + "motion_policy": {"modifiers": []}, + } + + +def _fork_join_program() -> dict[str, object]: + program = { + "schema_version": EXECUTION_PROGRAM_SCHEMA, + "task": "fork_join_demo", + "goal_description": "Move two objects in parallel, then finish.", + "start": "v_start", + "goal": "v_goal", + "nodes": [ + {"id": "v_start", "semantic": "ready"}, + {"id": "v_left", "semantic": "left branch active"}, + {"id": "v_right", "semantic": "right branch active"}, + {"id": "v_join", "semantic": "branches complete"}, + {"id": "v_goal", "semantic": "task complete"}, + ], + "edges": [ + { + "id": "e_left_pick", + "source": "v_start", + "target": "v_left", + "semantic_step_id": "s_left", + "actions": [_action("PickUp", "left_arm", "left_object")], + "depends_on": [], + "resources": ["arm:left_arm"], + }, + { + "id": "e_right_pick", + "source": "v_start", + "target": "v_right", + "semantic_step_id": "s_right", + "actions": [_action("PickUp", "right_arm", "right_object")], + "depends_on": [], + "resources": ["arm:right_arm"], + }, + { + "id": "e_left_join", + "source": "v_left", + "target": "v_join", + "semantic_step_id": "s_left", + "actions": [_action("MoveHeldObject", "left_arm", "left_object")], + "depends_on": ["e_left_pick"], + "resources": ["arm:left_arm"], + }, + { + "id": "e_right_join", + "source": "v_right", + "target": "v_join", + "semantic_step_id": "s_right", + "actions": [_action("MoveHeldObject", "right_arm", "right_object")], + # Cross-branch dependency not implied by state continuity, so + # the renderer must draw a visible dashed dependency arrow. + "depends_on": ["e_right_pick", "e_left_pick"], + "resources": ["arm:right_arm"], + }, + { + "id": "e_finish", + "source": "v_join", + "target": "v_goal", + "semantic_step_id": "s_finish", + "actions": [_action("MoveJoints", None, "home")], + "depends_on": ["e_left_join", "e_right_join"], + "resources": ["arm:auto"], + }, + ], + "semantic_steps": [ + { + "id": "s_left", + "parent_step_id": "s_left", + "operator": "place_relative", + "object": "left_object", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {"relation": "on"}, + "depends_on": [], + "postcondition": {"type": "semantic_goal"}, + "edge_ids": ["e_left_pick", "e_left_join"], + }, + { + "id": "s_right", + "parent_step_id": "s_right", + "operator": "place_relative", + "object": "right_object", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": {"relation": "on"}, + "depends_on": [], + "postcondition": {"type": "semantic_goal"}, + "edge_ids": ["e_right_pick", "e_right_join"], + }, + { + "id": "s_finish", + "parent_step_id": "s_finish", + "operator": "hold_hover", + "object": "home", + "actor": {"mode": "auto"}, + "goal": {}, + "depends_on": ["s_left", "s_right"], + "postcondition": {"type": "semantic_goal"}, + "edge_ids": ["e_finish"], + }, + ], + "allocation_groups": [ + { + "id": "g_parallel", + "semantic_step_ids": ["s_left", "s_right"], + "arm_constraint": "distinct_arms", + "execution_policy": "parallel_if_feasible", + "parallel_action_classes": ["PickUp"], + "workspace_policy": "shared_target_serial", + } + ], + "motion_policy_version": MOTION_POLICY_VERSION, + } + return validate_execution_program(program) + + +def test_seed_renderer_produces_a_compact_headless_png() -> None: + first = _image(render_seed_task_graph_png(_chain_program())) + second = _image(render_seed_task_graph_png(_chain_program())) + + assert first.size == second.size + assert first.width > first.height + assert first.height < 1_200 + + +def test_fork_join_layout_uses_actor_lanes_and_dependency_links() -> None: + program = _fork_join_program() + data = _graph_data(program, _RuntimeOverlay({}, {}, {})) + levels = _dag_levels(data.graph) + positions = _dag_positions( + data, + levels, + {"left": 2.6, "auto": 7.8, "right": 13.0}, + ) + + assert positions["v_left"][0] < 5.15 + assert positions["v_right"][0] > 10.45 + assert positions["v_start"][0] == pytest.approx(7.8) + assert positions["v_join"][0] == pytest.approx(7.8) + assert ("e_left_join", "e_finish") in _dependency_pairs(data) + assert ("e_right_join", "e_finish") in _dependency_pairs(data) + + image = _image(render_seed_task_graph_png(program)) + assert image.width > image.height + assert _contains_color(image, "#168A78") + assert _contains_color(image, "#D97706") + assert _contains_color(image, "#8A94A0") + + +def test_parallel_single_phase_edges_are_rendered_as_a_multigraph() -> None: + program = compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "parallel_press", + "goal": "Press both independent buttons.", + "semantic_steps": [ + { + "id": "s_left", + "operator": "press", + "object": "left_button", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + }, + { + "id": "s_right", + "operator": "press", + "object": "right_button", + "actor": {"mode": "required", "arm": "right_arm"}, + "goal": {}, + "depends_on": [], + }, + ], + } + ) + assert {(edge["source"], edge["target"]) for edge in program["edges"]} == { + ("v0_start", "v_goal") + } + + image = _image(render_seed_task_graph_png(program)) + + assert _contains_color(image, "#168A78") + assert _contains_color(image, "#D97706") + + +def test_runtime_renderer_overlays_observed_statuses() -> None: + program = _fork_join_program() + runtime = { + **program, + "runtime": { + "schema_version": "action_engine_runtime_record_v1", + "status": "failed", + "events": [ + { + "event": "edge", + "edge_id": "e_left_pick", + "arm": "left_arm", + "status": "executed", + }, + { + "event": "edge", + "edge_id": "e_right_pick", + "arm": "right_arm", + "status": "failed", + }, + ], + }, + } + + image = _image(render_task_graph_png(runtime)) + + assert _contains_color(image, "#25834B") + assert _contains_color(image, "#C43E3E") + + +def test_runtime_renderer_accepts_v2_seed_graph_envelope() -> None: + task_agent = { + "schema_version": TASK_AGENT_SCHEMA, + "task": "v2_runtime_overlay", + "goal": "Hold the cup.", + "semantic_steps": [ + { + "id": "hold", + "operator": "hold_hover", + "object": "cup", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ], + } + seed = compile_task_agent_v2(task_agent) + document = { + **seed, + "runtime": { + "schema_version": "action_engine_runtime_record_v2", + "status": "success", + "events": [], + }, + } + + image = _image(render_task_graph_png(document)) + + assert _contains_color(image, "#25834B") + + +def test_runtime_record_without_program_is_rejected() -> None: + with pytest.raises(ValueError, match="do not contain graph topology"): + render_task_graph_png( + { + "schema_version": "action_engine_runtime_record_v1", + "events": [], + } + ) diff --git a/tests/gen_sim/action_engine/test_motion_policy.py b/tests/gen_sim/action_engine/test_motion_policy.py new file mode 100644 index 000000000..aa253c6fb --- /dev/null +++ b/tests/gen_sim/action_engine/test_motion_policy.py @@ -0,0 +1,72 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.domain import motion_policy +from embodichain.gen_sim.action_engine.runtime.motion_policy import ( + resolve_motion_policy, +) + + +def test_upright_policy_matches_mature_runtime_across_robot_profiles() -> None: + upright = motion_policy(("orientation", "upright")) + franka = resolve_motion_policy("dual_franka", "PickUp", upright) + ur10 = resolve_motion_policy("dual_ur10", "PickUp", upright) + + assert franka["lift_height"] == pytest.approx(0.30) + assert ur10["lift_height"] == pytest.approx(0.30) + assert ur10["rotate_upright"] == pytest.approx(0.7853981633974483) + + +def test_policy_resolution_returns_detached_values() -> None: + upright = motion_policy(("orientation", "upright")) + first = resolve_motion_policy("ur10", "MoveHeldObject", upright) + first["surface_clearance"] = 123.0 + + second = resolve_motion_policy("dual_ur10", "MoveHeldObject", upright) + + assert second["surface_clearance"] == pytest.approx(0.05) + + +def test_unknown_action_base_is_rejected_instead_of_falling_back() -> None: + with pytest.raises(ValueError, match="Unknown Action Engine motion base"): + resolve_motion_policy("dual_franka", "TypoAction", motion_policy()) + + +def test_upright_and_handover_role_modifiers_compose_without_named_cross_product() -> ( + None +): + resolved = resolve_motion_policy( + "dual_franka", + "PickUp", + motion_policy( + ("orientation", "upright"), + ("handover_role", "transfer"), + ), + ) + + assert resolved["rotate_upright"] == pytest.approx(0.7853981633974483) + assert resolved["pick_object_part"] == "top" + assert "approach_direction_mode" not in resolved + assert resolved["sample_interval"] == 80 + + +def test_named_policy_strings_require_graph_regeneration() -> None: + with pytest.raises(ValueError, match="named string policies are no longer"): + resolve_motion_policy("dual_franka", "PickUp", "legacy_flat_name") diff --git a/tests/gen_sim/action_engine/test_orientation.py b/tests/gen_sim/action_engine/test_orientation.py new file mode 100644 index 000000000..57aa7ccd0 --- /dev/null +++ b/tests/gen_sim/action_engine/test_orientation.py @@ -0,0 +1,145 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest + +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + compile_orientation_constraint, +) +from embodichain.gen_sim.action_engine.protocol import TASK_SPEC_SCHEMA +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + + +def test_unspecified_orientation_has_no_hard_constraint() -> None: + constraint = compile_orientation_constraint({}) + + assert constraint.terms == () + assert constraint.planning_preference == "minimize_rotation_from_current" + assert not constraint.requires_reference + + +def test_explicit_preserve_compiles_to_rotation_match() -> None: + constraint = compile_orientation_constraint({"orientation_goal": "preserve"}) + + assert constraint.terms == (MatchRotationConstraint(reference="step_start"),) + assert constraint.requires_reference + + +def test_upright_compiles_to_directed_axis_when_requested() -> None: + constraint = compile_orientation_constraint( + { + "orientation_goal": "upright", + "upright_local_axis": "z", + "orientation_directed": True, + } + ) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="z", + target_axis="world_up", + directed=True, + ), + ) + + +def test_upright_rejects_non_boolean_directed_flag() -> None: + with pytest.raises(ValueError, match="orientation_directed must be a boolean"): + compile_orientation_constraint( + { + "orientation_goal": "upright", + "orientation_directed": "false", + } + ) + + +def test_legacy_long_axis_upright_remains_undirected() -> None: + constraint = compile_orientation_constraint( + { + "orientation_goal": "upright", + "upright_local_axis": "long_axis", + } + ) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="long_axis", + target_axis="world_up", + directed=False, + ), + ) + + +def test_serialized_constraint_keeps_term_local_tolerance() -> None: + constraint = compile_orientation_constraint( + { + "orientation_constraint": { + "terms": [ + { + "type": "align_axis", + "local_axis": "z", + "target_axis": "world_up", + "directed": True, + "tolerance": 0.1, + "scope": "terminal", + } + ] + } + } + ) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="z", + target_axis="world_up", + directed=True, + tolerance=0.1, + ), + ) + + +def test_new_placement_without_orientation_request_has_no_hard_constraint() -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "place_can", + "level": "L1", + "instruction": "Place the can beside the box.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "place", + "task_type": "E1", + "params": { + "object_role": "can", + "target_role": "box", + "relation": "left_of", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "semantic_goal"}, + "oracle": {}, + "metadata": {}, + } + + graph = instantiate_seed_graph(task, {"can": "can", "box": "box"}) + + assert graph["task_groups"][0]["goal"]["orientation_goal"] == "none" diff --git a/tests/gen_sim/action_engine/test_unbound.py b/tests/gen_sim/action_engine/test_unbound.py new file mode 100644 index 000000000..5aa3980dc --- /dev/null +++ b/tests/gen_sim/action_engine/test_unbound.py @@ -0,0 +1,111 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy + +import pytest + +from embodichain.gen_sim.action_engine.agent import ActionAgent +import embodichain.gen_sim.action_engine.agent as action_agent_module +from embodichain.gen_sim.action_engine.unbound import validate_unbound_action_plan + + +def _selector(reference: str) -> dict: + return { + "kind": "scene_ref", + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _candidate() -> dict: + return { + "candidate_id": "candidate_01", + "draft": { + "task_id": "place_can", + "instruction": "Place the can on the table.", + "steps": [ + { + "id": "place", + "task_type": "E1", + "object": _selector("the can"), + "target": _selector("the table"), + "depends_on": [], + } + ], + }, + } + + +def test_action_agent_drafts_without_scene_uids() -> None: + candidate = _candidate() + original = deepcopy(candidate) + + draft = ActionAgent(registry=object()).draft(candidate) + + assert draft["candidate_id"] == "candidate_01" + assert draft["steps"][0]["object"]["reference"] == "the can" + assert "uid" not in str(draft).lower() + assert candidate == original + + +def test_unbound_plan_rejects_noncanonical_action_recipe() -> None: + draft = ActionAgent(registry=object()).draft(_candidate()) + draft["steps"][0]["actions"] = ["UnknownAction"] + + with pytest.raises(ValueError, match="task contract"): + validate_unbound_action_plan(draft) + + +def test_action_agent_rejects_missing_atomic_action_during_draft() -> None: + class Registry: + def names(self): + return () + + def executable_names(self): + return () + + with pytest.raises(ValueError, match="AtomicAction is not registered"): + ActionAgent(registry=Registry()).draft(_candidate()) + + +def test_bind_and_plan_requires_the_exact_unbound_candidate( + monkeypatch: pytest.MonkeyPatch, +) -> None: + agent = ActionAgent(registry=object()) + unbound = agent.draft(_candidate()) + grounded = { + "selected_candidate_id": "candidate_01", + "task_draft": deepcopy(_candidate()["draft"]), + } + monkeypatch.setattr( + action_agent_module, + "_validate_grounded_plan", + lambda value: deepcopy(value), + ) + monkeypatch.setattr(agent, "plan", lambda value: {"task": value["task_draft"]}) + + graph = agent.bind_and_plan(unbound, grounded) + assert graph["task"] == grounded["task_draft"] + + altered = deepcopy(unbound) + altered["instruction"] = "A different instruction." + with pytest.raises(ValueError, match="does not match"): + agent.bind_and_plan(altered, grounded) From 9f28ef0d612d59e4b213fd02c4f15515a0321557 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Mon, 24 Aug 2026 16:01:40 +0800 Subject: [PATCH 64/85] feat(gen-sim): migrate task orchestration onto scene baseline --- embodichain/gen_sim/task_engine/__init__.py | 179 +++ embodichain/gen_sim/task_engine/__main__.py | 27 + .../gen_sim/task_engine/_bundle_runner.py | 197 +++ embodichain/gen_sim/task_engine/agent.py | 289 ++++ embodichain/gen_sim/task_engine/cli.py | 246 ++++ embodichain/gen_sim/task_engine/config.py | 190 +++ embodichain/gen_sim/task_engine/contracts.py | 410 ++++++ embodichain/gen_sim/task_engine/defaults.yaml | 33 + .../gen_sim/task_engine/interpretation.py | 1269 ++++++++++++++++ embodichain/gen_sim/task_engine/ontology.py | 221 +++ .../task_engine/orchestration/__init__.py | 109 ++ .../task_engine/orchestration/artifacts.py | 331 +++++ .../task_engine/orchestration/contracts.py | 647 +++++++++ .../task_engine/orchestration/coordinator.py | 865 +++++++++++ .../task_engine/orchestration/legacy_scene.py | 381 +++++ .../orchestration/scene_adapter.py | 1050 ++++++++++++++ .../task_engine/orchestration/scene_source.py | 281 ++++ .../gen_sim/task_engine/run_directory.py | 87 ++ .../gen_sim/task_engine/scene/__init__.py | 53 + .../task_engine/scene/conservative_graph.py | 247 ++++ .../gen_sim/task_engine/scene/contracts.py | 304 ++++ .../gen_sim/task_engine/scene/feasibility.py | 746 ++++++++++ .../task_engine/scene/final_inspection.py | 428 ++++++ .../task_engine/scene/scene_engine_v1.py | 242 ++++ .../gen_sim/task_engine/scene_backend.py | 474 ++++++ .../gen_sim/task_engine/state_machine.py | 294 ++++ embodichain/gen_sim/task_engine/workflow.py | 1275 +++++++++++++++++ .../gen_sim/task_engine/workflow_contracts.py | 179 +++ setup.py | 2 + tests/gen_sim/task_engine/__init__.py | 19 + .../task_engine/orchestration/__init__.py | 19 + .../orchestration/test_architecture.py | 90 ++ .../orchestration/test_coordinator_cli.py | 1102 ++++++++++++++ .../orchestration/test_legacy_scene.py | 164 +++ .../orchestration/test_scene_adapter.py | 780 ++++++++++ tests/gen_sim/task_engine/scene/__init__.py | 19 + .../scene/test_final_inspection.py | 118 ++ .../task_engine/scene/test_scene_boundary.py | 582 ++++++++ tests/gen_sim/task_engine/test_agent.py | 271 ++++ .../task_engine/test_interpretation.py | 96 ++ .../task_engine/test_parallel_workflow.py | 912 ++++++++++++ .../gen_sim/task_engine/test_run_directory.py | 58 + .../gen_sim/task_engine/test_scene_backend.py | 214 +++ tests/gen_sim/task_engine/test_workflow.py | 337 +++++ 44 files changed, 15837 insertions(+) create mode 100644 embodichain/gen_sim/task_engine/__init__.py create mode 100644 embodichain/gen_sim/task_engine/__main__.py create mode 100644 embodichain/gen_sim/task_engine/_bundle_runner.py create mode 100644 embodichain/gen_sim/task_engine/agent.py create mode 100644 embodichain/gen_sim/task_engine/cli.py create mode 100644 embodichain/gen_sim/task_engine/config.py create mode 100644 embodichain/gen_sim/task_engine/contracts.py create mode 100644 embodichain/gen_sim/task_engine/defaults.yaml create mode 100644 embodichain/gen_sim/task_engine/interpretation.py create mode 100644 embodichain/gen_sim/task_engine/ontology.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/__init__.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/artifacts.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/contracts.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/coordinator.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/legacy_scene.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/scene_adapter.py create mode 100644 embodichain/gen_sim/task_engine/orchestration/scene_source.py create mode 100644 embodichain/gen_sim/task_engine/run_directory.py create mode 100644 embodichain/gen_sim/task_engine/scene/__init__.py create mode 100644 embodichain/gen_sim/task_engine/scene/conservative_graph.py create mode 100644 embodichain/gen_sim/task_engine/scene/contracts.py create mode 100644 embodichain/gen_sim/task_engine/scene/feasibility.py create mode 100644 embodichain/gen_sim/task_engine/scene/final_inspection.py create mode 100644 embodichain/gen_sim/task_engine/scene/scene_engine_v1.py create mode 100644 embodichain/gen_sim/task_engine/scene_backend.py create mode 100644 embodichain/gen_sim/task_engine/state_machine.py create mode 100644 embodichain/gen_sim/task_engine/workflow.py create mode 100644 embodichain/gen_sim/task_engine/workflow_contracts.py create mode 100644 tests/gen_sim/task_engine/__init__.py create mode 100644 tests/gen_sim/task_engine/orchestration/__init__.py create mode 100644 tests/gen_sim/task_engine/orchestration/test_architecture.py create mode 100644 tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py create mode 100644 tests/gen_sim/task_engine/orchestration/test_legacy_scene.py create mode 100644 tests/gen_sim/task_engine/orchestration/test_scene_adapter.py create mode 100644 tests/gen_sim/task_engine/scene/__init__.py create mode 100644 tests/gen_sim/task_engine/scene/test_final_inspection.py create mode 100644 tests/gen_sim/task_engine/scene/test_scene_boundary.py create mode 100644 tests/gen_sim/task_engine/test_agent.py create mode 100644 tests/gen_sim/task_engine/test_interpretation.py create mode 100644 tests/gen_sim/task_engine/test_parallel_workflow.py create mode 100644 tests/gen_sim/task_engine/test_run_directory.py create mode 100644 tests/gen_sim/task_engine/test_scene_backend.py create mode 100644 tests/gen_sim/task_engine/test_workflow.py diff --git a/embodichain/gen_sim/task_engine/__init__.py b/embodichain/gen_sim/task_engine/__init__.py new file mode 100644 index 000000000..087983cc4 --- /dev/null +++ b/embodichain/gen_sim/task_engine/__init__.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scene-independent task interpretation and protocol ownership.""" + +from __future__ import annotations + +from typing import Any + +from .agent import ( + TaskAgent, + TaskGenerationError, + derive_scene_request, + derive_success_spec, +) +from .contracts import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + SceneRequest, + SuccessSpec, + TaskCandidate, + TaskCandidateSet, + TaskDraft, + canonical_hash, + validate_scene_request, + validate_success_spec, + validate_task_candidate, + validate_task_candidate_set, + validate_task_draft, +) +from .interpretation import ( + INSTRUCTION_INTENT_SCHEMA, + InstructionCaller, + InstructionDraftResult, + InstructionIntent, + interpret_instruction_draft, + validate_instruction_intent, +) +from .ontology import ( + RELATIONS, + TASK_CONTRACTS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, + TaskContract, + task_contract, + task_success_type, +) +from .config import ( + TASK_ENGINE_DEFAULTS_SCHEMA, + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, + load_task_engine_config, +) +from .state_machine import ( + StageStatus, + TaskEngineState, + WorkflowStage, + complete_stage, + fail_stage, + initial_state, + replay_events, + skip_stage, + start_stage, +) +from .workflow_contracts import ( + TASK_RUN_REQUEST_SCHEMA, + SceneInputKind, + TaskRunRequest, + scene_input_kind, + validate_scene_history_root, + validate_scene_output_separation, + validate_task_run_request, +) + +__all__ = [ + "INSTRUCTION_INTENT_SCHEMA", + "InstructionCaller", + "InstructionDraftResult", + "InstructionIntent", + "RELATIONS", + "SCENE_REQUEST_SCHEMA", + "SUCCESS_SPEC_SCHEMA", + "SceneRequest", + "SuccessSpec", + "TASK_CANDIDATE_SET_SCHEMA", + "TASK_CONTRACTS", + "TASK_DRAFT_SCHEMA", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "TaskAgent", + "TaskCandidate", + "TaskCandidateSet", + "TaskContract", + "TaskDraft", + "TaskGenerationError", + "TASK_RUN_REQUEST_SCHEMA", + "TASK_ENGINE_DEFAULTS_SCHEMA", + "SceneInputKind", + "SceneAnalysis", + "SceneEngineBackend", + "SceneRevision", + "StageStatus", + "TaskEngineState", + "TaskEngineExecutionCfg", + "TaskEnginePlanningCfg", + "TaskEngineWorkflowCfg", + "TaskEngineRunResult", + "TaskEngineWorkflow", + "TaskRunRequest", + "WorkflowStage", + "TASK_ENGINE_RUN_MANIFEST_SCHEMA", + "SubprocessActionExecutor", + "canonical_hash", + "derive_scene_request", + "derive_success_spec", + "complete_stage", + "fail_stage", + "initial_state", + "interpret_instruction_draft", + "load_task_engine_config", + "replay_events", + "task_contract", + "task_success_type", + "scene_input_kind", + "validate_scene_history_root", + "scene_blueprint_objects", + "skip_stage", + "start_stage", + "validate_instruction_intent", + "validate_scene_request", + "validate_scene_output_separation", + "validate_success_spec", + "validate_task_candidate", + "validate_task_candidate_set", + "validate_task_draft", + "validate_task_run_request", +] + +_SCENE_BACKEND_EXPORTS = { + "SceneAnalysis", + "SceneEngineBackend", + "SceneRevision", + "scene_blueprint_objects", +} +_WORKFLOW_EXPORTS = { + "TASK_ENGINE_RUN_MANIFEST_SCHEMA", + "SubprocessActionExecutor", + "TaskEngineRunResult", + "TaskEngineWorkflow", +} + + +def __getattr__(name: str) -> Any: + """Load orchestration entry points lazily to avoid engine import cycles.""" + if name in _SCENE_BACKEND_EXPORTS: + from . import scene_backend + + return getattr(scene_backend, name) + if name in _WORKFLOW_EXPORTS: + from . import workflow + + return getattr(workflow, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/embodichain/gen_sim/task_engine/__main__.py b/embodichain/gen_sim/task_engine/__main__.py new file mode 100644 index 000000000..9e4f06dbe --- /dev/null +++ b/embodichain/gen_sim/task_engine/__main__.py @@ -0,0 +1,27 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Module entry point for Task Engine workflows.""" + +from __future__ import annotations + +from .cli import main + +__all__ = ["main"] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/task_engine/_bundle_runner.py b/embodichain/gen_sim/task_engine/_bundle_runner.py new file mode 100644 index 000000000..5bc80bd9b --- /dev/null +++ b/embodichain/gen_sim/task_engine/_bundle_runner.py @@ -0,0 +1,197 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Private subprocess boundary for executing one prepared Task Engine bundle.""" + +from __future__ import annotations + +import argparse +from contextlib import contextmanager +import json +from pathlib import Path +import sys +from typing import Any, Iterator, Sequence + +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.protocol import ( + AGENT_CONFIG_FILENAME, + EXECUTION_PROGRAM_FILENAME, + FAST_GYM_CONFIG_FILENAME, +) +from embodichain.gen_sim.action_engine.runtime import ExecutionReport + +from .orchestration.artifacts import ( + GROUNDED_TASK_PLAN_FILENAME, + STATIC_SCENE_MANIFEST_FILENAME, + write_execution_report, +) +from .orchestration.contracts import validate_grounded_task_plan +from .orchestration.scene_source import verify_scene_source_fingerprint + +__all__ = ["execute_bundle", "main"] + + +def main(argv: Sequence[str] | None = None) -> int: + """Parse the private runner protocol and execute one bundle.""" + parser = argparse.ArgumentParser( + prog="embodichain.gen_sim.task_engine._bundle_runner" + ) + parser.add_argument("--bundle", required=True) + args, forwarded = parser.parse_known_args(argv) + return execute_bundle(args.bundle, forwarded) + + +def execute_bundle( + bundle: str | Path, + forwarded: Sequence[str] = (), +) -> int: + """Execute one prepared bundle through the existing Action Engine launcher.""" + root = Path(bundle).expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(f"Bundle directory does not exist: {root}") + agent_config = root / AGENT_CONFIG_FILENAME + gym_config = root / FAST_GYM_CONFIG_FILENAME + for path in (agent_config, gym_config): + if not path.is_file(): + raise FileNotFoundError(f"Bundle is missing required artifact: {path}") + task_id = _bundle_task_id(root, agent_config) + run_args = list(forwarded) + if run_args and run_args[0] == "--": + run_args.pop(0) + rejection = _preflight_bundle( + root, + agent_config=agent_config, + gym_config=gym_config, + forwarded=run_args, + ) + if rejection is not None: + write_execution_report(root, rejection) + _print_json(rejection.as_mapping()) + return 2 + legacy_argv = [ + "--task_name", + task_id, + "--gym_config", + str(gym_config), + "--agent_config", + str(agent_config), + "--task-engine-report", + *run_args, + ] + from embodichain.gen_sim.action_engine.cli import run_agent + + with _temporary_argv(["run_agent", *legacy_argv]): + return int(run_agent.cli() or 0) + + +def _preflight_bundle( + bundle: Path, + *, + agent_config: Path, + gym_config: Path, + forwarded: Sequence[str], +) -> ExecutionReport | None: + static_manifest_path = bundle / STATIC_SCENE_MANIFEST_FILENAME + if static_manifest_path.is_file(): + static_manifest = _read_json(static_manifest_path) + source = static_manifest.get("source", {}) + if isinstance(source, dict) and isinstance( + source.get("source_fingerprint"), dict + ): + verify_scene_source_fingerprint(source["source_fingerprint"]) + grounded_path = bundle / GROUNDED_TASK_PLAN_FILENAME + if not grounded_path.is_file(): + return None + grounded = validate_grounded_task_plan(_read_json(grounded_path)) + agent = _read_json(agent_config) + graph_value = agent.get("seed_task_graph", EXECUTION_PROGRAM_FILENAME) + if not isinstance(graph_value, str) or not graph_value: + raise ValueError("Bundle agent_config.seed_task_graph must be a path string.") + graph_path = Path(graph_value).expanduser() + if not graph_path.is_absolute(): + graph_path = (bundle / graph_path).resolve() + else: + graph_path = graph_path.resolve() + if graph_path != bundle and bundle not in graph_path.parents: + raise ValueError("Bundle SeedGraph path escapes the bundle directory.") + if not graph_path.is_file(): + raise FileNotFoundError(f"Bundle is missing SeedGraph: {graph_path}") + action_agent = ActionAgent() + try: + action_agent.preflight( + graph_path, + scene_manifest=grounded["scene_manifest"], + ) + except (TypeError, ValueError, OSError) as exc: + return action_agent.rejection_report( + graph_path, + exc, + grounded_plan=grounded, + environment_count=_environment_count(gym_config, forwarded), + ) + return None + + +def _environment_count(gym_config: Path, forwarded: Sequence[str]) -> int: + value: Any = _read_json(gym_config).get("num_envs", 1) + for index, argument in enumerate(forwarded): + if argument == "--num_envs" and index + 1 < len(forwarded): + value = forwarded[index + 1] + elif argument.startswith("--num_envs="): + value = argument.partition("=")[2] + try: + return max(1, int(value)) + except (TypeError, ValueError): + return 1 + + +def _bundle_task_id(bundle: Path, agent_config: Path) -> str: + grounded_path = bundle / GROUNDED_TASK_PLAN_FILENAME + if grounded_path.is_file(): + task_id = _read_json(grounded_path).get("task_id") + else: + task_id = _read_json(agent_config).get("task_name") + if not isinstance(task_id, str) or not task_id.strip(): + raise ValueError("Bundle does not declare a non-empty task ID.") + return task_id.strip() + + +def _read_json(path: Path) -> dict[str, Any]: + try: + value = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Unable to read JSON artifact {path}: {exc}") from exc + if not isinstance(value, dict): + raise ValueError(f"JSON artifact must contain an object: {path}") + return value + + +@contextmanager +def _temporary_argv(arguments: list[str]) -> Iterator[None]: + original = sys.argv + sys.argv = arguments + try: + yield + finally: + sys.argv = original + + +def _print_json(value: Any) -> None: + print(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/task_engine/agent.py b/embodichain/gen_sim/task_engine/agent.py new file mode 100644 index 000000000..f5ab773f0 --- /dev/null +++ b/embodichain/gen_sim/task_engine/agent.py @@ -0,0 +1,289 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scene-independent semantic candidate generation for Task Engine.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor, as_completed +from copy import deepcopy +from dataclasses import dataclass +from typing import Any + +from .contracts import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + TaskCandidate, + TaskCandidateSet, + canonical_hash, + validate_task_candidate, + validate_task_candidate_set, +) +from .interpretation import ( + InstructionCaller, + InstructionDraftResult, + interpret_instruction_draft, + validate_instruction_intent, +) +from .ontology import TASK_CONTRACTS, task_success_type + +__all__ = [ + "TaskAgent", + "TaskGenerationError", + "derive_scene_request", + "derive_success_spec", +] + +DraftInterpreter = Callable[..., InstructionDraftResult] + + +class TaskGenerationError(ValueError): + """Raised when every independently generated candidate fails validation.""" + + +@dataclass(frozen=True) +class _CandidateAttempt: + index: int + result: InstructionDraftResult | None = None + error: str = "" + + +class TaskAgent: + """Generate, validate, normalize, and vote on independent task drafts.""" + + def __init__( + self, + *, + caller: InstructionCaller | None = None, + interpreter: DraftInterpreter = interpret_instruction_draft, + ) -> None: + self._caller = caller + self._interpreter = interpreter + + def generate( + self, + task_id: str, + instruction: str, + model: str | None = None, + candidate_count: int = 3, + ) -> TaskCandidateSet: + """Generate candidates concurrently and retain votes after deduplication.""" + normalized_task_id = str(task_id).strip() + normalized_instruction = str(instruction).strip() + if not normalized_task_id or not normalized_instruction: + raise ValueError("task_id and instruction must be non-empty.") + if ( + isinstance(candidate_count, bool) + or not isinstance(candidate_count, int) + or candidate_count < 1 + ): + raise ValueError("candidate_count must be a positive integer.") + + attempts: list[_CandidateAttempt] = [] + with ThreadPoolExecutor( + max_workers=candidate_count, + thread_name_prefix="task-agent", + ) as executor: + futures = { + executor.submit( + self._interpreter, + normalized_instruction, + model=model, + caller=self._caller, + ): index + for index in range(candidate_count) + } + for future in as_completed(futures): + index = futures[future] + try: + attempts.append( + _CandidateAttempt(index=index, result=future.result()) + ) + except Exception as error: # Each candidate is an isolated vote. + attempts.append( + _CandidateAttempt( + index=index, + error=f"candidate_{index + 1:02d}: {type(error).__name__}: {error}", + ) + ) + attempts.sort(key=lambda item: item.index) + errors = [item.error for item in attempts if item.result is None] + unique: dict[str, TaskCandidate] = {} + valid_response_count = 0 + for attempt in attempts: + if attempt.result is None: + continue + assert attempt.result is not None + try: + canonical_intent = _canonicalize_intent(attempt.result.intent) + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": normalized_task_id, + "instruction": normalized_instruction, + "steps": canonical_intent["steps"], + } + semantic_hash = canonical_hash(draft["steps"]) + candidate_id = f"candidate_{len(unique) + 1:02d}" + candidate = validate_task_candidate( + { + "candidate_id": candidate_id, + "draft": draft, + "scene_request": derive_scene_request(draft), + "success_spec": derive_success_spec(draft), + "semantic_hash": semantic_hash, + "vote_count": 1, + "attempts": attempt.result.attempts, + "normalizations": deepcopy(list(attempt.result.normalizations)), + } + ) + existing = unique.get(semantic_hash) + if existing is not None: + existing["vote_count"] += 1 + existing["attempts"] = max( + existing["attempts"], attempt.result.attempts + ) + existing["normalizations"].extend(candidate["normalizations"]) + else: + unique[semantic_hash] = candidate + valid_response_count += 1 + except Exception as error: # Post-processing failures stay candidate-local. + errors.append( + f"candidate_{attempt.index + 1:02d}: " + f"{type(error).__name__}: {error}" + ) + + if not unique: + raise TaskGenerationError( + "All Task Agent candidates failed validation: " + "; ".join(errors) + ) + + return validate_task_candidate_set( + { + "schema_version": TASK_CANDIDATE_SET_SCHEMA, + "task_id": normalized_task_id, + "instruction": normalized_instruction, + "candidates": list(unique.values()), + "requested_candidate_count": candidate_count, + "valid_response_count": valid_response_count, + "errors": errors, + } + ) + + +def derive_scene_request(draft: Mapping[str, Any]) -> dict[str, Any]: + """Derive structural scene constraints without classifying reference text.""" + from .contracts import validate_scene_request, validate_task_draft + + normalized = validate_task_draft(draft) + references: list[dict[str, Any]] = [] + for step in normalized["steps"]: + task_type = str(step["task_type"]) + contract = TASK_CONTRACTS[task_type] + for role in ("object", "target"): + selector = step[role] + if selector["kind"] != "scene_ref": + continue + if role == "object": + structure = contract.source_structure + affordances = sorted(contract.scene_affordances) + initial_state = {"orientation": "fallen"} if task_type == "E2" else {} + attributes: dict[str, Any] = {} + else: + structure = _target_structure(task_type, str(step["relation"])) + affordances = _target_affordances(task_type, str(step["relation"])) + initial_state = {} + attributes = {} + references.append( + { + "reference_id": f"{step['id']}.{role}", + "step_id": step["id"], + "role": role, + "reference": selector["reference"], + "quantifier": selector["quantifier"], + "count": selector["count"], + "source_structure": structure, + "affordances": affordances, + "initial_state": initial_state, + "attributes": attributes, + } + ) + if not references: + raise ValueError("A TaskDraft must contain at least one scene_ref selector.") + return validate_scene_request( + { + "schema_version": SCENE_REQUEST_SCHEMA, + "task_id": normalized["task_id"], + "references": references, + } + ) + + +def derive_success_spec(draft: Mapping[str, Any]) -> dict[str, Any]: + """Derive every success term exclusively from the E-task ontology.""" + from .contracts import validate_success_spec, validate_task_draft + + normalized = validate_task_draft(draft) + return validate_success_spec( + { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": normalized["task_id"], + "op": "all", + "terms": [ + { + "step_id": step["id"], + "type": task_success_type(step["task_type"], step), + } + for step in normalized["steps"] + ], + }, + draft=normalized, + ) + + +def _canonicalize_intent(intent: Mapping[str, Any]) -> dict[str, Any]: + """Remove arbitrary model step IDs while preserving the explicit DAG order.""" + normalized = validate_instruction_intent(intent) + id_map = { + step["id"]: f"step_{index + 1:02d}" + for index, step in enumerate(normalized["steps"]) + } + steps = deepcopy(normalized["steps"]) + for step in steps: + old_id = step["id"] + step["id"] = id_map[old_id] + step["depends_on"] = [id_map[item] for item in step["depends_on"]] + for selector_name in ("object", "target"): + selector = step[selector_name] + if selector["kind"] == "step_result": + selector["step_id"] = id_map[selector["step_id"]] + return validate_instruction_intent({"steps": steps}) + + +def _target_affordances(task_type: str, relation: str) -> list[str]: + if task_type == "E3" or (task_type == "E1" and relation == "inside"): + return ["container"] + return [] + + +def _target_structure(task_type: str, relation: str) -> str: + if task_type == "E1" and relation == "on": + return "physical_entity" + if task_type == "E3" or (task_type == "E1" and relation == "inside"): + return "rigid_object" + return "spatial_reference" diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py new file mode 100644 index 000000000..fec6de86a --- /dev/null +++ b/embodichain/gen_sim/task_engine/cli.py @@ -0,0 +1,246 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Unified CLI for complete Task Engine workflows.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +import sys +from typing import Any, Final, Sequence + +from .config import load_task_engine_config +from .orchestration.scene_adapter import SceneAdapter +from .run_directory import reserve_run_directory +from .workflow import SubprocessActionExecutor, TaskEngineWorkflow +from .workflow_contracts import ( + TASK_RUN_REQUEST_SCHEMA, + validate_scene_history_root, + validate_scene_output_separation, +) + +__all__ = ["build_parser", "main"] + + +_ROBOT_PROFILES = ( + "ur5", + "ur10", + "dual_ur5", + "dual_ur10", + "franka", + "dual_franka", +) +_MODES: Final = ("image", "image-edit", "scene", "scene-edit") + + +def build_parser() -> argparse.ArgumentParser: + """Build the Task Engine parser.""" + parser = argparse.ArgumentParser( + prog="embodichain task-engine", + description="Prepare, run, or complete one Scene and Action workflow.", + ) + subparsers = parser.add_subparsers(dest="command", required=True) + prepare_parser = subparsers.add_parser( + "prepare", help="Prepare a bundle without simulator execution." + ) + _add_workflow_arguments(prepare_parser) + run_all_parser = subparsers.add_parser( + "run-all", help="Prepare and execute one complete workflow." + ) + _add_workflow_arguments(run_all_parser) + run_parser = subparsers.add_parser( + "run", help="Execute an already prepared Task Engine bundle." + ) + run_parser.add_argument("--bundle", required=True) + run_parser.add_argument("--output-root", required=True) + run_parser.add_argument("--config", default=None) + run_parser.add_argument("--seed", type=int, default=0) + run_parser.add_argument("--num-envs", type=int, default=None) + run_parser.add_argument("--dataset-saving", action="store_true") + return parser + + +def _add_workflow_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--mode", choices=_MODES, required=True) + parser.add_argument("--task-id", "--task_id", required=True) + instruction = parser.add_mutually_exclusive_group(required=True) + instruction.add_argument("--instruction") + instruction.add_argument("--task-file", "--task_file") + parser.add_argument("--image") + parser.add_argument("--scene") + parser.add_argument("--scene-edit", "--scene_edit", default=None) + parser.add_argument("--output-root", required=True) + parser.add_argument("--config", default=None) + parser.add_argument("--model", default=None) + parser.add_argument("--vlm-model", default=None) + parser.add_argument("--base-seed", type=int, default=0) + parser.add_argument( + "--dataset_saving", + action="store_true", + help="Opt in to the Gym project's dataset recorder during execution.", + ) + parser.add_argument( + "--robot-profile", + choices=_ROBOT_PROFILES, + default="franka", + ) + + +def main(argv: Sequence[str] | None = None) -> int: + """Dispatch one Task Engine workflow command.""" + parser = build_parser() + arguments = list(sys.argv[1:] if argv is None else argv) + if arguments and arguments[0] not in { + "prepare", + "run", + "run-all", + "-h", + "--help", + }: + arguments.insert(0, "run-all") + args = parser.parse_args(arguments) + if args.command == "run": + return _run_prepared_bundle(args) + return _run_workflow( + args, + execute=args.command == "run-all", + parser=parser, + ) + + +def _run_workflow( + args: argparse.Namespace, + *, + execute: bool, + parser: argparse.ArgumentParser, +) -> int: + try: + image, scene, edit = _mode_inputs(args) + except ValueError as exc: + parser.error(str(exc)) + if scene is not None: + validate_scene_history_root(scene, args.output_root) + instruction = _instruction(args) + adapter = SceneAdapter(model=args.model, robot_profile=args.robot_profile) + workflow = TaskEngineWorkflow(scene_adapter=adapter) + with reserve_run_directory(args.output_root) as allocation: + if scene is not None: + validate_scene_output_separation(scene, allocation.path) + result = workflow.run( + { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": args.task_id, + "task_instruction": instruction, + "image_path": image, + "gym_project": scene, + "scene_edit_prompt": edit, + "output_dir": allocation.path.as_posix(), + }, + config_path=args.config, + model=args.model, + vlm_model=args.vlm_model, + base_seed=args.base_seed, + dataset_saving=args.dataset_saving, + run_id=allocation.run_id, + created_at=allocation.created_at, + execute=execute, + ) + _print_json( + { + "run_id": allocation.run_id, + "status": result.status, + "failure_class": result.failure_class, + "output_dir": result.output_dir.as_posix(), + "manifest": result.manifest_path.as_posix(), + "final_bundle": ( + None if result.final_bundle is None else result.final_bundle.as_posix() + ), + } + ) + accepted = result.succeeded if execute else result.status == "prepared" + return 0 if accepted else 2 + + +def _run_prepared_bundle(args: argparse.Namespace) -> int: + _, _, execution_cfg = load_task_engine_config(args.config) + num_envs = execution_cfg.num_envs if args.num_envs is None else int(args.num_envs) + if num_envs < 1: + raise ValueError("num_envs must be positive.") + with reserve_run_directory(args.output_root) as allocation: + report = SubprocessActionExecutor()( + args.bundle, + allocation.path, + seed=int(args.seed), + num_envs=num_envs, + dataset_saving=bool(args.dataset_saving), + ) + environments = report.get("environments", ()) + successes = [ + bool(item.get("success")) for item in environments if isinstance(item, dict) + ] + accepted = ( + str(report.get("status")) not in {"rejected", "aborted"} + and len(successes) == num_envs + and sum(successes) >= execution_cfg.required_successes + ) + _print_json( + { + "run_id": allocation.run_id, + "status": "succeeded" if accepted else "failed", + "output_dir": allocation.path.as_posix(), + "execution_report": report, + } + ) + return 0 if accepted else 2 + + +def _instruction(args: argparse.Namespace) -> str: + instruction = ( + str(args.instruction).strip() + if args.instruction is not None + else Path(args.task_file).expanduser().read_text(encoding="utf-8").strip() + ) + if not instruction: + raise ValueError("Task instruction must not be empty.") + return instruction + + +def _mode_inputs(args: argparse.Namespace) -> tuple[str | None, str | None, str | None]: + image = None if args.image is None else str(args.image).strip() + scene = None if args.scene is None else str(args.scene).strip() + edit = None if args.scene_edit is None else str(args.scene_edit).strip() + expected = { + "image": (True, False, False), + "image-edit": (True, False, True), + "scene": (False, True, False), + "scene-edit": (False, True, True), + }[args.mode] + actual = (bool(image), bool(scene), bool(edit)) + if actual != expected: + raise ValueError( + f"mode={args.mode!r} requires image/scene/edit={expected}, got {actual}." + ) + return image, scene, edit + + +def _print_json(value: Any) -> None: + print(json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/embodichain/gen_sim/task_engine/config.py b/embodichain/gen_sim/task_engine/config.py new file mode 100644 index 000000000..60cda86ef --- /dev/null +++ b/embodichain/gen_sim/task_engine/config.py @@ -0,0 +1,190 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Configuration owned by Task Engine orchestration.""" + +from __future__ import annotations + +from collections.abc import Mapping +from importlib.resources import files +from pathlib import Path +from typing import Any, Final + +import yaml + +from embodichain.utils import configclass + +__all__ = [ + "TASK_ENGINE_DEFAULTS_SCHEMA", + "TaskEngineExecutionCfg", + "TaskEnginePlanningCfg", + "TaskEngineWorkflowCfg", + "load_task_engine_config", +] + +TASK_ENGINE_DEFAULTS_SCHEMA: Final = "embodichain.task-engine-defaults/v1" + + +@configclass +class TaskEngineExecutionCfg: + """Success policy for vectorized simulator execution.""" + + num_envs: int = 1 + success_policy: str = "any" + min_successful_envs: int = 1 + + def __post_init__(self) -> None: + if ( + isinstance(self.num_envs, bool) + or not isinstance(self.num_envs, int) + or self.num_envs < 1 + ): + raise ValueError("num_envs must be a positive integer.") + if self.success_policy not in {"any", "all", "at_least"}: + raise ValueError("success_policy must be any, all, or at_least.") + if ( + isinstance(self.min_successful_envs, bool) + or not isinstance(self.min_successful_envs, int) + or not 1 <= self.min_successful_envs <= self.num_envs + ): + raise ValueError("min_successful_envs must be in [1, num_envs].") + if self.success_policy == "any" and self.min_successful_envs != 1: + raise ValueError("success_policy=any requires min_successful_envs=1.") + if self.success_policy == "all" and self.min_successful_envs != self.num_envs: + raise ValueError( + "success_policy=all requires min_successful_envs=num_envs." + ) + + @property + def required_successes(self) -> int: + """Return the number of successful replicas required for acceptance.""" + if self.success_policy == "all": + return self.num_envs + if self.success_policy == "any": + return 1 + return self.min_successful_envs + + +@configclass +class TaskEngineWorkflowCfg: + """Conservative first-version orchestration limits. + + The packaged YAML owns retry limits so deployment testing can tune them + without changing the orchestration implementation. + """ + + max_parallel_workers: int = 2 + max_scene_attempts: int = 2 + max_action_attempts: int = 3 + + def __post_init__(self) -> None: + for field_name in ( + "max_parallel_workers", + "max_scene_attempts", + "max_action_attempts", + ): + value = getattr(self, field_name) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{field_name} must be a positive integer.") + + +@configclass +class TaskEnginePlanningCfg: + """Task interpretation and Action bundle generation defaults.""" + + candidate_count: int = 3 + planning_mode: str = "offline" + max_episodes: int = 1 + max_episode_steps: int = 4000 + + def __post_init__(self) -> None: + for field_name in ( + "candidate_count", + "max_episodes", + "max_episode_steps", + ): + value = getattr(self, field_name) + if isinstance(value, bool) or not isinstance(value, int) or value < 1: + raise ValueError(f"{field_name} must be a positive integer.") + if self.planning_mode not in {"offline", "ab"}: + raise ValueError("planning_mode must be offline or ab.") + + +def load_task_engine_config( + path: str | Path | None = None, +) -> tuple[ + TaskEngineWorkflowCfg, + TaskEnginePlanningCfg, + TaskEngineExecutionCfg, +]: + """Load strict Task Engine defaults from YAML. + + Args: + path: Optional override YAML. The packaged defaults are used when omitted. + + Returns: + Validated workflow, planning, and execution configurations. + + Raises: + TypeError: If a configuration section is not a mapping. + ValueError: If the YAML schema or fields are invalid. + """ + content = ( + Path(path).expanduser().resolve().read_text(encoding="utf-8") + if path is not None + else files(__package__).joinpath("defaults.yaml").read_text(encoding="utf-8") + ) + raw = yaml.safe_load(content) + if not isinstance(raw, Mapping): + raise TypeError("Task Engine configuration must be a mapping.") + expected = {"schema_version", "workflow", "planning", "execution"} + if set(raw) != expected: + raise ValueError("Task Engine configuration fields are invalid.") + if raw.get("schema_version") != TASK_ENGINE_DEFAULTS_SCHEMA: + raise ValueError("Task Engine configuration schema_version is invalid.") + workflow = _mapping(raw.get("workflow"), "workflow") + planning = _mapping(raw.get("planning"), "planning") + execution = _mapping(raw.get("execution"), "execution") + if set(workflow) != { + "max_parallel_workers", + "max_scene_attempts", + "max_action_attempts", + }: + raise ValueError("Task Engine workflow configuration fields are invalid.") + if set(planning) != { + "candidate_count", + "planning_mode", + "max_episodes", + "max_episode_steps", + }: + raise ValueError("Task Engine planning configuration fields are invalid.") + if set(execution) != { + "num_envs", + "success_policy", + "min_successful_envs", + }: + raise ValueError("Task Engine execution configuration fields are invalid.") + return ( + TaskEngineWorkflowCfg(**workflow), + TaskEnginePlanningCfg(**planning), + TaskEngineExecutionCfg(**execution), + ) + + +def _mapping(value: Any, field_name: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"Task Engine {field_name} configuration must be a mapping.") + return dict(value) diff --git a/embodichain/gen_sim/task_engine/contracts.py b/embodichain/gen_sim/task_engine/contracts.py new file mode 100644 index 000000000..237423e55 --- /dev/null +++ b/embodichain/gen_sim/task_engine/contracts.py @@ -0,0 +1,410 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict, JSON-safe public contracts owned by Task Engine.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import hashlib +import json +from typing import Any, TypeAlias + +from .interpretation import validate_instruction_intent +from .ontology import TASK_CONTRACTS, task_success_type + +__all__ = [ + "SCENE_REQUEST_SCHEMA", + "SUCCESS_SPEC_SCHEMA", + "TASK_CANDIDATE_SET_SCHEMA", + "TASK_DRAFT_SCHEMA", + "SceneRequest", + "SuccessSpec", + "TaskCandidate", + "TaskCandidateSet", + "TaskDraft", + "canonical_hash", + "validate_scene_request", + "validate_success_spec", + "validate_task_candidate", + "validate_task_candidate_set", + "validate_task_draft", +] + +TASK_DRAFT_SCHEMA = "action_engine_task_draft_v1" +SCENE_REQUEST_SCHEMA = "action_engine_scene_request_v1" +SUCCESS_SPEC_SCHEMA = "action_engine_success_spec_v1" +TASK_CANDIDATE_SET_SCHEMA = "action_engine_task_candidate_set_v1" + +TaskDraft: TypeAlias = dict[str, Any] +SceneRequest: TypeAlias = dict[str, Any] +SuccessSpec: TypeAlias = dict[str, Any] +TaskCandidate: TypeAlias = dict[str, Any] +TaskCandidateSet: TypeAlias = dict[str, Any] + +_SUCCESS_TYPES = frozenset( + {contract.success_type for contract in TASK_CONTRACTS.values()} | {"semantic_goal"} +) +_DRAFT_KEYS = frozenset({"schema_version", "task_id", "instruction", "steps"}) +_SCENE_REQUEST_KEYS = frozenset({"schema_version", "task_id", "references"}) +_REFERENCE_KEYS = frozenset( + { + "reference_id", + "step_id", + "role", + "reference", + "quantifier", + "count", + "source_structure", + "affordances", + "initial_state", + "attributes", + } +) +_SUCCESS_KEYS = frozenset({"schema_version", "task_id", "op", "terms"}) +_SUCCESS_TERM_KEYS = frozenset({"step_id", "type"}) +_CANDIDATE_KEYS = frozenset( + { + "candidate_id", + "draft", + "scene_request", + "success_spec", + "semantic_hash", + "vote_count", + "attempts", + "normalizations", + } +) +_CANDIDATE_SET_KEYS = frozenset( + { + "schema_version", + "task_id", + "instruction", + "candidates", + "requested_candidate_count", + "valid_response_count", + "errors", + } +) + + +def canonical_hash(value: Any) -> str: + """Return the stable SHA-256 of one JSON-safe protocol value.""" + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def validate_task_draft(value: Mapping[str, Any]) -> TaskDraft: + result = _mapping(value, "TaskDraft") + _keys(result, _DRAFT_KEYS, "TaskDraft") + _schema(result, TASK_DRAFT_SCHEMA, "TaskDraft") + result["task_id"] = _nonempty(result.get("task_id"), "TaskDraft.task_id") + result["instruction"] = _nonempty( + result.get("instruction"), "TaskDraft.instruction" + ) + intent = validate_instruction_intent({"steps": result.get("steps")}) + result["steps"] = intent["steps"] + return result + + +def validate_scene_request(value: Mapping[str, Any]) -> SceneRequest: + result = _mapping(value, "SceneRequest") + _keys(result, _SCENE_REQUEST_KEYS, "SceneRequest") + _schema(result, SCENE_REQUEST_SCHEMA, "SceneRequest") + task_id = _nonempty(result.get("task_id"), "SceneRequest.task_id") + references: list[dict[str, Any]] = [] + for index, raw in enumerate( + _sequence(result.get("references"), "SceneRequest.references") + ): + context = f"SceneRequest.references[{index}]" + reference = _mapping(raw, context) + _keys(reference, _REFERENCE_KEYS, context) + for key in ("reference_id", "step_id", "role", "reference", "source_structure"): + reference[key] = _nonempty(reference.get(key), f"{context}.{key}") + reference["role"] = _enum( + reference["role"], {"object", "target"}, f"{context}.role" + ) + reference["quantifier"] = _enum( + reference.get("quantifier"), + {"one", "all", "count"}, + f"{context}.quantifier", + ) + reference["count"] = _integer( + reference.get("count"), f"{context}.count", minimum=0 + ) + if reference["quantifier"] in {"one", "all"} and reference["count"] != 0: + raise ValueError( + f"{context} quantifier={reference['quantifier']} requires count=0." + ) + if reference["quantifier"] == "count" and reference["count"] < 1: + raise ValueError(f"{context} quantifier=count requires count>=1.") + reference["affordances"] = _strings( + reference.get("affordances"), f"{context}.affordances" + ) + reference["initial_state"] = _mapping( + reference.get("initial_state"), f"{context}.initial_state" + ) + reference["attributes"] = _mapping( + reference.get("attributes"), f"{context}.attributes" + ) + references.append(reference) + _unique([item["reference_id"] for item in references], "SceneRequest reference IDs") + result["task_id"] = task_id + result["references"] = references + _json_safe(result, "SceneRequest") + return result + + +def validate_success_spec( + value: Mapping[str, Any], + *, + draft: Mapping[str, Any] | None = None, +) -> SuccessSpec: + result = _mapping(value, "SuccessSpec") + _keys(result, _SUCCESS_KEYS, "SuccessSpec") + _schema(result, SUCCESS_SPEC_SCHEMA, "SuccessSpec") + task_id = _nonempty(result.get("task_id"), "SuccessSpec.task_id") + if result.get("op") != "all": + raise ValueError("SuccessSpec.op must be 'all'.") + terms: list[dict[str, str]] = [] + for index, raw in enumerate(_sequence(result.get("terms"), "SuccessSpec.terms")): + context = f"SuccessSpec.terms[{index}]" + term = _mapping(raw, context) + _keys(term, _SUCCESS_TERM_KEYS, context) + terms.append( + { + "step_id": _nonempty(term.get("step_id"), f"{context}.step_id"), + "type": _enum(term.get("type"), set(_SUCCESS_TYPES), f"{context}.type"), + } + ) + if not terms: + raise ValueError("SuccessSpec.terms must not be empty.") + _unique([term["step_id"] for term in terms], "SuccessSpec step IDs") + if draft is not None: + normalized_draft = validate_task_draft(draft) + if normalized_draft["task_id"] != task_id: + raise ValueError("SuccessSpec.task_id must match TaskDraft.task_id.") + expected = [ + { + "step_id": step["id"], + "type": task_success_type(step["task_type"], step), + } + for step in normalized_draft["steps"] + ] + if terms != expected: + raise ValueError( + "SuccessSpec terms must be ordered, complete, and derived from " + "task_success_type." + ) + result["task_id"] = task_id + result["terms"] = terms + return result + + +def validate_task_candidate(value: Mapping[str, Any]) -> TaskCandidate: + result = _mapping(value, "TaskCandidate") + _keys(result, _CANDIDATE_KEYS, "TaskCandidate") + result["candidate_id"] = _nonempty( + result.get("candidate_id"), "TaskCandidate.candidate_id" + ) + result["draft"] = validate_task_draft(result.get("draft")) + result["scene_request"] = validate_scene_request(result.get("scene_request")) + result["success_spec"] = validate_success_spec( + result.get("success_spec"), draft=result["draft"] + ) + for name in ("scene_request", "success_spec"): + if result[name]["task_id"] != result["draft"]["task_id"]: + raise ValueError(f"TaskCandidate {name}.task_id must match its draft.") + from .agent import derive_scene_request + + if result["scene_request"] != derive_scene_request(result["draft"]): + raise ValueError( + "TaskCandidate.scene_request must be derived exactly from its draft." + ) + result["semantic_hash"] = _digest( + result.get("semantic_hash"), "TaskCandidate.semantic_hash" + ) + if result["semantic_hash"] != canonical_hash(result["draft"]["steps"]): + raise ValueError( + "TaskCandidate.semantic_hash does not match its canonical steps." + ) + result["vote_count"] = _integer( + result.get("vote_count"), "TaskCandidate.vote_count", minimum=1 + ) + result["attempts"] = _integer( + result.get("attempts"), "TaskCandidate.attempts", minimum=1, maximum=2 + ) + result["normalizations"] = _mapping_sequence( + result.get("normalizations"), "TaskCandidate.normalizations" + ) + return result + + +def validate_task_candidate_set(value: Mapping[str, Any]) -> TaskCandidateSet: + result = _mapping(value, "TaskCandidateSet") + _keys(result, _CANDIDATE_SET_KEYS, "TaskCandidateSet") + _schema(result, TASK_CANDIDATE_SET_SCHEMA, "TaskCandidateSet") + task_id = _nonempty(result.get("task_id"), "TaskCandidateSet.task_id") + instruction = _nonempty(result.get("instruction"), "TaskCandidateSet.instruction") + requested = _integer( + result.get("requested_candidate_count"), + "TaskCandidateSet.requested_candidate_count", + minimum=1, + ) + valid = _integer( + result.get("valid_response_count"), + "TaskCandidateSet.valid_response_count", + minimum=1, + maximum=requested, + ) + candidates = [ + validate_task_candidate(item) + for item in _sequence(result.get("candidates"), "TaskCandidateSet.candidates") + ] + if not candidates: + raise ValueError("TaskCandidateSet.candidates must not be empty.") + _unique([item["candidate_id"] for item in candidates], "TaskCandidate IDs") + _unique( + [item["semantic_hash"] for item in candidates], "TaskCandidate semantic hashes" + ) + if sum(item["vote_count"] for item in candidates) != valid: + raise ValueError( + "TaskCandidate vote_count values must sum to valid_response_count." + ) + for candidate in candidates: + if ( + candidate["draft"]["task_id"] != task_id + or candidate["draft"]["instruction"] != instruction + ): + raise ValueError("Every TaskCandidate draft must match its candidate set.") + errors = _strings(result.get("errors"), "TaskCandidateSet.errors", allow_empty=True) + if valid + len(errors) != requested: + raise ValueError( + "Valid responses plus errors must equal requested_candidate_count." + ) + result.update( + { + "task_id": task_id, + "instruction": instruction, + "requested_candidate_count": requested, + "valid_response_count": valid, + "candidates": candidates, + "errors": errors, + } + ) + return result + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"{context} must be a list.") + return list(value) + + +def _keys( + value: Mapping[str, Any], expected: set[str] | frozenset[str], context: str +) -> None: + if set(value) != set(expected): + raise ValueError( + f"{context} requires exactly fields {sorted(expected)}; received {sorted(value)}." + ) + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{context} must be a string.") + return value.strip() + + +def _nonempty(value: Any, context: str) -> str: + result = _string(value, context) + if not result: + raise ValueError(f"{context} must not be empty.") + return result + + +def _enum(value: Any, choices: set[str], context: str) -> str: + result = _string(value, context) + if result not in choices: + raise ValueError(f"{context} must be one of {sorted(choices)}.") + return result + + +def _integer( + value: Any, context: str, *, minimum: int, maximum: int | None = None +) -> int: + if ( + not isinstance(value, int) + or isinstance(value, bool) + or value < minimum + or (maximum is not None and value > maximum) + ): + raise ValueError(f"{context} must be an integer in the allowed range.") + return value + + +def _strings(value: Any, context: str, *, allow_empty: bool = False) -> list[str]: + result = [_string(item, context) for item in _sequence(value, context)] + if not allow_empty and any(not item for item in result): + raise ValueError(f"{context} values must not be empty.") + if len(result) != len(set(result)): + raise ValueError(f"{context} values must be unique.") + return result + + +def _mapping_sequence(value: Any, context: str) -> list[dict[str, Any]]: + result = [_mapping(item, context) for item in _sequence(value, context)] + _json_safe(result, context) + return result + + +def _digest(value: Any, context: str) -> str: + result = _string(value, context) + if len(result) != 64 or any( + character not in "0123456789abcdef" for character in result + ): + raise ValueError(f"{context} must be a lowercase SHA-256 digest.") + return result + + +def _unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must be unique.") + + +def _json_safe(value: Any, context: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError) as error: + raise ValueError(f"{context} must be finite and JSON serializable.") from error diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml new file mode 100644 index 000000000..33169e461 --- /dev/null +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +schema_version: embodichain.task-engine-defaults/v1 + +workflow: + max_parallel_workers: 2 + max_scene_attempts: 2 + max_action_attempts: 3 + +planning: + candidate_count: 3 + planning_mode: offline + max_episodes: 1 + max_episode_steps: 4000 + +execution: + num_envs: 1 + success_policy: any + min_successful_envs: 1 diff --git a/embodichain/gen_sim/task_engine/interpretation.py b/embodichain/gen_sim/task_engine/interpretation.py new file mode 100644 index 000000000..2b971344c --- /dev/null +++ b/embodichain/gen_sim/task_engine/interpretation.py @@ -0,0 +1,1269 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scene-independent structured interpretation for Task Engine.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +import os +from pathlib import Path +import re +from time import perf_counter +from typing import Any, TypeAlias + +from .ontology import ( + RELATIONS, + TASK_CONTRACTS, + TERMINAL_BEHAVIORS, + TRANSPORT_DIRECTIONS, +) + +__all__ = [ + "INSTRUCTION_INTENT_SCHEMA", + "InstructionDraftResult", + "InstructionIntent", + "InstructionCaller", + "interpret_instruction_draft", + "validate_instruction_intent", +] + +InstructionCaller = Callable[..., Mapping[str, Any]] +InstructionIntent: TypeAlias = dict[str, Any] +TASK_TYPES = frozenset(TASK_CONTRACTS) + +_RELATIONS = RELATIONS +_ARMS = frozenset({"none", "auto", "left_arm", "right_arm"}) +_ORIENTATIONS = frozenset({"none", "preserve", "upright"}) +_TARGET_STATES = frozenset({"none", "open", "closed", "activated"}) +_LAYOUTS = frozenset({"none", "line"}) +_AXES = frozenset({"none", "world_x", "world_y"}) +_DIRECTIONS = TRANSPORT_DIRECTIONS +_TERMINAL_BEHAVIORS = TERMINAL_BEHAVIORS +_SELECTOR_KINDS = frozenset({"none", "scene_ref", "step_result"}) +_QUANTIFIERS = frozenset({"one", "all", "count"}) +_STEP_KEYS = frozenset( + { + "id", + "task_type", + "object", + "target", + "relation", + "required_arm", + "transfer_arm", + "receive_arm", + "orientation_goal", + "target_state", + "target_setting", + "layout", + "axis", + "direction", + "terminal_behavior", + "depends_on", + } +) +_INTENT_TASK_FIELD_REGISTRY = { + task_type: contract.applicable_intent_fields + for task_type, contract in TASK_CONTRACTS.items() +} +_INTENT_FIELD_DEFAULTS: dict[str, Any] = { + "target": None, + "relation": "none", + "required_arm": "none", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "none", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", +} +_SELECTOR_KEYS = frozenset( + { + "kind", + "step_id", + "reference", + "quantifier", + "count", + } +) +_FORBIDDEN_FIELDS = frozenset( + { + "atomic_action", + "atomic_actions", + "atomicaction", + "coordinates", + "bbox", + "bboxes", + "grasp_pose", + "keypoint", + "keypoints", + "joint_positions", + "joints", + "pose", + "position", + "qpos", + "rotation", + "target_pose", + "translation", + "trajectory", + "waypoints", + } +) +# MiMo's OpenAI-compatible endpoint can spend the whole completion budget in +# hidden reasoning when the request leaves thinking enabled. A sparse final +# JSON object then looks like a schema failure to the deterministic verifier. +# Keep the budget bounded and turn reasoning off for the text interpretation +# call; the parser must return an auditable object rather than a thought trace. +_MIMO_MAX_COMPLETION_TOKENS = 4096 +_GEN_SIM_DIR = Path(__file__).resolve().parents[1] +_GEN_SIM_ENV_PATH = _GEN_SIM_DIR / ".env" +_GEN_CONFIG_PATH = _GEN_SIM_DIR / "simready_pipeline" / "configs" / "gen_config.json" + + +class _MissingRequiredTargetError(ValueError): + """Identify a validation failure that receives targeted repair guidance.""" + + +class _MissingRequiredObjectError(ValueError): + """Identify a missing manipulated-object selector for targeted repair.""" + + +@dataclass(frozen=True) +class InstructionDraftResult: + """One validated, scene-independent interpretation and its audit metadata.""" + + intent: InstructionIntent + model: str + attempts: int + latency_seconds: float + normalizations: tuple[dict[str, Any], ...] + + +# Object semantics remain open natural-language references until the dedicated +# scene-grounding phase resolves them. All other values are strict protocol +# enums; non-canonical model output is repaired by the model, never guessed by +# a local language alias table. + +_SELECTOR_SCHEMA = { + "type": "object", + "additionalProperties": False, + "required": sorted(_SELECTOR_KEYS), + "properties": { + "kind": {"type": "string", "enum": sorted(_SELECTOR_KINDS)}, + "step_id": {"type": "string"}, + "reference": {"type": "string"}, + "quantifier": {"type": "string", "enum": sorted(_QUANTIFIERS)}, + "count": {"type": "integer", "minimum": 0}, + }, +} + +_INTENT_OUTPUT_SCHEMA = { + "title": "ActionEngineInstructionIntent", + "type": "object", + "additionalProperties": False, + "required": ["steps"], + "properties": { + "steps": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "additionalProperties": False, + "required": sorted(_STEP_KEYS), + "properties": { + "id": {"type": "string"}, + "task_type": {"type": "string", "enum": sorted(TASK_TYPES)}, + "object": _SELECTOR_SCHEMA, + "target": _SELECTOR_SCHEMA, + "relation": {"type": "string", "enum": sorted(_RELATIONS)}, + "required_arm": {"type": "string", "enum": sorted(_ARMS)}, + "transfer_arm": {"type": "string", "enum": sorted(_ARMS)}, + "receive_arm": {"type": "string", "enum": sorted(_ARMS)}, + "orientation_goal": { + "type": "string", + "enum": sorted(_ORIENTATIONS), + }, + "target_state": { + "type": "string", + "enum": sorted(_TARGET_STATES), + }, + "target_setting": {"type": "integer"}, + "layout": {"type": "string", "enum": sorted(_LAYOUTS)}, + "axis": {"type": "string", "enum": sorted(_AXES)}, + "direction": { + "type": "string", + "enum": sorted(_DIRECTIONS), + }, + "terminal_behavior": { + "type": "string", + "enum": sorted(_TERMINAL_BEHAVIORS), + }, + "depends_on": { + "type": "array", + "items": {"type": "string"}, + }, + }, + }, + } + }, +} + +# Keep a read-only-by-convention public copy for callers that need to configure +# a structured client. The schema is an input contract, not a persisted task +# graph; ``validate_instruction_intent`` remains the authoritative verifier. +INSTRUCTION_INTENT_SCHEMA = deepcopy(_INTENT_OUTPUT_SCHEMA) + + +def interpret_instruction_draft( + instruction: str, + *, + model: str | None = None, + caller: InstructionCaller | None = None, +) -> InstructionDraftResult: + """Interpret one instruction without reading or grounding a scene.""" + instruction_text = str(instruction).strip() + if not instruction_text: + raise ValueError("instruction must be non-empty.") + prompt = _instruction_prompt(instruction_text) + invoke = caller or _default_instruction_caller + # An injected caller owns its transport and does not need provider config. + selected_model = model if caller is not None else _instruction_model(model) + if caller is None and selected_model is None: + raise ValueError( + "A text LLM model is required through --llm-model, " + "ACTION_ENGINE_LLM_MODEL, or OPENAI_MODEL." + ) + started = perf_counter() + first_error: Exception | None = None + for attempt in range(2): + current_prompt = prompt + if first_error is not None: + current_prompt += ( + "\n\nREPAIR OVERRIDE: the previous JSON was invalid. Return a corrected " + "JSON object only; do not repeat the sparse response. Every step " + "must contain all 16 step keys and every selector all 5 selector " + "keys. Keep semantic fields explicit: E4 requires transfer_arm " + "and receive_arm, and E1/E3 require target plus relation (unless " + "E1 layout=line). Use canonical defaults only for fields that do " + "not apply. Validation error: " + f"{first_error}\n" + "Copy this complete shape before filling values (shape only; do " + "not copy its values or step count):\n" + f"{json.dumps(_instruction_shape_example(), ensure_ascii=False, sort_keys=True)}\n" + "Selector kind rules:\n" + f"{_instruction_selector_rules()}" + f"{_instruction_repair_guidance(first_error)}" + ) + try: + response = invoke( + prompt=current_prompt, + schema=deepcopy(INSTRUCTION_INTENT_SCHEMA), + model=selected_model, + ) + normalized, normalizations = _normalize_instruction_intent_fields( + _coerce_instruction_response(response) + ) + intent = validate_instruction_intent(normalized) + return InstructionDraftResult( + intent=intent, + model=selected_model or "injected_caller", + attempts=attempt + 1, + latency_seconds=perf_counter() - started, + normalizations=tuple(normalizations), + ) + except (TypeError, ValueError) as error: + if attempt: + raise ValueError( + "Instruction intent failed validation after one repair: " f"{error}" + ) from error + first_error = error + raise AssertionError("unreachable") + + +def _normalize_instruction_intent_fields( + value: Mapping[str, Any], +) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Canonicalize defaults and uniquely constrained cross-step continuity. + + The strict public validator deliberately remains unchanged. This pass is + confined to the LLM boundary, where weak JSON-mode providers sometimes + copy a meaningful value into an inapplicable slot such as E4.required_arm. + Required scene facts and ambiguous arm assignments are never inferred here + and still fail closed. + """ + result = deepcopy(dict(value)) + raw_steps = result.get("steps") + if not isinstance(raw_steps, list): + return result, [] + changes: list[dict[str, Any]] = [] + for index, raw_step in enumerate(raw_steps): + if not isinstance(raw_step, dict) or set(raw_step) != _STEP_KEYS: + continue + task_type = raw_step.get("task_type") + applicable = _INTENT_TASK_FIELD_REGISTRY.get(task_type) + if applicable is None: + continue + for field, configured_default in _INTENT_FIELD_DEFAULTS.items(): + field_applies = field in applicable + if task_type == "E1" and field in {"target", "relation"}: + field_applies = raw_step.get("layout") != "line" + if task_type == "E1" and field == "axis": + field_applies = raw_step.get("layout") == "line" + if field_applies: + continue + default = ( + _empty_selector() + if field == "target" and configured_default is None + else deepcopy(configured_default) + ) + if raw_step[field] == default: + continue + previous = deepcopy(raw_step[field]) + raw_step[field] = default + changes.append( + { + "path": f"steps[{index}].{field}", + "from": previous, + "to": deepcopy(default), + "reason": f"inapplicable_for_{task_type}", + } + ) + target = raw_step.get("target") + if ( + task_type == "E5" + and isinstance(target, Mapping) + and target.get("kind") == "none" + and raw_step.get("relation") == "none" + and raw_step.get("direction") == "none" + and raw_step.get("terminal_behavior") == "hold" + ): + raw_step["direction"] = "up" + changes.append( + { + "path": f"steps[{index}].direction", + "from": "none", + "to": "up", + "reason": "e5_hold_defaults_to_lift", + } + ) + _normalize_handover_arm_continuity(raw_steps, changes) + return result, changes + + +def _normalize_handover_arm_continuity( + steps: Sequence[Any], + changes: list[dict[str, Any]], +) -> None: + """Repair a same-arm E4 only when adjacent ownership fixes both roles.""" + by_id: dict[str, Mapping[str, Any]] = {} + for step in steps: + if not isinstance(step, dict) or set(step) != _STEP_KEYS: + return + step_id = step.get("id") + if not isinstance(step_id, str) or not step_id or step_id in by_id: + return + by_id[step_id] = step + + explicit_arms = {"left_arm", "right_arm"} + for index, step in enumerate(steps): + assert isinstance(step, dict) + transfer = step.get("transfer_arm") + receive = step.get("receive_arm") + if ( + step.get("task_type") != "E4" + or transfer not in explicit_arms + or transfer != receive + ): + continue + + object_key = _object_lineage_key(step, by_id) + upstream_arm: str | None = None + for producer in reversed(steps[:index]): + assert isinstance(producer, Mapping) + if object_key is None or _object_lineage_key(producer, by_id) != object_key: + continue + candidate = ( + producer.get("receive_arm") + if producer.get("task_type") == "E4" + else producer.get("required_arm") + ) + if candidate in explicit_arms: + upstream_arm = str(candidate) + break + + downstream_arm: str | None = None + for consumer in steps[index + 1 :]: + assert isinstance(consumer, Mapping) + if object_key is None or _object_lineage_key(consumer, by_id) != object_key: + continue + candidate = ( + consumer.get("transfer_arm") + if consumer.get("task_type") == "E4" + else consumer.get("required_arm") + ) + if candidate in explicit_arms: + downstream_arm = str(candidate) + break + + desired_transfer = upstream_arm or str(transfer) + desired_receive = downstream_arm or str(receive) + if desired_transfer == desired_receive: + continue + for field, desired in ( + ("transfer_arm", desired_transfer), + ("receive_arm", desired_receive), + ): + if step[field] == desired: + continue + previous = step[field] + step[field] = desired + changes.append( + { + "path": f"steps[{index}].{field}", + "from": previous, + "to": desired, + "reason": "handover_arm_continuity", + } + ) + + +def _object_lineage_key( + step: Mapping[str, Any], + by_id: Mapping[str, Mapping[str, Any]], + seen: frozenset[str] = frozenset(), +) -> tuple[str, str] | None: + """Resolve object identity only through explicit step-result lineage.""" + selector = step.get("object") + if not isinstance(selector, Mapping): + return None + kind = selector.get("kind") + if kind == "scene_ref": + step_id = step.get("id") + if not isinstance(step_id, str) or not step_id: + return None + return ("step_result", step_id) + if kind != "step_result": + return None + producer_id = selector.get("step_id") + if not isinstance(producer_id, str) or producer_id in seen: + return None + producer = by_id.get(producer_id) + if producer is None: + return None + return _object_lineage_key(producer, by_id, seen | {producer_id}) + + +def _empty_selector() -> dict[str, Any]: + """Return the canonical selector value for an inapplicable target.""" + return { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + + +def validate_instruction_intent(value: Mapping[str, Any]) -> dict[str, Any]: + """Validate the private, non-graph instruction interpretation contract.""" + if not isinstance(value, Mapping): + raise TypeError("Instruction intent must be a mapping.") + _reject_forbidden_fields(value) + if set(value) != {"steps"}: + raise ValueError("Instruction intent may contain only 'steps'.") + raw_steps = value.get("steps") + if not isinstance(raw_steps, Sequence) or isinstance(raw_steps, (str, bytes)): + raise ValueError("Instruction intent steps must be a list.") + if not raw_steps: + raise ValueError("Instruction intent steps must not be empty.") + steps = [] + ids: set[str] = set() + dependencies: dict[str, list[str]] = {} + for index, raw in enumerate(raw_steps): + context = f"InstructionIntent.steps[{index}]" + if not isinstance(raw, Mapping): + raise ValueError(f"{context} must be a mapping.") + if set(raw) != _STEP_KEYS: + raise ValueError( + f"{context} requires exactly fields {sorted(_STEP_KEYS)}; " + f"received {sorted(raw)}." + ) + step = deepcopy(dict(raw)) + step_id = _nonempty(step["id"], f"{context}.id") + if step_id in ids: + raise ValueError(f"Duplicate instruction step ID {step_id!r}.") + ids.add(step_id) + step["id"] = step_id + step["task_type"] = _choice( + step["task_type"], TASK_TYPES, f"{context}.task_type" + ) + step["object"] = _validate_selector(step["object"], f"{context}.object") + step["target"] = _validate_selector(step["target"], f"{context}.target") + step["relation"] = _canonical_relation(step["relation"], f"{context}.relation") + for key in ("required_arm", "transfer_arm", "receive_arm"): + step[key] = _canonical_arm(step[key], f"{context}.{key}") + step["orientation_goal"] = _canonical_orientation( + step["orientation_goal"], f"{context}.orientation_goal" + ) + step["target_state"] = _choice( + step["target_state"], _TARGET_STATES, f"{context}.target_state" + ) + if isinstance(step["target_setting"], bool) or not isinstance( + step["target_setting"], int + ): + raise ValueError(f"{context}.target_setting must be an integer.") + step["layout"] = _choice(step["layout"], _LAYOUTS, f"{context}.layout") + step["axis"] = _choice(step["axis"], _AXES, f"{context}.axis") + step["direction"] = _choice( + step["direction"], _DIRECTIONS, f"{context}.direction" + ) + step["terminal_behavior"] = _choice( + step["terminal_behavior"], + _TERMINAL_BEHAVIORS, + f"{context}.terminal_behavior", + ) + raw_depends = step["depends_on"] + if not isinstance(raw_depends, Sequence) or isinstance( + raw_depends, (str, bytes) + ): + raise ValueError(f"{context}.depends_on must be a list.") + step["depends_on"] = [ + _nonempty(item, f"{context}.depends_on") for item in raw_depends + ] + if step_id in step["depends_on"]: + raise ValueError(f"{context}.depends_on cannot contain its own ID.") + dependencies[step_id] = step["depends_on"] + _validate_task_fields(step, context) + steps.append(step) + positions = {str(step["id"]): index for index, step in enumerate(steps)} + for index, step in enumerate(steps): + for selector_name in ("object", "target"): + selector = step[selector_name] + if selector["kind"] != "step_result": + continue + reference = str(selector["step_id"]) + if reference not in positions: + raise ValueError( + f"Instruction step {step['id']!r} {selector_name} references " + f"unknown step {reference!r}." + ) + if positions[reference] >= index: + raise ValueError( + f"Instruction step {step['id']!r} {selector_name} must reference " + f"a preceding step, not {reference!r}." + ) + for step_id, depends_on in dependencies.items(): + unknown = set(depends_on) - ids + if unknown: + raise ValueError( + f"Instruction step {step_id!r} has unknown dependencies " + f"{sorted(unknown)}." + ) + _validate_dag(dependencies) + return {"steps": steps} + + +def _validate_selector(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + if set(value) != _SELECTOR_KEYS: + raise ValueError( + f"{context} requires exactly fields {sorted(_SELECTOR_KEYS)}; " + f"received {sorted(value)}." + ) + selector = deepcopy(dict(value)) + selector["kind"] = _choice(selector["kind"], _SELECTOR_KINDS, f"{context}.kind") + selector["step_id"] = _selector_string(selector["step_id"], f"{context}.step_id") + selector["reference"] = _selector_string( + selector["reference"], f"{context}.reference" + ) + selector["quantifier"] = _canonical_quantifier( + selector["quantifier"], f"{context}.quantifier" + ) + if isinstance(selector["count"], bool) or not isinstance(selector["count"], int): + raise ValueError(f"{context}.count must be an integer.") + if selector["count"] < 0: + raise ValueError(f"{context}.count must be non-negative.") + kind = selector["kind"] + if kind == "scene_ref" and not selector["reference"]: + raise ValueError(f"{context} scene_ref requires a reference.") + if kind == "step_result": + if not selector["step_id"]: + raise ValueError(f"{context} step_result requires step_id.") + if selector["reference"]: + raise ValueError( + f"{context} step_result may identify only a prior step_id." + ) + if selector["quantifier"] != "one" or selector["count"] != 0: + raise ValueError( + f"{context} step_result requires quantifier=one and count=0." + ) + if kind == "scene_ref" and selector["step_id"]: + raise ValueError(f"{context} scene_ref cannot carry step_id.") + if kind == "none" and (selector["step_id"] or selector["reference"]): + raise ValueError(f"{context} kind=none cannot carry constraints.") + if kind == "none" and (selector["quantifier"] != "one" or selector["count"] != 0): + raise ValueError(f"{context} kind=none requires quantifier=one and count=0.") + if selector["quantifier"] == "one" and selector["count"] != 0: + raise ValueError(f"{context} quantifier=one requires count=0.") + if selector["quantifier"] == "all" and selector["count"] != 0: + raise ValueError(f"{context} quantifier=all requires count=0.") + if selector["quantifier"] == "count" and selector["count"] < 1: + raise ValueError(f"{context} quantifier=count requires count>=1.") + return selector + + +def _validate_task_fields(step: Mapping[str, Any], context: str) -> None: + task_type = str(step["task_type"]) + if step["object"]["kind"] == "none": + raise _MissingRequiredObjectError( + f"{context} {task_type} requires an object selector." + ) + target_kind = str(step["target"]["kind"]) + if task_type not in {"E1", "E3", "E5"} and step["relation"] != "none": + raise ValueError(f"{context} {task_type} does not accept relation.") + if task_type == "E3" and step["relation"] != "above": + raise ValueError(f"{context} E3 relation must be above.") + target_setting = int(step["target_setting"]) + if task_type != "E8" and target_setting != 0: + raise ValueError(f"{context} target_setting is only valid for E8.") + if task_type != "E1" and step["axis"] != "none": + raise ValueError(f"{context} axis is only valid for E1 line arrangement.") + if task_type == "E1" and step["layout"] != "line" and step["axis"] != "none": + raise ValueError(f"{context} axis is only valid for E1 line arrangement.") + if task_type not in {"E6", "E7", "E9"} and step["target_state"] != "none": + raise ValueError(f"{context} target_state is not valid for {task_type}.") + if task_type != "E4" and step["transfer_arm"] != "none": + raise ValueError(f"{context} transfer_arm is only valid for E4.") + if task_type != "E4" and step["receive_arm"] != "none": + raise ValueError(f"{context} receive_arm is only valid for E4.") + orientation_goal = str(step["orientation_goal"]) + if task_type == "E2" and orientation_goal != "upright": + raise ValueError(f"{context} E2 orientation_goal must be upright.") + if task_type not in {"E1", "E2", "E4"} and orientation_goal != "none": + raise ValueError( + f"{context} orientation_goal is only valid for E1, E2, and E4." + ) + if task_type == "E1" and step["layout"] == "line": + if target_kind != "none": + raise ValueError(f"{context} E1 line arrangement cannot carry a target.") + if step["relation"] != "none": + raise ValueError(f"{context} E1 line arrangement cannot carry a relation.") + elif task_type in {"E1", "E3"}: + if target_kind == "none": + raise _MissingRequiredTargetError( + f"{context} {task_type} requires a target selector." + ) + if step["relation"] == "none" and task_type == "E3": + raise ValueError(f"{context} {task_type} requires a symbolic relation.") + elif task_type == "E5": + direction = str(step["direction"]) + terminal = str(step["terminal_behavior"]) + if terminal not in _TERMINAL_BEHAVIORS - {"none"}: + raise ValueError(f"{context} E5 requires terminal_behavior hold/place.") + if target_kind == "none": + if step["relation"] != "none": + raise ValueError(f"{context} E5 relation requires a target selector.") + if direction == "none" and terminal != "place": + raise ValueError( + f"{context} E5 requires a direction or target relation." + ) + else: + if step["relation"] == "none": + raise ValueError(f"{context} E5 target requires a relation.") + if direction != "none": + raise ValueError( + f"{context} E5 target relation cannot also carry direction." + ) + elif target_kind != "none": + raise ValueError(f"{context} {task_type} does not accept a target selector.") + if task_type != "E5": + if step["direction"] != "none": + raise ValueError(f"{context} direction is only valid for E5.") + if step["terminal_behavior"] != "none": + raise ValueError(f"{context} terminal_behavior is only valid for E5.") + if task_type == "E4": + transfer = str(step["transfer_arm"]) + receive = str(step["receive_arm"]) + if transfer not in {"left_arm", "right_arm"} or receive not in { + "left_arm", + "right_arm", + }: + raise ValueError(f"{context} E4 requires two explicit arms.") + if transfer == receive: + raise ValueError(f"{context} E4 transfer and receive arms must differ.") + if step["required_arm"] not in {"none", "auto"}: + raise ValueError( + f"{context} E4 uses transfer_arm/receive_arm, not required_arm." + ) + if task_type == "E5" and step["required_arm"] not in {"none", "auto"}: + raise ValueError(f"{context} E5 always uses both arms, not required_arm.") + if task_type == "E6" and step["target_state"] != "open": + raise ValueError(f"{context} E6 target_state must be open.") + if task_type == "E7" and step["target_state"] != "closed": + raise ValueError(f"{context} E7 target_state must be closed.") + if task_type == "E9" and step["target_state"] != "activated": + raise ValueError(f"{context} E9 target_state must be activated.") + if step["layout"] == "line" and task_type != "E1": + raise ValueError(f"{context} only E1 supports layout=line.") + + +def _instruction_prompt(instruction: str) -> str: + return ( + "Convert the user's explicit L1-L3 instruction into typed E1-E9 task " + "intent. Understand synonyms, ellipsis, and pronouns, but " + "do not invent missing objects. Use step_result for cross-step pronouns " + "and explicit references to the result of an earlier manipulation. Keep " + "an independently selected repeated noun as scene_ref; identical text " + "alone does not prove object identity. " + "Object directions are robot-relative; arm names are robot body sides. " + "Preserve each concrete object or target phrase from the instruction as " + "an open scene_ref.reference. Do not classify it or emit a scene UID. " + "Emit no AtomicAction, category label, affordance, coordinates, poses, " + "paths, or reasoning. Encode explicit ordering with depends_on; same-action set " + "members may remain independent. Use empty strings and 'none' for " + "inapplicable required fields. A request to retract the transfer arm " + "immediately after an E4 handover is a mandatory runtime retreat/home " + "barrier for that E4; do " + "not emit a separate task step for it. The exact output keys are steps -> id, " + "task_type, object, target, relation, required_arm, transfer_arm, " + "receive_arm, orientation_goal, target_state, target_setting, layout, " + "axis, direction, terminal_behavior, depends_on; each selector has kind, " + "step_id, reference, quantifier, count.\n\n" + "Use orientation_goal=none unless the instruction explicitly requests " + "upright orientation or preserving the original orientation. Spatial " + "placement and handover alone do not imply preserve. " + "Emptying, dumping, or pouring contents from one container into another " + "is exactly one E3 step: object selects the source container, target " + "selects the receiving container, and relation=above. Pickup and staging " + "are internal to E3; do not emit a separate E1 for an explicit grab. " + "Opening or pulling out a drawer is E6 with object selecting that drawer " + "and target_state=open. Closing or pushing in a drawer is E7 with object " + "selecting that drawer and target_state=closed. " + f"Instruction:\n{instruction}\n\n" + f"E1-E9 catalog:\n{json.dumps(_intent_capability_catalog(), ensure_ascii=False, sort_keys=True)}\n\n" + "Shape-only complete JSON example (do not copy its step count or values; " + "copy every key, including keys whose value is none/empty/0):\n" + f"{json.dumps(_instruction_shape_example(), ensure_ascii=False, sort_keys=True)}\n\n" + "Selector kind rules (these are not extra output fields):\n" + f"{_instruction_selector_rules()}\n\n" + "For E5, use target+relation for moving an object relative to another " + "object, or direction for a small robot-relative move. A dual-arm pick, " + "lift, raise, or hold request without another target uses direction=up " + "and terminal_behavior=hold. Use hold unless the instruction explicitly " + "says to put/release the object. For pick " + "and release at the original location, use direction=none and place. A dual-arm " + "pick/move/transport request is E5, not E1. Final checklist: every step " + "has all 16 step keys; every object and target " + "has all 5 selector keys. For an inapplicable field use the canonical " + "default shown in the example, never omit the field. E4 must explicitly " + "state transfer_arm and receive_arm. E1/E3 must explicitly state target " + "and relation (except E1 layout=line)." + ) + + +def _instruction_shape_example() -> dict[str, Any]: + """Return a compact field-complete example for providers with weak schemas.""" + selector = { + "kind": "scene_ref", + "step_id": "", + "reference": "example object A", + "quantifier": "one", + "count": 0, + } + empty_selector = { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + return { + "steps": [ + { + "id": "step_1", + "task_type": "E2", + "object": selector, + "target": empty_selector, + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + ] + } + + +def _instruction_selector_rules() -> str: + """Return the mutually exclusive selector encodings for model prompts.""" + step_result = { + "kind": "step_result", + "step_id": "step_1", + "reference": "", + "quantifier": "one", + "count": 0, + } + return ( + "- kind=none: step_id and reference are empty strings; " + "quantifier='one'; count=0.\n" + "- kind=scene_ref: step_id is empty and reference preserves the concrete " + "object phrase from the user's instruction. Repeated scene_ref text does " + "not establish cross-step identity.\n" + "- kind=step_result: use it only for a pronoun that means exactly one " + "object, or an explicit continuation of the result of an earlier " + "instruction step. Set step_id to that prior " + "step ID and set reference='', quantifier='one', count=0. Do not copy " + "the prior object's phrase into this selector. Replace step_1 in this " + f"complete shape with the actual prior step ID: {json.dumps(step_result, sort_keys=True)}\n" + "A step_result may identify only a prior step_id; it cannot carry any " + "other object constraint." + ) + + +def _instruction_repair_guidance(error: Exception) -> str: + """Add narrow semantic guidance for errors weak JSON-mode models repeat.""" + if "E4 transfer and receive arms must differ" in str(error): + return ( + "\nSame-arm handover repair rule: transfer_arm and receive_arm must " + "name different arms. Preserve the explicitly stated transfer arm. " + "When a later clause clearly continues with the handed object using " + "the other arm, use that arm as receive_arm. Resolve coreference from " + "the instruction semantics; identical scene_ref text alone does not " + "prove that two independently selected objects are the same.\n" + ) + if isinstance(error, _MissingRequiredObjectError): + return ( + "\nMissing-object repair rule: preserve the selected task_type and " + "set object to a scene_ref that preserves the explicit manipulated " + "object phrase from the instruction. For E6/E7 the drawer, door, or " + "other articulated part is the object selector; target remains none.\n" + ) + if not isinstance(error, _MissingRequiredTargetError): + return "" + if " E3 requires a target selector" in str(error): + return ( + "\nMissing-target repair rule for E3: keep task_type=E3. object is " + "the source container whose contents are poured, target is the " + "receiving container, and relation must be above. An explicit grab " + "is part of E3 and must not be reclassified as E1.\n" + ) + return ( + "\nMissing-target repair rule: for a non-line E1 placement, object is " + "the item being moved and target is the explicit reference object " + "after the spatial relation in the original instruction. For example, " + "in 'place it to the left of the striped pedestal', object is the earlier " + "step_result for 'it', while target selects the striped pedestal; target " + "must not use kind=none. Use target kind=step_result only when the " + "reference object itself is exactly the result of a prior step.\n" + ) + + +def _intent_capability_catalog() -> dict[str, dict[str, Any]]: + """Return the LLM's thin, import-safe E1-E9 capability view. + + Action Engine's online planning catalog also reports runtime availability + and therefore imports simulator action classes. Text interpretation only + needs symbolic E semantics and must remain testable before a simulator + backend is installed. + """ + return { + task_type: { + "semantics": contract.semantics, + "applicable_fields": sorted(_INTENT_TASK_FIELD_REGISTRY[task_type]), + } + for task_type, contract in TASK_CONTRACTS.items() + } + + +def _default_instruction_caller( + *, + prompt: str, + schema: Mapping[str, Any], + model: str | None, +) -> Mapping[str, Any]: + from langchain_core.messages import HumanMessage, SystemMessage + from langchain_openai import ChatOpenAI + + settings = _load_llm_settings(model=model) + kwargs: dict[str, Any] = { + "api_key": settings["api_key"], + "model": settings["model"], + "temperature": 0, + "http_socket_options": (), + } + for key in ("base_url", "default_query"): + if settings[key]: + kwargs[key] = settings[key] + if _is_mimo_compatible(settings): + # MiMo documents ``thinking`` as a provider extension carried in the + # OpenAI client's extra body. Disabling it is important here: hidden + # reasoning can consume the completion and leave only id/object/type. + kwargs.update( + { + "max_completion_tokens": _MIMO_MAX_COMPLETION_TOKENS, + "extra_body": {"thinking": {"type": "disabled"}}, + } + ) + client = ChatOpenAI(**kwargs) + # The full schema remains in the prompt and the local validator is still + # authoritative even when the provider only offers JSON mode. + structured = _structured_output_runnable( + client, + schema, + settings=settings, + ) + schema_prompt = ( + f"{prompt}\n\nReturn one JSON object conforming exactly to this JSON " + f"Schema:\n{json.dumps(schema, ensure_ascii=False, sort_keys=True)}" + ) + response = structured.invoke( + [ + SystemMessage( + content=( + "Return only the requested structured JSON response. Never " + "return reasoning, coordinates, or AtomicAction nodes." + ) + ), + HumanMessage(content=schema_prompt), + ] + ) + return _coerce_instruction_response(response) + + +def _instruction_model(explicit: str | None) -> str | None: + if isinstance(explicit, str) and explicit.strip(): + return explicit.strip() + # Keep model selection separate from credential loading. Reading the local + # dotenv file is side-effect free and gives generation the documented + # priority without leaking credentials into TaskSpec metadata. + for name in ("TASK_ENGINE_LLM_MODEL", "ACTION_ENGINE_LLM_MODEL", "OPENAI_MODEL"): + for source in ( + os.environ, + _load_local_env(), + ): + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _load_local_env() -> dict[str, str]: + """Read Task Engine model configuration without mutating the environment.""" + return _load_env_file(_GEN_SIM_ENV_PATH) + + +def _is_mimo_compatible(settings: Mapping[str, Any]) -> bool: + model = str(settings.get("model", "")).casefold() + base_url = str(settings.get("base_url", "")).casefold() + return "mimo" in model or "xiaomimimo.com" in base_url + + +def _structured_output_runnable( + client: Any, + schema: Mapping[str, Any], + *, + settings: Mapping[str, Any], +) -> Any: + if not hasattr(client, "with_structured_output"): + return client + method = "json_mode" if _is_mimo_compatible(settings) else "json_schema" + try: + return client.with_structured_output(schema, method=method) + except (TypeError, ValueError): + if method == "json_mode" and hasattr(client, "bind"): + from langchain_core.output_parsers import JsonOutputParser + + return ( + client.bind(response_format={"type": "json_object"}) + | JsonOutputParser() + ) + return client.with_structured_output(schema) + + +def _load_llm_settings(*, model: str | None) -> dict[str, Any]: + local_env = _load_local_env() + config: dict[str, Any] = {} + if _GEN_CONFIG_PATH.is_file(): + raw = json.loads(_GEN_CONFIG_PATH.read_text(encoding="utf-8")) + if isinstance(raw, Mapping): + llm = raw.get("llm", {}) + if isinstance(llm, Mapping): + configured = llm.get("openai_compatible", {}) + if isinstance(configured, Mapping): + config = dict(configured) + api_key, base_url = _resolve_transport_settings(local_env, config) + selected_model = ( + (model.strip() if isinstance(model, str) else "") + or _first_env_value( + local_env, + "TASK_ENGINE_LLM_MODEL", + "ACTION_ENGINE_LLM_MODEL", + "OPENAI_MODEL", + "LLM_MODEL", + ) + or str(config.get("model", "")).strip() + ) + default_query = config.get("default_query", {}) or {} + if not api_key: + raise ValueError( + "OPENAI_API_KEY is required for Task Engine interpretation. Set it " + f"in the process environment or {_GEN_SIM_ENV_PATH}." + ) + if not selected_model: + raise ValueError( + "A text LLM model is required through model=, TASK_ENGINE_LLM_MODEL, " + f"OPENAI_MODEL, or {_GEN_CONFIG_PATH}." + ) + if not isinstance(default_query, Mapping): + raise ValueError("LLM default_query must be a mapping.") + return { + "api_key": api_key, + "model": selected_model, + "base_url": base_url, + "default_query": dict(default_query), + } + + +def _resolve_transport_settings( + local_env: Mapping[str, str], + config: Mapping[str, Any], +) -> tuple[str, str]: + """Resolve an API key and endpoint from one configuration source.""" + transports = ( + ( + _mapping_value(os.environ, "OPENAI_API_KEY"), + _mapping_value( + os.environ, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(local_env, "OPENAI_API_KEY"), + _mapping_value( + local_env, + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + ), + ), + ( + _mapping_value(config, "api_key"), + _mapping_value(config, "base_url"), + ), + ) + for api_key, base_url in transports: + if api_key and base_url: + return api_key, base_url.rstrip("/") + for api_key, base_url in transports: + if api_key: + return api_key, base_url.rstrip("/") + return "", "" + + +def _mapping_value(source: Mapping[str, Any], *names: str) -> str: + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return "" + + +def _load_env_file(path: Path) -> dict[str, str]: + if not path.is_file(): + return {} + values: dict[str, str] = {} + for line_number, raw_line in enumerate( + path.read_text(encoding="utf-8").splitlines(), start=1 + ): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + if line.startswith("export "): + line = line[len("export ") :].lstrip() + if "=" not in line: + continue + key, raw_value = line.split("=", 1) + key = key.strip() + if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", key): + raise ValueError(f"Invalid dotenv key at {path}:{line_number}.") + value = raw_value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + value = value[1:-1] + elif " #" in value: + value = value.split(" #", 1)[0].rstrip() + values[key] = value + return values + + +def _first_env_value(local_env: Mapping[str, str], *names: str) -> str | None: + for source in (os.environ, local_env): + for name in names: + value = source.get(name) + if isinstance(value, str) and value.strip(): + return value.strip() + return None + + +def _choice(value: Any, allowed: set[str] | frozenset[str], context: str) -> str: + if not isinstance(value, str) or value not in allowed: + raise ValueError(f"{context} must be one of {sorted(allowed)}.") + return value + + +def _selector_string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{context} must be a string.") + return value.strip() + + +def _canonical_quantifier(value: Any, context: str) -> str: + return _choice(value, _QUANTIFIERS, context) + + +def _canonical_arm(value: Any, context: str) -> str: + return _choice(value, _ARMS, context) + + +def _canonical_relation(value: Any, context: str) -> str: + return _choice(value, _RELATIONS, context) + + +def _canonical_orientation(value: Any, context: str) -> str: + return _choice(value, _ORIENTATIONS, context) + + +def _nonempty(value: Any, context: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{context} must be a non-empty string.") + return value.strip() + + +def _topological_steps(steps: Sequence[Mapping[str, Any]]) -> list[dict[str, Any]]: + """Return a stable topological ordering for validated intent steps.""" + by_id = {str(step["id"]): dict(step) for step in steps} + effective_dependencies: dict[str, tuple[str, ...]] = {} + for step_id, step in by_id.items(): + deps = list(str(dep) for dep in step["depends_on"]) + for selector in (step["object"], step["target"]): + if selector["kind"] == "step_result": + reference = str(selector["step_id"]) + if reference not in deps: + deps.append(reference) + effective_dependencies[step_id] = tuple(deps) + pending = set(by_id) + ordered: list[dict[str, Any]] = [] + original = [str(step["id"]) for step in steps] + while pending: + ready = [ + step_id + for step_id in original + if step_id in pending + and all(str(dep) not in pending for dep in effective_dependencies[step_id]) + ] + if not ready: + raise ValueError("Instruction intent dependencies contain a cycle.") + for step_id in ready: + ordered.append(by_id[step_id]) + pending.remove(step_id) + return ordered + + +def _coerce_instruction_response(response: Any) -> Mapping[str, Any]: + """Coerce common structured-client response wrappers without accepting prose.""" + if isinstance(response, Mapping): + return dict(response) + model_dump = getattr(response, "model_dump", None) + if callable(model_dump): + dumped = model_dump() + if isinstance(dumped, Mapping): + return dict(dumped) + content = getattr(response, "content", response) + if isinstance(content, Mapping): + return dict(content) + if isinstance(content, list): + content = "\n".join( + str(item.get("text", "")) + for item in content + if isinstance(item, Mapping) and item.get("type") == "text" + ) + if not isinstance(content, str): + raise ValueError( + f"Instruction model output has unsupported type {type(content).__name__}." + ) + text = content.strip() + if text.startswith("```"): + lines = text.splitlines() + if lines: + lines = lines[1:] + if lines and lines[-1].strip().startswith("```"): + lines = lines[:-1] + text = "\n".join(lines).strip() + try: + parsed = json.loads(text) + except json.JSONDecodeError as exc: + raise ValueError(f"Instruction model output is not valid JSON: {exc}") from exc + if not isinstance(parsed, Mapping): + raise ValueError("Instruction model output must decode to a JSON object.") + return dict(parsed) + + +def _validate_dag(dependencies: Mapping[str, Sequence[str]]) -> None: + visiting: set[str] = set() + visited: set[str] = set() + + def visit(node: str) -> None: + if node in visiting: + raise ValueError("Instruction intent dependencies contain a cycle.") + if node in visited: + return + visiting.add(node) + for dependency in dependencies[node]: + visit(str(dependency)) + visiting.remove(node) + visited.add(node) + + for node in dependencies: + visit(node) + + +def _reject_forbidden_fields(value: Any) -> None: + if isinstance(value, Mapping): + forbidden = _FORBIDDEN_FIELDS & {str(key).strip().lower() for key in value} + if forbidden: + raise ValueError( + f"Instruction intent contains forbidden fields {sorted(forbidden)}." + ) + for item in value.values(): + _reject_forbidden_fields(item) + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + for item in value: + _reject_forbidden_fields(item) diff --git a/embodichain/gen_sim/task_engine/ontology.py b/embodichain/gen_sim/task_engine/ontology.py new file mode 100644 index 000000000..eb3d090a1 --- /dev/null +++ b/embodichain/gen_sim/task_engine/ontology.py @@ -0,0 +1,221 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Scene-independent semantic ontology for the canonical E1-E9 tasks.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any + +__all__ = [ + "RELATIONS", + "TASK_CONTRACTS", + "TERMINAL_BEHAVIORS", + "TRANSPORT_DIRECTIONS", + "TaskContract", + "task_contract", + "task_success_type", +] + + +# These are protocol values consumed by executable planners. They are not a +# vocabulary for matching words in user instructions. +RELATIONS = frozenset( + { + "none", + "on", + "inside", + "above", + "left_of", + "right_of", + "front_of", + "behind", + "front_left_of", + "front_right_of", + "back_left_of", + "back_right_of", + } +) +TRANSPORT_DIRECTIONS = frozenset( + { + "none", + "world_x", + "world_y", + "front", + "back", + "left", + "right", + "front_left", + "front_right", + "back_left", + "back_right", + "up", + "down", + } +) +TERMINAL_BEHAVIORS = frozenset({"none", "hold", "place"}) + + +@dataclass(frozen=True, slots=True) +class TaskContract: + """One scene-independent semantic E-task contract.""" + + task_type: str + semantics: str + applicable_intent_fields: frozenset[str] + source_structure: str + required_affordances: frozenset[str] + success_type: str + scene_affordances: frozenset[str] + + +def _contract( + task_type: str, + semantics: str, + applicable_intent_fields: frozenset[str], + source_structure: str, + required_affordances: frozenset[str], + success_type: str, + *, + scene_affordances: frozenset[str] | None = None, +) -> TaskContract: + return TaskContract( + task_type=task_type, + semantics=semantics, + applicable_intent_fields=applicable_intent_fields, + source_structure=source_structure, + required_affordances=required_affordances, + success_type=success_type, + scene_affordances=scene_affordances or required_affordances, + ) + + +TASK_CONTRACTS: Mapping[str, TaskContract] = MappingProxyType( + { + "E1": _contract( + "E1", + "Pick, move, and place one object at a symbolic relation.", + frozenset( + { + "target", + "relation", + "required_arm", + "orientation_goal", + "layout", + "axis", + } + ), + "rigid_object", + frozenset({"graspable", "placeable"}), + "semantic_goal", + ), + "E2": _contract( + "E2", + "Make one fallen object upright and place it stably.", + frozenset({"required_arm", "orientation_goal"}), + "rigid_object", + frozenset({"graspable", "orientable"}), + "object_upright", + ), + "E3": _contract( + "E3", + "Pick up when needed and pour contents from a source container " + "into a target container.", + frozenset({"target", "relation", "required_arm"}), + "rigid_object", + frozenset({"graspable", "pourable"}), + "poured", + ), + "E4": _contract( + "E4", + "Transfer one held object from one arm to the other.", + frozenset({"transfer_arm", "receive_arm", "orientation_goal"}), + "rigid_object", + frozenset({"graspable", "handover"}), + "handover_complete", + ), + "E5": _contract( + "E5", + "Use both arms to pick, move, and optionally release one shared rigid object.", + frozenset({"target", "relation", "direction", "terminal_behavior"}), + "rigid_object", + frozenset({"dual_graspable"}), + "held_by_both_grippers", + scene_affordances=frozenset({"dual_graspable", "rigid"}), + ), + "E6": _contract( + "E6", + "Pull an articulated part to its requested state.", + frozenset({"required_arm", "target_state"}), + "articulation", + frozenset({"pullable"}), + "articulation_joint_near", + scene_affordances=frozenset({"articulated", "pullable"}), + ), + "E7": _contract( + "E7", + "Push an articulated part to its requested state.", + frozenset({"required_arm", "target_state"}), + "articulation", + frozenset({"pushable"}), + "articulation_joint_near", + scene_affordances=frozenset({"articulated", "pushable"}), + ), + "E8": _contract( + "E8", + "Turn one knob to a requested setting.", + frozenset({"required_arm", "target_setting"}), + "articulation", + frozenset({"turnable"}), + "articulation_joint_near", + ), + "E9": _contract( + "E9", + "Press one button until its requested terminal state.", + frozenset({"required_arm", "target_state"}), + "articulation", + frozenset({"pressable"}), + "pressed", + ), + } +) + + +def task_contract(task_type: str) -> TaskContract: + """Return the canonical contract or reject an unknown E-task type.""" + try: + return TASK_CONTRACTS[str(task_type)] + except KeyError as exc: + raise ValueError(f"Unsupported task type {task_type!r}.") from exc + + +def task_success_type( + task_type: str, + params: Mapping[str, Any] | None = None, +) -> str: + """Resolve a TaskSpec success type, including E5's terminal behavior.""" + contract = task_contract(task_type) + if contract.task_type != "E5": + return contract.success_type + terminal_behavior = str((params or {}).get("terminal_behavior", "hold")) + if terminal_behavior == "hold": + return "held_by_both_grippers" + if terminal_behavior == "place": + return "semantic_goal" + raise ValueError("E5 terminal_behavior must be 'hold' or 'place'.") diff --git a/embodichain/gen_sim/task_engine/orchestration/__init__.py b/embodichain/gen_sim/task_engine/orchestration/__init__.py new file mode 100644 index 000000000..ed3a9539c --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/__init__.py @@ -0,0 +1,109 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Task-owned orchestration across task, scene, and action engines.""" + +from __future__ import annotations + +from embodichain.gen_sim.action_engine.agent import ActionAgent, ActionGraph +from embodichain.gen_sim.action_engine.runtime import ExecutionReport + +from .artifacts import ( + ArtifactTransaction, + CONSERVATIVE_SCENE_GRAPH_FILENAME, + TaskEngineArtifactPaths, + FEASIBILITY_REPORT_FILENAME, + PREPARATION_FAILURE_FILENAME, + STATIC_SCENE_MANIFEST_FILENAME, + task_engine_artifact_paths, + write_execution_report, + write_preparation_failure, +) +from .contracts import ( + BINDING_REPORT_SCHEMA, + EXECUTION_REPORT_SCHEMA, + GROUNDED_TASK_PLAN_SCHEMA, + ROLE_BINDINGS_SCHEMA, + SCENE_MANIFEST_SCHEMA, + BindingReport, + GroundedTaskPlan, + RoleBindings, + SceneManifest, +) +from .coordinator import ( + TaskEngineCoordinator, + PreparationResult, + build_grounded_task_plan, + lower_task_candidate, +) +from .scene_adapter import ( + CandidateSelection, + SceneAdaptation, + SceneAdapter, + SceneAdapterProtocolError, +) +from .scene_source import ( + SceneSourceFingerprint, + SceneSourceRef, + fingerprint_scene_source, + verify_scene_source_fingerprint, +) +from .legacy_scene import ( + LEGACY_SCENE_CONVERSION_SCHEMA, + LegacySceneRevision, + convert_legacy_gym_project, + restore_locked_scene_entities, +) + +__all__ = [ + "ActionAgent", + "ActionGraph", + "ArtifactTransaction", + "CONSERVATIVE_SCENE_GRAPH_FILENAME", + "BINDING_REPORT_SCHEMA", + "BindingReport", + "TaskEngineArtifactPaths", + "TaskEngineCoordinator", + "EXECUTION_REPORT_SCHEMA", + "ExecutionReport", + "FEASIBILITY_REPORT_FILENAME", + "GROUNDED_TASK_PLAN_SCHEMA", + "GroundedTaskPlan", + "PREPARATION_FAILURE_FILENAME", + "PreparationResult", + "ROLE_BINDINGS_SCHEMA", + "RoleBindings", + "SCENE_MANIFEST_SCHEMA", + "STATIC_SCENE_MANIFEST_FILENAME", + "SceneAdaptation", + "CandidateSelection", + "SceneAdapter", + "SceneAdapterProtocolError", + "SceneManifest", + "SceneSourceFingerprint", + "SceneSourceRef", + "build_grounded_task_plan", + "task_engine_artifact_paths", + "fingerprint_scene_source", + "LEGACY_SCENE_CONVERSION_SCHEMA", + "LegacySceneRevision", + "convert_legacy_gym_project", + "restore_locked_scene_entities", + "verify_scene_source_fingerprint", + "lower_task_candidate", + "write_execution_report", + "write_preparation_failure", +] diff --git a/embodichain/gen_sim/task_engine/orchestration/artifacts.py b/embodichain/gen_sim/task_engine/orchestration/artifacts.py new file mode 100644 index 000000000..b941c86dc --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/artifacts.py @@ -0,0 +1,331 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Transactional publication for Task Engine artifacts.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +import json +import os +from pathlib import Path +import shutil +import tempfile +from typing import Any + +from embodichain.gen_sim.action_engine.runtime import ( + EXECUTION_REPORT_FILENAME, + write_execution_report as _write_execution_report, +) + +__all__ = [ + "BINDING_REPORT_FILENAME", + "CONSERVATIVE_SCENE_GRAPH_FILENAME", + "EXECUTION_REPORT_FILENAME", + "GROUNDED_TASK_PLAN_FILENAME", + "FEASIBILITY_REPORT_FILENAME", + "FINAL_SCENE_INSPECTION_FILENAME", + "PREPARATION_FAILURE_FILENAME", + "ROLE_BINDINGS_FILENAME", + "SCENE_MANIFEST_FILENAME", + "STATIC_SCENE_MANIFEST_FILENAME", + "SUCCESS_SPEC_FILENAME", + "TASK_CANDIDATE_SET_FILENAME", + "TASK_DRAFT_FILENAME", + "SCENE_REQUEST_FILENAME", + "ArtifactTransaction", + "TaskEngineArtifactPaths", + "task_engine_artifact_paths", + "write_task_engine_artifacts", + "write_execution_report", + "write_preparation_failure", +] + + +TASK_CANDIDATE_SET_FILENAME = "task_candidate_set.json" +TASK_DRAFT_FILENAME = "task_draft.json" +SCENE_REQUEST_FILENAME = "scene_request.json" +SUCCESS_SPEC_FILENAME = "success_spec.json" +SCENE_MANIFEST_FILENAME = "scene_manifest.json" +STATIC_SCENE_MANIFEST_FILENAME = "static_scene_manifest.json" +CONSERVATIVE_SCENE_GRAPH_FILENAME = "conservative_scene_graph.json" +ROLE_BINDINGS_FILENAME = "role_bindings.json" +BINDING_REPORT_FILENAME = "binding_report.json" +FEASIBILITY_REPORT_FILENAME = "feasibility_report.json" +FINAL_SCENE_INSPECTION_FILENAME = "final_scene_inspection.json" +GROUNDED_TASK_PLAN_FILENAME = "grounded_task_plan.json" +PREPARATION_FAILURE_FILENAME = "preparation_failure.json" + + +@dataclass(frozen=True) +class TaskEngineArtifactPaths: + """Canonical Task Engine paths rooted at one published bundle.""" + + root: Path + task_candidate_set: Path + task_draft: Path + scene_request: Path + success_spec: Path + scene_manifest: Path + static_scene_manifest: Path + conservative_scene_graph: Path + role_bindings: Path + binding_report: Path + feasibility_report: Path + final_scene_inspection: Path + grounded_task_plan: Path + preparation_failure: Path + execution_report: Path + + +def task_engine_artifact_paths( + output_dir: str | Path, +) -> TaskEngineArtifactPaths: + """Return all Task Engine paths without creating the directory.""" + root = Path(output_dir).expanduser().resolve() + return TaskEngineArtifactPaths( + root=root, + task_candidate_set=root / TASK_CANDIDATE_SET_FILENAME, + task_draft=root / TASK_DRAFT_FILENAME, + scene_request=root / SCENE_REQUEST_FILENAME, + success_spec=root / SUCCESS_SPEC_FILENAME, + scene_manifest=root / SCENE_MANIFEST_FILENAME, + static_scene_manifest=root / STATIC_SCENE_MANIFEST_FILENAME, + conservative_scene_graph=root / CONSERVATIVE_SCENE_GRAPH_FILENAME, + role_bindings=root / ROLE_BINDINGS_FILENAME, + binding_report=root / BINDING_REPORT_FILENAME, + feasibility_report=root / FEASIBILITY_REPORT_FILENAME, + final_scene_inspection=root / FINAL_SCENE_INSPECTION_FILENAME, + grounded_task_plan=root / GROUNDED_TASK_PLAN_FILENAME, + preparation_failure=root / PREPARATION_FAILURE_FILENAME, + execution_report=root / EXECUTION_REPORT_FILENAME, + ) + + +class ArtifactTransaction: + """Build a complete bundle beside its destination and publish it by rename.""" + + def __init__(self, output_dir: str | Path, *, overwrite: bool = False) -> None: + raw = Path(output_dir).expanduser() + self.output_dir = ( + (Path.cwd() / raw).resolve() if not raw.is_absolute() else raw.resolve() + ) + self.overwrite = bool(overwrite) + self.staging_dir: Path | None = None + self._committed = False + + def __enter__(self) -> "ArtifactTransaction": + destination = self.output_dir + destination.parent.mkdir(parents=True, exist_ok=True) + if destination.exists() and not self.overwrite: + raise FileExistsError( + f"Output directory already exists: {destination}. " + "Pass overwrite=True to replace it." + ) + self.staging_dir = Path( + tempfile.mkdtemp( + prefix=f".{destination.name}.staging-", + dir=destination.parent, + ) + ) + return self + + def commit(self) -> Path: + """Rewrite staging-local absolute paths, then atomically publish.""" + if self.staging_dir is None: + raise RuntimeError("ArtifactTransaction has not been entered.") + if self._committed: + raise RuntimeError("ArtifactTransaction has already been committed.") + staging = self.staging_dir + destination = self.output_dir + _relocate_json_paths(staging, destination) + + backup: Path | None = None + if destination.exists(): + if not self.overwrite: + raise FileExistsError( + f"Output directory already exists: {destination}." + ) + backup = Path( + tempfile.mkdtemp( + prefix=f".{destination.name}.backup-", + dir=destination.parent, + ) + ) + backup.rmdir() + os.replace(destination, backup) + try: + os.replace(staging, destination) + except BaseException: + if backup is not None and backup.exists() and not destination.exists(): + os.replace(backup, destination) + raise + else: + self._committed = True + self.staging_dir = None + if backup is not None: + _remove_path(backup) + _fsync_directory(destination.parent) + return destination + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> bool: + if self.staging_dir is not None and self.staging_dir.exists(): + shutil.rmtree(self.staging_dir) + return False + + +def write_task_engine_artifacts( + output_dir: str | Path, + *, + candidate_set: Mapping[str, Any], + scene_manifest: Mapping[str, Any] | None, + role_bindings: Mapping[str, Any] | None, + binding_report: Mapping[str, Any], + grounded_task_plan: Mapping[str, Any] | None = None, + static_scene_manifest: Mapping[str, Any] | None = None, + conservative_scene_graph: Mapping[str, Any] | None = None, + feasibility_report: Mapping[str, Any] | None = None, + final_scene_inspection: Mapping[str, Any] | None = None, +) -> TaskEngineArtifactPaths: + """Write Task Engine protocols into an unpublished staging directory. + + An unsuccessful adaptation can omit SceneManifest and RoleBindings rather + than publishing protocol filenames whose payloads do not satisfy their + schemas. + """ + paths = task_engine_artifact_paths(output_dir) + paths.root.mkdir(parents=True, exist_ok=True) + _write_json(paths.task_candidate_set, candidate_set) + if scene_manifest is not None: + _write_json(paths.scene_manifest, scene_manifest) + if static_scene_manifest is not None: + _write_json(paths.static_scene_manifest, static_scene_manifest) + if conservative_scene_graph is not None: + _write_json(paths.conservative_scene_graph, conservative_scene_graph) + if role_bindings is not None: + _write_json(paths.role_bindings, role_bindings) + _write_json(paths.binding_report, binding_report) + if feasibility_report is not None: + _write_json(paths.feasibility_report, feasibility_report) + if final_scene_inspection is not None: + _write_json(paths.final_scene_inspection, final_scene_inspection) + + if grounded_task_plan is not None: + _write_json(paths.grounded_task_plan, grounded_task_plan) + _write_json(paths.task_draft, grounded_task_plan["task_draft"]) + candidate_id = grounded_task_plan["selected_candidate_id"] + selected = next( + candidate + for candidate in candidate_set["candidates"] + if candidate["candidate_id"] == candidate_id + ) + _write_json(paths.scene_request, selected["scene_request"]) + _write_json(paths.success_spec, grounded_task_plan["success_spec"]) + return paths + + +def write_execution_report(output_dir: str | Path, value: Any) -> Path: + """Publish through the Action Engine-owned report boundary.""" + return _write_execution_report(output_dir, value) + + +def write_preparation_failure(output_dir: str | Path, value: Any) -> Path: + """Write a strict-JSON audit for a failed candidate planning transaction.""" + path = task_engine_artifact_paths(output_dir).preparation_failure + path.parent.mkdir(parents=True, exist_ok=True) + _write_json(path, value) + return path + + +def _write_json(path: Path, value: Any) -> None: + try: + payload = ( + json.dumps( + value, + ensure_ascii=False, + indent=2, + sort_keys=False, + allow_nan=False, + ) + + "\n" + ).encode("utf-8") + except (TypeError, ValueError) as exc: + raise ValueError(f"Artifact {path.name} is not strict JSON data.") from exc + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as stream: + temporary = Path(stream.name) + stream.write(payload) + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary, path) + temporary = None + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def _relocate_json_paths(staging: Path, destination: Path) -> None: + """Replace staging-root absolute paths embedded by the legacy generator.""" + source_prefix = staging.resolve().as_posix() + destination_prefix = destination.resolve().as_posix() + + def relocate(value: Any) -> Any: + if isinstance(value, str): + if value == source_prefix: + return destination_prefix + if value.startswith(source_prefix + "/"): + return destination_prefix + value[len(source_prefix) :] + return value + if isinstance(value, list): + return [relocate(item) for item in value] + if isinstance(value, dict): + return {key: relocate(item) for key, item in value.items()} + return value + + for path in staging.rglob("*.json"): + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Generated artifact is invalid JSON: {path}") from exc + relocated = relocate(value) + if relocated != value: + _write_json(path, relocated) + + +def _remove_path(path: Path) -> None: + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + + +def _fsync_directory(path: Path) -> None: + try: + descriptor = os.open(path, os.O_RDONLY) + except OSError: + return + try: + os.fsync(descriptor) + finally: + os.close(descriptor) diff --git a/embodichain/gen_sim/task_engine/orchestration/contracts.py b/embodichain/gen_sim/task_engine/orchestration/contracts.py new file mode 100644 index 000000000..5c38f284c --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/contracts.py @@ -0,0 +1,647 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict cross-engine contracts for scene binding and orchestration.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +import math +from typing import Any, TypeAlias + +from embodichain.gen_sim.action_engine.domain import ( + validate_scene_requirements, + validate_task_spec, +) +from embodichain.gen_sim.action_engine.runtime import ( + EXECUTION_REPORT_SCHEMA, + validate_execution_report, +) +from embodichain.gen_sim.task_engine import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + SceneRequest, + SuccessSpec, + TaskCandidate, + TaskCandidateSet, + TaskDraft, + canonical_hash, + task_success_type, + validate_scene_request, + validate_success_spec, + validate_task_candidate, + validate_task_candidate_set, + validate_task_draft, +) + +__all__ = [ + "BINDING_REPORT_SCHEMA", + "EXECUTION_REPORT_SCHEMA", + "GROUNDED_TASK_PLAN_SCHEMA", + "ROLE_BINDINGS_SCHEMA", + "SCENE_MANIFEST_SCHEMA", + "SCENE_REQUEST_SCHEMA", + "SUCCESS_SPEC_SCHEMA", + "TASK_CANDIDATE_SET_SCHEMA", + "TASK_DRAFT_SCHEMA", + "BindingReport", + "ExecutionReport", + "GroundedTaskPlan", + "RoleBindings", + "SceneManifest", + "SceneRequest", + "SuccessSpec", + "TaskCandidate", + "TaskCandidateSet", + "TaskDraft", + "canonical_hash", + "validate_binding_report", + "validate_execution_report", + "validate_grounded_task_plan", + "validate_role_bindings", + "validate_scene_manifest", + "validate_scene_request", + "validate_success_spec", + "validate_task_candidate", + "validate_task_candidate_set", + "validate_task_draft", +] + +SCENE_MANIFEST_SCHEMA = "action_engine_scene_manifest_v1" +ROLE_BINDINGS_SCHEMA = "action_engine_role_bindings_v1" +BINDING_REPORT_SCHEMA = "action_engine_binding_report_v1" +GROUNDED_TASK_PLAN_SCHEMA = "action_engine_grounded_task_plan_v1" +SceneManifest: TypeAlias = dict[str, Any] +RoleBindings: TypeAlias = dict[str, Any] +BindingReport: TypeAlias = dict[str, Any] +GroundedTaskPlan: TypeAlias = dict[str, Any] +ExecutionReport: TypeAlias = dict[str, Any] + + +def validate_scene_manifest(value: Mapping[str, Any]) -> SceneManifest: + result = _mapping(value, "SceneManifest") + _keys( + result, + {"schema_version", "scene_id", "source_format", "robot_profile", "objects"}, + "SceneManifest", + ) + _schema(result, SCENE_MANIFEST_SCHEMA, "SceneManifest") + for key in ("scene_id", "source_format", "robot_profile"): + result[key] = _nonempty(result.get(key), f"SceneManifest.{key}") + object_keys = { + "uid", + "role", + "name", + "description", + "category", + "color", + "affordances", + "initial_state", + "attributes", + } + objects = [] + for index, raw in enumerate( + _sequence(result.get("objects"), "SceneManifest.objects") + ): + context = f"SceneManifest.objects[{index}]" + item = _mapping(raw, context) + _keys(item, object_keys, context) + item["uid"] = _nonempty(item.get("uid"), f"{context}.uid") + for key in ("role", "name", "description", "category"): + item[key] = _string(item.get(key), f"{context}.{key}") + if item.get("color") is not None: + item["color"] = _string(item.get("color"), f"{context}.color") + item["affordances"] = _strings( + item.get("affordances"), f"{context}.affordances" + ) + item["initial_state"] = _mapping( + item.get("initial_state"), f"{context}.initial_state" + ) + item["attributes"] = _mapping(item.get("attributes"), f"{context}.attributes") + objects.append(item) + _unique([item["uid"] for item in objects], "SceneManifest object UIDs") + result["objects"] = objects + _json_safe(result, "SceneManifest") + return result + + +def validate_role_bindings(value: Mapping[str, Any]) -> RoleBindings: + result = _mapping(value, "RoleBindings") + _keys( + result, + { + "schema_version", + "task_id", + "candidate_id", + "reference_bindings", + "role_bindings", + }, + "RoleBindings", + ) + _schema(result, ROLE_BINDINGS_SCHEMA, "RoleBindings") + for key in ("task_id", "candidate_id"): + result[key] = _nonempty(result.get(key), f"RoleBindings.{key}") + result["reference_bindings"] = _string_lists( + result.get("reference_bindings"), "RoleBindings.reference_bindings" + ) + if any(not uids for uids in result["reference_bindings"].values()): + raise ValueError("RoleBindings.reference_bindings values must not be empty.") + result["role_bindings"] = _string_map( + result.get("role_bindings"), "RoleBindings.role_bindings" + ) + return result + + +def validate_binding_report(value: Mapping[str, Any]) -> BindingReport: + result = _mapping(value, "BindingReport") + _keys( + result, + { + "schema_version", + "task_id", + "status", + "selected_candidate_id", + "selection_reason", + "candidates", + }, + "BindingReport", + ) + _schema(result, BINDING_REPORT_SCHEMA, "BindingReport") + result["task_id"] = _nonempty(result.get("task_id"), "BindingReport.task_id") + result["status"] = _enum( + result.get("status"), + {"bound", "ambiguous", "unsatisfied"}, + "BindingReport.status", + ) + result["selected_candidate_id"] = _string( + result.get("selected_candidate_id"), "BindingReport.selected_candidate_id" + ) + result["selection_reason"] = _string( + result.get("selection_reason"), "BindingReport.selection_reason" + ) + if result["status"] == "bound" and not result["selected_candidate_id"]: + raise ValueError("A bound BindingReport requires selected_candidate_id.") + candidate_keys = { + "candidate_id", + "semantic_hash", + "status", + "references", + "reasons", + } + reference_keys = { + "reference_id", + "status", + "confidence", + "candidate_uids", + "selected_uids", + "reasons", + } + candidates = [] + for index, raw in enumerate( + _sequence(result.get("candidates"), "BindingReport.candidates") + ): + context = f"BindingReport.candidates[{index}]" + candidate = _mapping(raw, context) + _keys(candidate, candidate_keys, context) + candidate["candidate_id"] = _nonempty( + candidate.get("candidate_id"), f"{context}.candidate_id" + ) + candidate["semantic_hash"] = _digest( + candidate.get("semantic_hash"), f"{context}.semantic_hash" + ) + candidate["status"] = _enum( + candidate.get("status"), + {"resolved", "ambiguous", "not_found", "incompatible"}, + f"{context}.status", + ) + references = [] + for ref_index, ref_raw in enumerate( + _sequence(candidate.get("references"), f"{context}.references") + ): + ref_context = f"{context}.references[{ref_index}]" + reference = _mapping(ref_raw, ref_context) + _keys(reference, reference_keys, ref_context) + reference["reference_id"] = _nonempty( + reference.get("reference_id"), f"{ref_context}.reference_id" + ) + reference["status"] = _enum( + reference.get("status"), + {"resolved", "ambiguous", "not_found", "incompatible"}, + f"{ref_context}.status", + ) + reference["confidence"] = _number( + reference.get("confidence"), + f"{ref_context}.confidence", + minimum=0.0, + maximum=1.0, + ) + reference["candidate_uids"] = _strings( + reference.get("candidate_uids"), + f"{ref_context}.candidate_uids", + allow_empty=True, + ) + reference["selected_uids"] = _strings( + reference.get("selected_uids"), + f"{ref_context}.selected_uids", + allow_empty=True, + ) + reference["reasons"] = _strings( + reference.get("reasons"), f"{ref_context}.reasons", allow_empty=True + ) + selected = set(reference["selected_uids"]) + candidates_for_reference = set(reference["candidate_uids"]) + if not selected <= candidates_for_reference: + raise ValueError( + f"{ref_context}.selected_uids must be a subset of candidate_uids." + ) + if reference["status"] == "resolved" and not selected: + raise ValueError( + f"{ref_context} status=resolved requires selected_uids." + ) + if reference["status"] != "resolved" and selected: + raise ValueError( + f"{ref_context} non-resolved status cannot select UIDs." + ) + if reference["status"] == "not_found" and candidates_for_reference: + raise ValueError( + f"{ref_context} status=not_found cannot carry candidate_uids." + ) + references.append(reference) + if not references: + raise ValueError(f"{context}.references must not be empty.") + _unique( + [item["reference_id"] for item in references], + f"{context} reference IDs", + ) + expected_status = _candidate_binding_status(references) + if candidate["status"] != expected_status: + raise ValueError( + f"{context}.status must be {expected_status!r} for its references." + ) + candidate["references"] = references + candidate["reasons"] = _strings( + candidate.get("reasons"), f"{context}.reasons", allow_empty=True + ) + candidates.append(candidate) + if not candidates: + raise ValueError("BindingReport.candidates must not be empty.") + _unique( + [item["candidate_id"] for item in candidates], "BindingReport candidate IDs" + ) + if result["selected_candidate_id"] and result["selected_candidate_id"] not in { + item["candidate_id"] for item in candidates + }: + raise ValueError("BindingReport.selected_candidate_id is unknown.") + if result["status"] != "bound" and result["selected_candidate_id"]: + raise ValueError( + "A non-bound BindingReport cannot carry selected_candidate_id." + ) + selected = next( + ( + candidate + for candidate in candidates + if candidate["candidate_id"] == result["selected_candidate_id"] + ), + None, + ) + if result["status"] == "bound" and ( + selected is None or selected["status"] != "resolved" + ): + raise ValueError( + "A bound BindingReport must select a resolved candidate audit." + ) + if result["status"] == "unsatisfied" and any( + candidate["status"] in {"resolved", "ambiguous"} for candidate in candidates + ): + raise ValueError( + "An unsatisfied BindingReport cannot contain resolved or ambiguous candidates." + ) + result["candidates"] = candidates + return result + + +def validate_grounded_task_plan(value: Mapping[str, Any]) -> GroundedTaskPlan: + result = _mapping(value, "GroundedTaskPlan") + keys = { + "schema_version", + "task_id", + "instruction", + "selected_candidate_id", + "task_draft", + "task_spec", + "scene_requirements", + "success_spec", + "scene_manifest", + "role_bindings", + "binding_report", + "hashes", + } + _keys(result, keys, "GroundedTaskPlan") + _schema(result, GROUNDED_TASK_PLAN_SCHEMA, "GroundedTaskPlan") + task_id = _nonempty(result.get("task_id"), "GroundedTaskPlan.task_id") + result["instruction"] = _nonempty( + result.get("instruction"), "GroundedTaskPlan.instruction" + ) + result["selected_candidate_id"] = _nonempty( + result.get("selected_candidate_id"), "GroundedTaskPlan.selected_candidate_id" + ) + result["task_draft"] = validate_task_draft(result.get("task_draft")) + # Candidate SuccessSpec terms use draft step IDs. Once count/all selectors + # are lowered, the grounded plan carries one term per concrete TaskSpec + # instance instead, and is checked against the v2 recipe below. + result["success_spec"] = validate_success_spec(result.get("success_spec")) + result["scene_manifest"] = validate_scene_manifest(result.get("scene_manifest")) + result["role_bindings"] = validate_role_bindings(result.get("role_bindings")) + result["binding_report"] = validate_binding_report(result.get("binding_report")) + result["task_spec"] = validate_task_spec( + _mapping(result.get("task_spec"), "GroundedTaskPlan.task_spec") + ) + result["scene_requirements"] = validate_scene_requirements( + _mapping( + result.get("scene_requirements"), + "GroundedTaskPlan.scene_requirements", + ) + ) + hashes = _mapping(result.get("hashes"), "GroundedTaskPlan.hashes") + _keys( + hashes, + {"task_draft", "task_spec", "scene_manifest", "role_bindings", "plan"}, + "GroundedTaskPlan.hashes", + ) + for key in hashes: + hashes[key] = _digest(hashes[key], f"GroundedTaskPlan.hashes.{key}") + if any( + part["task_id"] != task_id + for part in ( + result["task_draft"], + result["task_spec"], + result["scene_requirements"], + result["success_spec"], + result["role_bindings"], + result["binding_report"], + ) + ): + raise ValueError("GroundedTaskPlan task IDs must agree.") + if result["task_draft"]["instruction"] != result["instruction"]: + raise ValueError("GroundedTaskPlan instruction must match TaskDraft.") + if result["task_spec"]["instruction"] != result["instruction"]: + raise ValueError("GroundedTaskPlan instruction must match TaskSpec.") + if ( + result["role_bindings"]["candidate_id"] != result["selected_candidate_id"] + or result["binding_report"]["selected_candidate_id"] + != result["selected_candidate_id"] + ): + raise ValueError("GroundedTaskPlan selected candidate IDs must agree.") + if result["binding_report"]["status"] != "bound": + raise ValueError("GroundedTaskPlan requires a bound BindingReport.") + selected_audit = next( + candidate + for candidate in result["binding_report"]["candidates"] + if candidate["candidate_id"] == result["selected_candidate_id"] + ) + if selected_audit["semantic_hash"] != canonical_hash(result["task_draft"]["steps"]): + raise ValueError( + "GroundedTaskPlan selected candidate hash must match TaskDraft." + ) + task_metadata = result["task_spec"].get("metadata", {}) + task_oracle = result["task_spec"].get("oracle", {}) + serialized_bindings = ( + task_metadata.get("role_bindings") + if isinstance(task_metadata, Mapping) + and task_metadata.get("role_bindings") is not None + else ( + task_oracle.get("role_bindings") + if isinstance(task_oracle, Mapping) + else None + ) + ) + if serialized_bindings != result["role_bindings"]["role_bindings"]: + raise ValueError( + "GroundedTaskPlan RoleBindings must match the TaskSpec binding hand-off." + ) + requirement_roles = { + str(item["role_id"]) for item in result["scene_requirements"]["objects"] + } + if requirement_roles != set(result["role_bindings"]["role_bindings"]): + raise ValueError( + "GroundedTaskPlan SceneRequirements roles must match RoleBindings." + ) + manifest_uids = {item["uid"] for item in result["scene_manifest"]["objects"]} + missing_uids = sorted( + set(result["role_bindings"]["role_bindings"].values()) - manifest_uids + ) + if missing_uids: + raise ValueError( + f"GroundedTaskPlan RoleBindings reference unknown scene UIDs {missing_uids}." + ) + reference_ids = { + f"{step['id']}.{role}" + for step in result["task_draft"]["steps"] + for role in ("object", "target") + if step[role]["kind"] == "scene_ref" + } + reference_bindings = result["role_bindings"]["reference_bindings"] + if set(reference_bindings) != reference_ids: + raise ValueError( + "GroundedTaskPlan reference bindings must cover every draft scene_ref exactly." + ) + bound_uids = {uid for uids in reference_bindings.values() for uid in uids} | set( + result["role_bindings"]["role_bindings"].values() + ) + unknown_uids = sorted(bound_uids - manifest_uids) + if unknown_uids: + raise ValueError( + "GroundedTaskPlan bindings reference unknown SceneManifest UIDs: " + f"{unknown_uids}." + ) + audited_bindings = { + reference["reference_id"]: reference["selected_uids"] + for reference in selected_audit["references"] + } + if audited_bindings != reference_bindings: + raise ValueError( + "GroundedTaskPlan RoleBindings must match the selected candidate audit." + ) + ontology_success = [ + { + "step_id": instance["id"], + "type": task_success_type(instance["task_type"], instance["params"]), + } + for instance in result["task_spec"]["task_instances"] + ] + recipe_success = [ + { + "step_id": term["task_instance_id"], + "type": term["type"], + } + for term in result["task_spec"]["success"]["terms"] + ] + if recipe_success != ontology_success: + raise ValueError( + "GroundedTaskPlan TaskSpec success recipe must be derived from " + "task_success_type." + ) + if result["success_spec"]["terms"] != ontology_success: + raise ValueError( + "GroundedTaskPlan SuccessSpec must exactly match the lowered " + "TaskSpec success recipe." + ) + expected_hashes = { + "task_draft": canonical_hash(result["task_draft"]), + "task_spec": canonical_hash(result["task_spec"]), + "scene_manifest": canonical_hash(result["scene_manifest"]), + "role_bindings": canonical_hash(result["role_bindings"]), + } + base = {key: value for key, value in result.items() if key != "hashes"} + expected_hashes["plan"] = canonical_hash(base) + if hashes != expected_hashes: + raise ValueError("GroundedTaskPlan hashes do not match their contents.") + result["task_id"] = task_id + result["hashes"] = hashes + _json_safe(result, "GroundedTaskPlan") + return result + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise ValueError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _candidate_binding_status(references: Sequence[Mapping[str, Any]]) -> str: + statuses = {str(reference["status"]) for reference in references} + if statuses == {"resolved"}: + return "resolved" + if "ambiguous" in statuses: + return "ambiguous" + if "incompatible" in statuses: + return "incompatible" + return "not_found" + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise ValueError(f"{context} must be a list.") + return list(value) + + +def _keys( + value: Mapping[str, Any], expected: set[str] | frozenset[str], context: str +) -> None: + if set(value) != set(expected): + raise ValueError( + f"{context} requires exactly fields {sorted(expected)}; received {sorted(value)}." + ) + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise ValueError(f"{context} must be a string.") + return value.strip() + + +def _nonempty(value: Any, context: str) -> str: + result = _string(value, context) + if not result: + raise ValueError(f"{context} must not be empty.") + return result + + +def _enum(value: Any, choices: set[str], context: str) -> str: + result = _string(value, context) + if result not in choices: + raise ValueError(f"{context} must be one of {sorted(choices)}.") + return result + + +def _integer( + value: Any, context: str, *, minimum: int, maximum: int | None = None +) -> int: + if ( + not isinstance(value, int) + or isinstance(value, bool) + or value < minimum + or (maximum is not None and value > maximum) + ): + raise ValueError(f"{context} must be an integer in the allowed range.") + return value + + +def _number(value: Any, context: str, *, minimum: float, maximum: float) -> float: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or not minimum <= float(value) <= maximum + ): + raise ValueError( + f"{context} must be a finite number between {minimum} and {maximum}." + ) + return float(value) + + +def _strings(value: Any, context: str, *, allow_empty: bool = False) -> list[str]: + result = [_string(item, context) for item in _sequence(value, context)] + if not allow_empty and any(not item for item in result): + raise ValueError(f"{context} values must not be empty.") + if len(result) != len(set(result)): + raise ValueError(f"{context} values must be unique.") + return result + + +def _string_map(value: Any, context: str) -> dict[str, str]: + result = _mapping(value, context) + return { + _nonempty(key, context): _nonempty(item, context) + for key, item in result.items() + } + + +def _string_lists(value: Any, context: str) -> dict[str, list[str]]: + result = _mapping(value, context) + return { + _nonempty(key, context): _strings(item, context) for key, item in result.items() + } + + +def _digest(value: Any, context: str) -> str: + result = _string(value, context) + if len(result) != 64 or any( + character not in "0123456789abcdef" for character in result + ): + raise ValueError(f"{context} must be a lowercase SHA-256 digest.") + return result + + +def _unique(values: Sequence[str], context: str) -> None: + if len(values) != len(set(values)): + raise ValueError(f"{context} must be unique.") + + +def _json_safe(value: Any, context: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError) as error: + raise ValueError(f"{context} must be finite and JSON serializable.") from error diff --git a/embodichain/gen_sim/task_engine/orchestration/coordinator.py b/embodichain/gen_sim/task_engine/orchestration/coordinator.py new file mode 100644 index 000000000..dab074430 --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/coordinator.py @@ -0,0 +1,865 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""End-to-end Task Engine preparation for an existing scene source.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, replace +import json +from pathlib import Path +import shutil +from typing import Any + +from embodichain.gen_sim.action_engine.generation import ( + GeneratedConfigPaths, + generate_action_engine_config, +) +from embodichain.gen_sim.action_engine.generation.artifacts import artifact_paths +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError +from embodichain.gen_sim.action_engine.domain.task_contracts import ( + TASK_CONTRACTS as ACTION_TASK_CONTRACTS, +) +from embodichain.gen_sim.action_engine.tasks import ( + GroundedTaskSpec, + ground_instruction_draft, +) +from embodichain.gen_sim.task_engine import ( + TaskAgent, + TaskCandidate, + TaskCandidateSet, + validate_scene_output_separation, + validate_task_candidate, + validate_task_candidate_set, +) +from embodichain.gen_sim.task_engine.scene import FeasibilityBroker, FeasibilityReport + +from .artifacts import ( + ArtifactTransaction, + TaskEngineArtifactPaths, + task_engine_artifact_paths, + write_task_engine_artifacts, + write_preparation_failure, +) +from .contracts import ( + GROUNDED_TASK_PLAN_SCHEMA, + GroundedTaskPlan, + RoleBindings, + canonical_hash, + validate_grounded_task_plan, + validate_binding_report, + validate_role_bindings, +) +from .scene_adapter import SceneAdaptation, SceneAdapter +from .scene_source import SceneSourceRef + +__all__ = [ + "TaskEngineCoordinator", + "PreparationResult", + "build_grounded_task_plan", + "lower_task_candidate", +] + + +BundleGenerator = Callable[..., GeneratedConfigPaths] +_PREPARATION_FAILURE_SCHEMA = "action_engine_preparation_failure_v1" + + +def lower_task_candidate( + candidate: Mapping[str, Any], + reference_bindings: Mapping[str, Any], + scene_objects: Sequence[Mapping[str, Any]], + robot_profile: str, +) -> GroundedTaskSpec: + """Lower a selected TaskCandidate across the Task/Action boundary.""" + normalized = validate_task_candidate(candidate) + if reference_bindings.get("schema_version") is not None: + role_bindings = validate_role_bindings(reference_bindings) + if role_bindings["task_id"] != normalized["draft"]["task_id"]: + raise ValueError("RoleBindings.task_id must match the TaskCandidate.") + if role_bindings["candidate_id"] != normalized["candidate_id"]: + raise ValueError("RoleBindings.candidate_id must match the TaskCandidate.") + raw_bindings = role_bindings["reference_bindings"] + else: + raw_bindings = reference_bindings + bindings = { + str(reference_id): [str(uid) for uid in uids] + for reference_id, uids in raw_bindings.items() + } + grounded = ground_instruction_draft( + normalized["draft"]["task_id"], + normalized["draft"]["instruction"], + {"steps": normalized["draft"]["steps"]}, + scene_objects, + robot_profile=robot_profile, + reference_bindings=bindings, + ) + _validate_lowered_success(normalized, bindings, grounded) + return grounded + + +@dataclass(frozen=True) +class PreparationResult: + """Published result of one Task -> Scene -> Action preparation attempt.""" + + status: str + output_dir: Path + candidate_set: TaskCandidateSet + adaptation: SceneAdaptation + artifacts: TaskEngineArtifactPaths + grounded_task_plan: GroundedTaskPlan | None = None + action_graph: dict[str, Any] | None = None + generated_paths: GeneratedConfigPaths | None = None + feasibility_report: FeasibilityReport | None = None + planning_attempts: tuple[dict[str, Any], ...] = () + unbound_action_plan: dict[str, Any] | None = None + + @property + def bound(self) -> bool: + return self.status == "bound" + + @property + def selected_candidate_id(self) -> str | None: + return self.adaptation.selected_candidate_id + + +@dataclass(frozen=True) +class _PlannedCandidate: + adaptation: SceneAdaptation + selected: TaskCandidate + role_bindings: RoleBindings + feasibility_report: FeasibilityReport | None + grounded: GroundedTaskSpec + grounded_plan: GroundedTaskPlan + action_graph: dict[str, Any] + unbound_action_plan: dict[str, Any] | None + + +class TaskEngineCoordinator: + """Run Task Agent, Scene Adapter, and Action Agent as one transaction.""" + + def __init__( + self, + *, + task_agent: TaskAgent | None = None, + scene_adapter: SceneAdapter | None = None, + action_agent: ActionAgent | None = None, + bundle_generator: BundleGenerator = generate_action_engine_config, + feasibility_broker: FeasibilityBroker | None = None, + ) -> None: + self.task_agent = task_agent or TaskAgent() + self.scene_adapter = scene_adapter or SceneAdapter() + self.action_agent = action_agent or ActionAgent() + self.bundle_generator = bundle_generator + self.feasibility_broker = feasibility_broker or FeasibilityBroker() + + def prepare( + self, + task_id: str, + instruction: str, + source: SceneSourceRef | str | Path, + output_dir: str | Path, + *, + model: str | None = None, + candidate_count: int = 3, + overwrite: bool = False, + planning_mode: str = "offline", + vlm_model: str | None = None, + max_episodes: int | None = None, + max_episode_steps: int | None = None, + randomize_scene: bool = False, + randomize_table_material: bool = False, + candidate_set: TaskCandidateSet | Mapping[str, Any] | None = None, + force_most_likely: bool = False, + final_inspection: Mapping[str, Any] | None = None, + unbound_action_plan: Mapping[str, Any] | None = None, + ) -> PreparationResult: + """Prepare and atomically publish a Task Engine bundle. + + Ambiguous and unsatisfied scene adaptations are valid terminal results. + They publish the complete audit hand-off but never publish a TaskSpec, + SeedGraph, Gym configuration, or GroundedTaskPlan. + """ + normalized_source = self._coerce_source(source) + validate_scene_output_separation(normalized_source.path, output_dir) + with ArtifactTransaction(output_dir, overwrite=overwrite) as transaction: + staging_dir = transaction.staging_dir + assert staging_dir is not None + if candidate_set is None: + normalized_candidates = self.task_agent.generate( + task_id, + instruction, + model=model, + candidate_count=candidate_count, + ) + else: + normalized_candidates = validate_task_candidate_set(candidate_set) + if normalized_candidates["task_id"] != str(task_id).strip(): + raise ValueError("TaskCandidateSet.task_id must match task_id.") + if normalized_candidates["instruction"] != str(instruction).strip(): + raise ValueError( + "TaskCandidateSet.instruction must match instruction." + ) + candidate_set = normalized_candidates + adaptation_kwargs: dict[str, Any] = {"force_most_likely": force_most_likely} + if final_inspection is not None: + adaptation_kwargs["final_inspection"] = final_inspection + adaptation = self.scene_adapter.adapt( + candidate_set, + normalized_source, + **adaptation_kwargs, + ) + status = str(adaptation.binding_report["status"]) + + if status != "bound": + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=None, + role_bindings=None, + binding_report=adaptation.binding_report, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + final_scene_inspection=final_inspection, + ) + published = transaction.commit() + return PreparationResult( + status=status, + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + planning_attempts=(), + ) + + selected = adaptation.selected_candidate + raw_role_bindings = adaptation.role_bindings + if selected is None or raw_role_bindings is None: + raise ValueError( + "A bound SceneAdaptation must include a selected candidate " + "and RoleBindings." + ) + feasibility_report = self._assess_feasibility( + selected, + raw_role_bindings, + adaptation, + ) + if ( + feasibility_report is not None + and feasibility_report["status"] == "contradicted" + and feasibility_report["remediation_class"] != "action_capability" + ): + adaptation, selected, raw_role_bindings, feasibility_report = ( + self._fallback_feasible_candidate( + candidate_set, + adaptation, + selected, + raw_role_bindings, + feasibility_report, + ) + ) + if ( + feasibility_report is not None + and feasibility_report["status"] == "contradicted" + ): + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=adaptation.scene_manifest, + role_bindings=raw_role_bindings, + binding_report=adaptation.binding_report, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + feasibility_report=feasibility_report, + final_scene_inspection=final_inspection, + ) + published = transaction.commit() + return PreparationResult( + status="infeasible", + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + feasibility_report=deepcopy(feasibility_report), + planning_attempts=(), + ) + robot_profile = str(adaptation.scene_manifest["robot_profile"]) + planned, planning_failures = self._plan_with_candidate_fallback( + candidate_set, + adaptation, + selected, + raw_role_bindings, + feasibility_report, + robot_profile=robot_profile, + unbound_action_plan=unbound_action_plan, + ) + if planned is None: + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=adaptation.scene_manifest, + role_bindings=raw_role_bindings, + binding_report=adaptation.binding_report, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + feasibility_report=feasibility_report, + final_scene_inspection=final_inspection, + ) + write_preparation_failure( + staging_dir, + { + "schema_version": _PREPARATION_FAILURE_SCHEMA, + "task_id": str(candidate_set["task_id"]), + "status": "planning_failed", + "selected_candidate_id": str(selected["candidate_id"]), + "attempts": planning_failures, + }, + ) + published = transaction.commit() + return PreparationResult( + status="planning_failed", + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + feasibility_report=deepcopy(feasibility_report), + planning_attempts=tuple(deepcopy(planning_failures)), + ) + + adaptation = planned.adaptation + selected = planned.selected + role_bindings = planned.role_bindings + feasibility_report = planned.feasibility_report + grounded = planned.grounded + grounded_plan = planned.grounded_plan + action_graph = planned.action_graph + + generator_kwargs: dict[str, Any] = { + "task_name": grounded_plan["task_id"], + "task_spec": grounded_plan["task_spec"], + "robot_profile": robot_profile, + "source_scene_z_rotation_degrees": ( + adaptation.prepared_scene.z_rotation_degrees + ), + "source_scene_xy_translation": list( + adaptation.prepared_scene.source_scene_xy_translation + ), + "body_scale_policy": adaptation.prepared_scene.body_scale_policy, + "body_scale": adaptation.prepared_scene.body_scale, + "overwrite": False, + "randomize_scene": randomize_scene, + "randomize_table_material": randomize_table_material, + "planning_mode": planning_mode, + "vlm_model": vlm_model, + } + if max_episodes is not None: + generator_kwargs["max_episodes"] = max_episodes + if max_episode_steps is not None: + generator_kwargs["max_episode_steps"] = max_episode_steps + compatibility_input = staging_dir / ".task_engine_input" + compatibility_input.mkdir() + task_spec_path = compatibility_input / "task_spec.json" + requirements_path = compatibility_input / "scene_requirements.json" + _write_compatibility_input(task_spec_path, grounded.task_spec) + _write_compatibility_input( + requirements_path, + grounded.scene_requirements, + ) + generator_kwargs["task_spec"] = task_spec_path + try: + generated = self.bundle_generator( + adaptation.source_config_path, + staging_dir, + **generator_kwargs, + ) + finally: + shutil.rmtree(compatibility_input, ignore_errors=True) + _require_matching_generated_graph(generated, action_graph) + write_task_engine_artifacts( + staging_dir, + candidate_set=candidate_set, + scene_manifest=adaptation.scene_manifest, + role_bindings=role_bindings, + binding_report=adaptation.binding_report, + grounded_task_plan=grounded_plan, + static_scene_manifest=adaptation.static_scene_manifest, + conservative_scene_graph=adaptation.conservative_scene_graph, + feasibility_report=feasibility_report, + final_scene_inspection=final_inspection, + ) + published = transaction.commit() + return PreparationResult( + status="bound", + output_dir=published, + candidate_set=deepcopy(candidate_set), + adaptation=adaptation, + artifacts=task_engine_artifact_paths(published), + grounded_task_plan=grounded_plan, + action_graph=deepcopy(action_graph), + generated_paths=artifact_paths( + published, + planning_mode=planning_mode, + ), + feasibility_report=deepcopy(feasibility_report), + planning_attempts=tuple(deepcopy(planning_failures)), + unbound_action_plan=deepcopy(planned.unbound_action_plan), + ) + + def _plan_with_candidate_fallback( + self, + candidate_set: Mapping[str, Any], + adaptation: SceneAdaptation, + selected: TaskCandidate, + role_bindings: RoleBindings, + feasibility_report: FeasibilityReport | None, + *, + robot_profile: str, + unbound_action_plan: Mapping[str, Any] | None, + ) -> tuple[_PlannedCandidate | None, list[dict[str, Any]]]: + """Treat lowering and Action planning failures as candidate-local.""" + candidates = { + str(candidate["candidate_id"]): candidate + for candidate in candidate_set.get("candidates", ()) + if isinstance(candidate, Mapping) and candidate.get("candidate_id") + } + resolved = { + str(audit["candidate_id"]) + for audit in adaptation.binding_report["candidates"] + if audit["status"] == "resolved" + } + selected_id = str(selected["candidate_id"]) + ordered_ids = [selected_id] + [ + candidate_id + for candidate_id in candidates + if candidate_id != selected_id and candidate_id in resolved + ] + failures: list[dict[str, Any]] = [] + + for candidate_id in ordered_ids: + candidate = candidates.get(candidate_id) + raw_bindings = ( + role_bindings + if candidate_id == selected_id + else adaptation.candidate_bindings.get(candidate_id) + ) + if candidate is None or raw_bindings is None: + continue + report = ( + feasibility_report + if candidate_id == selected_id + else self._assess_feasibility(candidate, raw_bindings, adaptation) + ) + if report is not None and report["status"] == "contradicted": + failures.append( + _candidate_failure( + candidate, + raw_bindings, + stage="static_feasibility", + error_type="FeasibilityContradiction", + error_message="Static feasibility contradicted this candidate.", + feasibility_report=report, + ) + ) + continue + + candidate_adaptation = _select_candidate_adaptation( + adaptation, + candidate, + raw_bindings, + failures, + ) + grounded: GroundedTaskSpec | None = None + grounded_plan: GroundedTaskPlan | None = None + candidate_unbound: Mapping[str, Any] | None = None + action_graph: Mapping[str, Any] | None = None + stage = "lowering" + try: + grounded = lower_task_candidate( + candidate, + raw_bindings, + adaptation.prepared_scene.planner_objects, + robot_profile, + ) + canonical_bindings = validate_role_bindings( + { + **deepcopy(raw_bindings), + "role_bindings": deepcopy(grounded.role_bindings), + } + ) + candidate_adaptation = replace( + candidate_adaptation, + role_bindings=deepcopy(canonical_bindings), + ) + stage = "grounded_plan" + grounded_plan = build_grounded_task_plan( + candidate=candidate, + task_spec=grounded.task_spec, + scene_requirements=grounded.scene_requirements, + scene_manifest=adaptation.scene_manifest, + role_bindings=canonical_bindings, + binding_report=candidate_adaptation.binding_report, + ) + stage = "action_planning" + bind_and_plan = getattr(self.action_agent, "bind_and_plan", None) + if callable(bind_and_plan): + candidate_unbound = ( + unbound_action_plan + if unbound_action_plan is not None + and str(unbound_action_plan.get("candidate_id")) == candidate_id + else self.action_agent.draft(candidate) + ) + action_graph = bind_and_plan(candidate_unbound, grounded_plan) + else: + action_graph = self.action_agent.plan(grounded_plan) + stage = "preflight" + preflight = getattr(self.action_agent, "preflight", None) + if callable(preflight): + preflight( + action_graph, + scene_manifest=adaptation.scene_manifest, + ) + except ActionCapabilityError: + raise + except (TypeError, ValueError, OSError) as error: + failures.append( + _candidate_failure( + candidate, + raw_bindings, + stage=stage, + error_type=type(error).__name__, + error_message=str(error), + feasibility_report=report, + grounded_task_plan=grounded_plan, + unbound_action_plan=candidate_unbound, + action_graph=action_graph, + ) + ) + continue + + assert grounded is not None and grounded_plan is not None + return ( + _PlannedCandidate( + adaptation=candidate_adaptation, + selected=deepcopy(candidate), + role_bindings=canonical_bindings, + feasibility_report=deepcopy(report), + grounded=grounded, + grounded_plan=grounded_plan, + action_graph=deepcopy(action_graph), + unbound_action_plan=( + None + if candidate_unbound is None + else deepcopy(dict(candidate_unbound)) + ), + ), + failures, + ) + return None, failures + + def _fallback_feasible_candidate( + self, + candidate_set: Mapping[str, Any], + adaptation: SceneAdaptation, + selected: TaskCandidate, + role_bindings: RoleBindings, + report: FeasibilityReport, + ) -> tuple[ + SceneAdaptation, + TaskCandidate, + RoleBindings, + FeasibilityReport | None, + ]: + """Try other resolved semantic candidates after a static contradiction.""" + candidates = { + str(candidate["candidate_id"]): candidate + for candidate in candidate_set.get("candidates", ()) + if isinstance(candidate, Mapping) and candidate.get("candidate_id") + } + selected_id = str(selected["candidate_id"]) + for audit in adaptation.binding_report["candidates"]: + candidate_id = str(audit["candidate_id"]) + if candidate_id == selected_id or audit["status"] != "resolved": + continue + candidate = candidates.get(candidate_id) + alternative_bindings = adaptation.candidate_bindings.get(candidate_id) + if candidate is None or alternative_bindings is None: + continue + alternative_report = self._assess_feasibility( + candidate, + alternative_bindings, + adaptation, + ) + if ( + alternative_report is not None + and alternative_report["status"] == "contradicted" + ): + continue + binding_report = validate_binding_report( + { + **deepcopy(adaptation.binding_report), + "selected_candidate_id": candidate_id, + "selection_reason": ( + "Selected the next resolved candidate after static " + f"feasibility contradicted {selected_id}." + ), + } + ) + chosen = deepcopy(candidate) + updated = replace( + adaptation, + selected_candidate=chosen, + role_bindings=deepcopy(alternative_bindings), + binding_report=binding_report, + ) + return ( + updated, + chosen, + deepcopy(alternative_bindings), + alternative_report, + ) + return adaptation, selected, role_bindings, report + + def _assess_feasibility( + self, + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + adaptation: SceneAdaptation, + ) -> FeasibilityReport | None: + """Intersect task requirements with scene and Action Engine capabilities.""" + manifest = adaptation.static_scene_manifest + registry = getattr(self.action_agent, "registry", None) + if manifest is None or registry is None: + return None + catalog = getattr(registry, "catalog", None) + if not callable(catalog): + return None + return self.feasibility_broker.assess( + candidate, + role_bindings, + manifest, + capability_catalog=catalog(), + task_actions={ + task_type: contract.core_actions + for task_type, contract in ACTION_TASK_CONTRACTS.items() + }, + ) + + def _coerce_source( + self, + source: SceneSourceRef | str | Path, + ) -> SceneSourceRef: + if isinstance(source, SceneSourceRef): + return source + path = Path(source).expanduser() + return SceneSourceRef( + path, + robot_profile=self.scene_adapter.robot_profile, + ) + + +def _select_candidate_adaptation( + adaptation: SceneAdaptation, + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + prior_failures: Sequence[Mapping[str, Any]], +) -> SceneAdaptation: + candidate_id = str(candidate["candidate_id"]) + current_id = adaptation.selected_candidate_id + if candidate_id == current_id and not prior_failures: + return adaptation + failed = ", ".join( + f"{failure['candidate_id']} failed {failure['stage']}" + for failure in prior_failures + ) + reason = f"Selected {candidate_id} after {failed}." + binding_report = validate_binding_report( + { + **deepcopy(adaptation.binding_report), + "selected_candidate_id": candidate_id, + "selection_reason": reason, + } + ) + return replace( + adaptation, + selected_candidate=deepcopy(candidate), + role_bindings=deepcopy(role_bindings), + binding_report=binding_report, + ) + + +def _candidate_failure( + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + *, + stage: str, + error_type: str, + error_message: str, + feasibility_report: Mapping[str, Any] | None = None, + grounded_task_plan: Mapping[str, Any] | None = None, + unbound_action_plan: Mapping[str, Any] | None = None, + action_graph: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + return { + "candidate_id": str(candidate["candidate_id"]), + "stage": stage, + "draft": deepcopy(candidate["draft"]), + "bindings": deepcopy(dict(role_bindings)), + "grounded_task_plan": ( + None if grounded_task_plan is None else deepcopy(dict(grounded_task_plan)) + ), + "unbound_action_plan": ( + None if unbound_action_plan is None else deepcopy(dict(unbound_action_plan)) + ), + "action_graph": None if action_graph is None else deepcopy(dict(action_graph)), + "feasibility_report": ( + None if feasibility_report is None else deepcopy(dict(feasibility_report)) + ), + "error": {"type": error_type, "message": error_message}, + } + + +# Short public name used in the phase-one design document. +def build_grounded_task_plan( + *, + candidate: Mapping[str, Any], + task_spec: Mapping[str, Any], + scene_requirements: Mapping[str, Any], + scene_manifest: Mapping[str, Any], + role_bindings: RoleBindings, + binding_report: Mapping[str, Any], +) -> GroundedTaskPlan: + """Assemble a validated plan with hashes over every authoritative hand-off.""" + draft = deepcopy(candidate["draft"]) + # A scene-ref quantifier may expand one draft step into several concrete + # task instances. The grounded plan records the executable success terms, + # while the selected TaskCandidate retains the pre-grounding SuccessSpec. + success_spec = { + **deepcopy(candidate["success_spec"]), + "terms": [ + { + "step_id": str(term["task_instance_id"]), + "type": str(term["type"]), + } + for term in task_spec["success"]["terms"] + ], + } + task = deepcopy(dict(task_spec)) + requirements = deepcopy(dict(scene_requirements)) + manifest = deepcopy(dict(scene_manifest)) + bindings = deepcopy(dict(role_bindings)) + report = deepcopy(dict(binding_report)) + base = { + "schema_version": GROUNDED_TASK_PLAN_SCHEMA, + "task_id": draft["task_id"], + "instruction": draft["instruction"], + "selected_candidate_id": candidate["candidate_id"], + "task_draft": draft, + "task_spec": task, + "scene_requirements": requirements, + "success_spec": success_spec, + "scene_manifest": manifest, + "role_bindings": bindings, + "binding_report": report, + } + plan = { + **base, + "hashes": { + "task_draft": canonical_hash(draft), + "task_spec": canonical_hash(task), + "scene_manifest": canonical_hash(manifest), + "role_bindings": canonical_hash(bindings), + "plan": canonical_hash(base), + }, + } + return validate_grounded_task_plan(plan) + + +def _validate_lowered_success( + candidate: TaskCandidate, + bindings: Mapping[str, list[str]], + grounded: GroundedTaskSpec, +) -> None: + success_by_step = { + term["step_id"]: term["type"] for term in candidate["success_spec"]["terms"] + } + expected: list[str] = [] + multiplicity_by_step: dict[str, int] = {} + for step in _topological_steps(candidate["draft"]["steps"]): + selector = step["object"] + multiplicity = 1 + if selector["kind"] == "scene_ref": + multiplicity = len(bindings.get(f"{step['id']}.object", ())) + elif selector["kind"] == "step_result": + multiplicity = multiplicity_by_step[str(selector["step_id"])] + multiplicity_by_step[str(step["id"])] = multiplicity + expected.extend([success_by_step[step["id"]]] * multiplicity) + actual = [term.get("type") for term in grounded.task_spec["success"]["terms"]] + if actual != expected: + raise ValueError( + "Lowered TaskSpec success terms do not match the expanded SuccessSpec." + ) + + +def _topological_steps( + steps: list[Mapping[str, Any]], +) -> list[Mapping[str, Any]]: + positions = {str(step["id"]): index for index, step in enumerate(steps)} + pending = {str(step["id"]): set(step["depends_on"]) for step in steps} + result: list[Mapping[str, Any]] = [] + emitted: set[str] = set() + while len(result) < len(steps): + ready = [ + step + for step in steps + if step["id"] not in emitted and pending[str(step["id"])] <= emitted + ] + if not ready: + raise ValueError("TaskDraft step dependencies contain a cycle.") + step = min(ready, key=lambda item: positions[str(item["id"])]) + result.append(step) + emitted.add(str(step["id"])) + return result + + +def _require_matching_generated_graph( + generated: GeneratedConfigPaths, + expected: Mapping[str, Any], +) -> None: + """Catch a compatibility-generator drift before publishing the bundle.""" + graph_path = getattr(generated, "seed_task_graph", None) + if graph_path is None or not Path(graph_path).is_file(): + # Injected generators used by API consumers may publish by other means. + # Task Engine's independently planned graph remains authoritative. + return + try: + actual = json.loads(Path(graph_path).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Generated SeedGraph is unreadable: {graph_path}") from exc + if canonical_hash(actual) != canonical_hash(expected): + raise ValueError( + "Legacy bundle generation produced a SeedGraph different from " + "ActionAgent.plan." + ) + + +def _write_compatibility_input(path: Path, value: Mapping[str, Any]) -> None: + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/task_engine/orchestration/legacy_scene.py b/embodichain/gen_sim/task_engine/orchestration/legacy_scene.py new file mode 100644 index 000000000..4cb17ff7d --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/legacy_scene.py @@ -0,0 +1,381 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Read-only conversion of legacy Gym projects into editable scene revisions.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass +import json +from pathlib import Path +import shutil +from typing import Any, Final + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_source_scene, +) + +from .scene_source import ( + SceneSourceFingerprint, + fingerprint_scene_source, + verify_scene_source_fingerprint, +) + +__all__ = [ + "LEGACY_SCENE_CONVERSION_SCHEMA", + "LegacySceneRevision", + "convert_legacy_gym_project", + "restore_locked_scene_entities", +] + +LEGACY_SCENE_CONVERSION_SCHEMA: Final = "embodichain.legacy-scene-conversion/v1" +_CONVERSION_MANIFEST = "legacy_conversion.json" + + +@dataclass(frozen=True) +class LegacySceneRevision: + """A new editable revision derived without modifying its legacy source.""" + + output_root: Path + scene_config_path: Path + scene_graph_path: Path + manifest_path: Path + source_fingerprint: SceneSourceFingerprint + locked_entity_uids: tuple[str, ...] + + +def convert_legacy_gym_project( + source: str | Path, + output_root: str | Path, +) -> LegacySceneRevision: + """Convert a supported legacy Gym project into a Scene Engine revision. + + Args: + source: Legacy Gym project directory or explicit configuration path. + output_root: Empty destination owned by the new scene revision. + + Returns: + Paths and provenance for the converted revision. + + Raises: + ValueError: If the source is not legacy or the destination already exists. + FileNotFoundError: If a referenced source asset is missing. + """ + resolved = resolve_source_scene(source) + if resolved.source_format != "legacy_gym_config": + raise ValueError("Legacy conversion requires a legacy Gym configuration.") + destination = Path(output_root).expanduser().resolve() + if destination.exists(): + if not destination.is_dir() or any(destination.iterdir()): + raise ValueError("Legacy scene revision output_root must be empty.") + source_fingerprint = fingerprint_scene_source(source) + prepared = prepare_scene(source) + export_root = destination / "scene_export" + assets_root = export_root / "mesh_assets" + assets_root.mkdir(parents=True, exist_ok=True) + semantics = {str(item.get("uid")): item for item in prepared.planner_objects} + + background = [ + _editable_entry(item, semantics=semantics, assets_root=assets_root) + for item in prepared.background + ] + rigid_objects = [ + _editable_entry(item, semantics=semantics, assets_root=assets_root) + for item in prepared.rigid_objects + ] + articulations = [ + _locked_articulation( + item, + source_root=resolved.path.parent, + destination_root=export_root / "locked_assets", + ) + for item in prepared.articulations + ] + table = next((item for item in background if item.get("uid") == "table"), None) + if table is None: + raise ValueError("Legacy conversion requires one table support object.") + _measure_support_metadata(table, export_root=export_root) + for item in rigid_objects: + _measure_center(item, export_root=export_root) + + scene_config = { + "format": "embodichain.scene-export/v1", + "scene_id": f"legacy-revision-{source_fingerprint.config_sha256[:16]}", + "background": background, + "rigid_object": rigid_objects, + "articulation": articulations, + } + scene_config_path = export_root / "scene_config.json" + _write_json(scene_config_path, scene_config) + scene_graph = { + "nodes": [ + { + "object_id": "table", + "parent_id": None, + "parent_relation": None, + "table_region": None, + "orientation_state": None, + }, + *[ + { + "object_id": str(item["uid"]), + "parent_id": "table", + "parent_relation": "on", + "table_region": None, + "orientation_state": None, + } + for item in rigid_objects + ], + ], + "relations": [], + } + scene_graph_path = export_root / "scene_graph.json" + _write_json(scene_graph_path, scene_graph) + locked_uids = tuple( + sorted(str(item["uid"]) for item in [*background, *articulations]) + ) + manifest = { + "schema_version": LEGACY_SCENE_CONVERSION_SCHEMA, + "source": source_fingerprint.to_dict(), + "scene_config": scene_config_path.as_posix(), + "audit_hierarchy": "unknown", + "operational_hierarchy": "assumed_on_table", + "assumptions": [ + { + "uid": str(item["uid"]), + "relation": "on", + "parent_uid": "table", + "confidence": None, + "source": "operational_assumption", + } + for item in rigid_objects + ], + "locked_entity_uids": list(locked_uids), + "locked_articulations": deepcopy(articulations), + "locked_background": deepcopy( + [item for item in background if item.get("uid") != "table"] + ), + } + manifest_path = destination / _CONVERSION_MANIFEST + _write_json(manifest_path, manifest) + verify_scene_source_fingerprint(source_fingerprint.to_dict()) + return LegacySceneRevision( + output_root=destination, + scene_config_path=scene_config_path, + scene_graph_path=scene_graph_path, + manifest_path=manifest_path, + source_fingerprint=source_fingerprint, + locked_entity_uids=locked_uids, + ) + + +def restore_locked_scene_entities(revision_root: str | Path) -> Path: + """Restore collision-only legacy entities after Scene Engine export. + + Args: + revision_root: Converted revision root containing ``legacy_conversion.json``. + + Returns: + Updated scene configuration path. + + Raises: + FileNotFoundError: If the conversion manifest or scene config is absent. + ValueError: If a generated scene attempts to reuse a locked UID. + """ + root = Path(revision_root).expanduser().resolve() + manifest_path = root / _CONVERSION_MANIFEST + if not manifest_path.is_file(): + raise FileNotFoundError( + f"Legacy conversion manifest not found: {manifest_path}" + ) + manifest = _read_mapping(manifest_path) + if manifest.get("schema_version") != LEGACY_SCENE_CONVERSION_SCHEMA: + raise ValueError("Legacy conversion manifest schema is invalid.") + config_path = root / "scene_export" / "scene_config.json" + config = _read_mapping(config_path) + existing = { + str(item.get("uid")) + for section in ("background", "rigid_object", "articulation") + for item in config.get(section, ()) + if isinstance(item, Mapping) and item.get("uid") + } + for section, key in ( + ("background", "locked_background"), + ("articulation", "locked_articulations"), + ): + values = manifest.get(key, ()) + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise TypeError(f"Legacy conversion manifest {key} must be a sequence.") + target = list(config.get(section, ())) + for raw in values: + item = deepcopy(dict(raw)) + uid = str(item.get("uid", "")) + if uid in existing: + raise ValueError(f"Generated scene reused locked entity UID {uid!r}.") + existing.add(uid) + target.append(item) + config[section] = target + _write_json(config_path, config) + verify_scene_source_fingerprint(manifest["source"]) + return config_path + + +def _editable_entry( + value: Mapping[str, Any], + *, + semantics: Mapping[str, Mapping[str, Any]], + assets_root: Path, +) -> dict[str, Any]: + item = deepcopy(dict(value)) + uid = str(item.get("uid", "")).strip() + if not uid: + raise ValueError("Converted scene entities require a UID.") + semantic = semantics.get(uid, {}) + for key in ("category", "name", "description"): + item[key] = str(semantic.get(key) or item.get(key) or uid) + shape = item.get("shape") + if not isinstance(shape, Mapping): + raise ValueError(f"Legacy scene entity {uid!r} has no supported shape.") + destination = assets_root / uid / f"{uid}.glb" + destination.parent.mkdir(parents=True, exist_ok=True) + _shape_to_glb(shape, destination) + item["shape"] = { + "shape_type": "Mesh", + "fpath": destination.relative_to(assets_root.parent).as_posix(), + "compute_uv": False, + } + item.setdefault("body_scale", [1.0, 1.0, 1.0]) + item.setdefault("init_pos", [0.0, 0.0, 0.0]) + item.setdefault("init_rot", [0.0, 0.0, 0.0]) + item.setdefault("attrs", {"mass": 1.0}) + item.setdefault("body_type", "kinematic" if uid == "table" else "dynamic") + item.setdefault("max_convex_hull_num", 1 if uid == "table" else 16) + return item + + +def _shape_to_glb(shape: Mapping[str, Any], destination: Path) -> None: + shape_type = str(shape.get("shape_type", "")) + if shape_type == "Mesh": + source = Path(str(shape.get("fpath", ""))).expanduser().resolve() + if not source.is_file(): + raise FileNotFoundError(f"Legacy mesh asset not found: {source}") + mesh = trimesh.load(source, force="scene") + elif shape_type == "Cube": + size = _vector(shape.get("size", [1.0, 1.0, 1.0]), length=3) + mesh = trimesh.Scene(trimesh.creation.box(extents=size)) + elif shape_type == "Sphere": + radius = float(shape.get("radius", 1.0)) + if not np.isfinite(radius) or radius <= 0.0: + raise ValueError("Legacy sphere radius must be positive and finite.") + mesh = trimesh.Scene(trimesh.creation.icosphere(radius=radius)) + else: + raise ValueError(f"Unsupported legacy shape_type {shape_type!r}.") + mesh.export(destination, file_type="glb") + + +def _locked_articulation( + value: Mapping[str, Any], + *, + source_root: Path, + destination_root: Path, +) -> dict[str, Any]: + item = deepcopy(dict(value)) + uid = str(item.get("uid", "")).strip() + raw = Path(str(item.get("fpath", ""))).expanduser() + source = raw.resolve() if raw.is_absolute() else (source_root / raw).resolve() + if not source.is_file(): + raise FileNotFoundError(f"Legacy articulation asset not found: {source}") + target_root = destination_root / uid + shutil.copytree(source.parent, target_root, dirs_exist_ok=True) + copied = target_root / source.name + item["fpath"] = copied.resolve().as_posix() + return item + + +def _measure_support_metadata(entry: dict[str, Any], *, export_root: Path) -> None: + bounds = _world_bounds(entry, export_root=export_root) + entry["support_surface_z"] = float(bounds[1, 2]) + rectangle = [ + [float(bounds[0, 0]), float(bounds[0, 1])], + [float(bounds[1, 0]), float(bounds[0, 1])], + [float(bounds[1, 0]), float(bounds[1, 1])], + [float(bounds[0, 0]), float(bounds[1, 1])], + ] + entry["support_contour_xy"] = rectangle + entry["support_optimization_rect_xy"] = deepcopy(rectangle) + entry["center_xy"] = [ + float((bounds[0, 0] + bounds[1, 0]) / 2.0), + float((bounds[0, 1] + bounds[1, 1]) / 2.0), + ] + + +def _measure_center(entry: dict[str, Any], *, export_root: Path) -> None: + bounds = _world_bounds(entry, export_root=export_root) + entry["center_xy"] = [ + float((bounds[0, 0] + bounds[1, 0]) / 2.0), + float((bounds[0, 1] + bounds[1, 1]) / 2.0), + ] + + +def _world_bounds(entry: Mapping[str, Any], *, export_root: Path) -> np.ndarray: + shape = dict(entry["shape"]) + mesh_path = (export_root / str(shape["fpath"])).resolve() + loaded = trimesh.load(mesh_path, force="scene") + mesh = loaded.to_geometry() + scale = np.asarray(_vector(entry.get("body_scale", [1.0] * 3), length=3)) + mesh.apply_scale(scale) + transform = np.eye(4) + transform[:3, :3] = Rotation.from_euler( + "XYZ", + _vector(entry.get("init_rot", [0.0] * 3), length=3), + degrees=True, + ).as_matrix() + transform[:3, 3] = _vector(entry.get("init_pos", [0.0] * 3), length=3) + mesh.apply_transform(transform) + return np.asarray(mesh.bounds, dtype=float) + + +def _vector(value: Any, *, length: int) -> list[float]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError("Legacy scene vector must be a sequence.") + result = [float(item) for item in value] + if len(result) != length or not np.all(np.isfinite(result)): + raise ValueError(f"Legacy scene vector must contain {length} finite values.") + return result + + +def _read_mapping(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(path) + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, Mapping): + raise TypeError(f"JSON document must contain an object: {path}") + return dict(value) + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py b/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py new file mode 100644 index 000000000..8f101e23b --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py @@ -0,0 +1,1050 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Bind Task Agent candidates to a redacted, authoritative scene inventory.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +import hashlib +import json +from pathlib import Path +from typing import Any + +from embodichain.gen_sim.action_engine.generation.models import PreparedScene +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_source_scene, +) +from embodichain.gen_sim.action_engine.tasks.assembly import ( + SceneInventory, + validate_source_compatibility, + validate_target_compatibility, +) +from embodichain.gen_sim.action_engine.tasks.grounding import ( + GroundingCaller, + ground_scene_references, +) +from embodichain.gen_sim.task_engine import TaskCandidate, TaskCandidateSet +from embodichain.gen_sim.task_engine.interpretation import ( + _default_instruction_caller, +) +from embodichain.gen_sim.task_engine.scene import ( + ConservativeSceneGraph, + SceneEngineV1Adapter, + StaticSceneManifest, + build_conservative_scene_graph, + validate_static_scene_manifest, +) +from embodichain.gen_sim.task_engine.scene.final_inspection import ( + apply_final_inspection, + validate_final_scene_inspection, +) + +from .contracts import ( + BINDING_REPORT_SCHEMA, + ROLE_BINDINGS_SCHEMA, + SCENE_MANIFEST_SCHEMA, + BindingReport, + RoleBindings, + SceneManifest, + validate_binding_report, + validate_role_bindings, + validate_scene_manifest, + validate_task_candidate, + validate_task_candidate_set, +) +from .scene_source import ( + SceneSourceRef, + fingerprint_scene_source, + scene_revision_id, +) + +__all__ = [ + "Adjudicator", + "CandidateSelection", + "SceneAdaptation", + "SceneAdapter", + "SceneAdapterProtocolError", +] + + +Adjudicator = Callable[..., Mapping[str, Any]] + +_REDACTED_KEYS = frozenset( + { + "absolute_position", + "bbox", + "bboxes", + "bounding_box", + "center", + "centroid", + "coordinates", + "dimensions", + "extrinsics", + "grasp_pose", + "init_local_pose", + "init_pos", + "init_rot", + "intrinsics", + "joint_positions", + "joints", + "keypoint", + "keypoints", + "location", + "matrix", + "pose", + "position", + "position_xyz", + "qpos", + "quaternion", + "rotation", + "scale", + "target_pose", + "trajectory", + "transform", + "translation", + "waypoints", + "world_x", + "world_y", + "world_z", + "x", + "y", + "z", + } +) + + +class SceneAdapterProtocolError(ValueError): + """The grounding or adjudication transport violated its JSON protocol.""" + + +@dataclass(frozen=True) +class CandidateSelection: + """Candidate binding against semantic scene data before materialization.""" + + scene_manifest: SceneManifest + role_bindings: RoleBindings | None + binding_report: BindingReport + selected_candidate: TaskCandidate | None + candidate_bindings: dict[str, RoleBindings] = field(default_factory=dict) + + @property + def selected_candidate_id(self) -> str | None: + """Return the chosen candidate identifier, when one was bindable.""" + return ( + str(self.selected_candidate["candidate_id"]) + if self.selected_candidate is not None + else None + ) + + +@dataclass(frozen=True) +class SceneAdaptation: + """Complete Scene Adapter result, including the reusable prepared scene.""" + + scene_manifest: SceneManifest + role_bindings: RoleBindings | None + binding_report: BindingReport + selected_candidate: TaskCandidate | None + prepared_scene: PreparedScene + source_config_path: Path + conservative_scene_graph: ConservativeSceneGraph + static_scene_manifest: StaticSceneManifest | None = None + candidate_bindings: dict[str, RoleBindings] = field(default_factory=dict) + + @property + def selected_candidate_id(self) -> str | None: + return ( + str(self.selected_candidate["candidate_id"]) + if self.selected_candidate is not None + else None + ) + + @property + def reference_bindings(self) -> dict[str, list[str]]: + if self.role_bindings is None: + return {} + return deepcopy(self.role_bindings["reference_bindings"]) + + +class SceneAdapter: + """Adapt one existing or packaged scene to a set of task candidates.""" + + def __init__( + self, + *, + model: str | None = None, + grounding_caller: GroundingCaller | None = None, + adjudicator: Adjudicator | None = None, + robot_profile: str = "franka", + scene_engine_adapter: SceneEngineV1Adapter | None = None, + ) -> None: + self.model = model + self.grounding_caller = grounding_caller + self.adjudicator = adjudicator + self.robot_profile = robot_profile + self.scene_engine_adapter = scene_engine_adapter or SceneEngineV1Adapter() + + def adapt( + self, + candidate_set: TaskCandidateSet | Sequence[Mapping[str, Any]], + source: SceneSourceRef | str | Path, + *, + grounding_caller: GroundingCaller | None = None, + adjudicator: Adjudicator | None = None, + force_most_likely: bool = False, + final_inspection: Mapping[str, Any] | None = None, + ) -> SceneAdaptation: + """Ground all candidates, then deterministically choose a bindable one.""" + task_id, instruction, candidates = _coerce_candidates(candidate_set) + source_ref = self._resolve_source(source) + source_fingerprint = fingerprint_scene_source(source_ref) + prepared = prepare_scene( + source_ref.path, + z_rotation_degrees=source_ref.z_rotation_degrees, + body_scale_policy=source_ref.body_scale_policy, + body_scale=source_ref.body_scale, + ) + if final_inspection is not None: + normalized_inspection = validate_final_scene_inspection(final_inspection) + if normalized_inspection["scene_revision_id"] != scene_revision_id( + source_ref + ): + raise ValueError( + "FinalSceneInspection does not describe the adapted scene revision." + ) + prepared = apply_final_inspection(prepared, normalized_inspection) + inventory = SceneInventory( + prepared.planner_objects, + robot_profile=source_ref.robot_profile, + ) + resolved_source = resolve_source_scene(source_ref.path) + manifest = _build_manifest( + prepared, + inventory, + source_format=resolved_source.source_format, + ) + static_manifest = self.scene_engine_adapter.adapt_prepared_scene( + prepared, + source_format=resolved_source.source_format, + robot_profile=inventory.profile, + ) + static_manifest["source"]["source_fingerprint"] = source_fingerprint.to_dict() + static_manifest = validate_static_scene_manifest(static_manifest) + conservative_scene_graph = build_conservative_scene_graph( + prepared, + scene_id=static_manifest["scene_id"], + ) + if fingerprint_scene_source(source_ref) != source_fingerprint: + raise RuntimeError("Source Gym project changed while it was being adapted.") + + selection = self._select_candidates( + task_id, + instruction, + candidates, + manifest=manifest, + inventory=inventory, + scene_objects=prepared.planner_objects, + grounding_caller=grounding_caller, + adjudicator=adjudicator, + force_most_likely=force_most_likely, + ) + return SceneAdaptation( + scene_manifest=manifest, + role_bindings=selection.role_bindings, + binding_report=selection.binding_report, + selected_candidate=selection.selected_candidate, + prepared_scene=prepared, + source_config_path=prepared.source_config_path, + conservative_scene_graph=conservative_scene_graph, + static_scene_manifest=static_manifest, + candidate_bindings=selection.candidate_bindings, + ) + + def select_objects( + self, + candidate_set: TaskCandidateSet | Sequence[Mapping[str, Any]], + scene_objects: Sequence[Mapping[str, Any]], + *, + source_format: str = "embodichain.scene-blueprint/v1", + robot_profile: str | None = None, + grounding_caller: GroundingCaller | None = None, + adjudicator: Adjudicator | None = None, + force_most_likely: bool = False, + ) -> CandidateSelection: + """Bind candidates to semantic objects before assets are generated. + + Args: + candidate_set: Validated Task Engine candidate set. + scene_objects: Blueprint-level semantic object records. + source_format: Provenance label included in the semantic manifest. + robot_profile: Optional robot profile override. + grounding_caller: Optional structured grounding transport. + adjudicator: Optional candidate tie-breaker. + force_most_likely: Resolve ranked UID hypotheses instead of rejecting + low-confidence or ambiguous responses. + + Returns: + Audited candidate selection without requiring generated assets. + """ + task_id, instruction, candidates = _coerce_candidates(candidate_set) + inventory = SceneInventory( + scene_objects, + robot_profile=robot_profile or self.robot_profile, + ) + manifest = _build_semantic_manifest( + inventory, + source_format=source_format, + ) + return self._select_candidates( + task_id, + instruction, + candidates, + manifest=manifest, + inventory=inventory, + scene_objects=scene_objects, + grounding_caller=grounding_caller, + adjudicator=adjudicator, + force_most_likely=force_most_likely, + ) + + def _select_candidates( + self, + task_id: str, + instruction: str, + candidates: Sequence[TaskCandidate], + *, + manifest: SceneManifest, + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], + grounding_caller: GroundingCaller | None, + adjudicator: Adjudicator | None, + force_most_likely: bool, + ) -> CandidateSelection: + invoke = grounding_caller or self.grounding_caller + use_default_adjudicator = invoke is None + if invoke is None: + invoke = _default_grounding_caller() + choose = adjudicator or self.adjudicator + if choose is None and use_default_adjudicator: + choose = _default_adjudicator(model=self.model) + audits: list[dict[str, Any]] = [] + bindings_by_candidate: dict[str, dict[str, tuple[str, ...]]] = {} + for candidate in candidates: + audit, bindings = _ground_candidate( + candidate, + instruction=instruction, + inventory=inventory, + scene_objects=scene_objects, + model=self.model, + caller=invoke, + force_most_likely=force_most_likely, + ) + audits.append(audit) + if bindings is not None: + bindings_by_candidate[str(candidate["candidate_id"])] = bindings + + selected_id, status, reason = _select_candidate( + candidates, + audits, + manifest=manifest, + instruction=instruction, + adjudicator=choose, + ) + report = validate_binding_report( + { + "schema_version": BINDING_REPORT_SCHEMA, + "task_id": task_id, + "status": status, + "selected_candidate_id": selected_id or "", + "selection_reason": reason, + "candidates": audits, + } + ) + selected = next( + ( + deepcopy(candidate) + for candidate in candidates + if candidate["candidate_id"] == selected_id + ), + None, + ) + candidate_bindings = { + candidate_id: validate_role_bindings( + { + "schema_version": ROLE_BINDINGS_SCHEMA, + "task_id": task_id, + "candidate_id": candidate_id, + "reference_bindings": { + key: list(value) for key, value in sorted(raw_bindings.items()) + }, + "role_bindings": {}, + } + ) + for candidate_id, raw_bindings in bindings_by_candidate.items() + } + role_bindings = None if selected_id is None else candidate_bindings[selected_id] + return CandidateSelection( + scene_manifest=manifest, + role_bindings=role_bindings, + binding_report=report, + selected_candidate=selected, + candidate_bindings=candidate_bindings, + ) + + def _resolve_source( + self, + source: SceneSourceRef | str | Path, + ) -> SceneSourceRef: + if isinstance(source, SceneSourceRef): + return source + return SceneSourceRef(source, robot_profile=self.robot_profile) + + +def _coerce_candidates( + value: TaskCandidateSet | Sequence[Mapping[str, Any]], +) -> tuple[str, str, list[TaskCandidate]]: + if isinstance(value, Mapping): + normalized = validate_task_candidate_set(value) + return ( + normalized["task_id"], + normalized["instruction"], + normalized["candidates"], + ) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + candidates = [validate_task_candidate(candidate) for candidate in value] + if not candidates: + raise ValueError("SceneAdapter requires at least one TaskCandidate.") + task_ids = {candidate["draft"]["task_id"] for candidate in candidates} + instructions = {candidate["draft"]["instruction"] for candidate in candidates} + if len(task_ids) != 1 or len(instructions) != 1: + raise ValueError("All TaskCandidates must describe the same task.") + return task_ids.pop(), instructions.pop(), candidates + raise TypeError("candidate_set must be a TaskCandidateSet or candidate sequence.") + + +def _default_grounding_caller() -> GroundingCaller: + # Keep provider setup lazy so package import and offline tests never load an + # LLM client. This is the same structured transport used by interpretation. + return _default_instruction_caller + + +def _default_adjudicator(*, model: str | None) -> Adjudicator: + caller = _default_grounding_caller() + + def adjudicate(**kwargs: Any) -> Mapping[str, Any]: + candidates = [ + { + key: deepcopy(candidate[key]) + for key in ( + "candidate_id", + "draft", + "scene_request", + "success_spec", + "vote_count", + ) + } + for candidate in kwargs["candidates"] + ] + allowed = [str(candidate["candidate_id"]) for candidate in candidates] + schema = { + "title": "ActionEngineTaskAdjudication", + "type": "object", + "additionalProperties": False, + "required": ["candidate_id"], + "properties": { + "candidate_id": {"type": "string", "enum": allowed}, + }, + } + prompt = ( + "Select exactly one already verified, fully bindable task candidate " + "that best matches the instruction and redacted scene manifest. Do " + "not alter a candidate or invent a new interpretation. Return only " + "candidate_id.\n\n" + f"Instruction:\n{kwargs['instruction']}\n\n" + "Candidates:\n" + f"{json.dumps(candidates, ensure_ascii=False, sort_keys=True)}\n\n" + "Redacted scene manifest:\n" + f"{json.dumps(kwargs['scene_manifest'], ensure_ascii=False, sort_keys=True)}" + ) + try: + return caller(prompt=prompt, schema=schema, model=model) + except (TypeError, ValueError) as exc: + raise SceneAdapterProtocolError( + f"Task adjudication returned invalid structured output: {exc}" + ) from exc + + return adjudicate + + +def _ground_candidate( + candidate: TaskCandidate, + *, + instruction: str, + inventory: SceneInventory, + scene_objects: Sequence[Mapping[str, Any]], + model: str | None, + caller: GroundingCaller, + force_most_likely: bool, +) -> tuple[dict[str, Any], dict[str, tuple[str, ...]] | None]: + responses: list[Any] = [] + + def audited_caller(**kwargs: Any) -> Mapping[str, Any]: + call_kwargs = dict(kwargs) + if force_most_likely: + call_kwargs["prompt"] = ( + f"{kwargs['prompt']}\n\nFINAL BINDING OVERRIDE: do not return " + "ambiguous merely because confidence is low. Choose the most " + "likely existing UID that satisfies the supplied structured " + "role, affordance, state, and attribute metadata. Return " + "candidate UIDs in descending likelihood order. Do not invent, " + "add, delete, move, or modify any scene object. Use not_found " + "when no structurally compatible existing object is plausible." + ) + response = caller(**call_kwargs) + responses.append(deepcopy(response)) + if force_most_likely: + return _force_most_likely_response(response, candidate=candidate) + return response + + candidate_id = str(candidate["candidate_id"]) + try: + result = ground_scene_references( + instruction=instruction, + intent=candidate["draft"], + inventory=inventory, + scene_objects=scene_objects, + model=model, + caller=audited_caller, + ) + except (TypeError, ValueError) as exc: + if responses: + audits = _audit_unresolved_response( + responses[-1], + candidate=candidate, + inventory=inventory, + error=str(exc), + ) + status = _candidate_status(audits) + return ( + _candidate_audit(candidate, status, audits, [str(exc)]), + None, + ) + raise SceneAdapterProtocolError( + f"Grounding candidate {candidate_id!r} failed before returning JSON: {exc}" + ) from exc + + raw_bindings = result.bindings + response_by_id = _response_bindings(responses[-1], candidate=candidate) + self_reference_reasons = _self_reference_reasons(candidate["draft"], raw_bindings) + reference_audits = [] + incompatible: set[str] = set() + request_by_id = { + str(request["reference_id"]): request + for request in candidate["scene_request"]["references"] + } + for reference_id, uids in raw_bindings.items(): + compatibility_reasons = _compatibility_reasons( + request_by_id[reference_id], + uids, + inventory=inventory, + draft=candidate["draft"], + ) + compatibility_reasons.extend(self_reference_reasons.get(reference_id, ())) + compatibility_reasons = sorted(set(compatibility_reasons)) + if compatibility_reasons: + incompatible.add(reference_id) + response = response_by_id[reference_id] + audit_reasons = list(compatibility_reasons) + if ( + force_most_likely + and response.get("status") == "ambiguous" + and response.get("uids") + ): + audit_reasons.append( + "Forced the highest-ranked structurally compatible UID from an " + "ambiguous low-confidence response." + ) + reference_audits.append( + { + "reference_id": reference_id, + "status": ("incompatible" if compatibility_reasons else "resolved"), + "confidence": float(response["confidence"]), + "candidate_uids": list(response["uids"]), + "selected_uids": ([] if compatibility_reasons else list(uids)), + "reasons": audit_reasons, + } + ) + if incompatible: + reasons = [ + f"Reference {reference_id!r} conflicts with authoritative scene semantics." + for reference_id in sorted(incompatible) + ] + return ( + _candidate_audit(candidate, "incompatible", reference_audits, reasons), + None, + ) + return _candidate_audit(candidate, "resolved", reference_audits, []), dict( + raw_bindings + ) + + +def _force_most_likely_response( + response: Mapping[str, Any], + *, + candidate: TaskCandidate, +) -> Mapping[str, Any]: + """Turn ranked low-confidence UID hypotheses into explicit selections.""" + if not isinstance(response, Mapping) or set(response) != {"bindings"}: + return response + requests = { + str(item["reference_id"]): item + for item in candidate["scene_request"]["references"] + } + raw_bindings = response.get("bindings") + if not isinstance(raw_bindings, Sequence) or isinstance(raw_bindings, (str, bytes)): + return response + result = deepcopy(dict(response)) + values = [] + for raw in raw_bindings: + if not isinstance(raw, Mapping): + return response + item = deepcopy(dict(raw)) + request = requests.get(str(item.get("reference_id", ""))) + uids = item.get("uids") + if ( + request is not None + and item.get("status") in {"resolved", "ambiguous"} + and isinstance(uids, Sequence) + and not isinstance(uids, (str, bytes)) + and uids + ): + quantifier = str(request["quantifier"]) + count = int(request["count"]) + if quantifier == "one": + item["uids"] = list(uids[:1]) + elif quantifier == "count": + item["uids"] = list(uids[:count]) + item["status"] = "resolved" + confidence = item.get("confidence") + if isinstance(confidence, (int, float)) and not isinstance( + confidence, bool + ): + item["confidence"] = max(0.5, float(confidence)) + values.append(item) + result["bindings"] = values + return result + + +def _response_bindings( + response: Any, + *, + candidate: TaskCandidate, +) -> dict[str, Mapping[str, Any]]: + expected = { + str(request["reference_id"]) + for request in candidate["scene_request"]["references"] + } + if not isinstance(response, Mapping) or set(response) != {"bindings"}: + raise SceneAdapterProtocolError( + "Grounding response must contain only bindings." + ) + values = response["bindings"] + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise SceneAdapterProtocolError("Grounding response bindings must be a list.") + result: dict[str, Mapping[str, Any]] = {} + for raw in values: + if not isinstance(raw, Mapping): + raise SceneAdapterProtocolError( + "Every grounding binding must be a mapping." + ) + reference_id = raw.get("reference_id") + if not isinstance(reference_id, str) or reference_id not in expected: + raise SceneAdapterProtocolError( + "Grounding response contains an unknown reference ID." + ) + if reference_id in result: + raise SceneAdapterProtocolError( + "Grounding response contains duplicate reference IDs." + ) + result[reference_id] = raw + if set(result) != expected: + raise SceneAdapterProtocolError( + "Grounding response omitted requested reference IDs." + ) + return result + + +def _audit_unresolved_response( + response: Any, + *, + candidate: TaskCandidate, + inventory: SceneInventory, + error: str, +) -> list[dict[str, Any]]: + by_id = _response_bindings(response, candidate=candidate) + audits: list[dict[str, Any]] = [] + for request in candidate["scene_request"]["references"]: + reference_id = str(request["reference_id"]) + raw = by_id[reference_id] + if set(raw) != {"reference_id", "status", "uids", "confidence"}: + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has unsupported fields." + ) + status = raw["status"] + if status not in {"resolved", "ambiguous", "not_found"}: + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has invalid status." + ) + uids = raw["uids"] + confidence = raw["confidence"] + if ( + not isinstance(uids, Sequence) + or isinstance(uids, (str, bytes)) + or any( + not isinstance(uid, str) or uid not in inventory.by_uid for uid in uids + ) + or len(set(uids)) != len(uids) + ): + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has invalid candidate UIDs." + ) + if ( + isinstance(confidence, bool) + or not isinstance(confidence, (int, float)) + or not 0.0 <= float(confidence) <= 1.0 + ): + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} has invalid confidence." + ) + audit_status = status + reasons: list[str] = [] + if status == "not_found" and uids: + raise SceneAdapterProtocolError( + f"Grounding binding {reference_id!r} status=not_found requires no UIDs." + ) + if status == "resolved": + audit_status = "incompatible" + reasons.append(error) + else: + reasons.append(f"Grounding returned status={status}.") + audits.append( + { + "reference_id": reference_id, + "status": audit_status, + "confidence": float(confidence), + "candidate_uids": list(uids), + "selected_uids": [], + "reasons": reasons, + } + ) + return audits + + +def _compatibility_reasons( + request: Mapping[str, Any], + uids: Sequence[str], + *, + inventory: SceneInventory, + draft: Mapping[str, Any], +) -> list[str]: + entities = [inventory.by_uid[uid] for uid in uids] + reasons: list[str] = [] + role = str(request["role"]) + step = next(item for item in draft["steps"] if item["id"] == request["step_id"]) + try: + if role == "object": + validate_source_compatibility(str(step["task_type"]), entities) + else: + for entity in entities: + validate_target_compatibility( + str(step["task_type"]), + entity, + relation=str(step["relation"]), + ) + except ValueError as exc: + reasons.append(str(exc)) + + expected_structure = str(request["source_structure"]) + for entity in entities: + # Source structure is strict for manipulated objects. Target structure + # is relation-dependent and is already checked by + # validate_target_compatibility; a table support surface must not be + # rejected merely because it is passive rather than a rigid object. + if role == "object": + if expected_structure == "articulation" and entity.role != "articulation": + reasons.append( + f"UID {entity.uid!r} is not an articulation as requested." + ) + if expected_structure in { + "rigid_object", + "movable", + } and entity.role not in { + "object", + "rigid_object", + }: + reasons.append( + f"UID {entity.uid!r} is not a movable rigid object as requested." + ) + required_affordances = set(request["affordances"]) + if entity.affordances: + missing = required_affordances - set(entity.affordances) + if missing: + reasons.append( + f"UID {entity.uid!r} explicitly lacks affordances {sorted(missing)}." + ) + for key, expected in request["initial_state"].items(): + if key in entity.initial_state and entity.initial_state[key] != expected: + reasons.append( + f"UID {entity.uid!r} state {key!r} conflicts with the request." + ) + for key, expected in request["attributes"].items(): + if key in entity.attributes and entity.attributes[key] != expected: + reasons.append( + f"UID {entity.uid!r} attribute {key!r} conflicts with the request." + ) + return sorted(set(reasons)) + + +def _self_reference_reasons( + draft: Mapping[str, Any], + bindings: Mapping[str, Sequence[str]], +) -> dict[str, list[str]]: + """Reject object/target identity overlap, including step_result selectors.""" + objects_by_step: dict[str, tuple[str, ...]] = {} + reasons: dict[str, list[str]] = {} + for step in draft["steps"]: + step_id = str(step["id"]) + object_uids = _selector_uids( + step["object"], + reference_id=f"{step_id}.object", + bindings=bindings, + objects_by_step=objects_by_step, + ) + target_uids = _selector_uids( + step["target"], + reference_id=f"{step_id}.target", + bindings=bindings, + objects_by_step=objects_by_step, + ) + overlap = sorted(set(object_uids) & set(target_uids)) + if overlap: + reason = ( + f"Grounding step {step_id!r} uses the same UID as object and " + f"target: {overlap}." + ) + for role in ("object", "target"): + selector = step[role] + if selector["kind"] == "scene_ref": + reasons.setdefault(f"{step_id}.{role}", []).append(reason) + objects_by_step[step_id] = object_uids + return reasons + + +def _selector_uids( + selector: Mapping[str, Any], + *, + reference_id: str, + bindings: Mapping[str, Sequence[str]], + objects_by_step: Mapping[str, tuple[str, ...]], +) -> tuple[str, ...]: + kind = str(selector["kind"]) + if kind == "scene_ref": + return tuple(str(uid) for uid in bindings[reference_id]) + if kind == "step_result": + return objects_by_step[str(selector["step_id"])] + return () + + +def _candidate_audit( + candidate: TaskCandidate, + status: str, + references: Sequence[Mapping[str, Any]], + reasons: Sequence[str], +) -> dict[str, Any]: + return { + "candidate_id": candidate["candidate_id"], + "semantic_hash": candidate["semantic_hash"], + "status": status, + "references": [deepcopy(dict(reference)) for reference in references], + "reasons": list(reasons), + } + + +def _candidate_status(references: Sequence[Mapping[str, Any]]) -> str: + statuses = {str(reference["status"]) for reference in references} + if statuses == {"resolved"}: + return "resolved" + if "ambiguous" in statuses: + return "ambiguous" + if "incompatible" in statuses: + return "incompatible" + return "not_found" + + +def _select_candidate( + candidates: Sequence[TaskCandidate], + audits: Sequence[Mapping[str, Any]], + *, + manifest: SceneManifest, + instruction: str, + adjudicator: Adjudicator | None, +) -> tuple[str | None, str, str]: + audit_by_id = {str(audit["candidate_id"]): audit for audit in audits} + bound = [ + candidate + for candidate in candidates + if audit_by_id[str(candidate["candidate_id"])]["status"] == "resolved" + ] + majority = [candidate for candidate in bound if int(candidate["vote_count"]) >= 2] + if len(majority) == 1: + return str(majority[0]["candidate_id"]), "bound", "majority_bindable" + if not majority and len(bound) == 1: + return str(bound[0]["candidate_id"]), "bound", "unique_bindable" + + choices = majority if majority else bound + if len(choices) > 1: + if adjudicator is None: + return None, "ambiguous", "multiple_conflicting_bindable_candidates" + raw = adjudicator( + instruction=instruction, + candidates=deepcopy(list(choices)), + scene_manifest=deepcopy(manifest), + ) + if not isinstance(raw, Mapping) or set(raw) != {"candidate_id"}: + raise SceneAdapterProtocolError( + "Adjudicator response must contain only candidate_id." + ) + selected_id = raw["candidate_id"] + allowed = {str(candidate["candidate_id"]) for candidate in choices} + if not isinstance(selected_id, str) or selected_id not in allowed: + raise SceneAdapterProtocolError( + "Adjudicator must select the candidate_id of a verified bindable candidate." + ) + return selected_id, "bound", "adjudicated_bindable" + if any(audit["status"] == "ambiguous" for audit in audits): + return None, "ambiguous", "no_fully_bound_candidate" + return None, "unsatisfied", "no_fully_bound_candidate" + + +def _build_manifest( + prepared: PreparedScene, + inventory: SceneInventory, + *, + source_format: str, +) -> SceneManifest: + objects = [ + { + "uid": entity.uid, + "role": entity.role, + "name": entity.name, + "description": entity.description, + "category": entity.category, + "color": entity.color, + "affordances": sorted(entity.affordances), + "initial_state": _redact_semantics(entity.initial_state), + "attributes": _redact_semantics(entity.attributes), + } + for entity in sorted(inventory.entities, key=lambda item: item.uid) + ] + scene_id = _canonical_hash( + { + "source_format": source_format, + "objects": objects, + "asset_hashes": prepared.asset_hashes, + "rotation": prepared.z_rotation_degrees, + "xy_translation": list(prepared.source_scene_xy_translation), + "body_scale_policy": prepared.body_scale_policy, + "body_scale": prepared.body_scale, + } + ) + return validate_scene_manifest( + { + "schema_version": SCENE_MANIFEST_SCHEMA, + "scene_id": scene_id, + "source_format": source_format, + "robot_profile": inventory.profile, + "objects": objects, + } + ) + + +def _build_semantic_manifest( + inventory: SceneInventory, + *, + source_format: str, +) -> SceneManifest: + objects = [ + { + "uid": entity.uid, + "role": entity.role, + "name": entity.name, + "description": entity.description, + "category": entity.category, + "color": entity.color, + "affordances": sorted(entity.affordances), + "initial_state": _redact_semantics(entity.initial_state), + "attributes": _redact_semantics(entity.attributes), + } + for entity in sorted(inventory.entities, key=lambda item: item.uid) + ] + return validate_scene_manifest( + { + "schema_version": SCENE_MANIFEST_SCHEMA, + "scene_id": _canonical_hash( + {"source_format": source_format, "objects": objects} + ), + "source_format": source_format, + "robot_profile": inventory.profile, + "objects": objects, + } + ) + + +def _redact_semantics(value: Mapping[str, Any]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, child in value.items(): + name = str(key) + normalized = name.strip().lower().replace("-", "_") + if normalized in _REDACTED_KEYS: + continue + if isinstance(child, Mapping): + nested = _redact_semantics(child) + if nested: + result[name] = nested + elif isinstance(child, (str, int, float, bool)) and not isinstance( + child, complex + ): + result[name] = child + elif isinstance(child, Sequence) and not isinstance(child, (str, bytes)): + simple = [item for item in child if isinstance(item, (str, bool))] + if len(simple) == len(child): + result[name] = simple + return result + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() diff --git a/embodichain/gen_sim/task_engine/orchestration/scene_source.py b/embodichain/gen_sim/task_engine/orchestration/scene_source.py new file mode 100644 index 000000000..f48cefb48 --- /dev/null +++ b/embodichain/gen_sim/task_engine/orchestration/scene_source.py @@ -0,0 +1,281 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Read-only references and integrity checks for existing Gym projects.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +import hashlib +import json +import os +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse +import xml.etree.ElementTree as ET + +from embodichain.data import get_data_path +from embodichain.gen_sim.action_engine.generation.source_scene import ( + resolve_source_scene, +) + +__all__ = [ + "SceneSourceFingerprint", + "SceneSourceRef", + "fingerprint_scene_source", + "scene_revision_id", + "verify_scene_source_fingerprint", +] + +_SCENE_SECTIONS = ("background", "rigid_object", "articulation") + + +@dataclass(frozen=True) +class SceneSourceRef: + """Reference an existing scene without copying or owning its files.""" + + path: Path | str + robot_profile: str = "franka" + z_rotation_degrees: float | None = None + body_scale_policy: str = "preserve" + body_scale: tuple[float, float, float] = (1.0, 1.0, 1.0) + + def __post_init__(self) -> None: + object.__setattr__(self, "path", Path(self.path).expanduser()) + + +@dataclass(frozen=True) +class SceneSourceFingerprint: + """Content evidence for one externally owned scene source.""" + + source_format: str + config_path: Path + config_sha256: str + asset_sha256: dict[str, str] + + def to_dict(self) -> dict[str, Any]: + """Return a JSON-safe audit view.""" + return { + "source_format": self.source_format, + "config_path": self.config_path.as_posix(), + "config_sha256": self.config_sha256, + "asset_sha256": dict(sorted(self.asset_sha256.items())), + } + + +def fingerprint_scene_source( + source: SceneSourceRef | str | Path, +) -> SceneSourceFingerprint: + """Hash a source config and referenced assets without copying either.""" + source_path = source.path if isinstance(source, SceneSourceRef) else source + resolved = resolve_source_scene(source_path) + config_bytes = resolved.path.read_bytes() + try: + config = json.loads(config_bytes) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene config is not valid JSON: {resolved.path}") from exc + if not isinstance(config, Mapping): + raise ValueError(f"Scene config must contain an object: {resolved.path}") + + asset_hashes: dict[str, str] = {} + for section in _SCENE_SECTIONS: + entries = config.get(section, ()) + if not isinstance(entries, Sequence) or isinstance(entries, (str, bytes)): + continue + for index, entry in enumerate(entries): + if not isinstance(entry, Mapping): + continue + references: list[tuple[str, Any]] = [] + shape = entry.get("shape") + if isinstance(shape, Mapping) and shape.get("fpath"): + references.append(("shape.fpath", shape["fpath"])) + if section == "articulation" and entry.get("fpath"): + references.append(("fpath", entry["fpath"])) + for field_name, reference in references: + asset_path = Path(str(reference)).expanduser() + if not asset_path.is_absolute(): + asset_path = resolved.path.parent / asset_path + asset_path = asset_path.resolve() + if not asset_path.is_file() and not Path(str(reference)).is_absolute(): + asset_path = ( + Path(get_data_path(str(reference))).expanduser().resolve() + ) + if not asset_path.is_file(): + raise FileNotFoundError( + f"Scene asset does not exist: {asset_path} " + f"({section}[{index}].{field_name})." + ) + for dependency in _asset_dependency_files(asset_path): + asset_hashes[dependency.as_posix()] = _sha256( + dependency.read_bytes() + ) + return SceneSourceFingerprint( + source_format=resolved.source_format, + config_path=resolved.path, + config_sha256=_sha256(config_bytes), + asset_sha256=asset_hashes, + ) + + +def verify_scene_source_fingerprint(expected: Mapping[str, Any]) -> None: + """Raise when an externally owned source changed after preparation.""" + required = {"source_format", "config_path", "config_sha256", "asset_sha256"} + if set(expected) != required: + raise ValueError("Scene source fingerprint fields are invalid.") + actual = fingerprint_scene_source(str(expected["config_path"])).to_dict() + normalized = { + "source_format": str(expected["source_format"]), + "config_path": Path(str(expected["config_path"])).resolve().as_posix(), + "config_sha256": str(expected["config_sha256"]), + "asset_sha256": dict(expected["asset_sha256"]), + } + if actual != normalized: + raise RuntimeError( + "Source Gym project changed after Task Engine preparation; " + "prepare a new bundle before running it." + ) + + +def scene_revision_id(source: SceneSourceRef | str | Path) -> str: + """Return a location-independent content identity for one scene revision. + + Volatile exporter IDs and absolute asset paths are excluded. Referenced + asset content remains part of the identity through SHA-256 placeholders. + + Args: + source: Scene project, configuration path, or Task Engine source reference. + + Returns: + Stable SHA-256 identity of scene semantics and referenced asset content. + """ + source_path = source.path if isinstance(source, SceneSourceRef) else source + resolved = resolve_source_scene(source_path) + try: + config = json.loads(resolved.path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene config is not valid JSON: {resolved.path}") from exc + if not isinstance(config, Mapping): + raise ValueError(f"Scene config must contain an object: {resolved.path}") + normalized = _normalize_revision_value( + dict(config), + config_root=resolved.path.parent, + ) + normalized.pop("scene_id", None) + payload = json.dumps( + normalized, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return _sha256(payload) + + +def _normalize_revision_value(value: Any, *, config_root: Path) -> Any: + if isinstance(value, Mapping): + result = { + str(key): _normalize_revision_value(item, config_root=config_root) + for key, item in value.items() + } + for key in ("fpath",): + raw = result.get(key) + if not isinstance(raw, str) or not raw: + continue + path = Path(raw).expanduser() + if not path.is_absolute(): + path = config_root / path + path = path.resolve() + if path.is_file(): + files = _asset_dependency_files(path) + result[key] = { + "sha256": _sha256(path.read_bytes()), + "dependency_sha256": { + Path( + os.path.relpath(dependency, start=path.parent) + ).as_posix(): (_sha256(dependency.read_bytes())) + for dependency in files + if dependency != path + }, + } + return result + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [ + _normalize_revision_value(item, config_root=config_root) for item in value + ] + return value + + +def _asset_dependency_files(asset_path: Path) -> tuple[Path, ...]: + """Return one asset and every local XML-declared dependency transitively.""" + pending = [asset_path.resolve()] + visited: set[Path] = set() + while pending: + path = pending.pop() + if path in visited: + continue + if not path.is_file(): + raise FileNotFoundError(f"Scene asset dependency does not exist: {path}") + visited.add(path) + if path.suffix.lower() not in {".urdf", ".xml", ".mjcf", ".xacro"}: + continue + try: + root = ET.parse(path).getroot() + except ET.ParseError: + # Opaque articulation assets remain valid direct dependencies even + # when their extension suggests XML. + continue + for element in root.iter(): + tag = element.tag.rsplit("}", maxsplit=1)[-1] + if tag not in {"mesh", "texture", "include"}: + continue + for attribute in ("filename", "file", "url"): + reference = element.attrib.get(attribute) + if reference: + pending.append(_resolve_asset_reference(path, reference)) + return tuple(sorted(visited)) + + +def _resolve_asset_reference(owner: Path, reference: str) -> Path: + """Resolve a local filesystem or ROS package URI without global state.""" + parsed = urlparse(reference) + if parsed.scheme in {"http", "https", "data"}: + raise ValueError( + f"Remote scene asset dependencies cannot be integrity-hashed: {reference}" + ) + if parsed.scheme == "file": + return Path(unquote(parsed.path)).expanduser().resolve() + if parsed.scheme == "package": + package_name = parsed.netloc + relative = Path(unquote(parsed.path.lstrip("/"))) + candidates = [ + ancestor / package_name / relative + for ancestor in (owner.parent, *owner.parents) + ] + candidates.append(owner.parent / relative) + for candidate in candidates: + if candidate.is_file(): + return candidate.resolve() + raise FileNotFoundError( + f"Unable to resolve package asset {reference!r} from {owner}." + ) + if parsed.scheme: + raise ValueError(f"Unsupported scene asset URI scheme: {reference}") + return (owner.parent / unquote(reference)).expanduser().resolve() + + +def _sha256(value: bytes) -> str: + return hashlib.sha256(value).hexdigest() diff --git a/embodichain/gen_sim/task_engine/run_directory.py b/embodichain/gen_sim/task_engine/run_directory.py new file mode 100644 index 000000000..0092cffc0 --- /dev/null +++ b/embodichain/gen_sim/task_engine/run_directory.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Collision-safe allocation of human-readable Task Engine run directories.""" + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path +from typing import Iterator + +__all__ = ["RunDirectory", "reserve_run_directory"] + + +@dataclass(frozen=True) +class RunDirectory: + """One reserved run identifier and its not-yet-published destination.""" + + run_id: str + output_root: Path + path: Path + created_at: datetime + + +@contextmanager +def reserve_run_directory( + output_root: str | Path, + *, + now: datetime | None = None, +) -> Iterator[RunDirectory]: + """Reserve a timestamped child name without creating its destination. + + Args: + output_root: Persistent task-history directory. + now: Optional timezone-aware timestamp used by deterministic tests. + + Yields: + A run directory allocation safe to publish through ArtifactTransaction. + """ + created_at = now or datetime.now().astimezone() + if created_at.tzinfo is None or created_at.utcoffset() is None: + raise ValueError("Task Engine run timestamps must include a timezone.") + root = Path(output_root).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + if not root.is_dir(): + raise NotADirectoryError(root) + + base = created_at.strftime("%Y%m%d_%H%M%S") + for collision_index in range(10_000): + run_id = base if collision_index == 0 else f"{base}_{collision_index:02d}" + destination = root / run_id + reservation = root / f".{run_id}.reserve" + if destination.exists(): + continue + try: + reservation.mkdir() + except FileExistsError: + continue + if destination.exists(): + reservation.rmdir() + continue + try: + yield RunDirectory( + run_id=run_id, + output_root=root, + path=destination, + created_at=created_at, + ) + finally: + reservation.rmdir() + return + raise RuntimeError("Unable to reserve a Task Engine run directory.") diff --git a/embodichain/gen_sim/task_engine/scene/__init__.py b/embodichain/gen_sim/task_engine/scene/__init__.py new file mode 100644 index 000000000..40a218b57 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/__init__.py @@ -0,0 +1,53 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Task Engine ownership of scene adaptation and static feasibility.""" + +from __future__ import annotations + +from .contracts import ( + ASSESSMENT_STATUSES, + FEASIBILITY_REPORT_SCHEMA, + STATIC_SCENE_MANIFEST_SCHEMA, + FeasibilityReport, + StaticSceneManifest, + validate_feasibility_report, + validate_static_scene_manifest, +) +from .feasibility import FeasibilityBroker +from .scene_engine_v1 import SceneEngineV1Adapter +from .conservative_graph import ( + CONSERVATIVE_SCENE_GRAPH_SCHEMA, + ConservativeSceneGraph, + build_conservative_scene_graph, + validate_conservative_scene_graph, +) + +__all__ = [ + "ASSESSMENT_STATUSES", + "CONSERVATIVE_SCENE_GRAPH_SCHEMA", + "FEASIBILITY_REPORT_SCHEMA", + "STATIC_SCENE_MANIFEST_SCHEMA", + "FeasibilityBroker", + "FeasibilityReport", + "SceneEngineV1Adapter", + "StaticSceneManifest", + "ConservativeSceneGraph", + "build_conservative_scene_graph", + "validate_feasibility_report", + "validate_static_scene_manifest", + "validate_conservative_scene_graph", +] diff --git a/embodichain/gen_sim/task_engine/scene/conservative_graph.py b/embodichain/gen_sim/task_engine/scene/conservative_graph.py new file mode 100644 index 000000000..ac3f2a634 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/conservative_graph.py @@ -0,0 +1,247 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Conservative hierarchy evidence for imported scenes.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +from pathlib import Path +from typing import Any, Final, TypeAlias + +__all__ = [ + "CONSERVATIVE_SCENE_GRAPH_SCHEMA", + "ConservativeSceneGraph", + "build_conservative_scene_graph", + "validate_conservative_scene_graph", +] + +CONSERVATIVE_SCENE_GRAPH_SCHEMA: Final = "embodichain.conservative-scene-graph/v1" +ConservativeSceneGraph: TypeAlias = dict[str, Any] + + +def build_conservative_scene_graph( + prepared_scene: Any, + *, + scene_id: str, +) -> ConservativeSceneGraph: + """Use exported hierarchy when available and mark every gap as unknown.""" + source_path = Path(getattr(prepared_scene, "source_config_path")).resolve() + uid_map = dict(getattr(prepared_scene, "uid_map", {}) or {}) + exported = _read_exported_graph(source_path.with_name("scene_graph.json")) + operational_assumptions = _legacy_operational_assumption_uids(source_path) + exported_nodes = { + str(node.get("object_id")): node + for node in exported.get("nodes", ()) + if isinstance(node, Mapping) and node.get("object_id") + } + + nodes: list[dict[str, Any]] = [] + for raw in getattr(prepared_scene, "planner_objects"): + uid = str(raw.get("uid", "")) + source_uid = str(raw.get("source_uid", uid)) + known = exported_nodes.get(source_uid) or exported_nodes.get(uid) + attributes = raw.get("attributes", {}) + final_support = ( + attributes.get("final_support") if isinstance(attributes, Mapping) else None + ) + initial_state = raw.get("initial_state", {}) + final_orientation = ( + initial_state.get("orientation") + if isinstance(initial_state, Mapping) + else None + ) + if uid == "table": + node = { + "uid": uid, + "parent_uid": None, + "parent_relation": "root", + "orientation": "unknown", + "source": "structural_root", + } + elif isinstance(final_support, Mapping): + relation = str(final_support.get("relation", "unknown")) + parent_uid = final_support.get("parent_uid", "unknown") + node = { + "uid": uid, + "parent_uid": ( + str(parent_uid) + if isinstance(parent_uid, str) and parent_uid + else "unknown" + ), + "parent_relation": relation if relation == "on" else "unknown", + "orientation": ( + "standing" + if final_orientation == "upright" + else ("lying" if final_orientation == "fallen" else "unknown") + ), + "source": "final_inspection", + } + elif ( + known is None + or uid in operational_assumptions + or source_uid in operational_assumptions + ): + node = { + "uid": uid, + "parent_uid": "unknown", + "parent_relation": "unknown", + "orientation": "unknown", + "source": "conservative_import", + } + else: + raw_parent = known.get("parent_id") + parent_uid = ( + uid_map.get(str(raw_parent), str(raw_parent)) + if raw_parent is not None + else "unknown" + ) + relation = known.get("parent_relation") + orientation = known.get("orientation_state") + node = { + "uid": uid, + "parent_uid": parent_uid, + "parent_relation": relation if relation == "on" else "unknown", + "orientation": ( + orientation if orientation in {"standing", "lying"} else "unknown" + ), + "source": "scene_graph", + } + nodes.append(node) + + relations = [] + for raw in exported.get("relations", ()): + if not isinstance(raw, Mapping): + continue + source_uid = uid_map.get(str(raw.get("source_id")), str(raw.get("source_id"))) + target_uid = uid_map.get(str(raw.get("target_id")), str(raw.get("target_id"))) + relation = str(raw.get("relation", "")) + if source_uid and target_uid and relation: + relations.append( + { + "source_uid": source_uid, + "relation": relation, + "target_uid": target_uid, + "source": "scene_graph", + } + ) + return validate_conservative_scene_graph( + { + "schema_version": CONSERVATIVE_SCENE_GRAPH_SCHEMA, + "scene_id": str(scene_id), + "nodes": nodes, + "relations": relations, + } + ) + + +def _legacy_operational_assumption_uids(source_path: Path) -> set[str]: + manifest_path = source_path.parent.parent / "legacy_conversion.json" + if not manifest_path.is_file(): + return set() + try: + value = json.loads(manifest_path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError( + f"Legacy conversion manifest is invalid JSON: {manifest_path}" + ) from exc + if not isinstance(value, Mapping): + raise ValueError("Legacy conversion manifest must contain an object.") + assumptions = value.get("assumptions", ()) + if not isinstance(assumptions, Sequence) or isinstance(assumptions, (str, bytes)): + raise ValueError("Legacy conversion assumptions must be a sequence.") + return { + str(item["uid"]) + for item in assumptions + if isinstance(item, Mapping) and isinstance(item.get("uid"), str) + } + + +def validate_conservative_scene_graph( + value: Mapping[str, Any], +) -> ConservativeSceneGraph: + """Validate and detach one conservative graph.""" + if not isinstance(value, Mapping): + raise TypeError("ConservativeSceneGraph must be a mapping.") + result = deepcopy(dict(value)) + expected = {"schema_version", "scene_id", "nodes", "relations"} + if set(result) != expected: + raise ValueError("ConservativeSceneGraph fields are invalid.") + if result.get("schema_version") != CONSERVATIVE_SCENE_GRAPH_SCHEMA: + raise ValueError("ConservativeSceneGraph schema version is invalid.") + if not isinstance(result.get("scene_id"), str) or not result["scene_id"]: + raise ValueError("ConservativeSceneGraph.scene_id must not be empty.") + nodes = _sequence(result.get("nodes"), "nodes") + normalized_nodes = [] + for index, raw in enumerate(nodes): + if not isinstance(raw, Mapping): + raise TypeError(f"ConservativeSceneGraph.nodes[{index}] must be a mapping.") + node = dict(raw) + if set(node) != { + "uid", + "parent_uid", + "parent_relation", + "orientation", + "source", + }: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}] fields are invalid." + ) + if not isinstance(node["uid"], str) or not node["uid"]: + raise ValueError(f"ConservativeSceneGraph.nodes[{index}].uid is invalid.") + if node["parent_uid"] is not None and not isinstance(node["parent_uid"], str): + raise TypeError( + f"ConservativeSceneGraph.nodes[{index}].parent_uid is invalid." + ) + if node["parent_relation"] not in {"root", "on", "unknown"}: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}].parent_relation is invalid." + ) + if node["orientation"] not in {"standing", "lying", "unknown"}: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}].orientation is invalid." + ) + if not isinstance(node["source"], str) or not node["source"]: + raise ValueError( + f"ConservativeSceneGraph.nodes[{index}].source is invalid." + ) + normalized_nodes.append(node) + if len({node["uid"] for node in normalized_nodes}) != len(normalized_nodes): + raise ValueError("ConservativeSceneGraph node UIDs must be unique.") + result["nodes"] = normalized_nodes + result["relations"] = [ + dict(item) for item in _sequence(result.get("relations"), "relations") + ] + json.dumps(result, allow_nan=False) + return result + + +def _read_exported_graph(path: Path) -> dict[str, Any]: + if not path.is_file(): + return {} + try: + value = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise ValueError(f"Scene graph is not valid JSON: {path}") from exc + return dict(value) if isinstance(value, Mapping) else {} + + +def _sequence(value: Any, field_name: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"ConservativeSceneGraph.{field_name} must be a sequence.") + return list(value) diff --git a/embodichain/gen_sim/task_engine/scene/contracts.py b/embodichain/gen_sim/task_engine/scene/contracts.py new file mode 100644 index 000000000..974864785 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/contracts.py @@ -0,0 +1,304 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""JSON contracts owned by the Scene Engine anti-corruption boundary.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import json +import math +from typing import Any, TypeAlias + +__all__ = [ + "ASSESSMENT_STATUSES", + "FEASIBILITY_REPORT_SCHEMA", + "REMEDIATION_CLASSES", + "STATIC_SCENE_MANIFEST_SCHEMA", + "FeasibilityReport", + "StaticSceneManifest", + "validate_feasibility_report", + "validate_static_scene_manifest", +] + + +STATIC_SCENE_MANIFEST_SCHEMA = "embodichain.static-scene-manifest/v1" +FEASIBILITY_REPORT_SCHEMA = "embodichain.scene-action-feasibility/v2" +ASSESSMENT_STATUSES = frozenset({"proven", "runtime_probe", "unknown", "contradicted"}) +REMEDIATION_CLASSES = frozenset( + {"none", "scene_remediable", "action_capability", "input_conflict", "terminal"} +) +_EVIDENCE_STATUSES = frozenset({"declared", "inferred", "verified", "contradicted"}) + +StaticSceneManifest: TypeAlias = dict[str, Any] +FeasibilityReport: TypeAlias = dict[str, Any] + + +def validate_static_scene_manifest(value: Mapping[str, Any]) -> StaticSceneManifest: + """Validate and detach one static scene manifest.""" + result = _mapping(value, "StaticSceneManifest") + _exact_keys( + result, + { + "schema_version", + "scene_id", + "source_format", + "robot_profile", + "source", + "adapter_capabilities", + "objects", + }, + "StaticSceneManifest", + ) + _schema(result, STATIC_SCENE_MANIFEST_SCHEMA, "StaticSceneManifest") + for key in ("scene_id", "source_format", "robot_profile"): + result[key] = _nonempty(result.get(key), f"StaticSceneManifest.{key}") + result["source"] = _mapping(result.get("source"), "StaticSceneManifest.source") + result["adapter_capabilities"] = _bool_mapping( + result.get("adapter_capabilities"), + "StaticSceneManifest.adapter_capabilities", + ) + + objects: list[dict[str, Any]] = [] + for index, raw in enumerate(_sequence(result.get("objects"), "objects")): + context = f"StaticSceneManifest.objects[{index}]" + item = _mapping(raw, context) + _exact_keys( + item, + { + "uid", + "source_uid", + "role", + "name", + "description", + "category", + "color", + "geometry", + "initial_pose", + "physics", + "articulation", + "affordances", + "initial_state", + "attributes", + "provenance", + }, + context, + ) + item["uid"] = _nonempty(item.get("uid"), f"{context}.uid") + item["source_uid"] = _string(item.get("source_uid"), f"{context}.source_uid") + item["role"] = _nonempty(item.get("role"), f"{context}.role") + for key in ("name", "description", "category"): + item[key] = _string(item.get(key), f"{context}.{key}") + color = item.get("color") + if color is not None: + color = _string(color, f"{context}.color") + item["color"] = color + for key in ( + "geometry", + "initial_pose", + "physics", + "articulation", + "initial_state", + "attributes", + "provenance", + ): + item[key] = _mapping(item.get(key), f"{context}.{key}") + item["affordances"] = [ + _validate_affordance(evidence, f"{context}.affordances[{evidence_index}]") + for evidence_index, evidence in enumerate( + _sequence(item.get("affordances"), f"{context}.affordances") + ) + ] + objects.append(item) + uids = [item["uid"] for item in objects] + if len(set(uids)) != len(uids): + raise ValueError("StaticSceneManifest object UIDs must be unique.") + result["objects"] = objects + _json_safe(result, "StaticSceneManifest") + return result + + +def validate_feasibility_report(value: Mapping[str, Any]) -> FeasibilityReport: + """Validate and detach one scene/action feasibility report.""" + result = _mapping(value, "FeasibilityReport") + _exact_keys( + result, + { + "schema_version", + "task_id", + "candidate_id", + "scene_id", + "status", + "remediation_class", + "checks", + "blockers", + "summary", + }, + "FeasibilityReport", + ) + _schema(result, FEASIBILITY_REPORT_SCHEMA, "FeasibilityReport") + for key in ("task_id", "candidate_id", "scene_id"): + result[key] = _nonempty(result.get(key), f"FeasibilityReport.{key}") + result["status"] = _status(result.get("status"), "FeasibilityReport.status") + remediation_class = result.get("remediation_class") + if remediation_class not in REMEDIATION_CLASSES: + raise ValueError( + "FeasibilityReport.remediation_class must be one of " + f"{sorted(REMEDIATION_CLASSES)}." + ) + result["remediation_class"] = str(remediation_class) + if result["status"] != "contradicted" and remediation_class != "none": + raise ValueError( + "A non-contradicted FeasibilityReport requires remediation_class=none." + ) + checks: list[dict[str, Any]] = [] + for index, raw in enumerate(_sequence(result.get("checks"), "checks")): + context = f"FeasibilityReport.checks[{index}]" + item = _mapping(raw, context) + _exact_keys( + item, + {"kind", "subject", "status", "reason", "evidence"}, + context, + ) + item["kind"] = _nonempty(item.get("kind"), f"{context}.kind") + item["subject"] = _nonempty(item.get("subject"), f"{context}.subject") + item["status"] = _status(item.get("status"), f"{context}.status") + item["reason"] = _nonempty(item.get("reason"), f"{context}.reason") + item["evidence"] = _mapping(item.get("evidence"), f"{context}.evidence") + checks.append(item) + result["checks"] = checks + blockers = _sequence(result.get("blockers"), "FeasibilityReport.blockers") + if any(not isinstance(item, str) or not item for item in blockers): + raise ValueError("FeasibilityReport.blockers must contain non-empty strings.") + result["blockers"] = list(blockers) + summary = _mapping(result.get("summary"), "FeasibilityReport.summary") + expected = set(ASSESSMENT_STATUSES) + if set(summary) != expected: + raise ValueError( + "FeasibilityReport.summary must count every assessment status." + ) + if any( + isinstance(value, bool) or not isinstance(value, int) or value < 0 + for value in summary.values() + ): + raise ValueError( + "FeasibilityReport.summary counts must be non-negative integers." + ) + if sum(summary.values()) != len(checks): + raise ValueError("FeasibilityReport.summary must match the check count.") + result["summary"] = dict(summary) + _json_safe(result, "FeasibilityReport") + return result + + +def _validate_affordance(value: Any, context: str) -> dict[str, Any]: + item = _mapping(value, context) + _exact_keys( + item, + { + "type", + "status", + "confidence", + "source", + "link_uid", + "frame", + "parameters", + }, + context, + ) + item["type"] = _nonempty(item.get("type"), f"{context}.type") + status = item.get("status") + if status not in _EVIDENCE_STATUSES: + raise ValueError( + f"{context}.status must be one of {sorted(_EVIDENCE_STATUSES)}." + ) + confidence = item.get("confidence") + if confidence is not None: + if ( + isinstance(confidence, bool) + or not isinstance(confidence, (int, float)) + or not math.isfinite(float(confidence)) + or not 0.0 <= float(confidence) <= 1.0 + ): + raise ValueError(f"{context}.confidence must be null or in [0, 1].") + confidence = float(confidence) + item["confidence"] = confidence + item["source"] = _nonempty(item.get("source"), f"{context}.source") + item["link_uid"] = _string(item.get("link_uid"), f"{context}.link_uid") + item["frame"] = _mapping(item.get("frame"), f"{context}.frame") + item["parameters"] = _mapping(item.get("parameters"), f"{context}.parameters") + return item + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{context} must be a mapping.") + return deepcopy(dict(value)) + + +def _sequence(value: Any, context: str) -> list[Any]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"{context} must be a sequence.") + return list(value) + + +def _exact_keys(value: Mapping[str, Any], expected: set[str], context: str) -> None: + if set(value) != expected: + missing = sorted(expected - set(value)) + extra = sorted(set(value) - expected) + raise ValueError(f"{context} fields differ; missing={missing}, extra={extra}.") + + +def _schema(value: Mapping[str, Any], expected: str, context: str) -> None: + if value.get("schema_version") != expected: + raise ValueError(f"{context}.schema_version must be {expected!r}.") + + +def _nonempty(value: Any, context: str) -> str: + result = _string(value, context).strip() + if not result: + raise ValueError(f"{context} must not be empty.") + return result + + +def _string(value: Any, context: str) -> str: + if not isinstance(value, str): + raise TypeError(f"{context} must be a string.") + return value + + +def _status(value: Any, context: str) -> str: + if value not in ASSESSMENT_STATUSES: + raise ValueError(f"{context} must be one of {sorted(ASSESSMENT_STATUSES)}.") + return str(value) + + +def _bool_mapping(value: Any, context: str) -> dict[str, bool]: + result = _mapping(value, context) + if any( + not isinstance(key, str) or not isinstance(item, bool) + for key, item in result.items() + ): + raise TypeError(f"{context} must map strings to booleans.") + return result + + +def _json_safe(value: Any, context: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError(f"{context} must contain strict JSON data.") from exc diff --git a/embodichain/gen_sim/task_engine/scene/feasibility.py b/embodichain/gen_sim/task_engine/scene/feasibility.py new file mode 100644 index 000000000..363eb4be9 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/feasibility.py @@ -0,0 +1,746 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Deterministic task, scene, robot, and action-capability intersection.""" + +from __future__ import annotations + +from collections import Counter +from collections.abc import Mapping, Sequence +import math +from typing import Any + +from .contracts import ( + ASSESSMENT_STATUSES, + FEASIBILITY_REPORT_SCHEMA, + FeasibilityReport, + validate_feasibility_report, + validate_static_scene_manifest, +) + +__all__ = ["FeasibilityBroker"] + + +_STATUS_PRIORITY = { + "proven": 0, + "runtime_probe": 1, + "unknown": 2, + "contradicted": 3, +} + + +class FeasibilityBroker: + """Produce an auditable compatibility report without repairing inputs.""" + + def assess( + self, + candidate: Mapping[str, Any], + role_bindings: Mapping[str, Any], + scene_manifest: Mapping[str, Any], + *, + capability_catalog: Mapping[str, Mapping[str, Any]], + task_actions: Mapping[str, Sequence[str]], + ) -> FeasibilityReport: + """Assess one grounded candidate against static and runtime capabilities.""" + manifest = validate_static_scene_manifest(scene_manifest) + draft = _mapping(candidate.get("draft"), "candidate.draft") + scene_request = _mapping( + candidate.get("scene_request"), "candidate.scene_request" + ) + bindings = role_bindings.get("reference_bindings", role_bindings) + bindings = _mapping(bindings, "role_bindings.reference_bindings") + objects = {item["uid"]: item for item in manifest["objects"]} + steps = { + str(item["id"]): item + for item in _sequence(draft.get("steps"), "candidate.draft.steps") + } + checks: list[dict[str, Any]] = [] + + for step_id, step in steps.items(): + task_type = str(step.get("task_type", "")) + actions = task_actions.get(task_type) + if not actions: + checks.append( + _check( + "task_capability", + step_id, + "contradicted", + f"Task type {task_type!r} has no registered action recipe.", + ) + ) + continue + for action_name in actions: + capability = capability_catalog.get(str(action_name)) + if capability is None: + checks.append( + _check( + "atomic_capability", + f"{step_id}:{action_name}", + "contradicted", + f"AtomicAction {action_name!r} is not registered.", + ) + ) + elif not bool(capability.get("runtime_available", False)): + checks.append( + _check( + "atomic_capability", + f"{step_id}:{action_name}", + "contradicted", + str( + capability.get("unavailable_reason") + or "Action is planning-only." + ), + evidence={"action": str(action_name)}, + ) + ) + else: + checks.append( + _check( + "atomic_capability", + f"{step_id}:{action_name}", + "proven", + "AtomicAction is registered and executable.", + evidence={"action": str(action_name)}, + ) + ) + if task_type == "E3": + runtime_observation = bool( + manifest.get("adapter_capabilities", {}).get( + "runtime_scene_observation", False + ) + ) + checks.append( + _check( + "content_observation", + step_id, + "runtime_probe" if runtime_observation else "contradicted", + ( + "Runtime scene observation can verify independently " + "modeled contents after pouring." + if runtime_observation + else "E3 requires independently observable content " + "bodies or fluid state; contents baked into one source " + "mesh cannot prove physical transfer." + ), + evidence={ + "runtime_scene_observation": runtime_observation, + "required_evidence": "content_inside_target_container", + }, + ) + ) + if task_type == "E8": + reference_id = f"{step_id}.object" + raw_uids = bindings.get(reference_id, ()) + if not isinstance(raw_uids, Sequence) or isinstance( + raw_uids, (str, bytes) + ): + raw_uids = () + setting_maps = [] + for uid in raw_uids: + entity = objects.get(str(uid), {}) + attributes = entity.get("attributes", {}) + joint_settings = ( + attributes.get("joint_settings", {}) + if isinstance(attributes, Mapping) + else {} + ) + if isinstance(joint_settings, Mapping): + setting_maps.extend( + list(values) + for values in joint_settings.values() + if isinstance(values, Sequence) + and not isinstance(values, (str, bytes)) + and values + ) + for evidence in entity.get("affordances", ()): + if str(evidence.get("type")) != "turnable": + continue + parameters = evidence.get("parameters", {}) + values = ( + parameters.get("setting_values") + if isinstance(parameters, Mapping) + else None + ) + if ( + isinstance(values, Sequence) + and not isinstance(values, (str, bytes)) + and values + ): + setting_maps.append(list(values)) + checks.append( + _check( + "setting_mapping", + step_id, + "proven" if len(setting_maps) == 1 else "contradicted", + ( + "One explicit knob setting-to-angle map is available." + if len(setting_maps) == 1 + else "E8 requires exactly one explicit setting_values " + "map; ordinal knob settings cannot be inferred from " + "joint limits." + ), + evidence={"setting_map_count": len(setting_maps)}, + ) + ) + + for request in _sequence( + scene_request.get("references"), "candidate.scene_request.references" + ): + reference_id = str(request.get("reference_id", "")) + raw_uids = bindings.get(reference_id, ()) + if not isinstance(raw_uids, Sequence) or isinstance(raw_uids, (str, bytes)): + raw_uids = () + uids = [str(uid) for uid in raw_uids] + if not uids: + checks.append( + _check( + "binding", + reference_id, + "contradicted", + "Reference has no grounded scene entity.", + ) + ) + continue + for uid in uids: + entity = objects.get(uid) + if entity is None: + checks.append( + _check( + "binding", + f"{reference_id}:{uid}", + "contradicted", + "Binding references an entity absent from the static manifest.", + ) + ) + continue + checks.extend(self._entity_checks(request, entity, reference_id)) + + checks.extend(self._workspace_checks(steps, bindings, objects)) + + statuses = Counter(check["status"] for check in checks) + status = max( + (check["status"] for check in checks), + key=_STATUS_PRIORITY.__getitem__, + default="unknown", + ) + blockers = sorted( + { + f"{check['subject']}: {check['reason']}" + for check in checks + if check["status"] == "contradicted" + } + ) + return validate_feasibility_report( + { + "schema_version": FEASIBILITY_REPORT_SCHEMA, + "task_id": str(draft.get("task_id", "")), + "candidate_id": str(candidate.get("candidate_id", "")), + "scene_id": manifest["scene_id"], + "status": status, + "remediation_class": _remediation_class(checks), + "checks": checks, + "blockers": blockers, + "summary": { + name: int(statuses.get(name, 0)) + for name in sorted(ASSESSMENT_STATUSES) + }, + } + ) + + def _entity_checks( + self, + request: Mapping[str, Any], + entity: Mapping[str, Any], + reference_id: str, + ) -> list[dict[str, Any]]: + uid = str(entity["uid"]) + subject = f"{reference_id}:{uid}" + checks = [self._structure_check(request, entity, subject)] + evidence_by_type: dict[str, list[Mapping[str, Any]]] = {} + for evidence in entity["affordances"]: + evidence_by_type.setdefault(str(evidence["type"]), []).append(evidence) + for affordance in request.get("affordances", ()): + name = str(affordance) + checks.append( + self._affordance_check(name, evidence_by_type.get(name, ()), subject) + ) + for field_name in ("initial_state", "attributes"): + required = request.get(field_name, {}) + actual = entity.get(field_name, {}) + if isinstance(required, Mapping) and isinstance(actual, Mapping): + for key, expected in required.items(): + if key not in actual: + status = "unknown" + reason = f"Required {field_name} field {key!r} is not declared." + elif actual[key] != expected: + status = "contradicted" + reason = f"Required {field_name} field {key!r} conflicts with the scene." + else: + status = "proven" + reason = f"Required {field_name} field {key!r} matches." + checks.append( + _check( + field_name, + subject, + status, + reason, + evidence={"field": str(key)}, + ) + ) + if str(request.get("role")) == "object": + checks.append( + _check( + "runtime_reachability", + subject, + "runtime_probe", + "Reachability, collision, and grasp geometry require live planning.", + ) + ) + if ( + str(request.get("role")) == "target" + and str(request.get("source_structure")) == "physical_entity" + ): + checks.append( + _check( + "placement_support", + subject, + "runtime_probe", + "Support depends on the payload, candidate pose, live geometry, " + "and post-release stability.", + evidence={ + "runtime_obligations": [ + "placement_candidates", + "object_supported_by", + "stable_for", + "final_support_revalidation", + ] + }, + ) + ) + return checks + + def _workspace_checks( + self, + steps: Mapping[str, Mapping[str, Any]], + bindings: Mapping[str, Any], + objects: Mapping[str, Mapping[str, Any]], + ) -> list[dict[str, Any]]: + """Defer arm-side compatibility to the live robot frame.""" + checks: list[dict[str, Any]] = [] + object_uids_by_step: dict[str, tuple[str, ...]] = {} + phases: list[dict[str, Any]] = [] + for step_id, step in steps.items(): + object_uids = _step_selector_uids( + step_id, + "object", + step.get("object"), + bindings, + object_uids_by_step, + ) + object_uids_by_step[step_id] = object_uids + target_uids = _step_selector_uids( + step_id, + "target", + step.get("target"), + bindings, + object_uids_by_step, + ) + task_type = str(step.get("task_type", "")) + required_arm = str(step.get("required_arm", "auto")) + if task_type == "E4": + required_arm = str(step.get("transfer_arm", "none")) + if required_arm in {"left_arm", "right_arm"}: + for uid in object_uids: + entity = objects.get(uid) + position = ( + entity.get("initial_pose", {}).get("position", ()) + if isinstance(entity, Mapping) + and isinstance(entity.get("initial_pose"), Mapping) + else () + ) + if ( + not isinstance(position, Sequence) + or isinstance(position, (str, bytes, bytearray)) + or len(position) < 2 + ): + continue + checks.append( + _check( + "arm_layout_risk", + f"{step_id}:{uid}", + "runtime_probe", + "Arm-side compatibility requires live left/right arm-base " + "poses and workspace geometry.", + evidence={ + "required_arm": required_arm, + "object_world_position": [ + float(position[0]), + float(position[1]), + ], + "arm_side_frame": "live_robot", + "mismatch_risk": None, + "geometry_certificate": False, + }, + ) + ) + + phases.extend( + _workflow_phases( + step_id, + task_type, + object_uids, + target_uids, + transfer_arm=str(step.get("transfer_arm", "none")), + receive_arm=str(step.get("receive_arm", "none")), + ) + ) + if phases: + checks.append( + _check( + "task_workspace", + "task_workflow", + "runtime_probe", + "Scene layout must satisfy pickup, transfer, placement, and " + "safety-clearance phases across the complete task workflow.", + evidence={ + "arm_side_frame": "live_robot", + "phases": phases, + "geometry_certificate": False, + }, + ) + ) + return checks + + @staticmethod + def _structure_check( + request: Mapping[str, Any], + entity: Mapping[str, Any], + subject: str, + ) -> dict[str, Any]: + expected = str(request.get("source_structure", "")) + role = str(entity.get("role", "")) + if expected in {"scene_entity", "spatial_reference"}: + if role in {"camera", "light", "robot", "sensor"}: + return _check( + "structure", + subject, + "contradicted", + f"Scene entity role {role!r} cannot be a spatial action target.", + evidence={"static_pose": _has_static_pose(entity)}, + ) + if not _has_static_pose(entity): + return _check( + "structure", + subject, + "unknown", + "Static scene evidence does not provide a finite spatial pose.", + evidence={"static_pose": False}, + ) + if role == "articulation": + return _check( + "structure", + subject, + "runtime_probe", + "Articulation has a static pose, but live spatial target lookup " + "must be validated at runtime.", + evidence={ + "static_pose": True, + "runtime_entity_kind": "articulation", + }, + ) + has_runtime_body = bool(entity.get("physics")) + if ( + role + in { + "background", + "object", + "rigid_object", + "support_surface", + "table", + } + and has_runtime_body + ): + return _check( + "structure", + subject, + "proven", + "Scene entity has a static pose and a rigid runtime body.", + evidence={ + "static_pose": True, + "runtime_entity_kind": "rigid_object", + }, + ) + return _check( + "structure", + subject, + "runtime_probe", + "Scene entity has a static pose, but its live target interface is " + "not proven by the static manifest.", + evidence={"static_pose": True, "runtime_entity_kind": "unknown"}, + ) + if expected == "physical_entity": + geometry = entity.get("geometry", {}) + shape = geometry.get("shape", {}) if isinstance(geometry, Mapping) else {} + asset_sha256 = ( + geometry.get("asset_sha256", "") + if isinstance(geometry, Mapping) + else "" + ) + physics = entity.get("physics", {}) + articulation = entity.get("articulation", {}) + has_physical_geometry = bool(shape) or bool(asset_sha256) + has_runtime_body = bool(physics) or bool(articulation) + if role in {"camera", "light", "sensor"}: + return _check( + "structure", + subject, + "contradicted", + f"Scene entity role {role!r} is not a physical collision body.", + evidence={"physical_geometry": False, "runtime_body": False}, + ) + if role == "articulation" or bool(articulation): + return _check( + "structure", + subject, + "contradicted", + "Placement on an articulation requires a link-level target " + "interface that the current runtime does not provide.", + evidence={ + "physical_geometry": has_physical_geometry, + "runtime_body": bool(articulation), + "runtime_entity_kind": "articulation", + "runtime_target_interface": False, + }, + ) + if has_physical_geometry and has_runtime_body: + return _check( + "structure", + subject, + "proven", + "Scene entity has physical geometry and a runtime body.", + evidence={"physical_geometry": True, "runtime_body": True}, + ) + return _check( + "structure", + subject, + "unknown", + "Static scene evidence does not prove physical geometry and a " + "runtime body required for placement.", + evidence={ + "physical_geometry": has_physical_geometry, + "runtime_body": has_runtime_body, + }, + ) + accepted_by_structure = { + "articulation": {"articulation"}, + "rigid_object": {"object", "rigid_object"}, + "movable": {"object", "rigid_object"}, + "support_surface": {"background", "support_surface", "table"}, + } + accepted = accepted_by_structure.get(expected) + if accepted is None: + return _check( + "structure", + subject, + "unknown", + f"Structure contract {expected!r} is not recognized by the broker.", + evidence={"scene_role": role}, + ) + if role in accepted: + return _check( + "structure", + subject, + "proven", + f"Scene role {role!r} satisfies structure {expected!r}.", + ) + return _check( + "structure", + subject, + "contradicted", + f"Scene role {role!r} does not satisfy structure {expected!r}.", + ) + + @staticmethod + def _affordance_check( + name: str, + evidence: Sequence[Mapping[str, Any]], + subject: str, + ) -> dict[str, Any]: + if not evidence: + return _check( + "affordance", + subject, + "unknown", + f"Affordance {name!r} has no evidence.", + evidence={"affordance": name}, + ) + statuses = {str(item.get("status")) for item in evidence} + if statuses == {"contradicted"}: + status = "contradicted" + reason = f"Affordance {name!r} is explicitly contradicted." + elif "verified" in statuses: + status = "proven" + reason = f"Affordance {name!r} has verified evidence." + else: + status = "runtime_probe" + reason = ( + f"Affordance {name!r} is declared but requires physical validation." + ) + return _check( + "affordance", + subject, + status, + reason, + evidence={ + "affordance": name, + "sources": sorted({str(item.get("source")) for item in evidence}), + }, + ) + + +def _has_static_pose(entity: Mapping[str, Any]) -> bool: + pose = entity.get("initial_pose") + if not isinstance(pose, Mapping): + return False + position = pose.get("position") + if ( + not isinstance(position, Sequence) + or isinstance(position, (str, bytes, bytearray)) + or len(position) != 3 + ): + return False + return all( + not isinstance(value, bool) + and isinstance(value, (int, float)) + and math.isfinite(float(value)) + for value in position + ) + + +def _step_selector_uids( + step_id: str, + role: str, + selector: Any, + bindings: Mapping[str, Any], + object_uids_by_step: Mapping[str, tuple[str, ...]], +) -> tuple[str, ...]: + """Resolve direct and prior-step selectors for static workspace advice.""" + raw = bindings.get(f"{step_id}.{role}", ()) + if isinstance(raw, Sequence) and not isinstance(raw, (str, bytes, bytearray)): + direct = tuple(str(uid) for uid in raw if str(uid)) + if direct: + return direct + if not isinstance(selector, Mapping) or selector.get("kind") != "step_result": + return () + source_step = str(selector.get("step_id", "")) + return tuple(object_uids_by_step.get(source_step, ())) + + +def _workflow_phases( + step_id: str, + task_type: str, + object_uids: Sequence[str], + target_uids: Sequence[str], + *, + transfer_arm: str, + receive_arm: str, +) -> list[dict[str, Any]]: + """Describe whole-task layout anchors without inventing geometry bounds.""" + phases: list[dict[str, Any]] = [] + if object_uids: + phases.append( + { + "step_id": step_id, + "phase": "pickup", + "object_uids": list(object_uids), + } + ) + if task_type == "E4": + phases.append( + { + "step_id": step_id, + "phase": "handover_shared_workspace", + "object_uids": list(object_uids), + "transfer_arm": transfer_arm, + "receive_arm": receive_arm, + } + ) + if target_uids: + phases.append( + { + "step_id": step_id, + "phase": "target_interaction", + "object_uids": list(object_uids), + "target_uids": list(target_uids), + } + ) + if task_type in {"E1", "E2", "E3", "E4", "E5"}: + phases.append( + { + "step_id": step_id, + "phase": "safety_clearance", + "object_uids": list(object_uids), + } + ) + return phases + + +def _check( + kind: str, + subject: str, + status: str, + reason: str, + *, + evidence: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + return { + "kind": kind, + "subject": subject, + "status": status, + "reason": reason, + "evidence": dict(evidence or {}), + } + + +def _remediation_class(checks: Sequence[Mapping[str, Any]]) -> str: + """Classify contradictions by the subsystem capable of changing them.""" + contradicted = [check for check in checks if check.get("status") == "contradicted"] + if not contradicted: + return "none" + kinds = {str(check.get("kind", "")) for check in contradicted} + if kinds.intersection({"task_capability", "atomic_capability"}): + return "action_capability" + # A new materialization seed can change observed pose/orientation, but it + # cannot change task semantics, entity roles, bindings, or declared affordances. + if kinds <= {"initial_state"}: + return "scene_remediable" + if kinds.intersection({"binding", "structure", "affordance", "attributes"}): + return "input_conflict" + return "terminal" + + +def _mapping(value: Any, context: str) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"{context} must be a mapping.") + return dict(value) + + +def _sequence(value: Any, context: str) -> list[Mapping[str, Any]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + raise TypeError(f"{context} must be a sequence.") + if any(not isinstance(item, Mapping) for item in value): + raise TypeError(f"{context} must contain mappings.") + return list(value) diff --git a/embodichain/gen_sim/task_engine/scene/final_inspection.py b/embodichain/gen_sim/task_engine/scene/final_inspection.py new file mode 100644 index 000000000..f7bf3b68f --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/final_inspection.py @@ -0,0 +1,428 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Geometry-derived evidence from one completed scene revision.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import replace +import json +from pathlib import Path +from typing import Any, Final, TypeAlias + +import numpy as np +from scipy.spatial.transform import Rotation +import trimesh + +from embodichain.gen_sim.action_engine.generation.models import PreparedScene +from embodichain.gen_sim.action_engine.generation.source_scene import ( + prepare_scene, + resolve_source_scene, +) + +__all__ = [ + "FINAL_SCENE_INSPECTION_SCHEMA", + "FinalSceneInspection", + "apply_final_inspection", + "inspect_final_scene", + "validate_final_scene_inspection", +] + +FINAL_SCENE_INSPECTION_SCHEMA: Final = "embodichain.final-scene-inspection/v1" +FinalSceneInspection: TypeAlias = dict[str, Any] + +_Y_UP_TO_Z_UP = np.array( + [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]], + dtype=float, +) + + +def inspect_final_scene( + source: str | Path, + *, + revision_id: str, + contact_tolerance_m: float = 0.03, +) -> FinalSceneInspection: + """Measure final AABBs, orientation, and support from exported geometry. + + Args: + source: Completed scene project or configuration path. + revision_id: Content identity already assigned to the completed revision. + contact_tolerance_m: Maximum support-surface contact gap in meters. + + Returns: + Strict geometry-derived final inspection document. + """ + tolerance = float(contact_tolerance_m) + if not np.isfinite(tolerance) or tolerance <= 0.0: + raise ValueError("contact_tolerance_m must be positive and finite.") + normalized_revision_id = str(revision_id) + if len(normalized_revision_id) != 64: + raise ValueError("revision_id must be a SHA-256 hexadecimal digest.") + try: + int(normalized_revision_id, 16) + except ValueError as exc: + raise ValueError("revision_id must be a SHA-256 hexadecimal digest.") from exc + resolved = resolve_source_scene(source) + prepared = prepare_scene(source) + runtime = { + str(item.get("uid")): item + for item in ( + *prepared.background, + *prepared.rigid_objects, + *prepared.articulations, + ) + if isinstance(item, Mapping) and item.get("uid") + } + measured: dict[str, dict[str, Any]] = {} + for raw in prepared.planner_objects: + uid = str(raw.get("uid", "")) + role = str(raw.get("role", "")) + geometry = _measure_geometry( + runtime.get(uid, raw), + convert_y_up=resolved.is_prompt2scene, + ) + measured[uid] = { + "uid": uid, + "role": role, + "orientation": _orientation(geometry), + "support": { + "parent_uid": None if uid == "table" else "unknown", + "relation": "root" if uid == "table" else "unknown", + "confidence": 1.0 if uid == "table" else None, + "gap_m": None, + "xy_overlap_ratio": None, + }, + "world_aabb": ( + None + if geometry is None + else { + "min": geometry["bounds"][0].tolist(), + "max": geometry["bounds"][1].tolist(), + } + ), + "evidence": { + "source": "final_geometry" if geometry is not None else "unmeasured", + "method": "world_aabb_and_dominant_axis", + }, + } + + for uid, item in measured.items(): + child_geometry = _geometry_from_record(item) + if uid == "table" or child_geometry is None: + continue + support = _support_for( + uid, + child_geometry, + measured, + tolerance=tolerance, + ) + if support is not None: + item["support"] = support + + return validate_final_scene_inspection( + { + "schema_version": FINAL_SCENE_INSPECTION_SCHEMA, + "scene_revision_id": normalized_revision_id, + "source_config_path": prepared.source_config_path.as_posix(), + "contact_tolerance_m": tolerance, + "objects": [measured[uid] for uid in sorted(measured)], + } + ) + + +def apply_final_inspection( + prepared_scene: PreparedScene, + inspection: Mapping[str, Any], +) -> PreparedScene: + """Return a detached PreparedScene enriched with measured final evidence. + + Args: + prepared_scene: Normalized scene to enrich without mutation. + inspection: Validated or raw final inspection mapping. + + Returns: + Prepared scene whose semantic state reflects measured final geometry. + """ + normalized = validate_final_scene_inspection(inspection) + by_uid = {str(item["uid"]): item for item in normalized["objects"]} + planner_objects = [] + for raw in prepared_scene.planner_objects: + item = deepcopy(raw) + evidence = by_uid.get(str(item.get("uid"))) + if evidence is not None: + initial_state = deepcopy(dict(item.get("initial_state", {}))) + initial_state.pop("orientation", None) + if evidence["orientation"] == "standing": + initial_state["orientation"] = "upright" + elif evidence["orientation"] == "lying": + initial_state["orientation"] = "fallen" + attributes = deepcopy(dict(item.get("attributes", {}))) + attributes["final_support"] = deepcopy(evidence["support"]) + attributes["final_world_aabb"] = deepcopy(evidence["world_aabb"]) + item["initial_state"] = initial_state + item["attributes"] = attributes + planner_objects.append(item) + return replace(prepared_scene, planner_objects=tuple(planner_objects)) + + +def validate_final_scene_inspection( + value: Mapping[str, Any], +) -> FinalSceneInspection: + """Validate and detach one final scene inspection document. + + Args: + value: Inspection mapping to validate. + + Returns: + Detached, normalized inspection document. + """ + if not isinstance(value, Mapping): + raise TypeError("FinalSceneInspection must be a mapping.") + result = deepcopy(dict(value)) + expected = { + "schema_version", + "scene_revision_id", + "source_config_path", + "contact_tolerance_m", + "objects", + } + if set(result) != expected: + raise ValueError("FinalSceneInspection fields are invalid.") + if result.get("schema_version") != FINAL_SCENE_INSPECTION_SCHEMA: + raise ValueError("FinalSceneInspection schema version is invalid.") + revision_id = result.get("scene_revision_id") + if not isinstance(revision_id, str) or len(revision_id) != 64: + raise ValueError("FinalSceneInspection.scene_revision_id is invalid.") + try: + int(revision_id, 16) + except ValueError as exc: + raise ValueError("FinalSceneInspection.scene_revision_id is invalid.") from exc + source_path = result.get("source_config_path") + if not isinstance(source_path, str) or not source_path: + raise ValueError("FinalSceneInspection.source_config_path is invalid.") + tolerance = result.get("contact_tolerance_m") + if ( + isinstance(tolerance, bool) + or not isinstance(tolerance, (int, float)) + or not np.isfinite(float(tolerance)) + or float(tolerance) <= 0.0 + ): + raise ValueError("FinalSceneInspection.contact_tolerance_m is invalid.") + objects = result.get("objects") + if not isinstance(objects, Sequence) or isinstance(objects, (str, bytes)): + raise TypeError("FinalSceneInspection.objects must be a sequence.") + normalized = [_validate_object(item, index) for index, item in enumerate(objects)] + if len({item["uid"] for item in normalized}) != len(normalized): + raise ValueError("FinalSceneInspection object UIDs must be unique.") + result["objects"] = normalized + result["contact_tolerance_m"] = float(tolerance) + json.dumps(result, ensure_ascii=False, allow_nan=False) + return result + + +def _validate_object(value: Any, index: int) -> dict[str, Any]: + if not isinstance(value, Mapping): + raise TypeError(f"FinalSceneInspection.objects[{index}] must be a mapping.") + item = deepcopy(dict(value)) + expected = {"uid", "role", "orientation", "support", "world_aabb", "evidence"} + if set(item) != expected: + raise ValueError(f"FinalSceneInspection.objects[{index}] fields are invalid.") + if not isinstance(item["uid"], str) or not item["uid"]: + raise ValueError(f"FinalSceneInspection.objects[{index}].uid is invalid.") + if not isinstance(item["role"], str) or not item["role"]: + raise ValueError(f"FinalSceneInspection.objects[{index}].role is invalid.") + if item["orientation"] not in {"standing", "lying", "unknown"}: + raise ValueError( + f"FinalSceneInspection.objects[{index}].orientation is invalid." + ) + if not isinstance(item["support"], Mapping) or not isinstance( + item["evidence"], Mapping + ): + raise TypeError("FinalSceneInspection support and evidence must be mappings.") + support = deepcopy(dict(item["support"])) + if set(support) != { + "parent_uid", + "relation", + "confidence", + "gap_m", + "xy_overlap_ratio", + }: + raise ValueError("FinalSceneInspection support fields are invalid.") + if support["parent_uid"] is not None and not isinstance(support["parent_uid"], str): + raise TypeError("FinalSceneInspection support parent_uid is invalid.") + if support["relation"] not in {"root", "on", "unknown"}: + raise ValueError("FinalSceneInspection support relation is invalid.") + for field_name in ("confidence", "gap_m", "xy_overlap_ratio"): + field_value = support[field_name] + if field_value is not None and ( + isinstance(field_value, bool) + or not isinstance(field_value, (int, float)) + or not np.isfinite(float(field_value)) + ): + raise ValueError(f"FinalSceneInspection support {field_name} is invalid.") + if ( + support["confidence"] is not None + and not 0.0 <= float(support["confidence"]) <= 1.0 + ): + raise ValueError("FinalSceneInspection support confidence is invalid.") + if ( + support["xy_overlap_ratio"] is not None + and not 0.0 <= float(support["xy_overlap_ratio"]) <= 1.0 + 1.0e-6 + ): + raise ValueError("FinalSceneInspection support overlap is invalid.") + item["support"] = support + aabb = item["world_aabb"] + if aabb is not None: + if not isinstance(aabb, Mapping) or set(aabb) != {"min", "max"}: + raise ValueError("FinalSceneInspection world_aabb is invalid.") + if aabb["min"] is None or aabb["max"] is None: + raise ValueError("FinalSceneInspection world_aabb vectors are invalid.") + minimum = _vector(aabb["min"], default=(0.0, 0.0, 0.0)) + maximum = _vector(aabb["max"], default=(0.0, 0.0, 0.0)) + if np.any(np.asarray(maximum) < np.asarray(minimum)): + raise ValueError("FinalSceneInspection world_aabb bounds are inverted.") + item["world_aabb"] = {"min": minimum, "max": maximum} + evidence = deepcopy(dict(item["evidence"])) + if set(evidence) != {"source", "method"} or any( + not isinstance(evidence[key], str) or not evidence[key] for key in evidence + ): + raise ValueError("FinalSceneInspection evidence is invalid.") + item["evidence"] = evidence + return item + + +def _measure_geometry( + entry: Mapping[str, Any], + *, + convert_y_up: bool, +) -> dict[str, Any] | None: + shape = entry.get("shape") + if not isinstance(shape, Mapping): + return None + shape_type = str(shape.get("shape_type", "")) + if shape_type == "Mesh": + path = Path(str(shape.get("fpath", ""))).expanduser().resolve() + if not path.is_file(): + return None + loaded = trimesh.load(path, force="scene") + mesh = loaded.to_geometry() + elif shape_type == "Cube": + mesh = trimesh.creation.box( + extents=_vector(shape.get("size"), default=(1, 1, 1)) + ) + elif shape_type == "Sphere": + radius = float(shape.get("radius", 1.0)) + mesh = trimesh.creation.icosphere(radius=radius) + else: + return None + scale = np.asarray(_vector(entry.get("body_scale"), default=(1, 1, 1))) + mesh.apply_scale(scale) + local_extents = np.asarray(mesh.extents, dtype=float) + conversion = _Y_UP_TO_Z_UP if convert_y_up else np.eye(3) + rotation = Rotation.from_euler( + "XYZ", + _vector(entry.get("init_rot"), default=(0, 0, 0)), + degrees=True, + ).as_matrix() + transform = np.eye(4) + transform[:3, :3] = rotation @ conversion + transform[:3, 3] = _vector(entry.get("init_pos"), default=(0, 0, 0)) + mesh.apply_transform(transform) + return { + "bounds": np.asarray(mesh.bounds, dtype=float), + "local_extents": local_extents, + "axis_transform": transform[:3, :3], + "shape_type": shape_type, + } + + +def _orientation(geometry: Mapping[str, Any] | None) -> str: + if geometry is None or geometry["shape_type"] == "Sphere": + return "unknown" + extents = np.asarray(geometry["local_extents"], dtype=float) + ordered = np.sort(extents) + if ordered[-1] <= 0.0 or ordered[-1] / max(ordered[-2], 1.0e-9) < 1.2: + return "unknown" + dominant = int(np.argmax(extents)) + axis = np.asarray(geometry["axis_transform"], dtype=float)[:, dominant] + vertical = abs(float(axis[2])) / max(float(np.linalg.norm(axis)), 1.0e-9) + if vertical >= 0.75: + return "standing" + if vertical <= 0.35: + return "lying" + return "unknown" + + +def _geometry_from_record(item: Mapping[str, Any]) -> np.ndarray | None: + aabb = item.get("world_aabb") + if not isinstance(aabb, Mapping): + return None + return np.asarray([aabb["min"], aabb["max"]], dtype=float) + + +def _support_for( + uid: str, + child: np.ndarray, + objects: Mapping[str, Mapping[str, Any]], + *, + tolerance: float, +) -> dict[str, Any] | None: + child_bottom = float(child[0, 2]) + child_area = max( + float((child[1, 0] - child[0, 0]) * (child[1, 1] - child[0, 1])), + 1.0e-9, + ) + candidates = [] + for parent_uid, parent_item in objects.items(): + if parent_uid == uid: + continue + parent = _geometry_from_record(parent_item) + if parent is None: + continue + overlap_x = max( + 0.0, min(child[1, 0], parent[1, 0]) - max(child[0, 0], parent[0, 0]) + ) + overlap_y = max( + 0.0, min(child[1, 1], parent[1, 1]) - max(child[0, 1], parent[0, 1]) + ) + overlap_ratio = float(overlap_x * overlap_y / child_area) + gap = child_bottom - float(parent[1, 2]) + if overlap_ratio >= 0.1 and -tolerance <= gap <= tolerance: + candidates.append((overlap_ratio, -abs(gap), parent_uid, gap)) + if not candidates: + return None + overlap_ratio, _, parent_uid, gap = max(candidates) + confidence = min(1.0, overlap_ratio * max(0.0, 1.0 - abs(gap) / tolerance)) + return { + "parent_uid": parent_uid, + "relation": "on", + "confidence": float(confidence), + "gap_m": float(gap), + "xy_overlap_ratio": float(overlap_ratio), + } + + +def _vector(value: Any, *, default: tuple[float, float, float]) -> list[float]: + raw = default if value is None else value + if not isinstance(raw, Sequence) or isinstance(raw, (str, bytes)): + raise TypeError("Scene geometry vectors must be sequences.") + result = [float(item) for item in raw] + if len(result) != 3 or not np.all(np.isfinite(result)): + raise ValueError("Scene geometry vectors must contain three finite values.") + return result diff --git a/embodichain/gen_sim/task_engine/scene/scene_engine_v1.py b/embodichain/gen_sim/task_engine/scene/scene_engine_v1.py new file mode 100644 index 000000000..af0c4a47c --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene/scene_engine_v1.py @@ -0,0 +1,242 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Adapt existing Scene Engine exports without changing their source schema.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +import hashlib +import json +from pathlib import Path +from typing import Any + +from .contracts import ( + STATIC_SCENE_MANIFEST_SCHEMA, + StaticSceneManifest, + validate_static_scene_manifest, +) + +__all__ = ["SceneEngineV1Adapter"] + + +class SceneEngineV1Adapter: + """Convert a normalized Scene Engine v1 export to the neutral manifest.""" + + def adapt_prepared_scene( + self, + prepared_scene: Any, + *, + source_format: str, + robot_profile: str, + ) -> StaticSceneManifest: + """Adapt the existing prepared-scene view through a duck-typed boundary.""" + planner_objects = tuple(getattr(prepared_scene, "planner_objects")) + runtime_objects = ( + tuple(getattr(prepared_scene, "background", ())) + + tuple(getattr(prepared_scene, "rigid_objects", ())) + + tuple(getattr(prepared_scene, "articulations", ())) + ) + runtime_by_uid = { + str(item.get("uid")): item + for item in runtime_objects + if isinstance(item, Mapping) and item.get("uid") + } + asset_hashes = dict(getattr(prepared_scene, "asset_hashes", {}) or {}) + objects = [ + self._object_manifest( + raw, + runtime=runtime_by_uid.get(str(raw.get("uid")), {}), + asset_sha256=str(asset_hashes.get(str(raw.get("uid")), "")), + ) + for raw in planner_objects + ] + identity = { + "source_format": str(source_format), + "robot_profile": str(robot_profile), + "objects": [_identity_object(item) for item in objects], + } + source_path = Path(getattr(prepared_scene, "source_config_path")) + return validate_static_scene_manifest( + { + "schema_version": STATIC_SCENE_MANIFEST_SCHEMA, + "scene_id": _canonical_hash(identity), + "source_format": str(source_format), + "robot_profile": str(robot_profile), + "source": { + "adapter": f"{type(self).__module__}.{type(self).__qualname__}", + "config_path": source_path.expanduser().resolve().as_posix(), + "config_sha256": _file_hash(source_path), + "asset_hashes": asset_hashes, + }, + "adapter_capabilities": { + "task_conditioned_generation": False, + "structured_affordances": any( + bool(item["affordances"]) for item in objects + ), + "articulation_instances": any( + item["role"] == "articulation" for item in objects + ), + "runtime_scene_observation": False, + }, + "objects": objects, + } + ) + + def _object_manifest( + self, + raw: Mapping[str, Any], + *, + runtime: Mapping[str, Any], + asset_sha256: str, + ) -> dict[str, Any]: + uid = str(raw.get("uid", "")).strip() + role = str(raw.get("role", "")).strip() + shape = raw.get("shape", runtime.get("shape", {})) + shape = deepcopy(dict(shape)) if isinstance(shape, Mapping) else {} + physics_keys = ("attrs", "body_type", "max_convex_hull_num") + physics = { + key: deepcopy(runtime[key]) for key in physics_keys if key in runtime + } + articulation = deepcopy(dict(runtime)) if role == "articulation" else {} + affordances = _affordance_evidence(raw.get("affordances", ())) + if role in {"background", "table", "support_surface"}: + affordances = _with_structural_evidence( + affordances, + "support_surface", + ) + if role in {"object", "rigid_object"}: + affordances = _with_structural_evidence(affordances, "rigid") + return { + "uid": uid, + "source_uid": str(raw.get("source_uid", "")), + "role": role, + "name": str(raw.get("name", "")), + "description": str(raw.get("description", "")), + "category": str(raw.get("category", "")), + "color": raw.get("color") if isinstance(raw.get("color"), str) else None, + "geometry": { + "shape": shape, + "asset_sha256": asset_sha256, + }, + "initial_pose": { + "position": deepcopy(list(raw.get("init_pos", ()))), + "rotation": deepcopy(list(raw.get("init_rot", ()))), + "scale": deepcopy(list(raw.get("body_scale", ()))), + }, + "physics": physics, + "articulation": articulation, + "affordances": affordances, + "initial_state": _mapping_or_empty(raw.get("initial_state")), + "attributes": _mapping_or_empty(raw.get("attributes")), + "provenance": { + "semantic_source": "scene_export", + "geometry_source": "prepared_scene", + "physics_source": "prepared_scene_runtime", + }, + } + + +def _affordance_evidence(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + return [] + result: list[dict[str, Any]] = [] + for raw in value: + if isinstance(raw, str) and raw.strip(): + result.append(_evidence(raw.strip(), status="declared")) + continue + if not isinstance(raw, Mapping): + continue + affordance_type = str(raw.get("type", raw.get("name", ""))).strip() + if not affordance_type: + continue + status = str(raw.get("status", "declared")) + result.append( + { + "type": affordance_type, + "status": status, + "confidence": raw.get("confidence"), + "source": str(raw.get("source", "scene_export")), + "link_uid": str(raw.get("link_uid", "")), + "frame": _mapping_or_empty(raw.get("frame")), + "parameters": _mapping_or_empty(raw.get("parameters")), + } + ) + return sorted(result, key=lambda item: (item["type"], item["source"])) + + +def _with_structural_evidence( + evidence: list[dict[str, Any]], affordance_type: str +) -> list[dict[str, Any]]: + if any(item["type"] == affordance_type for item in evidence): + return evidence + return sorted( + [ + *evidence, + _evidence(affordance_type, status="verified", source="adapter_structure"), + ], + key=lambda item: (item["type"], item["source"]), + ) + + +def _evidence( + affordance_type: str, + *, + status: str, + source: str = "scene_export", +) -> dict[str, Any]: + return { + "type": affordance_type, + "status": status, + "confidence": None, + "source": source, + "link_uid": "", + "frame": {}, + "parameters": {}, + } + + +def _mapping_or_empty(value: Any) -> dict[str, Any]: + return deepcopy(dict(value)) if isinstance(value, Mapping) else {} + + +def _identity_object(value: Mapping[str, Any]) -> dict[str, Any]: + result = deepcopy(dict(value)) + geometry = result.get("geometry") + if isinstance(geometry, dict): + shape = geometry.get("shape") + if isinstance(shape, dict) and geometry.get("asset_sha256"): + shape.pop("fpath", None) + return result + + +def _canonical_hash(value: Any) -> str: + payload = json.dumps( + value, + ensure_ascii=False, + sort_keys=True, + separators=(",", ":"), + allow_nan=False, + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _file_hash(path: Path) -> str: + resolved = path.expanduser().resolve() + if not resolved.is_file(): + return "" + return hashlib.sha256(resolved.read_bytes()).hexdigest() diff --git a/embodichain/gen_sim/task_engine/scene_backend.py b/embodichain/gen_sim/task_engine/scene_backend.py new file mode 100644 index 000000000..dd4a665c5 --- /dev/null +++ b/embodichain/gen_sim/task_engine/scene_backend.py @@ -0,0 +1,474 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Task-owned adapter for Scene Engine analysis, revisions, and edits.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import dataclass, replace +from pathlib import Path +import json +import shutil +from typing import Any + +from embodichain.gen_sim.action_engine.generation.source_scene import ( + resolve_source_scene, +) +from embodichain.gen_sim.scene_engine.pipeline import ( + SceneBlueprintPackage, + SceneMaterialization, + analyze_edit, + analyze_image, + materialize_blueprint, + materialize_edit, +) + +from .orchestration.legacy_scene import ( + convert_legacy_gym_project, + restore_locked_scene_entities, +) +from .orchestration.scene_adapter import CandidateSelection, SceneAdapter +from .orchestration.scene_source import ( + SceneSourceFingerprint, + fingerprint_scene_source, + scene_revision_id, + verify_scene_source_fingerprint, +) +from .scene.final_inspection import FinalSceneInspection, inspect_final_scene +from .workflow_contracts import TaskRunRequest, scene_input_kind + +__all__ = [ + "SceneRemediableError", + "SceneAnalysis", + "SceneEngineBackend", + "SceneRevision", + "scene_blueprint_objects", +] + +_LOCKED_SCENE_MANIFEST = "locked_scene_entities.json" + + +class SceneRemediableError(RuntimeError): + """A Scene output failure that permits a fresh materialization attempt.""" + + +@dataclass(frozen=True) +class SceneAnalysis: + """Scene semantics available before asset materialization.""" + + input_kind: str + source: Path + blueprint: SceneBlueprintPackage | None + source_fingerprint: SceneSourceFingerprint | None + + +@dataclass(frozen=True) +class SceneRevision: + """One immutable scene source selected for final Action preparation.""" + + source: Path + output_root: Path | None + revision_id: str + seed: int + edit_plan: dict[str, Any] | None + source_fingerprint: SceneSourceFingerprint | None + + +class SceneEngineBackend: + """Expose Scene Engine stages without giving it workflow ownership.""" + + def analyze( + self, + request: TaskRunRequest, + output_root: str | Path, + ) -> SceneAnalysis: + """Analyze an image or fingerprint an existing read-only project. + + Args: + request: Validated Task Engine run request. + output_root: Directory for image-understanding artifacts. + + Returns: + Scene semantics and immutable source provenance. + """ + root = Path(output_root).expanduser().resolve() + if scene_input_kind(request) == "image": + image_path = Path(str(request["image_path"])).resolve() + blueprint = analyze_image(image_path, root) + return SceneAnalysis( + input_kind="image", + source=image_path, + blueprint=blueprint, + source_fingerprint=None, + ) + source = Path(str(request["gym_project"])).resolve() + return SceneAnalysis( + input_kind="gym_project", + source=source, + blueprint=None, + source_fingerprint=fingerprint_scene_source(source), + ) + + def select( + self, + analysis: SceneAnalysis, + candidate_set: Mapping[str, Any], + scene_adapter: SceneAdapter, + *, + force_most_likely: bool, + ) -> CandidateSelection: + """Select a task candidate from blueprint or existing-scene semantics. + + Args: + analysis: Pre-materialization scene analysis. + candidate_set: Task candidates to ground and vote. + scene_adapter: Task-owned semantic binding adapter. + force_most_likely: Whether ranked UID hypotheses must be resolved. + + Returns: + Audited initial candidate selection. + """ + if analysis.blueprint is not None: + return scene_adapter.select_objects( + candidate_set, + scene_blueprint_objects(analysis.blueprint), + force_most_likely=force_most_likely, + ) + adaptation = scene_adapter.adapt( + candidate_set, + analysis.source, + force_most_likely=force_most_likely, + ) + return CandidateSelection( + scene_manifest=adaptation.scene_manifest, + role_bindings=adaptation.role_bindings, + binding_report=adaptation.binding_report, + selected_candidate=adaptation.selected_candidate, + candidate_bindings=adaptation.candidate_bindings, + ) + + def materialize( + self, + analysis: SceneAnalysis, + request: TaskRunRequest, + output_root: str | Path, + *, + seed: int, + ) -> SceneRevision: + """Produce a new revision, or return the untouched existing source. + + Args: + analysis: Pre-materialization scene analysis. + request: Validated Task Engine run request. + output_root: Fresh directory for this scene attempt. + seed: Attempt seed recorded for recovery audit. + + Returns: + Final scene source for binding and Action Engine generation. + """ + root = Path(output_root).expanduser().resolve() + edit_prompt = request["scene_edit_prompt"] + if analysis.input_kind == "image": + assert analysis.blueprint is not None + root.mkdir(parents=True, exist_ok=False) + blueprint = replace(analysis.blueprint, output_root=root) + materialization = materialize_blueprint(blueprint, seed=seed) + edit_plan = None + if edit_prompt is not None: + edit_blueprint = analyze_edit( + output_root=root, + edit_prompt=str(edit_prompt), + ) + edit_plan = edit_blueprint.scene_edit_plan.to_dict() + materialization = materialize_edit(edit_blueprint, seed=seed) + revision = _revision(materialization, seed=seed, edit_plan=edit_plan) + _write_revision_audit( + root, + revision_id=revision.revision_id, + seed=seed, + edit_plan=edit_plan, + ) + return revision + + fingerprint = analysis.source_fingerprint + assert fingerprint is not None + if edit_prompt is None: + verify_scene_source_fingerprint(fingerprint.to_dict()) + return SceneRevision( + source=analysis.source, + output_root=None, + revision_id=scene_revision_id(analysis.source), + seed=seed, + edit_plan=None, + source_fingerprint=fingerprint, + ) + + resolved = resolve_source_scene(analysis.source) + if resolved.source_format == "legacy_gym_config": + converted = convert_legacy_gym_project(analysis.source, root) + editable_root = converted.output_root + else: + editable_root = _copy_scene_export_revision(resolved.path, root) + edit_blueprint = analyze_edit( + output_root=editable_root, + edit_prompt=str(edit_prompt), + ) + edit_plan = edit_blueprint.scene_edit_plan.to_dict() + materialization = materialize_edit(edit_blueprint, seed=seed) + if resolved.source_format == "legacy_gym_config": + restore_locked_scene_entities(editable_root) + else: + _restore_scene_export_locked_entities(editable_root) + verify_scene_source_fingerprint(fingerprint.to_dict()) + _write_revision_audit( + editable_root, + revision_id=scene_revision_id(materialization.scene_config_path), + seed=seed, + edit_plan=edit_plan, + source_fingerprint=fingerprint, + ) + return SceneRevision( + source=materialization.scene_config_path, + output_root=editable_root, + revision_id=scene_revision_id(materialization.scene_config_path), + seed=seed, + edit_plan=edit_plan, + source_fingerprint=fingerprint, + ) + + def inspect( + self, + revision: SceneRevision, + output_path: str | Path, + ) -> FinalSceneInspection: + """Inspect final geometry and publish support/orientation evidence. + + Args: + revision: Completed immutable scene revision. + output_path: JSON path receiving the inspection document. + + Returns: + Validated final scene inspection. + """ + try: + actual_revision_id = scene_revision_id(revision.source) + except (OSError, TypeError, ValueError) as exc: + raise SceneRemediableError( + f"Final scene content could not be hashed: {exc}" + ) from exc + if actual_revision_id != revision.revision_id: + raise RuntimeError("Final scene changed before geometry inspection.") + try: + inspection = inspect_final_scene( + revision.source, + revision_id=actual_revision_id, + ) + except (OSError, TypeError, ValueError) as exc: + raise SceneRemediableError( + f"Final scene assets could not be inspected: {exc}" + ) from exc + path = Path(output_path).expanduser().resolve() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(inspection, ensure_ascii=False, indent=2, allow_nan=False) + + "\n", + encoding="utf-8", + ) + return inspection + + +def scene_blueprint_objects(blueprint: SceneBlueprintPackage) -> list[dict[str, Any]]: + """Convert image semantics into the redacted grounding inventory shape. + + Args: + blueprint: Scene Engine image-understanding package. + + Returns: + Semantic objects with unknown physical fields represented conservatively. + """ + nodes = blueprint.scene_graph.node_by_id() + result = [] + for item in blueprint.scene.objects: + node = nodes.get(item.id) + orientation = None if node is None else node.orientation_state + initial_state = {} + if orientation == "lying": + initial_state["orientation"] = "fallen" + elif orientation == "standing": + initial_state["orientation"] = "upright" + result.append( + { + "uid": item.id, + "source_uid": item.id, + "role": "table" if item.kind == "table" else "rigid_object", + "name": item.name, + "description": item.description, + "category": item.category, + "color": None, + "init_pos": [0.0, 0.0, 0.0], + "affordances": [], + "initial_state": initial_state, + "attributes": {}, + } + ) + return result + + +def _copy_scene_export_revision(source_config: Path, output_root: Path) -> Path: + if output_root.exists(): + if not output_root.is_dir() or any(output_root.iterdir()): + raise ValueError("Scene revision output_root must be empty.") + source_root = source_config.parent + destination = output_root / "scene_export" + shutil.copytree(source_root, destination) + config_path = destination / "scene_config.json" + config = _read_json_mapping(config_path) + background = list(config.get("background", ())) + rigid_objects = list(config.get("rigid_object", ())) + articulations = list(config.get("articulation", ())) + editable_rigid = [ + item + for item in rigid_objects + if isinstance(item, Mapping) and _scene_editable_rigid(item) + ] + locked_rigid = [item for item in rigid_objects if item not in editable_rigid] + table = [ + item + for item in background + if isinstance(item, Mapping) and item.get("uid") == "table" + ] + if len(table) != 1: + raise ValueError("Scene export revision requires exactly one table.") + locked = { + "schema_version": "embodichain.locked-scene-entities/v1", + "background": [item for item in background if item not in table], + "rigid_object": locked_rigid, + "articulation": articulations, + } + config["background"] = table + config["rigid_object"] = editable_rigid + config["articulation"] = [] + _write_json_mapping(config_path, config) + _write_json_mapping(output_root / _LOCKED_SCENE_MANIFEST, locked) + graph_path = destination / "scene_graph.json" + if graph_path.is_file(): + graph = _read_json_mapping(graph_path) + editable_uids = {str(item.get("uid")) for item in [*table, *editable_rigid]} + graph["nodes"] = [ + item + for item in graph.get("nodes", ()) + if isinstance(item, Mapping) and item.get("object_id") in editable_uids + ] + graph["relations"] = [ + item + for item in graph.get("relations", ()) + if isinstance(item, Mapping) + and item.get("source_id") in editable_uids + and item.get("target_id") in editable_uids + ] + _write_json_mapping(graph_path, graph) + return output_root + + +def _restore_scene_export_locked_entities(output_root: Path) -> None: + manifest = _read_json_mapping(output_root / _LOCKED_SCENE_MANIFEST) + if manifest.get("schema_version") != "embodichain.locked-scene-entities/v1": + raise ValueError("Locked scene entity manifest schema is invalid.") + config_path = output_root / "scene_export" / "scene_config.json" + config = _read_json_mapping(config_path) + existing = { + str(item.get("uid")) + for section in ("background", "rigid_object", "articulation") + for item in config.get(section, ()) + if isinstance(item, Mapping) and item.get("uid") + } + for section in ("background", "rigid_object", "articulation"): + target = list(config.get(section, ())) + for raw in manifest.get(section, ()): + item = deepcopy(dict(raw)) + uid = str(item.get("uid", "")) + if not uid or uid in existing: + raise ValueError(f"Scene edit reused locked entity UID {uid!r}.") + target.append(item) + existing.add(uid) + config[section] = target + _write_json_mapping(config_path, config) + + +def _scene_editable_rigid(value: Mapping[str, Any]) -> bool: + shape = value.get("shape") + return ( + isinstance(shape, Mapping) + and shape.get("shape_type") == "Mesh" + and isinstance(shape.get("fpath"), str) + and bool(shape["fpath"]) + ) + + +def _read_json_mapping(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, Mapping): + raise TypeError(f"JSON artifact must contain an object: {path}") + return dict(value) + + +def _write_json_mapping(path: Path, value: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) + + +def _revision( + value: SceneMaterialization, + *, + seed: int, + edit_plan: dict[str, Any] | None, +) -> SceneRevision: + return SceneRevision( + source=value.scene_config_path, + output_root=value.output_root, + revision_id=scene_revision_id(value.scene_config_path), + seed=seed, + edit_plan=edit_plan, + source_fingerprint=None, + ) + + +def _write_revision_audit( + output_root: Path, + *, + revision_id: str, + seed: int, + edit_plan: Mapping[str, Any] | None, + source_fingerprint: SceneSourceFingerprint | None = None, +) -> None: + payload = { + "schema_version": "embodichain.scene-revision-attempt/v1", + "revision_id": revision_id, + "seed": int(seed), + "edit_plan": None if edit_plan is None else dict(edit_plan), + "source_fingerprint": ( + None if source_fingerprint is None else source_fingerprint.to_dict() + ), + } + (output_root / "scene_revision_attempt.json").write_text( + json.dumps(payload, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/task_engine/state_machine.py b/embodichain/gen_sim/task_engine/state_machine.py new file mode 100644 index 000000000..ffdaaba15 --- /dev/null +++ b/embodichain/gen_sim/task_engine/state_machine.py @@ -0,0 +1,294 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Typed, replayable state transitions for cross-engine orchestration.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from copy import deepcopy +from dataclasses import dataclass, field +from enum import Enum +from types import MappingProxyType +from typing import Any + +from .workflow_contracts import TaskRunRequest, validate_task_run_request + +__all__ = [ + "StageStatus", + "TaskEngineState", + "WorkflowStage", + "complete_stage", + "fail_stage", + "initial_state", + "replay_events", + "skip_stage", + "start_stage", +] + + +class WorkflowStage(str, Enum): + """Stable stages shared by all four supported input combinations.""" + + INPUT = "input" + TASK_CANDIDATES = "task_candidates" + SCENE_PREPARATION = "scene_preparation" + SCENE_EDIT = "scene_edit" + CANDIDATE_SELECTION = "candidate_selection" + SCENE_FINALIZATION = "scene_finalization" + UNBOUND_ACTION = "unbound_action" + FINAL_INSPECTION = "final_inspection" + FINAL_BINDING = "final_binding" + STATIC_FEASIBILITY = "static_feasibility" + GROUNDED_ACTION = "grounded_action" + EXECUTION = "execution" + + +class StageStatus(str, Enum): + """Lifecycle of one independently schedulable workflow stage.""" + + PENDING = "pending" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + SKIPPED = "skipped" + + +_DEPENDENCIES: dict[WorkflowStage, frozenset[WorkflowStage]] = { + WorkflowStage.INPUT: frozenset(), + WorkflowStage.TASK_CANDIDATES: frozenset({WorkflowStage.INPUT}), + WorkflowStage.SCENE_PREPARATION: frozenset({WorkflowStage.INPUT}), + WorkflowStage.SCENE_EDIT: frozenset({WorkflowStage.SCENE_PREPARATION}), + WorkflowStage.CANDIDATE_SELECTION: frozenset( + { + WorkflowStage.TASK_CANDIDATES, + WorkflowStage.SCENE_PREPARATION, + } + ), + WorkflowStage.SCENE_FINALIZATION: frozenset( + {WorkflowStage.CANDIDATE_SELECTION, WorkflowStage.SCENE_EDIT} + ), + WorkflowStage.UNBOUND_ACTION: frozenset({WorkflowStage.CANDIDATE_SELECTION}), + WorkflowStage.FINAL_INSPECTION: frozenset({WorkflowStage.SCENE_FINALIZATION}), + WorkflowStage.FINAL_BINDING: frozenset( + {WorkflowStage.FINAL_INSPECTION, WorkflowStage.UNBOUND_ACTION} + ), + WorkflowStage.STATIC_FEASIBILITY: frozenset({WorkflowStage.FINAL_BINDING}), + WorkflowStage.GROUNDED_ACTION: frozenset({WorkflowStage.STATIC_FEASIBILITY}), + WorkflowStage.EXECUTION: frozenset({WorkflowStage.GROUNDED_ACTION}), +} + +_SKIPPABLE_STAGES = frozenset({WorkflowStage.SCENE_EDIT}) + + +@dataclass(frozen=True) +class TaskEngineState: + """Immutable state snapshot plus an append-only transition audit.""" + + request: Mapping[str, Any] + stages: Mapping[WorkflowStage, StageStatus] + events: tuple[Mapping[str, Any], ...] = field(default_factory=tuple) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "request", + MappingProxyType(deepcopy(dict(self.request))), + ) + object.__setattr__( + self, + "stages", + MappingProxyType(dict(self.stages)), + ) + object.__setattr__( + self, + "events", + tuple(MappingProxyType(deepcopy(dict(event))) for event in self.events), + ) + + @property + def terminal(self) -> bool: + """Return whether execution succeeded or any stage failed.""" + return ( + self.stages[WorkflowStage.EXECUTION] == StageStatus.SUCCEEDED + or StageStatus.FAILED in self.stages.values() + ) + + def to_dict(self) -> dict[str, Any]: + """Return one JSON-safe audit snapshot.""" + return { + "request": deepcopy(dict(self.request)), + "stages": { + stage.value: self.stages[stage].value for stage in WorkflowStage + }, + "events": deepcopy([dict(event) for event in self.events]), + } + + +def initial_state(request: TaskRunRequest) -> TaskEngineState: + """Create a validated state with the optional edit stage resolved.""" + normalized = validate_task_run_request(request) + stages = {stage: StageStatus.PENDING for stage in WorkflowStage} + stages[WorkflowStage.INPUT] = StageStatus.SUCCEEDED + events = ( + { + "sequence": 1, + "stage": WorkflowStage.INPUT.value, + "from": StageStatus.PENDING.value, + "to": StageStatus.SUCCEEDED.value, + }, + ) + state = TaskEngineState(request=normalized, stages=stages, events=events) + if normalized["scene_edit_prompt"] is None: + state = skip_stage(state, WorkflowStage.SCENE_EDIT) + return state + + +def start_stage(state: TaskEngineState, stage: WorkflowStage) -> TaskEngineState: + """Start a pending stage only after every dependency has completed.""" + if state.terminal: + raise ValueError("A terminal TaskEngineState cannot start another stage.") + if state.stages[stage] != StageStatus.PENDING: + raise ValueError(f"Stage {stage.value!r} is not pending.") + incomplete = [ + dependency.value + for dependency in _DEPENDENCIES[stage] + if state.stages[dependency] not in {StageStatus.SUCCEEDED, StageStatus.SKIPPED} + ] + if incomplete: + raise ValueError( + f"Stage {stage.value!r} has incomplete dependencies: {incomplete}." + ) + return _transition(state, stage, StageStatus.RUNNING) + + +def complete_stage(state: TaskEngineState, stage: WorkflowStage) -> TaskEngineState: + """Complete one running stage.""" + if state.stages[stage] != StageStatus.RUNNING: + raise ValueError(f"Stage {stage.value!r} is not running.") + return _transition(state, stage, StageStatus.SUCCEEDED) + + +def fail_stage( + state: TaskEngineState, + stage: WorkflowStage, + *, + reason: str, +) -> TaskEngineState: + """Fail a stage, including a later retry of a previously successful stage.""" + if state.terminal: + raise ValueError("A terminal TaskEngineState cannot fail another stage.") + if state.stages[stage] not in { + StageStatus.PENDING, + StageStatus.RUNNING, + StageStatus.SUCCEEDED, + }: + raise ValueError(f"Stage {stage.value!r} cannot be failed now.") + normalized_reason = str(reason).strip() + if not normalized_reason: + raise ValueError("A failed stage requires a non-empty reason.") + return _transition( + state, + stage, + StageStatus.FAILED, + details={"reason": normalized_reason}, + ) + + +def skip_stage(state: TaskEngineState, stage: WorkflowStage) -> TaskEngineState: + """Skip one optional pending stage.""" + if stage not in _SKIPPABLE_STAGES: + raise ValueError("Only the optional scene_edit stage can be skipped.") + if state.stages[stage] != StageStatus.PENDING: + raise ValueError(f"Stage {stage.value!r} is not pending.") + return _transition(state, stage, StageStatus.SKIPPED) + + +def replay_events( + request: TaskRunRequest, + events: Sequence[Mapping[str, Any]], +) -> TaskEngineState: + """Rebuild a state by validating and applying its transition audit. + + Args: + request: Original workflow request used to create the state. + events: Complete ordered event audit to validate and replay. + + Returns: + The immutable state reconstructed from the supplied audit. + + Raises: + TypeError: If the audit is not a sequence of event mappings. + ValueError: If any event is missing, altered, or not a valid transition. + """ + if not isinstance(events, Sequence) or isinstance(events, (str, bytes)): + raise TypeError("Task Engine events must be a sequence of mappings.") + recorded = [] + for event in events: + if not isinstance(event, Mapping): + raise TypeError("Each Task Engine event must be a mapping.") + recorded.append(deepcopy(dict(event))) + + state = initial_state(request) + initial_events = [dict(event) for event in state.events] + if recorded[: len(initial_events)] != initial_events: + raise ValueError("Replay event does not match the canonical initial state.") + + for expected in recorded[len(initial_events) :]: + try: + stage = WorkflowStage(expected["stage"]) + target = StageStatus(expected["to"]) + if target == StageStatus.RUNNING: + replayed = start_stage(state, stage) + elif target == StageStatus.SUCCEEDED: + replayed = complete_stage(state, stage) + elif target == StageStatus.FAILED: + replayed = fail_stage(state, stage, reason=expected["reason"]) + elif target == StageStatus.SKIPPED: + replayed = skip_stage(state, stage) + else: + raise ValueError(f"Unsupported replay target: {target.value!r}.") + except (KeyError, TypeError, ValueError) as exc: + raise ValueError("Replay event does not match a valid transition.") from exc + if dict(replayed.events[-1]) != expected: + raise ValueError("Replay event does not match the generated transition.") + state = replayed + return state + + +def _transition( + state: TaskEngineState, + stage: WorkflowStage, + status: StageStatus, + *, + details: dict[str, Any] | None = None, +) -> TaskEngineState: + previous = state.stages[stage] + stages = dict(state.stages) + stages[stage] = status + event = { + "sequence": len(state.events) + 1, + "stage": stage.value, + "from": previous.value, + "to": status.value, + } + if details: + event.update(deepcopy(details)) + return TaskEngineState( + request=dict(state.request), + stages=stages, + events=(*state.events, event), + ) diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py new file mode 100644 index 000000000..fca9778c4 --- /dev/null +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -0,0 +1,1275 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Parallel Task Engine workflow with bounded, fully audited recovery.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from copy import deepcopy +from dataclasses import dataclass, replace +from datetime import datetime +import json +from pathlib import Path +import shutil +import subprocess +import sys +from typing import Any, Final + +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError +from embodichain.gen_sim.action_engine.runtime import ( + EXECUTION_REPORT_FILENAME, + validate_execution_report, +) +from embodichain.gen_sim.scene_engine.errors import SceneServiceError + +from .agent import TaskAgent +from .config import ( + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, + load_task_engine_config, +) +from .contracts import canonical_hash +from .orchestration.artifacts import ArtifactTransaction +from .orchestration.coordinator import PreparationResult, TaskEngineCoordinator +from .orchestration.scene_adapter import CandidateSelection, SceneAdapter +from .orchestration.scene_source import SceneSourceRef +from .scene_backend import ( + SceneAnalysis, + SceneEngineBackend, + SceneRemediableError, + SceneRevision, +) +from .state_machine import ( + TaskEngineState, + WorkflowStage, + complete_stage, + fail_stage, + initial_state, + start_stage, +) +from .workflow_contracts import TaskRunRequest, validate_task_run_request + +__all__ = [ + "TASK_ENGINE_RUN_MANIFEST_SCHEMA", + "ActionExecutor", + "SubprocessActionExecutor", + "TaskEngineRunResult", + "TaskEngineWorkflow", +] + +TASK_ENGINE_RUN_MANIFEST_SCHEMA: Final = "embodichain.task-engine-run/v1" +ActionExecutor = Callable[..., Mapping[str, Any]] + + +@dataclass(frozen=True) +class TaskEngineRunResult: + """Published outcome of one isolated cross-engine workflow run.""" + + status: str + output_dir: Path + manifest_path: Path + state_path: Path + final_bundle: Path | None + failure_class: str | None = None + + @property + def succeeded(self) -> bool: + """Return whether real simulator execution met the configured policy.""" + return self.status == "succeeded" + + +class SubprocessActionExecutor: + """Execute a prepared bundle through Task Engine's private runner.""" + + def __call__( + self, + bundle: str | Path, + output_root: str | Path, + *, + seed: int, + num_envs: int, + dataset_saving: bool = False, + ) -> Mapping[str, Any]: + """Run one simulator attempt and preserve its report and trajectory. + + Args: + bundle: Prepared Action Engine bundle. + output_root: Fresh directory for this execution attempt. + seed: Action Engine random seed. + num_envs: Number of vectorized scene replicas. + dataset_saving: Whether to enable the Gym project's dataset recorder. + + Returns: + Validated Action Engine execution report. + """ + bundle_root = Path(bundle).expanduser().resolve() + attempt_root = Path(output_root).expanduser().resolve() + attempt_root.mkdir(parents=True, exist_ok=False) + command = [ + sys.executable, + "-u", + "-m", + "embodichain.gen_sim.task_engine._bundle_runner", + "--bundle", + bundle_root.as_posix(), + "--num_envs", + str(num_envs), + "--seed", + str(seed), + "--headless", + ] + if not dataset_saving: + command.append("--filter_dataset_saving") + log_path = attempt_root / "action.log" + print( + "[Task Engine] Starting " + f"{attempt_root.name}: seed={seed}, num_envs={num_envs}, " + f"dataset_saving={dataset_saving}", + flush=True, + ) + completed = _run_streaming_process(command, log_path) + print( + f"[Task Engine] Completed {attempt_root.name}: " + f"returncode={completed.returncode}", + flush=True, + ) + report_path = bundle_root / EXECUTION_REPORT_FILENAME + process_record = { + "command": command, + "returncode": completed.returncode, + "combined_log": log_path.name, + "stdout": completed.stdout, + "stderr": completed.stderr, + } + _write_json(attempt_root / "process.json", process_record) + if not report_path.is_file(): + raise RuntimeError( + "Action execution did not publish execution_report.json; " + f"returncode={completed.returncode}." + ) + report = validate_execution_report(_read_json(report_path)) + shutil.copy2(report_path, attempt_root / EXECUTION_REPORT_FILENAME) + trajectory_copy = _copy_trajectory_record(report, attempt_root) + if report["action_count"] > 0 and trajectory_copy is None: + raise RuntimeError( + "Action execution report did not expose a readable trajectory record." + ) + _write_json( + attempt_root / "execution_attempt.json", + { + "seed": seed, + "num_envs": num_envs, + "dataset_saving": dataset_saving, + "returncode": completed.returncode, + "trajectory_copy": trajectory_copy, + "report": report, + }, + ) + return report + + +def _run_streaming_process( + command: list[str], + log_path: str | Path, +) -> subprocess.CompletedProcess[str]: + """Run a child while teeing its combined output to the terminal and disk. + + Args: + command: Argument vector passed directly to the child process. + log_path: File receiving the exact combined stdout and stderr bytes. + + Returns: + Completed process metadata with a decoded copy of the combined output. + """ + resolved_log = Path(log_path).expanduser().resolve() + resolved_log.parent.mkdir(parents=True, exist_ok=True) + captured = bytearray() + process: subprocess.Popen[bytes] | None = None + try: + with resolved_log.open("wb") as log_stream: + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=0, + ) + assert process.stdout is not None + while True: + chunk = process.stdout.read(64 * 1024) + if not chunk: + break + captured.extend(chunk) + log_stream.write(chunk) + log_stream.flush() + _write_terminal_chunk(chunk) + returncode = process.wait() + except BaseException: + if process is not None and process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + raise + output = captured.decode("utf-8", errors="replace") + return subprocess.CompletedProcess( + args=command, + returncode=returncode, + stdout=output, + stderr="", + ) + + +def _write_terminal_chunk(chunk: bytes) -> None: + """Best-effort write of raw child output to the parent terminal.""" + try: + stream = getattr(sys.stdout, "buffer", None) + if stream is not None: + stream.write(chunk) + stream.flush() + return + sys.stdout.write(chunk.decode("utf-8", errors="replace")) + sys.stdout.flush() + except (BrokenPipeError, OSError, ValueError): + return + + +class TaskEngineWorkflow: + """Run Scene and Action work concurrently under Task Engine ownership.""" + + def __init__( + self, + *, + task_agent: TaskAgent | None = None, + scene_adapter: SceneAdapter | None = None, + action_agent: ActionAgent | None = None, + coordinator: TaskEngineCoordinator | None = None, + scene_backend: SceneEngineBackend | None = None, + action_executor: ActionExecutor | None = None, + ) -> None: + self.task_agent = task_agent or TaskAgent() + self.scene_adapter = scene_adapter or SceneAdapter() + self.action_agent = action_agent or ActionAgent() + self.coordinator = coordinator or TaskEngineCoordinator( + task_agent=self.task_agent, + scene_adapter=self.scene_adapter, + action_agent=self.action_agent, + ) + self.scene_backend = scene_backend or SceneEngineBackend() + self.action_executor = action_executor or SubprocessActionExecutor() + + def run( + self, + request: TaskRunRequest | Mapping[str, Any], + *, + workflow_cfg: TaskEngineWorkflowCfg | None = None, + planning_cfg: TaskEnginePlanningCfg | None = None, + execution_cfg: TaskEngineExecutionCfg | None = None, + config_path: str | Path | None = None, + model: str | None = None, + vlm_model: str | None = None, + base_seed: int = 0, + dataset_saving: bool = False, + run_id: str | None = None, + created_at: datetime | None = None, + overwrite: bool = False, + execute: bool = True, + ) -> TaskEngineRunResult: + """Run all stages and publish success only after simulator acceptance. + + Args: + request: One of the four image/project plus optional-edit inputs. + workflow_cfg: Optional retry and concurrency configuration. + planning_cfg: Optional interpretation and bundle generation defaults. + execution_cfg: Optional vectorized success policy. + config_path: YAML used for omitted workflow or execution config. + model: Optional Task and grounding model override. + vlm_model: Optional Action Engine VLM override. + base_seed: First audited scene and action attempt seed. + dataset_saving: Whether Action attempts may initialize dataset recording. + run_id: Optional externally allocated run identifier. + created_at: Optional timezone-aware run creation timestamp. + overwrite: Whether to atomically replace an existing run directory. + execute: Whether to execute the prepared bundle in the simulator. + + Returns: + Published run status, manifest, state audit, and final bundle path. + """ + normalized = validate_task_run_request(request) + if not isinstance(dataset_saving, bool): + raise TypeError("dataset_saving must be a boolean.") + if workflow_cfg is None or planning_cfg is None or execution_cfg is None: + loaded_workflow, loaded_planning, loaded_execution = ( + load_task_engine_config(config_path) + ) + workflow_cfg = workflow_cfg or loaded_workflow + planning_cfg = planning_cfg or loaded_planning + execution_cfg = execution_cfg or loaded_execution + effective_candidate_count = planning_cfg.candidate_count + effective_run_id = str(run_id or Path(normalized["output_dir"]).name).strip() + if not effective_run_id or Path(effective_run_id).name != effective_run_id: + raise ValueError("run_id must be one non-empty path component.") + effective_created_at = created_at or datetime.now().astimezone() + if ( + effective_created_at.tzinfo is None + or effective_created_at.utcoffset() is None + ): + raise ValueError("created_at must include a timezone.") + run_metadata = { + "run_id": effective_run_id, + "created_at": effective_created_at.isoformat(), + "dataset_saving": bool(dataset_saving), + } + state = initial_state(normalized) + attempts: list[dict[str, Any]] = [] + output_dir = Path(normalized["output_dir"]) + + with ArtifactTransaction(output_dir, overwrite=overwrite) as transaction: + staging = transaction.staging_dir + assert staging is not None + analysis_root = staging / "scene_analysis" + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + with ThreadPoolExecutor( + max_workers=workflow_cfg.max_parallel_workers, + thread_name_prefix="task-engine-input", + ) as executor: + candidate_future = executor.submit( + self.task_agent.generate, + normalized["task_id"], + normalized["task_instruction"], + model, + effective_candidate_count, + ) + analysis_future = executor.submit( + self.scene_backend.analyze, + normalized, + analysis_root, + ) + try: + candidate_set = candidate_future.result() + except Exception as exc: + analysis_future.cancel() + state = fail_stage( + state, + WorkflowStage.TASK_CANDIDATES, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="task_generation", + ) + state = complete_stage(state, WorkflowStage.TASK_CANDIDATES) + try: + analysis = analysis_future.result() + except Exception as exc: + state = fail_stage( + state, + WorkflowStage.SCENE_PREPARATION, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="scene_analysis", + ) + state = complete_stage(state, WorkflowStage.SCENE_PREPARATION) + + state = start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + try: + selection = self.scene_backend.select( + analysis, + candidate_set, + self.scene_adapter, + force_most_likely=True, + ) + except Exception as exc: + state = fail_stage( + state, + WorkflowStage.CANDIDATE_SELECTION, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="input_conflict", + failure_class="candidate_selection", + ) + _write_json( + staging / "initial_binding_report.json", selection.binding_report + ) + if selection.selected_candidate is None: + if normalized["scene_edit_prompt"] is None: + state = fail_stage( + state, + WorkflowStage.CANDIDATE_SELECTION, + reason=str(selection.binding_report["selection_reason"]), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="input_conflict", + failure_class="unbound_scene_reference", + ) + provisional = _highest_vote_candidate(candidate_set) + selection = replace( + selection, + selected_candidate=deepcopy(provisional), + ) + _write_json( + staging / "provisional_candidate.json", + { + "candidate_id": provisional["candidate_id"], + "reason": "explicit_scene_edit_may_materialize_missing_reference", + "binding_status": selection.binding_report["status"], + }, + ) + state = complete_stage(state, WorkflowStage.CANDIDATE_SELECTION) + + state = start_stage(state, WorkflowStage.UNBOUND_ACTION) + if normalized["scene_edit_prompt"] is not None: + state = start_stage(state, WorkflowStage.SCENE_EDIT) + else: + state = start_stage(state, WorkflowStage.SCENE_FINALIZATION) + + unbound_plan: Mapping[str, Any] | None = None + unbound_failures: list[dict[str, Any]] = [] + unbound_error: Exception | None = None + scene_error: Exception | None = None + inspection_error = False + preparation_error: Exception | None = None + preparation: PreparationResult | None = None + scene_attempt_limit = ( + 1 + if analysis.input_kind == "gym_project" + and normalized["scene_edit_prompt"] is None + else workflow_cfg.max_scene_attempts + ) + for scene_index in range(1, scene_attempt_limit + 1): + inspection_error = False + scene_seed = int(base_seed) + scene_index - 1 + attempt_root = staging / "attempts" / f"scene_{scene_index:04d}" + attempt_root.mkdir(parents=True) + attempt = { + "scene_attempt": scene_index, + "scene_seed": scene_seed, + "status": "running", + "scene_revision": None, + "final_inspection": None, + "unbound_action_plan": None, + "final_unbound_action_plan": None, + "unbound_transition": None, + "unbound_failures": [], + "preparation": None, + "planning_attempts": [], + "action_attempts": [], + "parallel_errors": [], + "error": None, + } + attempts.append(attempt) + revision: SceneRevision | None = None + try: + if unbound_plan is None: + with ThreadPoolExecutor( + max_workers=workflow_cfg.max_parallel_workers, + thread_name_prefix="task-engine-parallel", + ) as executor: + scene_future = executor.submit( + self.scene_backend.materialize, + analysis, + normalized, + attempt_root / "scene_revision", + seed=scene_seed, + ) + draft_future = executor.submit( + self._draft_with_fallback, + candidate_set, + selection, + ) + try: + revision = scene_future.result() + except Exception as exc: + scene_error = exc + revision = None + try: + unbound_plan, unbound_failures = draft_future.result() + except Exception as exc: + unbound_error = exc + if unbound_error is not None: + raise unbound_error + state = complete_stage(state, WorkflowStage.UNBOUND_ACTION) + if scene_error is not None: + raise scene_error + assert revision is not None + else: + revision = self.scene_backend.materialize( + analysis, + normalized, + attempt_root / "scene_revision", + seed=scene_seed, + ) + scene_error = None + except Exception as exc: + if revision is not None: + attempt["scene_revision"] = _revision_record(revision) + state = _complete_materialized_scene( + state, + has_edit=normalized["scene_edit_prompt"] is not None, + ) + if unbound_error is not None: + attempt["status"] = "unbound_action_failed" + attempt["error"] = _error_record(unbound_error) + if scene_error is not None: + attempt["parallel_errors"].append( + { + "branch": "scene", + **_error_record(scene_error), + } + ) + _write_json(attempt_root / "attempt.json", attempt) + break + scene_error = exc + if unbound_plan is not None: + attempt["unbound_action_plan"] = deepcopy(dict(unbound_plan)) + attempt["unbound_failures"] = deepcopy(unbound_failures) + _write_json( + attempt_root / "unbound_action_plan.json", unbound_plan + ) + attempt["status"] = "scene_failed" + attempt["error"] = _error_record(exc) + _write_json(attempt_root / "attempt.json", attempt) + if ( + scene_index < scene_attempt_limit + and _is_scene_remediable_error(exc) + ): + continue + break + + attempt["scene_revision"] = _revision_record(revision) + attempt["unbound_action_plan"] = deepcopy(dict(unbound_plan)) + attempt["unbound_failures"] = deepcopy(unbound_failures) + _write_json(attempt_root / "unbound_action_plan.json", unbound_plan) + state = _complete_materialized_scene( + state, + has_edit=normalized["scene_edit_prompt"] is not None, + ) + if state.stages[WorkflowStage.FINAL_INSPECTION].value == "pending": + state = start_stage(state, WorkflowStage.FINAL_INSPECTION) + try: + final_inspection = self.scene_backend.inspect( + revision, + attempt_root / "final_scene_inspection.json", + ) + except Exception as exc: + scene_error = exc + inspection_error = True + attempt["status"] = "scene_inspection_failed" + attempt["error"] = _error_record(exc) + _write_json(attempt_root / "attempt.json", attempt) + if ( + scene_index < scene_attempt_limit + and _is_scene_remediable_error(exc) + ): + continue + break + attempt["final_inspection"] = deepcopy(dict(final_inspection)) + if state.stages[WorkflowStage.FINAL_INSPECTION].value == "running": + state = complete_stage(state, WorkflowStage.FINAL_INSPECTION) + + bundle_root = attempt_root / "bundle" + try: + preparation = self.coordinator.prepare( + normalized["task_id"], + normalized["task_instruction"], + SceneSourceRef( + revision.source, + robot_profile=self.scene_adapter.robot_profile, + ), + bundle_root, + model=model, + candidate_count=effective_candidate_count, + planning_mode=planning_cfg.planning_mode, + vlm_model=vlm_model, + max_episodes=planning_cfg.max_episodes, + max_episode_steps=planning_cfg.max_episode_steps, + candidate_set=candidate_set, + force_most_likely=True, + final_inspection=final_inspection, + unbound_action_plan=unbound_plan, + ) + except Exception as exc: + preparation_error = exc + attempt["status"] = "preparation_error" + attempt["error"] = _error_record(exc) + _write_json(attempt_root / "attempt.json", attempt) + break + attempt["preparation"] = preparation.status + attempt["planning_attempts"] = deepcopy( + list(preparation.planning_attempts) + ) + if preparation.status == "bound": + attempt["status"] = "prepared" + _write_json(attempt_root / "attempt.json", attempt) + break + attempt["status"] = "preparation_failed" + attempt["error"] = { + "type": "PreparationFailure", + "message": preparation.status, + } + _write_json(attempt_root / "attempt.json", attempt) + if not _scene_remediable( + preparation, + analysis=analysis, + request=normalized, + ): + break + + if preparation is None or preparation.status != "bound": + failure_class = ( + "action_capability" + if unbound_error is not None + or isinstance(preparation_error, ActionCapabilityError) + else ( + "preparation_error" + if preparation_error is not None + else _preparation_failure_class( + preparation, + scene_error=scene_error, + analysis=analysis, + request=normalized, + ) + ) + ) + failed_stage = ( + WorkflowStage.FINAL_INSPECTION + if inspection_error + else ( + WorkflowStage.UNBOUND_ACTION + if unbound_error is not None + else _failure_stage(failure_class, normalized) + ) + ) + if state.stages[failed_stage].value in { + "pending", + "running", + "succeeded", + }: + state = fail_stage( + state, + failed_stage, + reason=( + str(unbound_error) + if unbound_error is not None + else ( + str(preparation_error) + if preparation_error is not None + else ( + str(scene_error) + if scene_error is not None + else ( + preparation.status + if preparation is not None + else failure_class + ) + ) + ) + ), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status=( + "input_conflict" + if failure_class == "input_conflict" + else "failed" + ), + failure_class=failure_class, + ) + + final_candidate_id = preparation.selected_candidate_id + if not isinstance(final_candidate_id, str) or not final_candidate_id: + raise ValueError( + "A bound preparation must select one non-empty candidate ID." + ) + selected_attempt = attempts[-1] + final_unbound = getattr(preparation, "unbound_action_plan", None) + if ( + final_unbound is None + and final_candidate_id != unbound_plan["candidate_id"] + ): + final_candidate = next( + item + for item in candidate_set["candidates"] + if item["candidate_id"] == final_candidate_id + ) + final_unbound = self.action_agent.draft(final_candidate) + elif final_unbound is None: + final_unbound = unbound_plan + if str(final_unbound.get("candidate_id")) != final_candidate_id: + raise ValueError( + "Final UnboundActionPlan candidate does not match preparation." + ) + selected_attempt["final_unbound_action_plan"] = deepcopy( + dict(final_unbound) + ) + selected_attempt["unbound_transition"] = { + "initial_candidate_id": str(unbound_plan["candidate_id"]), + "initial_hash": canonical_hash(unbound_plan), + "final_candidate_id": final_candidate_id, + "final_hash": canonical_hash(final_unbound), + "changed": final_unbound != unbound_plan, + } + _write_json( + preparation.output_dir.parent / "final_unbound_action_plan.json", + final_unbound, + ) + _write_json( + preparation.output_dir.parent / "attempt.json", + selected_attempt, + ) + + for stage in ( + WorkflowStage.FINAL_BINDING, + WorkflowStage.STATIC_FEASIBILITY, + WorkflowStage.GROUNDED_ACTION, + ): + state = start_stage(state, stage) + state = complete_stage(state, stage) + if not execute: + final_root = staging / "final" + final_bundle = final_root / "bundle" + final_root.mkdir() + shutil.copytree(preparation.output_dir, final_bundle) + selected_attempt["status"] = "prepared" + _write_json( + preparation.output_dir.parent / "attempt.json", + selected_attempt, + ) + _write_json( + final_root / "selection.json", + { + "scene_attempt": selected_attempt["scene_attempt"], + "candidate_id": final_candidate_id, + "action_attempt": None, + "execution_report": None, + }, + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="prepared", + failure_class=None, + final_bundle=final_bundle, + ) + state = start_stage(state, WorkflowStage.EXECUTION) + + successful_report: Mapping[str, Any] | None = None + successful_action_root: Path | None = None + success_terms = _bundle_success_terms(preparation.output_dir) + for action_index in range(1, workflow_cfg.max_action_attempts + 1): + action_seed = int(base_seed) + action_index - 1 + action_root = ( + preparation.output_dir.parent + / "action_attempts" + / f"action_{action_index:04d}" + ) + action_record: dict[str, Any] = { + "action_attempt": action_index, + "seed": action_seed, + "status": "running", + "successful_environments": 0, + "required_successes": execution_cfg.required_successes, + "error": None, + } + try: + report = self.action_executor( + preparation.output_dir, + action_root, + seed=action_seed, + num_envs=execution_cfg.num_envs, + dataset_saving=bool(dataset_saving), + ) + successes = _environment_successes( + report, + required_semantic_steps=success_terms, + ) + if len(successes) != execution_cfg.num_envs: + raise ValueError( + "Execution report environment count does not match " + "TaskEngineExecutionCfg.num_envs." + ) + action_record["successful_environments"] = sum(successes) + accepted = ( + str(report.get("status")) not in {"rejected", "aborted"} + and sum(successes) >= execution_cfg.required_successes + ) + action_record["status"] = "succeeded" if accepted else "failed" + _write_json(action_root / "task_engine_attempt.json", action_record) + selected_attempt["action_attempts"].append(deepcopy(action_record)) + if accepted: + successful_report = deepcopy(dict(report)) + successful_action_root = action_root + break + except Exception as exc: + action_record["status"] = "failed" + action_record["error"] = _error_record(exc) + action_root.mkdir(parents=True, exist_ok=True) + _write_json(action_root / "task_engine_attempt.json", action_record) + selected_attempt["action_attempts"].append(deepcopy(action_record)) + _write_json( + preparation.output_dir.parent / "attempt.json", selected_attempt + ) + + if successful_report is None: + selected_attempt["status"] = "execution_failed" + _write_json( + preparation.output_dir.parent / "attempt.json", selected_attempt + ) + state = fail_stage( + state, + WorkflowStage.EXECUTION, + reason="All bounded Action Engine execution attempts failed.", + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="action_execution", + ) + + selected_attempt["status"] = "succeeded" + _write_json( + preparation.output_dir.parent / "attempt.json", selected_attempt + ) + state = complete_stage(state, WorkflowStage.EXECUTION) + final_root = staging / "final" + final_bundle = final_root / "bundle" + final_root.mkdir() + shutil.copytree(preparation.output_dir, final_bundle) + _write_json( + final_root / "selection.json", + { + "scene_attempt": selected_attempt["scene_attempt"], + "candidate_id": final_candidate_id, + "action_attempt": int(successful_action_root.name.split("_")[-1]), + "execution_report": successful_report, + "success_spec_steps": list(success_terms), + }, + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="succeeded", + failure_class=None, + final_bundle=final_bundle, + ) + + def _draft_with_fallback( + self, + candidate_set: Mapping[str, Any], + selection: CandidateSelection, + ) -> tuple[Mapping[str, Any], list[dict[str, Any]]]: + selected_id = selection.selected_candidate_id + resolved_ids = { + str(item["candidate_id"]) + for item in selection.binding_report["candidates"] + if item["status"] == "resolved" + } + ordered = [selected_id] + [ + str(item["candidate_id"]) + for item in candidate_set["candidates"] + if item["candidate_id"] != selected_id + and item["candidate_id"] in resolved_ids + ] + failures = [] + for candidate_id in ordered: + candidate = next( + item + for item in candidate_set["candidates"] + if item["candidate_id"] == candidate_id + ) + try: + return self.action_agent.draft(candidate), failures + except ActionCapabilityError: + raise + except (TypeError, ValueError) as exc: + failures.append( + { + "candidate_id": candidate_id, + "stage": "unbound_action", + "draft": deepcopy(candidate["draft"]), + "error": _error_record(exc), + } + ) + raise ValueError( + "No selected task candidate can be represented by Action Engine." + ) + + @staticmethod + def _publish( + transaction: ArtifactTransaction, + staging: Path, + request: Mapping[str, Any], + workflow_cfg: TaskEngineWorkflowCfg, + planning_cfg: TaskEnginePlanningCfg, + execution_cfg: TaskEngineExecutionCfg, + run_metadata: Mapping[str, Any], + state: TaskEngineState, + attempts: Sequence[Mapping[str, Any]], + *, + status: str, + failure_class: str | None, + final_bundle: Path | None = None, + ) -> TaskEngineRunResult: + state_path = staging / "workflow_state.json" + manifest_path = staging / "run_manifest.json" + _write_json(state_path, state.to_dict()) + _write_json( + manifest_path, + { + "schema_version": TASK_ENGINE_RUN_MANIFEST_SCHEMA, + "run_id": run_metadata["run_id"], + "created_at": run_metadata["created_at"], + "output_root": Path(request["output_dir"]).parent.as_posix(), + "run_dir": Path(request["output_dir"]).as_posix(), + "status": status, + "failure_class": failure_class, + "request": deepcopy(dict(request)), + "configuration": { + "workflow": { + "max_parallel_workers": workflow_cfg.max_parallel_workers, + "max_scene_attempts": workflow_cfg.max_scene_attempts, + "max_action_attempts": workflow_cfg.max_action_attempts, + }, + "planning": { + "candidate_count": planning_cfg.candidate_count, + "planning_mode": planning_cfg.planning_mode, + "max_episodes": planning_cfg.max_episodes, + "max_episode_steps": planning_cfg.max_episode_steps, + }, + "execution": { + "num_envs": execution_cfg.num_envs, + "success_policy": execution_cfg.success_policy, + "min_successful_envs": execution_cfg.min_successful_envs, + "dataset_saving": bool(run_metadata["dataset_saving"]), + }, + }, + "attempts": deepcopy(list(attempts)), + "final_bundle": ( + None if final_bundle is None else final_bundle.as_posix() + ), + }, + ) + published = transaction.commit() + return TaskEngineRunResult( + status=status, + output_dir=published, + manifest_path=published / manifest_path.name, + state_path=published / state_path.name, + final_bundle=( + None if final_bundle is None else published / "final" / "bundle" + ), + failure_class=failure_class, + ) + + +def _complete_materialized_scene( + state: TaskEngineState, + *, + has_edit: bool, +) -> TaskEngineState: + if has_edit and state.stages[WorkflowStage.SCENE_EDIT].value == "running": + state = complete_stage(state, WorkflowStage.SCENE_EDIT) + state = start_stage(state, WorkflowStage.SCENE_FINALIZATION) + if state.stages[WorkflowStage.SCENE_FINALIZATION].value == "running": + state = complete_stage(state, WorkflowStage.SCENE_FINALIZATION) + return state + + +def _scene_remediable( + preparation: PreparationResult, + *, + analysis: SceneAnalysis, + request: Mapping[str, Any], +) -> bool: + if preparation.status != "infeasible": + return False + report = preparation.feasibility_report + if not isinstance(report, Mapping) or report.get("remediation_class") != ( + "scene_remediable" + ): + return False + if analysis.input_kind == "image": + return True + return request["scene_edit_prompt"] is not None + + +def _is_scene_remediable_error(error: Exception) -> bool: + """Return whether one typed Scene failure may create a new attempt.""" + return isinstance(error, (SceneRemediableError, SceneServiceError)) + + +def _preparation_failure_class( + preparation: PreparationResult | None, + *, + scene_error: Exception | None, + analysis: SceneAnalysis, + request: Mapping[str, Any], +) -> str: + if scene_error is not None: + return "scene_materialization" + if preparation is None: + return "scene_materialization" + if preparation.status == "planning_failed": + return "action_capability" + if preparation.status in {"ambiguous", "unsatisfied"}: + return "input_conflict" + if preparation.status == "infeasible": + report = preparation.feasibility_report + remediation = ( + str(report.get("remediation_class")) + if isinstance(report, Mapping) + else "terminal" + ) + if remediation == "action_capability": + return "action_capability" + if remediation == "input_conflict": + return "input_conflict" + if remediation != "scene_remediable": + return "terminal_feasibility" + if ( + analysis.input_kind == "gym_project" + and request["scene_edit_prompt"] is None + ): + return "read_only_scene_infeasible" + return "scene_infeasible" + return "preparation" + + +def _failure_stage( + failure_class: str, + request: Mapping[str, Any], +) -> WorkflowStage: + if failure_class == "action_capability": + return WorkflowStage.GROUNDED_ACTION + if failure_class == "preparation_error": + return WorkflowStage.FINAL_BINDING + if failure_class == "input_conflict": + return WorkflowStage.FINAL_BINDING + if failure_class in { + "scene_infeasible", + "read_only_scene_infeasible", + "terminal_feasibility", + }: + return WorkflowStage.STATIC_FEASIBILITY + if failure_class == "scene_materialization": + return ( + WorkflowStage.SCENE_EDIT + if request["scene_edit_prompt"] is not None + else WorkflowStage.SCENE_FINALIZATION + ) + return WorkflowStage.GROUNDED_ACTION + + +def _environment_successes( + report: Mapping[str, Any], + *, + required_semantic_steps: Sequence[str] = (), +) -> list[bool]: + environments = report.get("environments") + if not isinstance(environments, Sequence) or isinstance(environments, (str, bytes)): + raise ValueError("Execution report environments must be a sequence.") + values = [] + for item in environments: + if not isinstance(item, Mapping) or not isinstance(item.get("success"), bool): + raise ValueError("Every execution environment requires boolean success.") + success = bool(item["success"]) + if required_semantic_steps: + semantics = item.get("semantic_success") + if not isinstance(semantics, Mapping): + success = False + else: + success = success and all( + semantics.get(step_id) is True + for step_id in required_semantic_steps + ) + values.append(success) + if not values: + raise ValueError("Execution report must contain at least one environment.") + return values + + +def _bundle_success_terms(bundle: Path) -> tuple[str, ...]: + path = bundle / "grounded_task_plan.json" + if not path.is_file(): + return () + try: + value = _read_json(path) + success_spec = value.get("success_spec") + terms = success_spec.get("terms") if isinstance(success_spec, Mapping) else None + strict = isinstance(value.get("schema_version"), str) + if not isinstance(terms, Sequence) or isinstance(terms, (str, bytes)): + if strict: + raise ValueError("GroundedTaskPlan has no valid SuccessSpec terms.") + return () + result = tuple( + str(item["step_id"]) + for item in terms + if isinstance(item, Mapping) and isinstance(item.get("step_id"), str) + ) + if len(result) != len(terms) or (strict and not result): + if strict: + raise ValueError("GroundedTaskPlan SuccessSpec terms are invalid.") + return () + return result + except OSError: + return () + + +def _highest_vote_candidate(candidate_set: Mapping[str, Any]) -> Mapping[str, Any]: + candidates = candidate_set.get("candidates") + if not isinstance(candidates, Sequence) or isinstance(candidates, (str, bytes)): + raise TypeError("TaskCandidateSet.candidates must be a sequence.") + values = [item for item in candidates if isinstance(item, Mapping)] + if not values: + raise ValueError("TaskCandidateSet requires at least one candidate.") + return max( + values, + key=lambda item: ( + int(item.get("vote_count", 0)), + str(item.get("candidate_id", "")), + ), + ) + + +def _copy_trajectory_record(report: Mapping[str, Any], output_root: Path) -> str | None: + raw = report.get("record_dir") + if not isinstance(raw, str) or not raw: + return None + source = Path(raw).expanduser().resolve() + if not source.is_dir(): + return None + destination = output_root / "trajectory" + if source == destination or destination in source.parents: + return source.as_posix() + shutil.copytree(source, destination) + return destination.as_posix() + + +def _revision_record(revision: SceneRevision) -> dict[str, Any]: + return { + "source": revision.source.as_posix(), + "output_root": ( + None if revision.output_root is None else revision.output_root.as_posix() + ), + "revision_id": revision.revision_id, + "seed": revision.seed, + "edit_plan": deepcopy(revision.edit_plan), + "source_fingerprint": ( + None + if revision.source_fingerprint is None + else revision.source_fingerprint.to_dict() + ), + } + + +def _error_record(error: Exception) -> dict[str, str]: + return { + "type": type(error).__name__, + "failure_type": ( + "scene_remediable" if _is_scene_remediable_error(error) else "terminal" + ), + "message": str(error), + } + + +def _read_json(path: Path) -> dict[str, Any]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"JSON artifact must contain an object: {path}") + return value + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(value, ensure_ascii=False, indent=2, allow_nan=False) + "\n", + encoding="utf-8", + ) diff --git a/embodichain/gen_sim/task_engine/workflow_contracts.py b/embodichain/gen_sim/task_engine/workflow_contracts.py new file mode 100644 index 000000000..a6b61b04b --- /dev/null +++ b/embodichain/gen_sim/task_engine/workflow_contracts.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Strict inputs for Task Engine cross-engine workflows.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +import json +from pathlib import Path +from typing import Any, Final, Literal, TypeAlias + +__all__ = [ + "TASK_RUN_REQUEST_SCHEMA", + "SceneInputKind", + "TaskRunRequest", + "scene_input_kind", + "validate_scene_history_root", + "validate_scene_output_separation", + "validate_task_run_request", +] + +TASK_RUN_REQUEST_SCHEMA: Final = "embodichain.task-engine-run-request/v1" +TaskRunRequest: TypeAlias = dict[str, Any] +SceneInputKind = Literal["image", "gym_project"] + +_REQUEST_KEYS = frozenset( + { + "schema_version", + "task_id", + "task_instruction", + "image_path", + "gym_project", + "scene_edit_prompt", + "output_dir", + } +) + + +def validate_task_run_request(value: Mapping[str, Any]) -> TaskRunRequest: + """Validate and detach one Task Engine run request. + + Version 1 deliberately has no ``scene_generation_prompt``. Image workflows + use the image-only Scene Engine generation behavior and may apply one + optional edit after that initial scene has been generated. + """ + if not isinstance(value, Mapping): + raise TypeError("TaskRunRequest must be a mapping.") + result = deepcopy(dict(value)) + if set(result) != _REQUEST_KEYS: + missing = sorted(_REQUEST_KEYS - set(result)) + extra = sorted(set(result) - _REQUEST_KEYS) + raise ValueError( + f"TaskRunRequest fields differ; missing={missing}, extra={extra}." + ) + if result.get("schema_version") != TASK_RUN_REQUEST_SCHEMA: + raise ValueError( + "TaskRunRequest.schema_version must be " f"{TASK_RUN_REQUEST_SCHEMA!r}." + ) + result["task_id"] = _nonempty(result.get("task_id"), "task_id") + result["task_instruction"] = _nonempty( + result.get("task_instruction"), "task_instruction" + ) + result["output_dir"] = _path(result.get("output_dir"), "output_dir") + + image_path = _optional_path(result.get("image_path"), "image_path") + gym_project = _optional_path(result.get("gym_project"), "gym_project") + if (image_path is None) == (gym_project is None): + raise ValueError( + "TaskRunRequest requires exactly one of image_path or gym_project." + ) + result["image_path"] = image_path + result["gym_project"] = gym_project + if gym_project is not None: + validate_scene_output_separation(gym_project, result["output_dir"]) + + edit_prompt = result.get("scene_edit_prompt") + if edit_prompt is not None: + edit_prompt = _nonempty(edit_prompt, "scene_edit_prompt") + result["scene_edit_prompt"] = edit_prompt + _json_safe(result) + return result + + +def scene_input_kind(request: Mapping[str, Any]) -> SceneInputKind: + """Return the selected scene input kind after validating ``request``.""" + normalized = validate_task_run_request(request) + return "image" if normalized["image_path"] is not None else "gym_project" + + +def validate_scene_output_separation( + gym_project: str | Path, + output_dir: str | Path, +) -> None: + """Reject output paths that could replace or modify a read-only source. + + Args: + gym_project: Existing Gym project directory or configuration path. + output_dir: Transactional output directory for the Task Engine run. + + Raises: + ValueError: If either path contains the other or both paths are equal. + """ + source = Path(gym_project).expanduser().resolve() + output = Path(output_dir).expanduser().resolve() + if source == output or source in output.parents or output in source.parents: + raise ValueError( + "Task Engine output_dir and source Gym project must not overlap." + ) + + +def validate_scene_history_root( + gym_project: str | Path, + output_root: str | Path, +) -> None: + """Protect a source project before reserving a history-directory child. + + A prior run may live below the same history root because every new run is + published to a distinct timestamped child. The inverse remains unsafe: + creating the history root at or below the source project would write a + reservation and output artifacts into the read-only source tree. + + Args: + gym_project: Existing Gym project directory or configuration path. + output_root: Parent directory under which a new run will be reserved. + + Raises: + ValueError: If the history root is equal to or contained by the source + project boundary. + """ + source = Path(gym_project).expanduser().resolve() + protected_root = source.parent if source.is_file() else source + history_root = Path(output_root).expanduser().resolve() + if protected_root == history_root or protected_root in history_root.parents: + raise ValueError( + "Task Engine output_root must not be inside the read-only source " + "Gym project." + ) + + +def _path(value: Any, field_name: str) -> str: + text = _nonempty(value, field_name) + return Path(text).expanduser().resolve().as_posix() + + +def _optional_path(value: Any, field_name: str) -> str | None: + if value is None: + return None + return _path(value, field_name) + + +def _nonempty(value: Any, field_name: str) -> str: + if not isinstance(value, str): + raise TypeError(f"TaskRunRequest.{field_name} must be a string.") + result = value.strip() + if not result: + raise ValueError(f"TaskRunRequest.{field_name} must not be empty.") + return result + + +def _json_safe(value: Any) -> None: + try: + json.dumps(value, ensure_ascii=False, allow_nan=False) + except (TypeError, ValueError) as exc: + raise ValueError("TaskRunRequest must contain strict JSON data.") from exc diff --git a/setup.py b/setup.py index ac131d484..06f687a3a 100644 --- a/setup.py +++ b/setup.py @@ -133,7 +133,9 @@ def main(): package_dir=get_package_dir(), package_data={ "embodichain": ["VERSION"], + "embodichain.gen_sim.action_engine.config": ["*.yaml"], "embodichain.gen_sim.simready_pipeline.configs": ["*.json"], + "embodichain.gen_sim.task_engine": ["*.yaml"], "embodichain_tasks.configs": ["**/*.json", "**/*.yaml", "**/*.yml"], }, cmdclass=cmdclass, diff --git a/tests/gen_sim/task_engine/__init__.py b/tests/gen_sim/task_engine/__init__.py new file mode 100644 index 000000000..b201491d8 --- /dev/null +++ b/tests/gen_sim/task_engine/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for Task Engine semantics and orchestration.""" + +from __future__ import annotations diff --git a/tests/gen_sim/task_engine/orchestration/__init__.py b/tests/gen_sim/task_engine/orchestration/__init__.py new file mode 100644 index 000000000..8256d7018 --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for Task Engine cross-engine orchestration.""" + +from __future__ import annotations diff --git a/tests/gen_sim/task_engine/orchestration/test_architecture.py b/tests/gen_sim/task_engine/orchestration/test_architecture.py new file mode 100644 index 000000000..ad3593c0a --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_architecture.py @@ -0,0 +1,90 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import ast +from pathlib import Path + +import embodichain.gen_sim as gen_sim_package +from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.task_engine import TaskAgent +from embodichain.gen_sim.task_engine import __main__ as task_engine_main +from embodichain.gen_sim.task_engine import cli as task_engine_cli +from embodichain.gen_sim.task_engine.orchestration.coordinator import ( + TaskEngineCoordinator, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import SceneAdapter + +_GEN_SIM_ROOT = Path(gen_sim_package.__file__).resolve().parent +_PURE_TASK_MODULES = ( + "agent.py", + "config.py", + "contracts.py", + "interpretation.py", + "ontology.py", + "state_machine.py", + "workflow_contracts.py", +) + + +def test_task_semantic_core_does_not_import_scene_action_or_orchestration() -> None: + forbidden = { + "embodichain.gen_sim.action_engine", + "embodichain.gen_sim.scene_engine", + "embodichain.gen_sim.task_engine.orchestration", + "embodichain.gen_sim.task_engine.scene", + } + offenders: list[str] = [] + for filename in _PURE_TASK_MODULES: + path = _GEN_SIM_ROOT / "task_engine" / filename + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + modules = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + modules = [node.module or ""] + else: + continue + if any( + module == prefix or module.startswith(prefix + ".") + for module in modules + for prefix in forbidden + ): + offenders.append(filename) + break + assert offenders == [] + + +def test_cross_engine_owners_are_explicit() -> None: + assert TaskAgent.__module__ == "embodichain.gen_sim.task_engine.agent" + assert ActionAgent.__module__ == "embodichain.gen_sim.action_engine.agent" + assert SceneAdapter.__module__.startswith( + "embodichain.gen_sim.task_engine.orchestration" + ) + assert TaskEngineCoordinator.__module__.startswith( + "embodichain.gen_sim.task_engine.orchestration" + ) + + +def test_task_engine_owns_its_module_entry_point() -> None: + assert task_engine_main.main is task_engine_cli.main + + +def test_legacy_cross_engine_packages_are_deleted() -> None: + assert not (_GEN_SIM_ROOT / "scene_bridge").exists() + assert not (_GEN_SIM_ROOT / "collaboration").exists() + assert not (_GEN_SIM_ROOT / "action_engine" / "collaboration").exists() diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py new file mode 100644 index 000000000..053b3e93f --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -0,0 +1,1102 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import replace +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.task_engine import cli +from embodichain.gen_sim.task_engine import _bundle_runner as bundle_runner +from embodichain.gen_sim.task_engine.orchestration.artifacts import ( + ArtifactTransaction, +) +from embodichain.gen_sim.task_engine.orchestration.contracts import ( + BINDING_REPORT_SCHEMA, + ROLE_BINDINGS_SCHEMA, + SCENE_MANIFEST_SCHEMA, + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + canonical_hash, +) +from embodichain.gen_sim.task_engine.orchestration.coordinator import ( + TaskEngineCoordinator, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( + SceneAdaptation, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import SceneSourceRef +from embodichain.gen_sim.action_engine.generation.artifacts import artifact_paths +from embodichain.gen_sim.action_engine.generation.models import PreparedScene +from embodichain.gen_sim.action_engine.protocol import ( + AGENT_CONFIG_FILENAME, + FAST_GYM_CONFIG_FILENAME, +) +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + build_execution_provenance, +) +from embodichain.gen_sim.task_engine.scene import SceneEngineV1Adapter + +_UPRIGHT_CAN_INSTRUCTION = "test-instruction" + + +def _candidate_set() -> dict: + selector = { + "kind": "scene_ref", + "step_id": "", + "reference": "red can", + "quantifier": "one", + "count": 0, + } + none_selector = { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + step = { + "id": "upright", + "task_type": "E2", + "object": selector, + "target": none_selector, + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "steps": [step], + } + candidate = { + "candidate_id": "candidate_01", + "draft": draft, + "scene_request": { + "schema_version": SCENE_REQUEST_SCHEMA, + "task_id": "upright_can", + "references": [ + { + "reference_id": "upright.object", + "step_id": "upright", + "role": "object", + "reference": "red can", + "quantifier": "one", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + }, + "success_spec": { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": "upright_can", + "op": "all", + "terms": [{"step_id": "upright", "type": "object_upright"}], + }, + "semantic_hash": canonical_hash([step]), + "vote_count": 1, + "attempts": 1, + "normalizations": [], + } + return { + "schema_version": TASK_CANDIDATE_SET_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "candidates": [candidate], + "requested_candidate_count": 1, + "valid_response_count": 1, + "errors": [], + } + + +def _candidate_set_with_alternative() -> dict: + candidates = _candidate_set() + alternative = deepcopy(candidates["candidates"][0]) + alternative["candidate_id"] = "candidate_02" + alternative["draft"]["steps"][0]["required_arm"] = "left_arm" + alternative["semantic_hash"] = canonical_hash(alternative["draft"]["steps"]) + candidates["candidates"].append(alternative) + candidates["requested_candidate_count"] = 2 + candidates["valid_response_count"] = 2 + return candidates + + +def _prepared_scene(tmp_path: Path) -> PreparedScene: + scene_path = tmp_path / "scene_config.json" + scene_path.write_text("{}", encoding="utf-8") + scene_object = { + "uid": "red_can", + "source_uid": "red_can", + "role": "rigid_object", + "name": "red can", + "description": "A red can.", + "category": "can", + "color": "red", + "position": [0.0, 0.0, 0.5], + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + return PreparedScene( + source_config_path=scene_path, + scene_dir=tmp_path, + planner_objects=(scene_object,), + background=(), + rigid_objects=(), + articulations=(), + uid_map={"red_can": "red_can"}, + table_top_z=None, + z_rotation_degrees=0.0, + body_scale_policy="preserve", + body_scale=(1.0, 1.0, 1.0), + asset_hashes={}, + ) + + +def _adaptation(tmp_path: Path, *, status: str = "bound") -> SceneAdaptation: + candidates = _candidate_set() + candidate = candidates["candidates"][0] + selected_id = candidate["candidate_id"] if status == "bound" else "" + role_bindings = ( + { + "schema_version": ROLE_BINDINGS_SCHEMA, + "task_id": "upright_can", + "candidate_id": "candidate_01", + "reference_bindings": {"upright.object": ["red_can"]}, + "role_bindings": {}, + } + if status == "bound" + else None + ) + return SceneAdaptation( + scene_manifest={ + "schema_version": SCENE_MANIFEST_SCHEMA, + "scene_id": "scene", + "source_format": "test", + "robot_profile": "dual_franka", + "objects": [ + { + "uid": "red_can", + "role": "rigid_object", + "name": "red can", + "description": "A red can.", + "category": "can", + "color": "red", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + }, + role_bindings=role_bindings, + binding_report={ + "schema_version": BINDING_REPORT_SCHEMA, + "task_id": "upright_can", + "status": status, + "selected_candidate_id": selected_id, + "selection_reason": "test", + "candidates": [ + { + "candidate_id": "candidate_01", + "semantic_hash": candidate["semantic_hash"], + "status": "resolved" if status == "bound" else status, + "references": [ + { + "reference_id": "upright.object", + "status": ( + "resolved" if status == "bound" else "ambiguous" + ), + "confidence": 1.0, + "candidate_uids": ["red_can"], + "selected_uids": (["red_can"] if status == "bound" else []), + "reasons": [], + } + ], + "reasons": [], + } + ], + }, + selected_candidate=deepcopy(candidate) if status == "bound" else None, + prepared_scene=_prepared_scene(tmp_path), + source_config_path=tmp_path / "scene_config.json", + conservative_scene_graph={ + "schema_version": "embodichain.conservative-scene-graph/v1", + "scene_id": "scene", + "nodes": [ + { + "uid": "red_can", + "parent_uid": "unknown", + "parent_relation": "unknown", + "orientation": "unknown", + "source": "test", + } + ], + "relations": [], + }, + ) + + +def _adaptation_with_alternative(tmp_path: Path) -> SceneAdaptation: + candidate_set = _candidate_set_with_alternative() + adaptation = _adaptation(tmp_path) + alternative = candidate_set["candidates"][1] + alternative_audit = deepcopy(adaptation.binding_report["candidates"][0]) + alternative_audit["candidate_id"] = "candidate_02" + alternative_audit["semantic_hash"] = alternative["semantic_hash"] + alternative_bindings = { + **deepcopy(adaptation.role_bindings), + "candidate_id": "candidate_02", + } + return replace( + adaptation, + binding_report={ + **deepcopy(adaptation.binding_report), + "candidates": [ + *deepcopy(adaptation.binding_report["candidates"]), + alternative_audit, + ], + }, + candidate_bindings={"candidate_02": alternative_bindings}, + ) + + +def test_artifact_transaction_rolls_back_and_preserves_existing_output( + tmp_path: Path, +) -> None: + output = tmp_path / "bundle" + output.mkdir() + (output / "kept.txt").write_text("old", encoding="utf-8") + + with pytest.raises(RuntimeError, match="fail"): + with ArtifactTransaction(output, overwrite=True) as transaction: + assert transaction.staging_dir is not None + (transaction.staging_dir / "partial.txt").write_text( + "partial", encoding="utf-8" + ) + raise RuntimeError("fail before commit") + + assert (output / "kept.txt").read_text(encoding="utf-8") == "old" + assert not (output / "partial.txt").exists() + + +def test_prepare_rejects_output_overlapping_read_only_source(tmp_path: Path) -> None: + source = tmp_path / "gym_project" + source.mkdir() + coordinator = TaskEngineCoordinator( + task_agent=object(), + scene_adapter=SimpleNamespace(robot_profile="franka"), + action_agent=object(), + feasibility_broker=object(), + ) + + with pytest.raises(ValueError, match="must not overlap"): + coordinator.prepare( + "task", + "Pick up the object.", + source, + source / "task_run", + overwrite=True, + ) + + +def test_unbound_prepare_publishes_only_audit_artifacts(tmp_path: Path) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path, status="ambiguous") + task_agent = SimpleNamespace(generate=lambda *args, **kwargs: candidates) + scene_adapter = SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ) + action_agent = SimpleNamespace( + plan=lambda *_args, **_kwargs: pytest.fail("Action Agent must not run") + ) + coordinator = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=scene_adapter, + action_agent=action_agent, + bundle_generator=lambda *_args, **_kwargs: pytest.fail( + "legacy generator must not run" + ), + ) + + result = coordinator.prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "bundle", + candidate_count=1, + ) + + assert result.status == "ambiguous" + assert (result.output_dir / "task_candidate_set.json").is_file() + assert (result.output_dir / "binding_report.json").is_file() + assert not (result.output_dir / "scene_manifest.json").exists() + assert not (result.output_dir / "role_bindings.json").exists() + assert not (result.output_dir / "grounded_task_plan.json").exists() + assert not (result.output_dir / FAST_GYM_CONFIG_FILENAME).exists() + + +def test_prepare_reuses_precomputed_candidates_without_rerunning_task_agent( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path, status="ambiguous") + task_agent = SimpleNamespace( + generate=lambda *args, **kwargs: pytest.fail("Task Agent must not rerun") + ) + coordinator = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ), + action_agent=object(), + ) + + result = coordinator.prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "candidate-reuse", + candidate_set=candidates, + force_most_likely=True, + ) + + assert result.status == "ambiguous" + assert result.candidate_set == candidates + + +def test_prepare_inherits_adapter_robot_profile_for_raw_scene_path( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path, status="ambiguous") + captured: dict[str, object] = {} + + def adapt(_candidates, source, **_kwargs): + captured["source"] = source + return adaptation + + coordinator = TaskEngineCoordinator( + task_agent=SimpleNamespace( + generate=lambda *args, **kwargs: pytest.fail("Task Agent must not rerun") + ), + scene_adapter=SimpleNamespace(robot_profile="ur10", adapt=adapt), + action_agent=object(), + ) + + result = coordinator.prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "ur10-bundle", + candidate_set=candidates, + ) + + assert result.status == "ambiguous" + assert isinstance(captured["source"], SceneSourceRef) + assert captured["source"].robot_profile == "ur10" + + +def test_contradicted_feasibility_publishes_audit_without_planning( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path) + static_manifest = SceneEngineV1Adapter().adapt_prepared_scene( + adaptation.prepared_scene, + source_format="test", + robot_profile="dual_franka", + ) + adaptation = replace( + adaptation, + static_scene_manifest=static_manifest, + ) + task_agent = SimpleNamespace(generate=lambda *args, **kwargs: candidates) + scene_adapter = SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ) + registry = SimpleNamespace( + catalog=lambda: { + name: { + "runtime_available": name != "PickUp", + "unavailable_reason": ( + "PickUp disabled for test." if name == "PickUp" else None + ), + } + for name in ("PickUp", "MoveHeldObject", "Place") + } + ) + action_agent = SimpleNamespace( + registry=registry, + plan=lambda *_args, **_kwargs: pytest.fail("Action Agent must not plan"), + ) + + result = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=scene_adapter, + action_agent=action_agent, + bundle_generator=lambda *_args, **_kwargs: pytest.fail( + "legacy generator must not run" + ), + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "infeasible-bundle", + candidate_count=1, + ) + + assert result.status == "infeasible" + assert result.feasibility_report is not None + assert result.feasibility_report["status"] == "contradicted" + assert result.feasibility_report["remediation_class"] == "action_capability" + assert result.artifacts.static_scene_manifest.is_file() + assert result.artifacts.feasibility_report.is_file() + assert not result.artifacts.grounded_task_plan.exists() + + +def test_feasibility_contradiction_falls_back_to_next_resolved_candidate( + tmp_path: Path, +) -> None: + candidate_set = _candidate_set() + first = candidate_set["candidates"][0] + second = deepcopy(first) + second["candidate_id"] = "candidate_02" + second["semantic_hash"] = "b" * 64 + candidate_set["candidates"].append(second) + adaptation = _adaptation(tmp_path) + second_audit = deepcopy(adaptation.binding_report["candidates"][0]) + second_audit["candidate_id"] = "candidate_02" + second_audit["semantic_hash"] = "b" * 64 + binding_report = deepcopy(adaptation.binding_report) + binding_report["candidates"].append(second_audit) + second_bindings = { + **deepcopy(adaptation.role_bindings), + "candidate_id": "candidate_02", + } + adaptation = replace( + adaptation, + binding_report=binding_report, + candidate_bindings={"candidate_02": second_bindings}, + static_scene_manifest={}, + ) + + class _Broker: + @staticmethod + def assess(candidate, *_args, **_kwargs): + return { + "status": ( + "runtime_probe" + if candidate["candidate_id"] == "candidate_02" + else "contradicted" + ) + } + + registry = SimpleNamespace(catalog=lambda: {}) + coordinator = TaskEngineCoordinator( + action_agent=SimpleNamespace(registry=registry), + feasibility_broker=_Broker(), + ) + + updated, selected, bindings, report = coordinator._fallback_feasible_candidate( + candidate_set, + adaptation, + first, + adaptation.role_bindings, + {"status": "contradicted"}, + ) + + assert selected["candidate_id"] == "candidate_02" + assert bindings["candidate_id"] == "candidate_02" + assert report["status"] == "runtime_probe" + assert updated.binding_report["selected_candidate_id"] == "candidate_02" + assert ( + "static feasibility contradicted candidate_01" + in updated.binding_report["selection_reason"] + ) + + +def test_bound_prepare_uses_sidecar_and_publishes_complete_bundle( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + adaptation = _adaptation(tmp_path) + task_agent = SimpleNamespace(generate=lambda *args, **kwargs: candidates) + scene_adapter = SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ) + graph = {"graph": "planned"} + action_agent = SimpleNamespace(plan=lambda _plan: deepcopy(graph)) + generator_calls = [] + + def generator(_scene, output, **kwargs): + generator_calls.append(kwargs) + task_spec_path = Path(kwargs["task_spec"]) + assert task_spec_path.is_file() + assert (task_spec_path.parent / "scene_requirements.json").is_file() + paths = artifact_paths(output) + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.seed_task_graph, + ): + path.parent.mkdir(parents=True, exist_ok=True) + value = graph if path == paths.seed_task_graph else {} + path.write_text(json.dumps(value), encoding="utf-8") + paths.seed_task_graph_png.write_bytes(b"png") + return paths + + result = TaskEngineCoordinator( + task_agent=task_agent, + scene_adapter=scene_adapter, + action_agent=action_agent, + bundle_generator=generator, + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "bundle", + candidate_count=1, + ) + + assert result.bound + assert generator_calls + assert not (result.output_dir / ".task_engine_input").exists() + grounded = json.loads( + (result.output_dir / "grounded_task_plan.json").read_text(encoding="utf-8") + ) + assert grounded["success_spec"]["terms"] == [ + {"step_id": "task_01", "type": "object_upright"} + ] + assert (result.output_dir / "seed_task_graph.json").is_file() + + +def test_prepare_falls_back_after_candidate_action_planning_failure( + tmp_path: Path, +) -> None: + candidates = _candidate_set_with_alternative() + adaptation = _adaptation_with_alternative(tmp_path) + planned_candidates: list[str] = [] + graph = {"graph": "planned"} + + def plan(grounded_plan): + candidate_id = grounded_plan["selected_candidate_id"] + planned_candidates.append(candidate_id) + if candidate_id == "candidate_01": + raise ValueError( + "SeedGraph TaskGroup 'task_04' requires unavailable state " + "{'predicate': 'arm_free', 'arm': 'right_arm'}." + ) + return deepcopy(graph) + + def generator(_scene, output, **_kwargs): + paths = artifact_paths(output) + for path in ( + paths.gym_config, + paths.agent_config, + paths.task_spec, + paths.scene_requirements, + paths.seed_task_graph, + ): + path.parent.mkdir(parents=True, exist_ok=True) + value = graph if path == paths.seed_task_graph else {} + path.write_text(json.dumps(value), encoding="utf-8") + paths.seed_task_graph_png.write_bytes(b"png") + return paths + + result = TaskEngineCoordinator( + task_agent=SimpleNamespace(generate=lambda *args, **kwargs: candidates), + scene_adapter=SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ), + action_agent=SimpleNamespace(plan=plan), + bundle_generator=generator, + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + tmp_path / "fallback-bundle", + candidate_count=2, + ) + + assert result.bound + assert result.selected_candidate_id == "candidate_02" + assert planned_candidates == ["candidate_01", "candidate_02"] + assert ( + "candidate_01 failed action_planning" + in result.adaptation.binding_report["selection_reason"] + ) + assert not result.artifacts.preparation_failure.exists() + + +def test_prepare_publishes_failure_context_when_all_candidates_fail_planning( + tmp_path: Path, +) -> None: + candidates = _candidate_set_with_alternative() + adaptation = _adaptation_with_alternative(tmp_path) + output = tmp_path / "failed-bundle" + output.mkdir() + (output / "stale.txt").write_text("old", encoding="utf-8") + + result = TaskEngineCoordinator( + task_agent=SimpleNamespace(generate=lambda *args, **kwargs: candidates), + scene_adapter=SimpleNamespace( + robot_profile="franka", + adapt=lambda *args, **kwargs: adaptation, + ), + action_agent=SimpleNamespace( + plan=lambda _plan: (_ for _ in ()).throw( + ValueError( + "SeedGraph TaskGroup 'task_04' requires unavailable state " + "{'predicate': 'arm_free', 'arm': 'right_arm'}." + ) + ) + ), + bundle_generator=lambda *_args, **_kwargs: pytest.fail( + "bundle generation must not run" + ), + ).prepare( + "upright_can", + _UPRIGHT_CAN_INSTRUCTION, + tmp_path / "scene_config.json", + output, + candidate_count=2, + overwrite=True, + ) + + assert result.status == "planning_failed" + assert not result.bound + assert result.artifacts.preparation_failure.is_file() + assert not (result.output_dir / "stale.txt").exists() + failure = json.loads( + result.artifacts.preparation_failure.read_text(encoding="utf-8") + ) + assert failure["schema_version"] == "action_engine_preparation_failure_v1" + assert failure["task_id"] == "upright_can" + assert failure["selected_candidate_id"] == "candidate_01" + assert [attempt["candidate_id"] for attempt in failure["attempts"]] == [ + "candidate_01", + "candidate_02", + ] + for index, attempt in enumerate(failure["attempts"]): + candidate_id = f"candidate_{index + 1:02d}" + assert attempt["stage"] == "action_planning" + assert attempt["draft"] == candidates["candidates"][index]["draft"] + assert attempt["bindings"]["candidate_id"] == candidate_id + assert attempt["grounded_task_plan"]["selected_candidate_id"] == candidate_id + assert "unbound_action_plan" in attempt + assert "action_graph" in attempt + assert attempt["error"]["type"] == "ValueError" + assert "arm_free" in attempt["error"]["message"] + + +def test_private_bundle_runner_forwards_arguments_without_leaking_sys_argv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + (bundle / AGENT_CONFIG_FILENAME).write_text( + json.dumps({"task_name": "task"}), encoding="utf-8" + ) + (bundle / FAST_GYM_CONFIG_FILENAME).write_text("{}", encoding="utf-8") + captured = [] + + def fake_cli() -> None: + import sys + + captured.append(list(sys.argv)) + + import embodichain.gen_sim.action_engine.cli as legacy_cli + + monkeypatch.setattr( + legacy_cli, + "run_agent", + SimpleNamespace(cli=fake_cli), + raising=False, + ) + import sys + + original = sys.argv + assert bundle_runner.main(["--bundle", str(bundle), "--seed", "7"]) == 0 + + assert sys.argv is original + assert captured[0][-2:] == ["--seed", "7"] + assert str(bundle / AGENT_CONFIG_FILENAME) in captured[0] + + +@pytest.mark.parametrize( + ("mode", "image", "scene", "edit"), + [ + ("image", "input.png", None, None), + ("image-edit", "input.png", None, "move the cup left"), + ("scene", None, "gym_project", None), + ("scene-edit", None, "gym_project", "move the cup left"), + ], +) +def test_unified_cli_accepts_exactly_four_modes( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + mode: str, + image: str | None, + scene: str | None, + edit: str | None, +) -> None: + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **kwargs): + captured["request"] = request + captured["kwargs"] = kwargs + output = Path(request["output_dir"]) + return SimpleNamespace( + status="succeeded", + succeeded=True, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + arguments = [ + "--mode", + mode, + "--task-id", + "task", + "--instruction", + "place the cup", + "--output-root", + str(tmp_path / "history"), + "--base-seed", + "9", + ] + if image is not None: + arguments.extend(["--image", str(tmp_path / image)]) + if scene is not None: + arguments.extend(["--scene", str(tmp_path / scene)]) + if edit is not None: + arguments.extend(["--scene-edit", edit]) + if mode == "image": + arguments.append("--dataset_saving") + + assert cli.main(arguments) == 0 + + request = captured["request"] + assert request["image_path"] == (None if image is None else str(tmp_path / image)) + assert request["gym_project"] == (None if scene is None else str(tmp_path / scene)) + assert request["scene_edit_prompt"] == edit + assert captured["kwargs"]["base_seed"] == 9 + assert captured["kwargs"]["dataset_saving"] is (mode == "image") + assert captured["kwargs"]["execute"] is True + payload = json.loads(capsys.readouterr().out) + assert payload["status"] == "succeeded" + assert payload["run_id"].replace("_", "").isdigit() + assert len(payload["run_id"]) == 15 + assert Path(payload["output_dir"]).parent == tmp_path / "history" + + +def test_unified_cli_reuses_history_root_without_modifying_prior_scene( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + history = tmp_path / "task1008" + source = ( + history + / "20260820_105939" + / "attempts" + / "scene_0001" + / "scene_revision" + / "scene_export" + ) + source.mkdir(parents=True) + marker = source / "scene_config.json" + marker.write_text('{"source": "unchanged"}\n', encoding="utf-8") + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **_kwargs): + captured["request"] = request + output = Path(request["output_dir"]) + return SimpleNamespace( + status="succeeded", + succeeded=True, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + + assert ( + cli.main( + [ + "--mode", + "scene", + "--task-id", + "task1008", + "--scene", + str(source), + "--instruction", + "place the cup on the book", + "--output-root", + str(history), + ] + ) + == 0 + ) + + output_dir = Path(captured["request"]["output_dir"]) + assert output_dir.parent == history + assert output_dir != source + assert marker.read_text(encoding="utf-8") == '{"source": "unchanged"}\n' + assert list(history.glob(".*.reserve")) == [] + + +def test_unified_cli_rejects_history_root_inside_source_before_reservation( + tmp_path: Path, +) -> None: + source = tmp_path / "scene_export" + source.mkdir() + output_root = source / "new_runs" + + with pytest.raises(ValueError, match="read-only source"): + cli.main( + [ + "--mode", + "scene", + "--task-id", + "task", + "--scene", + str(source), + "--instruction", + "place the cup", + "--output-root", + str(output_root), + ] + ) + + assert not output_root.exists() + + +def test_unified_cli_rejects_mode_input_mismatch(tmp_path: Path) -> None: + with pytest.raises(SystemExit, match="2"): + cli.main( + [ + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + str(tmp_path / "input.png"), + "--scene", + str(tmp_path / "scene"), + "--output-root", + str(tmp_path / "history"), + ] + ) + + +def test_public_cli_exposes_prepare_run_and_run_all_modes() -> None: + parser = cli.build_parser() + help_text = parser.format_help() + assert "prepare" in help_text + assert "run-all" in help_text + assert "run" in help_text + assert "--overwrite" not in help_text + assert "--run-after-prepare" not in help_text + arguments = parser.parse_args( + [ + "prepare", + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + "input.png", + "--output-root", + "history", + "--dataset_saving", + ] + ) + assert arguments.command == "prepare" + assert arguments.dataset_saving is True + + +def test_prepare_cli_stops_before_simulator_execution( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **kwargs): + captured.update(kwargs) + output = Path(request["output_dir"]) + return SimpleNamespace( + status="prepared", + succeeded=False, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + + result = cli.main( + [ + "prepare", + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + str(tmp_path / "input.png"), + "--output-root", + str(tmp_path / "history"), + ] + ) + + assert result == 0 + assert captured["execute"] is False + + +def test_run_cli_executes_an_existing_bundle( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + + class Executor: + def __call__(self, _bundle, output, **kwargs): + Path(output).mkdir() + assert kwargs["num_envs"] == 2 + return { + "status": "failed", + "environments": [ + {"success": True}, + {"success": False}, + ], + } + + monkeypatch.setattr(cli, "SubprocessActionExecutor", Executor) + + result = cli.main( + [ + "run", + "--bundle", + str(bundle), + "--output-root", + str(tmp_path / "history"), + "--num-envs", + "2", + ] + ) + + assert result == 0 + + +def test_private_bundle_runner_publishes_rejected_preflight_report( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + (bundle / AGENT_CONFIG_FILENAME).write_text( + json.dumps({"task_name": "task"}), encoding="utf-8" + ) + (bundle / FAST_GYM_CONFIG_FILENAME).write_text("{}", encoding="utf-8") + report = ExecutionReport( + task_id="task", + plan_hash="0" * 64, + action_graph_hash="1" * 64, + status="rejected", + run_id="preflight", + episode_id="0", + provenance=build_execution_provenance(), + environments=( + { + "env_id": "0", + "success": False, + "semantic_success": {}, + "action_count": 0, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + }, + ), + error="ValueError: planning-only action", + ) + monkeypatch.setattr( + bundle_runner, + "_preflight_bundle", + lambda *args, **kwargs: report, + ) + + assert bundle_runner.main(["--bundle", str(bundle)]) == 2 + payload = json.loads((bundle / "execution_report.json").read_text(encoding="utf-8")) + assert payload["status"] == "rejected" + assert payload["action_count"] == 0 diff --git a/tests/gen_sim/task_engine/orchestration/test_legacy_scene.py b/tests/gen_sim/task_engine/orchestration/test_legacy_scene.py new file mode 100644 index 000000000..286441fb0 --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_legacy_scene.py @@ -0,0 +1,164 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path + +import trimesh + +from embodichain.gen_sim.action_engine.generation.source_scene import prepare_scene +from embodichain.gen_sim.task_engine.orchestration.legacy_scene import ( + convert_legacy_gym_project, + restore_locked_scene_entities, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import ( + fingerprint_scene_source, + scene_revision_id, +) +from embodichain.gen_sim.task_engine.scene import build_conservative_scene_graph + + +def _legacy_project(tmp_path: Path) -> Path: + project = tmp_path / "legacy" + assets = project / "assets" + assets.mkdir(parents=True) + trimesh.creation.box(extents=[1.0, 1.0, 0.1]).export( + assets / "table.glb", file_type="glb" + ) + trimesh.creation.cylinder(radius=0.03, height=0.12).export( + assets / "can.glb", file_type="glb" + ) + trimesh.creation.box(extents=[0.3, 0.2, 0.4]).export( + assets / "cabinet.glb", file_type="glb" + ) + (assets / "cabinet.urdf").write_text( + '' + '\n', + encoding="utf-8", + ) + config = { + "background": [ + { + "uid": "table_0", + "name": "table", + "description": "A work table.", + "category": "table", + "shape": {"shape_type": "Mesh", "fpath": "assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "can_0", + "name": "red can", + "description": "A red can.", + "category": "can", + "shape": {"shape_type": "Mesh", "fpath": "assets/can.glb"}, + "init_pos": [0.0, 0.1, 0.2], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "articulation": [ + { + "uid": "cabinet_0", + "name": "cabinet", + "description": "A fixed cabinet.", + "category": "cabinet", + "fpath": "assets/cabinet.urdf", + "init_pos": [0.5, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + } + (project / "gym_config.json").write_text(json.dumps(config), encoding="utf-8") + return project + + +def test_legacy_conversion_is_read_only_and_restores_locked_articulation( + tmp_path: Path, +) -> None: + project = _legacy_project(tmp_path) + original = fingerprint_scene_source(project) + + revision = convert_legacy_gym_project(project, tmp_path / "revision") + converted = json.loads(revision.scene_config_path.read_text(encoding="utf-8")) + manifest = json.loads(revision.manifest_path.read_text(encoding="utf-8")) + + assert fingerprint_scene_source(project) == original + assert converted["format"] == "embodichain.scene-export/v1" + assert converted["background"][0]["uid"] == "table" + assert converted["rigid_object"][0]["uid"] == "can" + assert converted["articulation"][0]["uid"] == "cabinet" + assert manifest["audit_hierarchy"] == "unknown" + assert manifest["operational_hierarchy"] == "assumed_on_table" + assert set(revision.locked_entity_uids) == {"table", "cabinet"} + + converted["articulation"] = [] + revision.scene_config_path.write_text(json.dumps(converted), encoding="utf-8") + restore_locked_scene_entities(revision.output_root) + restored = json.loads(revision.scene_config_path.read_text(encoding="utf-8")) + + assert restored["articulation"][0]["uid"] == "cabinet" + assert Path(restored["articulation"][0]["fpath"]).is_file() + assert fingerprint_scene_source(project) == original + + +def test_legacy_conversion_separates_audit_and_operational_hierarchy( + tmp_path: Path, +) -> None: + revision = convert_legacy_gym_project( + _legacy_project(tmp_path), + tmp_path / "revision", + ) + + operational = json.loads(revision.scene_graph_path.read_text(encoding="utf-8")) + conservative = build_conservative_scene_graph( + prepare_scene(revision.scene_config_path), + scene_id="legacy-scene", + ) + + operational_can = next( + node for node in operational["nodes"] if node["object_id"] == "can" + ) + conservative_can = next( + node for node in conservative["nodes"] if node["uid"] == "can" + ) + assert operational_can["parent_id"] == "table" + assert operational_can["parent_relation"] == "on" + assert conservative_can["parent_uid"] == "unknown" + assert conservative_can["parent_relation"] == "unknown" + assert conservative_can["source"] == "conservative_import" + + +def test_scene_identity_covers_transitive_urdf_meshes(tmp_path: Path) -> None: + project = _legacy_project(tmp_path) + original_fingerprint = fingerprint_scene_source(project) + original_revision = scene_revision_id(project) + + trimesh.creation.box(extents=[0.6, 0.2, 0.4]).export( + project / "assets" / "cabinet.glb", + file_type="glb", + ) + + changed_fingerprint = fingerprint_scene_source(project) + assert changed_fingerprint.asset_sha256 != original_fingerprint.asset_sha256 + assert scene_revision_id(project) != original_revision diff --git a/tests/gen_sim/task_engine/orchestration/test_scene_adapter.py b/tests/gen_sim/task_engine/orchestration/test_scene_adapter.py new file mode 100644 index 000000000..2ae3b749f --- /dev/null +++ b/tests/gen_sim/task_engine/orchestration/test_scene_adapter.py @@ -0,0 +1,780 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +import json +from pathlib import Path + +import pytest + +import embodichain.gen_sim.task_engine.orchestration.scene_adapter as scene_adapter_module +from embodichain.gen_sim.task_engine.contracts import ( + SCENE_REQUEST_SCHEMA, + SUCCESS_SPEC_SCHEMA, + TASK_CANDIDATE_SET_SCHEMA, + TASK_DRAFT_SCHEMA, + canonical_hash, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( + SceneAdapter, + SceneAdapterProtocolError, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import ( + SceneSourceRef, + fingerprint_scene_source, + verify_scene_source_fingerprint, +) +from embodichain.gen_sim.task_engine.agent import ( + derive_scene_request, + derive_success_spec, +) + +_UPRIGHT_CAN_INSTRUCTION = "test-instruction" + + +@pytest.fixture +def scene_export(tmp_path: Path) -> Path: + export = tmp_path / "scene_export" + assets = export / "meshes" + assets.mkdir(parents=True) + for name in ("table", "red_can", "blue_can"): + (assets / f"{name}.glb").write_bytes(f"mesh:{name}".encode()) + config = { + "format": "embodichain.scene-export/v1", + "scene_id": "2026-03-18T10:20:30Z", + "background": [ + { + "uid": "table", + "name": "table", + "description": "A work table.", + "category": "table", + "affordances": ["support_surface"], + "shape": {"shape_type": "Mesh", "fpath": "meshes/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": f"{color}_can", + "name": f"{color} can", + "description": f"A {color} soda can.", + "category": "can", + "attributes": { + "color": color, + "geometry": {"position": [1.0, 2.0, 3.0]}, + }, + "affordances": ["graspable", "orientable", "placeable"], + "initial_state": {"orientation": "fallen"}, + "shape": { + "shape_type": "Mesh", + "fpath": f"meshes/{color}_can.glb", + }, + "init_pos": [0.0, offset, 0.7], + "init_rot": [0.0, 0.0, 90.0], + "body_scale": [1.0, 1.0, 1.0], + } + for color, offset in (("red", 0.2), ("blue", -0.2)) + ], + } + (export / "scene_config.json").write_text(json.dumps(config), encoding="utf-8") + return export + + +def _legacy_gym_project(tmp_path: Path, filename: str) -> Path: + project = tmp_path / filename.removesuffix(".json") + assets = project / "assets" + assets.mkdir(parents=True) + for name in ("table.glb", "red_can.glb", "cabinet.urdf"): + (assets / name).write_bytes(f"asset:{name}".encode()) + config = { + "background": [ + { + "uid": "table_0", + "name": "table", + "description": "A work table.", + "category": "table", + "shape": {"shape_type": "Mesh", "fpath": "assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "red_can_0", + "name": "red can", + "description": "A red soda can.", + "category": "can", + "affordances": ["graspable", "orientable", "placeable"], + "initial_state": {"orientation": "fallen"}, + "shape": { + "shape_type": "Mesh", + "fpath": "assets/red_can.glb", + }, + "init_pos": [0.0, 0.2, 0.7], + "init_rot": [0.0, 0.0, 90.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "articulation": [ + { + "uid": "cabinet_0", + "name": "cabinet", + "description": "A fixed articulated cabinet.", + "category": "cabinet", + "fpath": "assets/cabinet.urdf", + "init_pos": [0.4, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + } + (project / filename).write_text(json.dumps(config), encoding="utf-8") + return project + + +def _selector(reference: str) -> dict: + return { + "kind": "scene_ref", + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _none_selector() -> dict: + return { + "kind": "none", + "step_id": "", + "reference": "", + "quantifier": "one", + "count": 0, + } + + +def _candidate(candidate_id: str, reference: str, *, votes: int = 1) -> dict: + step = { + "id": "upright", + "task_type": "E2", + "object": _selector(reference), + "target": _none_selector(), + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "steps": [step], + } + return { + "candidate_id": candidate_id, + "draft": draft, + "scene_request": { + "schema_version": SCENE_REQUEST_SCHEMA, + "task_id": "upright_can", + "references": [ + { + "reference_id": "upright.object", + "step_id": "upright", + "role": "object", + "reference": reference, + "quantifier": "one", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ], + }, + "success_spec": { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": "upright_can", + "op": "all", + "terms": [{"step_id": "upright", "type": "object_upright"}], + }, + "semantic_hash": canonical_hash([step]), + "vote_count": votes, + "attempts": 1, + "normalizations": [], + } + + +def _candidate_set(candidates: list[dict]) -> dict: + return { + "schema_version": TASK_CANDIDATE_SET_SCHEMA, + "task_id": "upright_can", + "instruction": _UPRIGHT_CAN_INSTRUCTION, + "candidates": candidates, + "requested_candidate_count": sum(item["vote_count"] for item in candidates), + "valid_response_count": sum(item["vote_count"] for item in candidates), + "errors": [], + } + + +def _placement_candidate(candidate_id: str = "place") -> dict: + candidate = _candidate(candidate_id, "red can") + step = candidate["draft"]["steps"][0] + step.update( + { + "task_type": "E1", + "target": _selector("table"), + "relation": "on", + "orientation_goal": "preserve", + } + ) + candidate["scene_request"]["references"] = [ + { + "reference_id": "upright.object", + "step_id": "upright", + "role": "object", + "reference": "red can", + "quantifier": "one", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "placeable"], + "initial_state": {}, + "attributes": {}, + }, + { + "reference_id": "upright.target", + "step_id": "upright", + "role": "target", + "reference": "table", + "quantifier": "one", + "count": 0, + "source_structure": "physical_entity", + "affordances": [], + "initial_state": {}, + "attributes": {}, + }, + ] + candidate["success_spec"]["terms"] = [ + {"step_id": "upright", "type": "semantic_goal"} + ] + candidate["semantic_hash"] = canonical_hash([step]) + return candidate + + +def _grounder(**kwargs) -> dict: + prompt = kwargs["prompt"] + uid = "blue_can" if '"reference": "blue can"' in prompt else "red_can" + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": [uid], + "confidence": 0.95, + } + ] + } + + +def test_scene_source_fingerprint_reads_without_copying(scene_export: Path) -> None: + before = sorted(path.relative_to(scene_export) for path in scene_export.rglob("*")) + fingerprint = fingerprint_scene_source(SceneSourceRef(scene_export)) + after = sorted(path.relative_to(scene_export) for path in scene_export.rglob("*")) + + assert fingerprint.config_path == scene_export / "scene_config.json" + assert len(fingerprint.config_sha256) == 64 + assert len(fingerprint.asset_sha256) == 3 + assert after == before + + +@pytest.mark.parametrize("filename", ["gym_config.json", "gym_config_merged.json"]) +def test_scene_adapter_supports_legacy_gym_configs( + tmp_path: Path, + filename: str, +) -> None: + project = _legacy_gym_project(tmp_path, filename) + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([_candidate("legacy", "red can")]), + project, + ) + + assert result.selected_candidate_id == "legacy" + assert result.static_scene_manifest["source_format"] == "legacy_gym_config" + assert any( + item["role"] == "articulation" + for item in result.static_scene_manifest["objects"] + ) + assert ( + result.prepared_scene.articulations[0]["fpath"] + == (project / "assets" / "cabinet.urdf").resolve().as_posix() + ) + + +def test_scene_source_fingerprint_covers_articulation_fpath(tmp_path: Path) -> None: + project = _legacy_gym_project(tmp_path, "gym_config.json") + original = fingerprint_scene_source(project) + articulation_path = project / "assets" / "cabinet.urdf" + + articulation_path.write_bytes(b"changed articulation") + changed = fingerprint_scene_source(project) + + assert articulation_path.resolve().as_posix() in original.asset_sha256 + assert changed.asset_sha256 != original.asset_sha256 + with pytest.raises(RuntimeError, match="changed after Task Engine preparation"): + verify_scene_source_fingerprint(original.to_dict()) + + +def test_scene_adapter_selects_bindable_majority_and_redacts_manifest( + scene_export: Path, +) -> None: + red = _candidate("red-majority", "red can", votes=2) + blue = _candidate("blue-minority", "blue can") + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([red, blue]), + scene_export, + ) + + hierarchy_by_uid = { + node["uid"]: node for node in result.conservative_scene_graph["nodes"] + } + assert hierarchy_by_uid["red_can"]["parent_uid"] == "unknown" + assert hierarchy_by_uid["red_can"]["parent_relation"] == "unknown" + + assert result.binding_report["status"] == "bound" + assert result.binding_report["candidates"][0]["status"] == "resolved" + assert result.selected_candidate_id == "red-majority" + assert result.reference_bindings == {"upright.object": ["red_can"]} + assert result.role_bindings["role_bindings"] == {} + red_manifest = next( + item for item in result.scene_manifest["objects"] if item["uid"] == "red_can" + ) + assert "position" not in json.dumps(red_manifest) + assert ( + result.prepared_scene.source_config_path == scene_export / "scene_config.json" + ) + assert result.static_scene_manifest is not None + static_by_uid = { + item["uid"]: item for item in result.static_scene_manifest["objects"] + } + assert static_by_uid["red_can"]["physics"]["body_type"] == "dynamic" + assert static_by_uid["red_can"]["geometry"]["asset_sha256"] + + +def test_scene_adapter_returns_report_for_business_level_non_binding( + scene_export: Path, +) -> None: + candidate = _candidate("missing", "green can") + + def not_found(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "not_found", + "uids": [], + "confidence": 0.0, + } + ] + } + + result = SceneAdapter(grounding_caller=not_found).adapt( + _candidate_set([candidate]), + scene_export, + ) + + assert result.selected_candidate is None + assert result.role_bindings is None + assert result.binding_report["status"] == "unsatisfied" + assert ( + result.binding_report["candidates"][0]["references"][0]["status"] == "not_found" + ) + + +def test_semantic_blueprint_selection_forces_ranked_low_confidence_uid() -> None: + candidate = _candidate("likely", "the can") + scene_objects = [ + { + "uid": "table", + "role": "table", + "name": "table", + "description": "A work table.", + "category": "table", + "init_pos": [0.0, 0.0, 0.0], + "affordances": ["support_surface"], + "initial_state": {}, + "attributes": {}, + }, + *[ + { + "uid": f"{color}_can", + "role": "rigid_object", + "name": f"{color} can", + "description": f"A {color} can.", + "category": "can", + "init_pos": [0.0, offset, 0.1], + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {"color": color}, + } + for color, offset in (("red", -0.1), ("blue", 0.1)) + ], + ] + + def ambiguous(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "ambiguous", + "uids": ["red_can", "blue_can"], + "confidence": 0.2, + } + ] + } + + result = SceneAdapter(grounding_caller=ambiguous).select_objects( + _candidate_set([candidate]), + scene_objects, + force_most_likely=True, + ) + + assert result.selected_candidate_id == "likely" + assert result.role_bindings["reference_bindings"] == {"upright.object": ["red_can"]} + reference = result.binding_report["candidates"][0]["references"][0] + assert reference["confidence"] == 0.2 + assert reference["candidate_uids"] == ["red_can", "blue_can"] + assert reference["selected_uids"] == ["red_can"] + assert reference["reasons"] == [ + "Forced the highest-ranked structurally compatible UID from an " + "ambiguous low-confidence response." + ] + + +def test_scene_adapter_uses_unique_bindable_then_injected_adjudication( + scene_export: Path, +) -> None: + red = _candidate("red", "red can") + blue = _candidate("blue", "blue can") + + def one_missing(**kwargs): + if '"reference": "red can"' in kwargs["prompt"]: + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "not_found", + "uids": [], + "confidence": 0.0, + } + ] + } + return _grounder(**kwargs) + + unique = SceneAdapter(grounding_caller=one_missing).adapt( + _candidate_set([red, blue]), + scene_export, + ) + assert unique.selected_candidate_id == "blue" + assert unique.binding_report["selection_reason"] == "unique_bindable" + + ambiguous = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([red, blue]), + scene_export, + ) + assert ambiguous.binding_report["status"] == "ambiguous" + + adjudicated = SceneAdapter( + grounding_caller=_grounder, + adjudicator=lambda **_kwargs: {"candidate_id": "blue"}, + ).adapt(_candidate_set([red, blue]), scene_export) + assert adjudicated.selected_candidate_id == "blue" + assert adjudicated.binding_report["selection_reason"] == "adjudicated_bindable" + + +def test_scene_adapter_runs_one_default_structured_adjudication( + scene_export: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + adjudications = 0 + + def caller(**kwargs): + nonlocal adjudications + if kwargs["schema"]["title"] == "ActionEngineTaskAdjudication": + adjudications += 1 + return {"candidate_id": "blue"} + return _grounder(**kwargs) + + monkeypatch.setattr( + scene_adapter_module, "_default_grounding_caller", lambda: caller + ) + result = SceneAdapter().adapt( + _candidate_set([_candidate("red", "red can"), _candidate("blue", "blue can")]), + scene_export, + ) + + assert result.selected_candidate_id == "blue" + assert result.binding_report["selection_reason"] == "adjudicated_bindable" + assert adjudications == 1 + + +def test_scene_adapter_accepts_direct_source_and_rejects_bad_protocol( + scene_export: Path, +) -> None: + candidate = _candidate("red", "red can") + direct = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([candidate]), + scene_export, + ) + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([candidate]), + SceneSourceRef(scene_export), + ) + assert result.scene_manifest == direct.scene_manifest + assert result.role_bindings == direct.role_bindings + + with pytest.raises(SceneAdapterProtocolError, match="unsupported fields"): + SceneAdapter( + grounding_caller=lambda **_kwargs: { + "bindings": [ + { + "reference_id": "upright.object", + "status": "not_found", + "uids": [], + "confidence": 0.0, + "invented": True, + } + ] + } + ).adapt(_candidate_set([candidate]), scene_export) + + +def test_explicit_scene_semantic_conflict_is_incompatible( + scene_export: Path, +) -> None: + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["initial_state"]["orientation"] = "upright" + config_path.write_text(json.dumps(config), encoding="utf-8") + + result = SceneAdapter(grounding_caller=_grounder).adapt( + _candidate_set([_candidate("red", "red can")]), + scene_export, + ) + reference = result.binding_report["candidates"][0]["references"][0] + assert result.binding_report["status"] == "unsatisfied" + assert reference["status"] == "incompatible" + assert result.binding_report["candidates"][0]["status"] == "incompatible" + assert "state 'orientation' conflicts" in reference["reasons"][0] + + +def test_scene_adapter_accepts_passive_support_target_and_rejects_self_reference( + scene_export: Path, +) -> None: + candidate = _placement_candidate() + + def place_on_table(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + }, + { + "reference_id": "upright.target", + "status": "resolved", + "uids": ["table"], + "confidence": 0.95, + }, + ] + } + + bound = SceneAdapter(grounding_caller=place_on_table).adapt( + _candidate_set([candidate]), scene_export + ) + assert bound.binding_report["status"] == "bound" + assert bound.reference_bindings["upright.target"] == ["table"] + + def self_reference(**_kwargs): + response = place_on_table() + response["bindings"][1]["uids"] = ["red_can"] + return response + + incompatible = SceneAdapter(grounding_caller=self_reference).adapt( + _candidate_set([candidate]), scene_export + ) + assert incompatible.binding_report["status"] == "unsatisfied" + assert incompatible.binding_report["candidates"][0]["status"] == "incompatible" + + +def test_scene_adapter_enforces_count_cardinality_in_audit( + scene_export: Path, +) -> None: + candidate = _candidate("two", "cans") + selector = candidate["draft"]["steps"][0]["object"] + selector.update(quantifier="count", count=2) + request = candidate["scene_request"]["references"][0] + request.update(reference="cans", quantifier="count", count=2) + candidate["semantic_hash"] = canonical_hash(candidate["draft"]["steps"]) + + def one_only(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + } + ] + } + + result = SceneAdapter(grounding_caller=one_only).adapt( + _candidate_set([candidate]), scene_export + ) + + assert result.binding_report["status"] == "unsatisfied" + audit = result.binding_report["candidates"][0] + assert audit["status"] == "incompatible" + assert "requires exactly 2 UIDs" in audit["references"][0]["reasons"][0] + + +def test_scene_adapter_binds_all_matching_uids( + scene_export: Path, +) -> None: + candidate = _candidate("all", "all cans") + candidate["draft"]["steps"][0]["object"].update(quantifier="all") + candidate["scene_request"] = derive_scene_request(candidate["draft"]) + candidate["semantic_hash"] = canonical_hash(candidate["draft"]["steps"]) + + def all_cans(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can", "blue_can"], + "confidence": 0.95, + } + ] + } + + result = SceneAdapter(grounding_caller=all_cans).adapt( + _candidate_set([candidate]), scene_export + ) + + assert result.binding_report["status"] == "bound" + assert result.reference_bindings == {"upright.object": ["red_can", "blue_can"]} + assert result.candidate_bindings[candidate["candidate_id"]][ + "reference_bindings" + ] == {"upright.object": ["red_can", "blue_can"]} + + +def test_scene_adapter_rejects_step_result_object_matching_same_step_target( + scene_export: Path, +) -> None: + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["affordances"].append("support_surface") + config_path.write_text(json.dumps(config), encoding="utf-8") + + candidate = _candidate("self-reference", "red can") + second = deepcopy(candidate["draft"]["steps"][0]) + second.update( + { + "id": "place_again", + "task_type": "E1", + "object": { + "kind": "step_result", + "step_id": "upright", + "reference": "", + "quantifier": "one", + "count": 0, + }, + "target": _selector("red can"), + "relation": "on", + "orientation_goal": "preserve", + "depends_on": ["upright"], + } + ) + candidate["draft"]["steps"].append(second) + candidate["scene_request"] = derive_scene_request(candidate["draft"]) + candidate["success_spec"] = derive_success_spec(candidate["draft"]) + candidate["semantic_hash"] = canonical_hash(candidate["draft"]["steps"]) + + def same_uid(**_kwargs): + return { + "bindings": [ + { + "reference_id": "upright.object", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + }, + { + "reference_id": "place_again.target", + "status": "resolved", + "uids": ["red_can"], + "confidence": 0.95, + }, + ] + } + + result = SceneAdapter(grounding_caller=same_uid).adapt( + _candidate_set([candidate]), scene_export + ) + + assert result.binding_report["status"] == "unsatisfied" + target_audit = result.binding_report["candidates"][0]["references"][1] + assert target_audit["status"] == "incompatible" + assert "same UID as object and target" in target_audit["reasons"][0] + + +def test_scene_source_fingerprint_covers_assets_and_config(scene_export: Path) -> None: + original = fingerprint_scene_source(scene_export) + + asset_path = scene_export / "meshes" / "red_can.glb" + asset_path.write_bytes(b"changed asset") + changed_asset = fingerprint_scene_source(scene_export) + assert changed_asset.asset_sha256 != original.asset_sha256 + + config_path = scene_export / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["body_scale"] = [1.1, 1.0, 1.0] + config["rigid_object"][0]["physics"] = {"mass": 0.25} + config_path.write_text(json.dumps(config), encoding="utf-8") + changed_config = fingerprint_scene_source(scene_export) + assert changed_config.config_sha256 != changed_asset.config_sha256 + + +def test_scene_source_verification_rejects_later_mutation(scene_export: Path) -> None: + expected = fingerprint_scene_source(scene_export).to_dict() + (scene_export / "meshes" / "red_can.glb").write_bytes(b"changed later") + + with pytest.raises(RuntimeError, match="changed after Task Engine preparation"): + verify_scene_source_fingerprint(expected) diff --git a/tests/gen_sim/task_engine/scene/__init__.py b/tests/gen_sim/task_engine/scene/__init__.py new file mode 100644 index 000000000..96ca57709 --- /dev/null +++ b/tests/gen_sim/task_engine/scene/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for Task Engine scene adaptation boundaries.""" + +from __future__ import annotations diff --git a/tests/gen_sim/task_engine/scene/test_final_inspection.py b/tests/gen_sim/task_engine/scene/test_final_inspection.py new file mode 100644 index 000000000..937a075d9 --- /dev/null +++ b/tests/gen_sim/task_engine/scene/test_final_inspection.py @@ -0,0 +1,118 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import trimesh + +from embodichain.gen_sim.task_engine.orchestration.scene_source import ( + scene_revision_id, +) +from embodichain.gen_sim.task_engine.scene.final_inspection import ( + inspect_final_scene, +) + + +def _scene_export(root: Path, *, scene_id: str, can_rotation: list[float]) -> Path: + export = root / "scene_export" + assets = export / "mesh_assets" + assets.mkdir(parents=True) + trimesh.creation.box(extents=[1.0, 0.1, 1.0]).export( + assets / "table.glb", file_type="glb" + ) + can = trimesh.creation.cylinder(radius=0.04, height=0.2) + can.apply_transform( + trimesh.transformations.rotation_matrix(np.pi / 2.0, [1.0, 0.0, 0.0]) + ) + can.export(assets / "can.glb", file_type="glb") + config = { + "format": "embodichain.scene-export/v1", + "scene_id": scene_id, + "background": [ + { + "uid": "table", + "name": "table", + "description": "A support table.", + "category": "table", + "shape": {"shape_type": "Mesh", "fpath": "mesh_assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + } + ], + "rigid_object": [ + { + "uid": "can", + "name": "red can", + "description": "A red can.", + "category": "can", + "shape": {"shape_type": "Mesh", "fpath": "mesh_assets/can.glb"}, + "init_pos": [0.0, 0.0, 0.15], + "init_rot": can_rotation, + "body_scale": [1.0, 1.0, 1.0], + } + ], + } + path = export / "scene_config.json" + path.write_text(json.dumps(config), encoding="utf-8") + return path + + +def test_scene_revision_id_ignores_exporter_timestamp_and_location( + tmp_path: Path, +) -> None: + first = _scene_export( + tmp_path / "first", scene_id="scene-100", can_rotation=[0, 0, 0] + ) + second = _scene_export( + tmp_path / "second", scene_id="scene-200", can_rotation=[0, 0, 0] + ) + + assert scene_revision_id(first) == scene_revision_id(second) + + value = json.loads(second.read_text(encoding="utf-8")) + value["rigid_object"][0]["init_pos"][0] = 0.25 + second.write_text(json.dumps(value), encoding="utf-8") + assert scene_revision_id(first) != scene_revision_id(second) + + +def test_final_inspection_recomputes_support_and_orientation(tmp_path: Path) -> None: + source = _scene_export( + tmp_path / "standing", scene_id="scene", can_rotation=[0.0, 0.0, 0.0] + ) + + inspection = inspect_final_scene(source, revision_id=scene_revision_id(source)) + + can = next(item for item in inspection["objects"] if item["uid"] == "can") + assert can["orientation"] == "standing" + assert can["support"]["parent_uid"] == "table" + assert can["support"]["relation"] == "on" + assert can["support"]["xy_overlap_ratio"] > 0.9 + + +def test_final_inspection_detects_lying_rotation(tmp_path: Path) -> None: + source = _scene_export( + tmp_path / "lying", scene_id="scene", can_rotation=[90.0, 0.0, 0.0] + ) + + inspection = inspect_final_scene(source, revision_id=scene_revision_id(source)) + + can = next(item for item in inspection["objects"] if item["uid"] == "can") + assert can["orientation"] == "lying" diff --git a/tests/gen_sim/task_engine/scene/test_scene_boundary.py b/tests/gen_sim/task_engine/scene/test_scene_boundary.py new file mode 100644 index 000000000..69b81af89 --- /dev/null +++ b/tests/gen_sim/task_engine/scene/test_scene_boundary.py @@ -0,0 +1,582 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.task_engine.scene import ( + FeasibilityBroker, + SceneEngineV1Adapter, +) +from embodichain.gen_sim.task_engine.agent import derive_scene_request +from embodichain.gen_sim.task_engine.contracts import TASK_DRAFT_SCHEMA + + +def _prepared_scene(tmp_path: Path) -> SimpleNamespace: + table = { + "uid": "table", + "source_uid": "table_0", + "role": "background", + "name": "table", + "description": "A support table.", + "category": "table", + "color": "brown", + "shape": {"shape_type": "Mesh", "fpath": "/assets/table.glb"}, + "init_pos": [0.0, 0.0, 0.0], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "attributes": {}, + "initial_state": {}, + "affordances": [], + } + can = { + "uid": "red_can", + "source_uid": "red_can_0", + "role": "rigid_object", + "name": "red can", + "description": "A fallen red can.", + "category": "can", + "color": "red", + "shape": {"shape_type": "Mesh", "fpath": "/assets/can.glb"}, + "init_pos": [0.1, 0.0, 0.7], + "init_rot": [90.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "attributes": {}, + "initial_state": {"orientation": "fallen"}, + "affordances": ["graspable", "orientable", "placeable"], + } + runtime_table = { + "uid": "table", + "shape": table["shape"], + "attrs": {"mass": 10.0}, + "body_type": "kinematic", + } + runtime_can = { + "uid": "red_can", + "shape": can["shape"], + "attrs": {"mass": 0.1}, + "body_type": "dynamic", + } + return SimpleNamespace( + source_config_path=tmp_path / "scene_config.json", + planner_objects=(table, can), + background=(runtime_table,), + rigid_objects=(runtime_can,), + articulations=(), + asset_hashes={"table": "a" * 64, "red_can": "b" * 64}, + ) + + +def _candidate(task_type: str, affordances: list[str]) -> dict: + return { + "candidate_id": "candidate_01", + "draft": { + "task_id": "task", + "steps": [{"id": "step_01", "task_type": task_type}], + }, + "scene_request": { + "references": [ + { + "reference_id": "step_01.object", + "role": "object", + "source_structure": "rigid_object", + "affordances": affordances, + "initial_state": ( + {"orientation": "fallen"} if task_type == "E2" else {} + ), + "attributes": {}, + } + ] + }, + } + + +def _catalog(*, pour_available: bool = False) -> dict[str, dict]: + return { + name: {"runtime_available": True, "unavailable_reason": None} + for name in ("PickUp", "MoveHeldObject", "Place", "TurnKnob") + } | { + "Pour": { + "runtime_available": pour_available, + "unavailable_reason": None if pour_available else "Pour is planning-only.", + } + } + + +def _selector(kind: str, *, reference: str = "") -> dict[str, object]: + return { + "kind": kind, + "step_id": "", + "reference": reference, + "quantifier": "one", + "count": 0, + } + + +def _relation_candidate(relation: str) -> dict: + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "place_relative", + "instruction": "place the can relative to the target", + "steps": [ + { + "id": "step_01", + "task_type": "E1", + "object": _selector("scene_ref", reference="red can"), + "target": _selector("scene_ref", reference="target"), + "relation": relation, + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "preserve", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + ], + } + return { + "candidate_id": "candidate_01", + "draft": draft, + "scene_request": derive_scene_request(draft), + } + + +def _manifest_with_target_kinds(tmp_path: Path) -> dict: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + by_uid = {item["uid"]: item for item in manifest["objects"]} + by_uid["red_can"]["affordances"].append( + { + "type": "container", + "status": "declared", + "confidence": None, + "source": "test", + "link_uid": "", + "frame": {}, + "parameters": {}, + } + ) + articulation = deepcopy(by_uid["red_can"]) + articulation.update( + uid="cabinet", + source_uid="cabinet_0", + role="articulation", + name="cabinet", + category="cabinet", + physics={}, + articulation={"runtime_uid": "cabinet"}, + affordances=[], + ) + manifest["objects"].append(articulation) + return manifest + + +def test_scene_engine_v1_adapter_preserves_static_execution_evidence( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="embodichain.scene-export/v1", + robot_profile="dual_franka", + ) + + by_uid = {item["uid"]: item for item in manifest["objects"]} + assert manifest["adapter_capabilities"]["task_conditioned_generation"] is False + assert by_uid["red_can"]["geometry"]["asset_sha256"] == "b" * 64 + assert by_uid["red_can"]["physics"]["body_type"] == "dynamic" + assert {item["type"] for item in by_uid["table"]["affordances"]} == { + "support_surface" + } + assert ( + next( + item + for item in by_uid["red_can"]["affordances"] + if item["type"] == "graspable" + )["status"] + == "declared" + ) + + +def test_e2_feasibility_requires_runtime_probe_for_geometry(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E2", ["graspable", "orientable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + assert report["status"] == "runtime_probe" + assert report["remediation_class"] == "none" + assert report["blockers"] == [] + assert report["summary"]["proven"] > 0 + assert report["summary"]["runtime_probe"] > 0 + + +def test_planning_only_action_is_reported_as_contradicted(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E3", ["graspable", "pourable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E3": ("Pour",)}, + ) + + assert report["status"] == "contradicted" + assert report["remediation_class"] == "action_capability" + assert any("planning-only" in blocker for blocker in report["blockers"]) + + +def test_e3_requires_runtime_content_observation_before_execution( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E3", ["graspable", "pourable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(pour_available=True), + task_actions={"E3": ("PickUp", "MoveHeldObject", "Pour")}, + ) + + assert report["status"] == "contradicted" + assert report["remediation_class"] == "terminal" + assert any( + check["kind"] == "content_observation" and check["status"] == "contradicted" + for check in report["checks"] + ) + assert any( + "baked into one source mesh" in blocker for blocker in report["blockers"] + ) + + +def test_e8_requires_explicit_setting_to_angle_mapping(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E8", ["turnable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E8": ("TurnKnob",)}, + ) + + assert any( + check["kind"] == "setting_mapping" and check["status"] == "contradicted" + for check in report["checks"] + ) + assert any("setting_values" in blocker for blocker in report["blockers"]) + + +def test_final_orientation_conflict_is_scene_remediable(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + can = next(item for item in manifest["objects"] if item["uid"] == "red_can") + can["initial_state"]["orientation"] = "upright" + + report = FeasibilityBroker().assess( + _candidate("E2", ["graspable", "orientable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + assert report["status"] == "contradicted" + assert report["remediation_class"] == "scene_remediable" + + +def test_missing_affordance_remains_unknown_instead_of_becoming_supported( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E1", ["graspable", "liquid_safe"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + assert report["status"] == "unknown" + assert any( + check["status"] == "unknown" and "liquid_safe" in check["reason"] + for check in report["checks"] + ) + + +def test_physical_object_can_be_a_runtime_support_without_support_affordance( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + candidate = _candidate("E1", ["graspable", "placeable"]) + candidate["draft"]["steps"][0].update( + target={"kind": "scene_ref"}, + relation="on", + ) + candidate["scene_request"]["references"].append( + { + "reference_id": "step_01.target", + "role": "target", + "source_structure": "physical_entity", + "affordances": [], + "initial_state": {}, + "attributes": {}, + } + ) + + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": ["red_can"], + }, + manifest, + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" and check["subject"] == "step_01.target:red_can" + ) + assert structure["status"] == "proven" + support_probe = next( + check for check in report["checks"] if check["kind"] == "placement_support" + ) + assert support_probe["status"] == "runtime_probe" + assert support_probe["evidence"]["runtime_obligations"] == [ + "placement_candidates", + "object_supported_by", + "stable_for", + "final_support_revalidation", + ] + assert report["blockers"] == [] + + +@pytest.mark.parametrize( + ("relation", "target_uid", "expected_structure", "expected_status"), + [ + ("on", "red_can", "physical_entity", "proven"), + ("on", "table", "physical_entity", "proven"), + ("on", "cabinet", "physical_entity", "contradicted"), + ("inside", "red_can", "rigid_object", "proven"), + ("inside", "table", "rigid_object", "contradicted"), + ("inside", "cabinet", "rigid_object", "contradicted"), + ("behind", "red_can", "spatial_reference", "proven"), + ("behind", "table", "spatial_reference", "proven"), + ("behind", "cabinet", "spatial_reference", "runtime_probe"), + ("front_of", "red_can", "spatial_reference", "proven"), + ("front_of", "table", "spatial_reference", "proven"), + ("front_of", "cabinet", "spatial_reference", "runtime_probe"), + ("left_of", "red_can", "spatial_reference", "proven"), + ("left_of", "table", "spatial_reference", "proven"), + ("left_of", "cabinet", "spatial_reference", "runtime_probe"), + ("right_of", "red_can", "spatial_reference", "proven"), + ("right_of", "table", "spatial_reference", "proven"), + ("right_of", "cabinet", "spatial_reference", "runtime_probe"), + ], +) +def test_relation_target_structure_matrix_uses_capability_semantics( + tmp_path: Path, + relation: str, + target_uid: str, + expected_structure: str, + expected_status: str, +) -> None: + candidate = _relation_candidate(relation) + target_request = next( + item + for item in candidate["scene_request"]["references"] + if item["role"] == "target" + ) + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": [target_uid], + }, + _manifest_with_target_kinds(tmp_path), + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" + and check["subject"] == f"step_01.target:{target_uid}" + ) + assert target_request["source_structure"] == expected_structure + assert structure["status"] == expected_status + + +def test_legacy_scene_entity_target_is_treated_as_an_abstract_structure( + tmp_path: Path, +) -> None: + candidate = _relation_candidate("behind") + target_request = next( + item + for item in candidate["scene_request"]["references"] + if item["role"] == "target" + ) + target_request["source_structure"] = "scene_entity" + + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": ["red_can"], + }, + _manifest_with_target_kinds(tmp_path), + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" and check["subject"] == "step_01.target:red_can" + ) + assert structure["status"] == "proven" + assert report["blockers"] == [] + + +def test_unknown_structure_contract_is_not_a_scene_contradiction( + tmp_path: Path, +) -> None: + candidate = _relation_candidate("behind") + target_request = next( + item + for item in candidate["scene_request"]["references"] + if item["role"] == "target" + ) + target_request["source_structure"] = "future_spatial_capability" + + report = FeasibilityBroker().assess( + candidate, + { + "step_01.object": ["red_can"], + "step_01.target": ["red_can"], + }, + _manifest_with_target_kinds(tmp_path), + capability_catalog=_catalog(), + task_actions={"E1": ("PickUp", "MoveHeldObject", "Place")}, + ) + + structure = next( + check + for check in report["checks"] + if check["kind"] == "structure" and check["subject"] == "step_01.target:red_can" + ) + assert structure["status"] == "unknown" + assert not any("future_spatial_capability" in item for item in report["blockers"]) + + +def test_required_arm_side_requires_the_live_robot_frame( + tmp_path: Path, +) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + red_can = next(item for item in manifest["objects"] if item["uid"] == "red_can") + red_can["initial_pose"]["position"][1] = -0.20 + candidate = _candidate("E2", ["graspable", "orientable"]) + candidate["draft"]["steps"][0]["required_arm"] = "right_arm" + + report = FeasibilityBroker().assess( + candidate, + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + probe = next( + check for check in report["checks"] if check["kind"] == "arm_layout_risk" + ) + assert probe["status"] == "runtime_probe" + assert probe["evidence"]["arm_side_frame"] == "live_robot" + assert probe["evidence"]["mismatch_risk"] is None + assert "expected_arm" not in probe["evidence"] + assert probe["evidence"]["geometry_certificate"] is False + assert report["blockers"] == [] + + +def test_workspace_report_covers_complete_task_phases(tmp_path: Path) -> None: + manifest = SceneEngineV1Adapter().adapt_prepared_scene( + _prepared_scene(tmp_path), + source_format="test", + robot_profile="dual_franka", + ) + report = FeasibilityBroker().assess( + _candidate("E2", ["graspable", "orientable"]), + {"step_01.object": ["red_can"]}, + manifest, + capability_catalog=_catalog(), + task_actions={"E2": ("PickUp", "MoveHeldObject", "Place")}, + ) + + workflow = next( + check for check in report["checks"] if check["kind"] == "task_workspace" + ) + phases = {item["phase"] for item in workflow["evidence"]["phases"]} + assert phases == {"pickup", "safety_clearance"} + assert workflow["status"] == "runtime_probe" diff --git a/tests/gen_sim/task_engine/test_agent.py b/tests/gen_sim/task_engine/test_agent.py new file mode 100644 index 000000000..cb6ab7cb2 --- /dev/null +++ b/tests/gen_sim/task_engine/test_agent.py @@ -0,0 +1,271 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +import threading +from time import sleep + +import pytest + +from embodichain.gen_sim.task_engine.contracts import ( + SUCCESS_SPEC_SCHEMA, + TASK_DRAFT_SCHEMA, + validate_success_spec, + validate_task_candidate, + validate_task_draft, +) +from embodichain.gen_sim.task_engine.agent import ( + TaskAgent, + TaskGenerationError, + derive_scene_request, + derive_success_spec, +) +from embodichain.gen_sim.task_engine.interpretation import InstructionDraftResult +from embodichain.gen_sim.task_engine.orchestration.coordinator import ( + lower_task_candidate, +) + +_TEST_INSTRUCTION = "test-instruction" + + +def _selector(kind="none", *, step_id="", reference="", quantifier="one", count=0): + return { + "kind": kind, + "step_id": step_id, + "reference": reference, + "quantifier": quantifier, + "count": count, + } + + +def _step(step_id="orient", reference="purple can"): + return { + "id": step_id, + "task_type": "E2", + "object": _selector("scene_ref", reference=reference), + "target": _selector(), + "relation": "none", + "required_arm": "auto", + "transfer_arm": "none", + "receive_arm": "none", + "orientation_goal": "upright", + "target_state": "none", + "target_setting": 0, + "layout": "none", + "axis": "none", + "direction": "none", + "terminal_behavior": "none", + "depends_on": [], + } + + +def _result(step): + return InstructionDraftResult( + intent={"steps": [deepcopy(step)]}, + model="injected_caller", + attempts=1, + latency_seconds=0.01, + normalizations=(), + ) + + +def test_task_agent_generates_concurrently_deduplicates_and_counts_votes(): + barrier = threading.Barrier(3) + lock = threading.Lock() + assigned = 0 + + def interpreter(_instruction, **_kwargs): + nonlocal assigned + with lock: + index = assigned + assigned += 1 + barrier.wait(timeout=2) + sleep(0.01) + if index < 2: + return _result(_step(step_id=f"arbitrary_{index}")) + return _result(_step(step_id="different", reference="orange can")) + + result = TaskAgent(interpreter=interpreter).generate("task", _TEST_INSTRUCTION) + + assert result["requested_candidate_count"] == 3 + assert result["valid_response_count"] == 3 + assert len(result["candidates"]) == 2 + assert sorted(item["vote_count"] for item in result["candidates"]) == [1, 2] + assert {item["draft"]["steps"][0]["id"] for item in result["candidates"]} == { + "step_01" + } + + +def test_scene_request_and_success_are_deterministic_contract_derivations(): + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "upright", + "instruction": _TEST_INSTRUCTION, + "steps": [_step(reference="all cans")], + } + draft["steps"][0]["object"].update(quantifier="all") + + request = derive_scene_request(draft) + success = derive_success_spec(draft) + + assert request["references"] == [ + { + "reference_id": "orient.object", + "step_id": "orient", + "role": "object", + "reference": "all cans", + "quantifier": "all", + "count": 0, + "source_structure": "rigid_object", + "affordances": ["graspable", "orientable"], + "initial_state": {"orientation": "fallen"}, + "attributes": {}, + } + ] + assert success["terms"] == [{"step_id": "orient", "type": "object_upright"}] + + +@pytest.mark.parametrize( + ("relation", "expected_structure", "expected_affordances"), + [ + ("on", "physical_entity", []), + ("inside", "rigid_object", ["container"]), + ("behind", "spatial_reference", []), + ("front_of", "spatial_reference", []), + ("left_of", "spatial_reference", []), + ("right_of", "spatial_reference", []), + ], +) +def test_target_requirements_describe_capabilities_not_concrete_roles( + relation: str, + expected_structure: str, + expected_affordances: list[str], +) -> None: + step = _step(step_id="place", reference="green can") + step.update( + task_type="E1", + target=_selector("scene_ref", reference="red can"), + relation=relation, + orientation_goal="preserve", + ) + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "stack", + "instruction": _TEST_INSTRUCTION, + "steps": [step], + } + + request = derive_scene_request(draft) + + target = next( + reference + for reference in request["references"] + if reference["role"] == "target" + ) + assert target["source_structure"] == expected_structure + assert target["affordances"] == expected_affordances + + +def test_lower_task_candidate_expands_success_for_all_binding(): + def interpreter(_instruction, **_kwargs): + step = _step(reference="all cans") + step["object"].update(quantifier="all") + return _result(step) + + candidate = TaskAgent(interpreter=interpreter).generate( + "upright", _TEST_INSTRUCTION, candidate_count=1 + )["candidates"][0] + grounded = lower_task_candidate( + candidate, + {"step_01.object": ["can_a", "can_b"]}, + [ + {"uid": "can_a", "role": "rigid_object", "description": "A can."}, + {"uid": "can_b", "role": "rigid_object", "description": "A can."}, + ], + "dual_franka", + ) + + assert grounded.task_spec["level"] == "L2" + assert [term["type"] for term in grounded.task_spec["success"]["terms"]] == [ + "object_upright", + "object_upright", + ] + + +def test_draft_rejects_grounded_fields_and_task_agent_fails_closed(): + draft = { + "schema_version": TASK_DRAFT_SCHEMA, + "task_id": "bad", + "instruction": "bad", + "steps": [_step()], + } + draft["steps"][0]["object"]["uid"] = "scene_uid" + with pytest.raises(ValueError, match="forbidden|exactly fields"): + validate_task_draft(draft) + + def invalid(_instruction, **_kwargs): + raise ValueError("invalid draft after repair") + + with pytest.raises(TaskGenerationError, match="All Task Agent candidates"): + TaskAgent(interpreter=invalid).generate("bad", "bad") + + +def test_task_candidate_rejects_scene_constraints_not_derived_from_draft(): + candidate = TaskAgent( + interpreter=lambda *_args, **_kwargs: _result(_step()) + ).generate("upright", _TEST_INSTRUCTION, candidate_count=1)["candidates"][0] + candidate["scene_request"]["references"][0]["affordances"] = [] + + with pytest.raises(ValueError, match="derived exactly"): + validate_task_candidate(candidate) + + +def test_success_spec_rejects_types_outside_task_ontology(): + with pytest.raises(ValueError, match="must be one of"): + validate_success_spec( + { + "schema_version": SUCCESS_SPEC_SCHEMA, + "task_id": "bad_success", + "op": "all", + "terms": [{"step_id": "step_01", "type": "looks_good"}], + } + ) + + +def test_task_agent_isolates_invalid_interpreter_results(): + lock = threading.Lock() + calls = 0 + + def interpreter(_instruction, **_kwargs): + nonlocal calls + with lock: + index = calls + calls += 1 + if index == 0: + invalid = _step() + invalid["object"]["uid"] = "red_can" + return _result(invalid) + return _result(_step()) + + result = TaskAgent(interpreter=interpreter).generate( + "upright", _TEST_INSTRUCTION, candidate_count=2 + ) + + assert result["valid_response_count"] == 1 + assert len(result["errors"]) == 1 + assert len(result["candidates"]) == 1 diff --git a/tests/gen_sim/task_engine/test_interpretation.py b/tests/gen_sim/task_engine/test_interpretation.py new file mode 100644 index 000000000..009362e87 --- /dev/null +++ b/tests/gen_sim/task_engine/test_interpretation.py @@ -0,0 +1,96 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from embodichain.gen_sim.task_engine import interpretation as interpretation_module + + +def _write_dotenv(path: Path) -> None: + path.write_text( + "\n".join( + ( + "OPENAI_API_KEY=dotenv-key", + "OPENAI_BASE_URL=https://dotenv.example/v1", + "OPENAI_MODEL=dotenv-model", + ) + ), + encoding="utf-8", + ) + + +def _clear_process_provider(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ( + "OPENAI_API_KEY", + "OPENAI_BASE_URL", + "OPENAI_API_BASE", + "LLM_URL", + "TASK_ENGINE_LLM_MODEL", + "ACTION_ENGINE_LLM_MODEL", + "OPENAI_MODEL", + "LLM_MODEL", + ): + monkeypatch.delenv(name, raising=False) + + +def test_partial_process_transport_does_not_mix_with_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + _write_dotenv(env_path) + monkeypatch.setattr(interpretation_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr( + interpretation_module, + "_GEN_CONFIG_PATH", + tmp_path / "missing.json", + ) + _clear_process_provider(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "unrelated-process-key") + + settings = interpretation_module._load_llm_settings(model=None) + + assert settings["api_key"] == "dotenv-key" + assert settings["base_url"] == "https://dotenv.example/v1" + assert settings["model"] == "dotenv-model" + + +def test_complete_process_transport_overrides_dotenv( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + env_path = tmp_path / ".env" + _write_dotenv(env_path) + monkeypatch.setattr(interpretation_module, "_GEN_SIM_ENV_PATH", env_path) + monkeypatch.setattr( + interpretation_module, + "_GEN_CONFIG_PATH", + tmp_path / "missing.json", + ) + _clear_process_provider(monkeypatch) + monkeypatch.setenv("OPENAI_API_KEY", "process-key") + monkeypatch.setenv("OPENAI_BASE_URL", "https://process.example/v1/") + monkeypatch.setenv("TASK_ENGINE_LLM_MODEL", "process-model") + + settings = interpretation_module._load_llm_settings(model=None) + + assert settings["api_key"] == "process-key" + assert settings["base_url"] == "https://process.example/v1" + assert settings["model"] == "process-model" diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py new file mode 100644 index 000000000..23aaab655 --- /dev/null +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -0,0 +1,912 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +import json +from pathlib import Path +import sys +from types import SimpleNamespace +from threading import Barrier + +import pytest + +from embodichain.gen_sim.action_engine.unbound import build_unbound_action_plan +from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError +from embodichain.gen_sim.action_engine.runtime import ( + ExecutionReport, + build_execution_provenance, +) +from embodichain.gen_sim.scene_engine.errors import SceneServiceError +from embodichain.gen_sim.task_engine.config import ( + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, +) +from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( + CandidateSelection, +) +from embodichain.gen_sim.task_engine.orchestration.scene_source import SceneSourceRef +from embodichain.gen_sim.task_engine.scene_backend import SceneAnalysis, SceneRevision +from embodichain.gen_sim.task_engine.workflow import ( + SubprocessActionExecutor, + TaskEngineWorkflow, + _environment_successes, + _run_streaming_process, +) +from embodichain.gen_sim.task_engine.workflow_contracts import TASK_RUN_REQUEST_SCHEMA + + +def _candidate_set() -> dict: + candidate = { + "candidate_id": "candidate_01", + "draft": { + "task_id": "place_can", + "instruction": "Place the can on the table.", + "steps": [ + { + "id": "place", + "task_type": "E1", + "object": { + "kind": "scene_ref", + "step_id": "", + "reference": "the can", + "quantifier": "one", + "count": 0, + }, + "target": { + "kind": "scene_ref", + "step_id": "", + "reference": "the table", + "quantifier": "one", + "count": 0, + }, + "depends_on": [], + } + ], + }, + } + return { + "task_id": "place_can", + "instruction": "Place the can on the table.", + "candidates": [candidate], + } + + +def _selection(candidate_set: Mapping[str, object]) -> CandidateSelection: + candidate = candidate_set["candidates"][0] + return CandidateSelection( + scene_manifest={}, + role_bindings={}, + binding_report={ + "status": "bound", + "selection_reason": "test", + "candidates": [{"candidate_id": "candidate_01", "status": "resolved"}], + }, + selected_candidate=candidate, + candidate_bindings={"candidate_01": {}}, + ) + + +def _request(tmp_path: Path, *, existing: bool = False, edit: bool = False) -> dict: + return { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": "place_can", + "task_instruction": "Place the can on the table.", + "image_path": None if existing else str(tmp_path / "input.png"), + "gym_project": str(tmp_path / "project") if existing else None, + "scene_edit_prompt": "Move the can left." if edit else None, + "output_dir": str(tmp_path / "run"), + } + + +class _TaskAgent: + def __init__(self, candidates: dict, barrier: Barrier | None = None) -> None: + self.candidates = candidates + self.barrier = barrier + + def generate(self, *_args, **_kwargs) -> dict: + if self.barrier is not None: + self.barrier.wait(timeout=2) + return self.candidates + + +class _ActionAgent: + def __init__(self, barrier: Barrier | None = None) -> None: + self.barrier = barrier + + def draft(self, candidate: Mapping[str, object]) -> dict: + if self.barrier is not None: + self.barrier.wait(timeout=2) + return build_unbound_action_plan(candidate) + + +class _FailingActionAgent: + def draft(self, _candidate: Mapping[str, object]) -> dict: + raise ActionCapabilityError("missing AtomicAction") + + +class _SceneBackend: + def __init__( + self, + selection: CandidateSelection, + *, + input_kind: str = "image", + input_barrier: Barrier | None = None, + materialize_barrier: Barrier | None = None, + materialize_failures: int = 0, + ) -> None: + self.selection = selection + self.input_kind = input_kind + self.input_barrier = input_barrier + self.materialize_barrier = materialize_barrier + self.materialize_failures = materialize_failures + self.seeds: list[int] = [] + + def analyze(self, request, output_root) -> SceneAnalysis: + if self.input_barrier is not None: + self.input_barrier.wait(timeout=2) + return SceneAnalysis( + input_kind=self.input_kind, + source=Path(request["image_path"] or request["gym_project"]), + blueprint=None, + source_fingerprint=None, + ) + + def select(self, *_args, **_kwargs) -> CandidateSelection: + return self.selection + + def materialize( + self, _analysis, _request, output_root, *, seed: int + ) -> SceneRevision: + if self.materialize_barrier is not None: + self.materialize_barrier.wait(timeout=2) + root = Path(output_root) + root.mkdir(parents=True) + self.seeds.append(seed) + if len(self.seeds) <= self.materialize_failures: + raise SceneServiceError("scene service failed") + source = root / "scene_config.json" + source.write_text("{}\n", encoding="utf-8") + return SceneRevision( + source=source, + output_root=root, + revision_id="0" * 64, + seed=seed, + edit_plan=None, + source_fingerprint=None, + ) + + def inspect(self, revision, output_path): + value = { + "schema_version": "embodichain.final-scene-inspection/v1", + "scene_revision_id": revision.revision_id, + "source_config_path": revision.source.as_posix(), + "contact_tolerance_m": 0.03, + "objects": [], + } + path = Path(output_path) + path.write_text(json.dumps(value), encoding="utf-8") + return value + + +class _Coordinator: + def __init__( + self, + statuses: list[str], + *, + infeasible_remediation: str = "scene_remediable", + ) -> None: + self.statuses = list(statuses) + self.infeasible_remediation = infeasible_remediation + self.calls = 0 + self.kwargs: list[dict] = [] + self.sources: list[object] = [] + + def prepare(self, _task_id, _instruction, _source, output_dir, **_kwargs): + status = self.statuses[min(self.calls, len(self.statuses) - 1)] + self.calls += 1 + self.kwargs.append(dict(_kwargs)) + self.sources.append(_source) + root = Path(output_dir) + root.mkdir(parents=True) + for name in ( + "conservative_scene_graph.json", + "seed_task_graph.json", + "grounded_task_plan.json", + ): + (root / name).write_text("{}\n", encoding="utf-8") + return SimpleNamespace( + status=status, + output_dir=root, + planning_attempts=(), + feasibility_report=( + {"remediation_class": self.infeasible_remediation} + if status == "infeasible" + else None + ), + selected_candidate_id="candidate_01" if status == "bound" else None, + ) + + +class _FailingCoordinator: + def prepare(self, *_args, **_kwargs): + raise RuntimeError("grounding service unavailable") + + +class _RebindingCoordinator(_Coordinator): + def __init__(self, final_candidate: Mapping[str, object]) -> None: + super().__init__(["bound"]) + self.final_candidate = final_candidate + + def prepare(self, *args, **kwargs): + result = super().prepare(*args, **kwargs) + result.selected_candidate_id = str(self.final_candidate["candidate_id"]) + result.unbound_action_plan = build_unbound_action_plan(self.final_candidate) + return result + + +class _InvalidSceneBackend(_SceneBackend): + def materialize(self, *_args, **kwargs): + self.seeds.append(int(kwargs["seed"])) + raise ValueError("invalid deterministic scene input") + + +class _Executor: + def __init__( + self, + successes: list[list[bool]], + *, + expected_dataset_saving: bool = False, + ) -> None: + self.successes = successes + self.expected_dataset_saving = expected_dataset_saving + self.calls = 0 + + def __call__( + self, + _bundle, + _output_root, + *, + seed: int, + num_envs: int, + dataset_saving: bool = False, + ): + values = self.successes[min(self.calls, len(self.successes) - 1)] + self.calls += 1 + assert len(values) == num_envs + assert dataset_saving is self.expected_dataset_saving + return { + "status": "succeeded" if all(values) else "failed", + "seed": seed, + "environments": [ + {"env_id": str(index), "success": success} + for index, success in enumerate(values) + ], + } + + +@pytest.mark.parametrize("existing", [False, True]) +@pytest.mark.parametrize("edit", [False, True]) +def test_parallel_workflow_supports_all_four_scene_inputs( + tmp_path: Path, + *, + existing: bool, + edit: bool, +) -> None: + candidates = _candidate_set() + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend( + _selection(candidates), + input_kind="gym_project" if existing else "image", + ), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor( + [[True, False, False, False]], + expected_dataset_saving=True, + ), + ) + + result = workflow.run( + _request(tmp_path, existing=existing, edit=edit), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + dataset_saving=True, + ) + + assert result.succeeded + + +def test_parallel_workflow_preserves_requested_robot_profile(tmp_path: Path) -> None: + candidates = _candidate_set() + coordinator = _Coordinator(["bound"]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_adapter=SimpleNamespace(robot_profile="ur10"), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=coordinator, + action_executor=_Executor([[True]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=1), + ) + + assert result.succeeded + assert isinstance(coordinator.sources[0], SceneSourceRef) + assert coordinator.sources[0].robot_profile == "ur10" + + +def test_parallel_workflow_accepts_one_success_and_publishes_all_graphs( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + input_barrier = Barrier(2) + work_barrier = Barrier(2) + scene = _SceneBackend( + _selection(candidates), + input_barrier=input_barrier, + materialize_barrier=work_barrier, + ) + coordinator = _Coordinator(["bound"]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates, input_barrier), + scene_backend=scene, + action_agent=_ActionAgent(work_barrier), + coordinator=coordinator, + action_executor=_Executor([[False, True, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + planning_cfg=TaskEnginePlanningCfg( + candidate_count=3, + planning_mode="offline", + max_episodes=1, + max_episode_steps=4000, + ), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + base_seed=11, + run_id="20260820_072436", + ) + + assert result.succeeded + assert scene.seeds == [11] + assert result.final_bundle is not None + assert (result.final_bundle / "conservative_scene_graph.json").is_file() + assert (result.final_bundle / "seed_task_graph.json").is_file() + assert (result.final_bundle / "grounded_task_plan.json").is_file() + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["run_id"] == "20260820_072436" + assert manifest["configuration"]["planning"] == { + "candidate_count": 3, + "planning_mode": "offline", + "max_episodes": 1, + "max_episode_steps": 4000, + } + assert manifest["configuration"]["execution"]["dataset_saving"] is False + assert coordinator.kwargs[0]["max_episode_steps"] == 4000 + assert coordinator.kwargs[0]["final_inspection"]["scene_revision_id"] == "0" * 64 + assert ( + coordinator.kwargs[0]["unbound_action_plan"]["candidate_id"] == "candidate_01" + ) + assert manifest["attempts"][0]["action_attempts"][0]["status"] == "succeeded" + assert manifest["attempts"][0]["final_unbound_action_plan"]["candidate_id"] == ( + "candidate_01" + ) + assert manifest["attempts"][0]["unbound_transition"]["changed"] is False + state = json.loads(result.state_path.read_text(encoding="utf-8")) + succeeded = [ + event["stage"] for event in state["events"] if event["to"] == "succeeded" + ] + assert ( + succeeded.index("scene_finalization") + < succeeded.index("final_inspection") + < succeeded.index("final_binding") + ) + + +def test_prepare_only_publishes_bundle_without_action_execution( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + + def fail_execution(*_args, **_kwargs): + pytest.fail("prepare-only workflow must not execute Action Engine") + + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=fail_execution, + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(), + execute=False, + ) + + assert result.status == "prepared" + assert result.final_bundle is not None + assert result.final_bundle.is_dir() + + +def test_final_candidate_rebinding_updates_attempt_unbound_audit( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + final_candidate = deepcopy(candidates["candidates"][0]) + final_candidate["candidate_id"] = "candidate_02" + candidates["candidates"].append(final_candidate) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_RebindingCoordinator(final_candidate), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + attempt = manifest["attempts"][0] + assert attempt["unbound_action_plan"]["candidate_id"] == "candidate_01" + assert attempt["final_unbound_action_plan"]["candidate_id"] == "candidate_02" + assert attempt["unbound_transition"]["changed"] is True + + +@pytest.mark.parametrize( + ("dataset_saving", "expects_filter"), + [(False, True), (True, False)], +) +def test_subprocess_executor_controls_dataset_saving_and_copies_trajectory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + dataset_saving: bool, + expects_filter: bool, +) -> None: + bundle = tmp_path / "bundle" + bundle.mkdir() + trajectory = tmp_path / "trajectory-source" + trajectory.mkdir() + (trajectory / "episode.json").write_text("{}\n", encoding="utf-8") + captured = {} + provenance = build_execution_provenance(episode_seed=7) + + def fake_run(command, log_path): + captured["command"] = command + captured["log_path"] = Path(log_path) + Path(log_path).write_text("child output\n", encoding="utf-8") + report = ExecutionReport( + task_id="place_can", + plan_hash="0" * 64, + action_graph_hash="1" * 64, + status="succeeded", + run_id="run", + episode_id="0", + provenance=provenance, + environments=tuple( + { + "env_id": str(index), + "success": True, + "semantic_success": {}, + "action_count": 1, + "retry_count": 0, + "recovery_count": 0, + "revision_count": 0, + "failures": [], + } + for index in range(4) + ), + action_count=4, + record_dir=trajectory.as_posix(), + ) + (bundle / "execution_report.json").write_text( + json.dumps(report.as_mapping()), encoding="utf-8" + ) + return SimpleNamespace(returncode=0, stdout="ok", stderr="") + + monkeypatch.setattr( + "embodichain.gen_sim.task_engine.workflow._run_streaming_process", + fake_run, + ) + attempt = tmp_path / "attempt" + + report = SubprocessActionExecutor()( + bundle, + attempt, + seed=7, + num_envs=4, + dataset_saving=dataset_saving, + ) + + assert report["status"] == "succeeded" + assert captured["command"][1:5] == [ + "-u", + "-m", + "embodichain.gen_sim.task_engine._bundle_runner", + "--bundle", + ] + assert " prepare" not in " ".join(captured["command"]) + assert " workflow" not in " ".join(captured["command"]) + assert ("--filter_dataset_saving" in captured["command"]) is expects_filter + assert captured["log_path"] == attempt / "action.log" + assert (attempt / "action.log").read_text(encoding="utf-8") == "child output\n" + assert (attempt / "trajectory" / "episode.json").is_file() + process = json.loads((attempt / "process.json").read_text(encoding="utf-8")) + assert process["combined_log"] == "action.log" + assert process["stdout"] == "ok" + assert process["stderr"] == "" + + +def test_streaming_process_tees_combined_binary_output( + tmp_path: Path, + capfd: pytest.CaptureFixture[str], +) -> None: + log_path = tmp_path / "action.log" + script = ( + "import os; " + "os.write(1, b'stdout\\x00'); " + "os.write(2, b'stderr\\rprogress\\n'); " + "raise SystemExit(7)" + ) + + completed = _run_streaming_process( + [sys.executable, "-c", script], + log_path, + ) + + expected = b"stdout\x00stderr\rprogress\n" + assert completed.returncode == 7 + assert completed.stdout.encode("utf-8") == expected + assert completed.stderr == "" + assert log_path.read_bytes() == expected + terminal = capfd.readouterr().out + assert "stdout\x00" in terminal + assert "stderr\rprogress" in terminal + + +def test_scene_remediation_changes_seed_before_action_execution(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + coordinator = _Coordinator(["infeasible", "bound"]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=coordinator, + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + base_seed=20, + ) + + assert result.succeeded + assert scene.seeds == [20, 21] + assert coordinator.calls == 2 + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert [item["status"] for item in manifest["attempts"]] == [ + "preparation_failed", + "succeeded", + ] + + +def test_input_conflict_feasibility_does_not_regenerate_scene(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator( + ["infeasible", "bound"], + infeasible_remediation="input_conflict", + ), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.failure_class == "input_conflict" + assert scene.seeds == [0] + + +def test_scene_service_retry_keeps_completed_unbound_plan(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates), materialize_failures=1) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + base_seed=30, + ) + + assert result.succeeded + assert scene.seeds == [30, 31] + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["attempts"][0]["unbound_action_plan"] is not None + + +def test_nonremediable_scene_error_does_not_change_scene_attempt_seed( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _InvalidSceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=3), + execution_cfg=TaskEngineExecutionCfg(), + base_seed=9, + ) + + assert not result.succeeded + assert scene.seeds == [9] + + +def test_execution_acceptance_requires_every_success_spec_term() -> None: + report = { + "environments": [ + { + "success": True, + "semantic_success": {"step_01": True, "step_02": False}, + }, + { + "success": True, + "semantic_success": {"step_01": True, "step_02": True}, + }, + ] + } + + assert _environment_successes( + report, + required_semantic_steps=("step_01", "step_02"), + ) == [False, True] + + +def test_unbound_failure_retains_completed_parallel_scene(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_FailingActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert not result.succeeded + assert result.failure_class == "action_capability" + assert scene.seeds == [0] + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["attempts"][0]["scene_revision"] is not None + state = json.loads(result.state_path.read_text(encoding="utf-8")) + assert state["stages"]["scene_finalization"] == "succeeded" + assert state["stages"]["unbound_action"] == "failed" + + +def test_preparation_exception_is_published_as_audited_failure(tmp_path: Path) -> None: + candidates = _candidate_set() + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_FailingCoordinator(), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert not result.succeeded + assert result.failure_class == "preparation_error" + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert manifest["attempts"][0]["status"] == "preparation_error" + assert manifest["attempts"][0]["error"]["type"] == "RuntimeError" + + +def test_explicit_edit_may_materialize_initially_missing_reference( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + unresolved = CandidateSelection( + scene_manifest={}, + role_bindings={}, + binding_report={ + "status": "unsatisfied", + "selection_reason": "the can is not visible before the explicit edit", + "candidates": [{"candidate_id": "candidate_01", "status": "unsatisfied"}], + }, + selected_candidate=None, + candidate_bindings={"candidate_01": {}}, + ) + scene = _SceneBackend(unresolved, input_kind="gym_project") + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path, existing=True, edit=True), + workflow_cfg=TaskEngineWorkflowCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.succeeded + provisional = json.loads( + (result.output_dir / "provisional_candidate.json").read_text(encoding="utf-8") + ) + assert provisional == { + "binding_status": "unsatisfied", + "candidate_id": "candidate_01", + "reason": "explicit_scene_edit_may_materialize_missing_reference", + } + + +def test_action_failure_retries_action_only_and_retains_attempts( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + executor = _Executor([[False, False, False, False]]) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=executor, + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_action_attempts=3), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert not result.succeeded + assert result.failure_class == "action_execution" + assert scene.seeds == [0] + assert executor.calls == 3 + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert len(manifest["attempts"][0]["action_attempts"]) == 3 + + +def test_action_retry_stops_after_first_success(tmp_path: Path) -> None: + candidates = _candidate_set() + executor = _Executor( + [ + [False, False, False, False], + [True, True, True, True], + [True, True, True, True], + ] + ) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=executor, + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_action_attempts=3), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.succeeded + assert executor.calls == 2 + manifest = json.loads(result.manifest_path.read_text(encoding="utf-8")) + assert [item["status"] for item in manifest["attempts"][0]["action_attempts"]] == [ + "failed", + "succeeded", + ] + + +def test_existing_edit_binding_conflict_does_not_invent_scene_repair( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates), input_kind="gym_project") + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["unsatisfied"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path, existing=True, edit=True), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.status == "input_conflict" + assert result.failure_class == "input_conflict" + assert scene.seeds == [0] + + +def test_image_binding_conflict_does_not_regenerate_scene(tmp_path: Path) -> None: + candidates = _candidate_set() + scene = _SceneBackend(_selection(candidates)) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["unsatisfied"]), + action_executor=_Executor([[True, False, False, False]]), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(max_scene_attempts=2), + execution_cfg=TaskEngineExecutionCfg(num_envs=4), + ) + + assert result.status == "input_conflict" + assert result.failure_class == "input_conflict" + assert scene.seeds == [0] diff --git a/tests/gen_sim/task_engine/test_run_directory.py b/tests/gen_sim/task_engine/test_run_directory.py new file mode 100644 index 000000000..e0305d59e --- /dev/null +++ b/tests/gen_sim/task_engine/test_run_directory.py @@ -0,0 +1,58 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from datetime import datetime, timedelta, timezone +from pathlib import Path + +import pytest + +from embodichain.gen_sim.task_engine.run_directory import reserve_run_directory + +_NOW = datetime(2026, 8, 20, 7, 24, 36, tzinfo=timezone(timedelta(hours=8))) + + +def test_run_directory_uses_local_second_timestamp(tmp_path: Path) -> None: + root = tmp_path / "task2_2" + + with reserve_run_directory(root, now=_NOW) as allocation: + assert allocation.run_id == "20260820_072436" + assert allocation.path == root / "20260820_072436" + assert not allocation.path.exists() + allocation.path.mkdir() + + assert allocation.path.is_dir() + assert not (root / ".20260820_072436.reserve").exists() + + +def test_run_directory_adds_suffix_for_same_second_runs(tmp_path: Path) -> None: + root = tmp_path / "task2_2" + (root / "20260820_072436").mkdir(parents=True) + + with reserve_run_directory(root, now=_NOW) as first: + with reserve_run_directory(root, now=_NOW) as second: + assert first.run_id == "20260820_072436_01" + assert second.run_id == "20260820_072436_02" + + +def test_run_directory_rejects_naive_timestamp(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="timezone"): + with reserve_run_directory( + tmp_path, + now=datetime(2026, 8, 20, 7, 24, 36), + ): + pass diff --git a/tests/gen_sim/task_engine/test_scene_backend.py b/tests/gen_sim/task_engine/test_scene_backend.py new file mode 100644 index 000000000..cc6f85bd9 --- /dev/null +++ b/tests/gen_sim/task_engine/test_scene_backend.py @@ -0,0 +1,214 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.scene_engine.core.scene import Scene +from embodichain.gen_sim.scene_engine.core.scene_graph import ( + SceneGraph, + SceneGraphNode, +) +from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.pipeline.api import ( + SceneBlueprintPackage, + SceneMaterialization, +) +import embodichain.gen_sim.task_engine.scene_backend as scene_backend_module +from embodichain.gen_sim.task_engine.scene_backend import ( + SceneEngineBackend, + scene_blueprint_objects, +) +from embodichain.gen_sim.task_engine.workflow_contracts import TASK_RUN_REQUEST_SCHEMA + + +def _request(tmp_path: Path, project: Path, *, edit: str | None) -> dict: + return { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": "task", + "task_instruction": "Move the cup.", + "image_path": None, + "gym_project": project.as_posix(), + "scene_edit_prompt": edit, + "output_dir": (tmp_path / "run").as_posix(), + } + + +def _scene_export(tmp_path: Path) -> Path: + export = tmp_path / "project" / "scene_export" + assets = export / "mesh_assets" + assets.mkdir(parents=True) + (assets / "table.glb").write_bytes(b"glTF-table") + (assets / "cup.glb").write_bytes(b"glTF-cup") + (export / "scene_config.json").write_text( + json.dumps( + { + "format": "embodichain.scene-export/v1", + "scene_id": "scene", + "background": [ + { + "uid": "table", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/table.glb", + }, + } + ], + "rigid_object": [ + { + "uid": "cup", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh_assets/cup.glb", + }, + } + ], + } + ), + encoding="utf-8", + ) + return export.parent + + +def test_blueprint_objects_preserve_semantics_without_geometry(tmp_path: Path) -> None: + scene = Scene( + objects=[ + SceneObject("table", "table", "table", "table", "A table."), + SceneObject("cup", "asset", "cup", "red cup", "A red cup."), + ] + ) + graph = SceneGraph( + nodes=[ + SceneGraphNode("table", None), + SceneGraphNode("cup", "table", "on", orientation_state="lying"), + ] + ) + package = SceneBlueprintPackage( + blueprint_id="blueprint", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=tmp_path / "scene_blueprint.json", + scene=scene, + scene_graph=graph, + ) + + objects = scene_blueprint_objects(package) + + cup = next(item for item in objects if item["uid"] == "cup") + assert cup["description"] == "A red cup." + assert cup["initial_state"] == {"orientation": "fallen"} + assert cup["affordances"] == [] + assert cup["init_pos"] == [0.0, 0.0, 0.0] + + +def test_existing_scene_edit_creates_revision_and_never_writes_source( + tmp_path: Path, + monkeypatch, +) -> None: + project = _scene_export(tmp_path) + source_config = project / "scene_export" / "scene_config.json" + source_value = json.loads(source_config.read_text(encoding="utf-8")) + articulation_path = project / "scene_export" / "cabinet.urdf" + articulation_path.write_text( + '\n', + encoding="utf-8", + ) + source_value["articulation"] = [ + { + "uid": "cabinet", + "name": "cabinet", + "description": "A fixed cabinet.", + "category": "cabinet", + "fpath": "cabinet.urdf", + } + ] + source_config.write_text(json.dumps(source_value), encoding="utf-8") + original = source_config.read_bytes() + prompts: list[str] = [] + + def fake_analyze_edit(*, output_root, edit_prompt): + prompts.append(edit_prompt) + return SimpleNamespace( + output_root=Path(output_root), + scene_edit_plan=SimpleNamespace( + to_dict=lambda: {"operations": [{"op": "move", "object_id": "cup"}]} + ), + ) + + def fake_materialize_edit(blueprint, *, seed=None): + assert seed == 7 + return SceneMaterialization( + scene=Scene(), + scene_graph=SceneGraph(nodes=[SceneGraphNode("table", None)]), + output_root=blueprint.output_root, + scene_config_path=blueprint.output_root + / "scene_export" + / "scene_config.json", + ) + + monkeypatch.setattr(scene_backend_module, "analyze_edit", fake_analyze_edit) + monkeypatch.setattr(scene_backend_module, "materialize_edit", fake_materialize_edit) + backend = SceneEngineBackend() + request = _request(tmp_path, project, edit="Move the cup left.") + analysis = backend.analyze(request, tmp_path / "analysis") + + revision = backend.materialize( + analysis, + request, + tmp_path / "revision", + seed=7, + ) + + assert prompts == ["Move the cup left."] + assert revision.source != source_config + assert revision.source.is_file() + assert len(revision.revision_id) == 64 + assert revision.edit_plan == {"operations": [{"op": "move", "object_id": "cup"}]} + assert source_config.read_bytes() == original + revision_config = json.loads(revision.source.read_text(encoding="utf-8")) + assert revision_config["articulation"][0]["uid"] == "cabinet" + audit = json.loads( + (tmp_path / "revision" / "scene_revision_attempt.json").read_text( + encoding="utf-8" + ) + ) + assert audit["seed"] == 7 + assert audit["revision_id"] == revision.revision_id + assert audit["edit_plan"] == revision.edit_plan + + +def test_final_inspection_rejects_scene_changed_after_revision(tmp_path: Path) -> None: + project = _scene_export(tmp_path) + backend = SceneEngineBackend() + request = _request(tmp_path, project, edit=None) + revision = backend.materialize( + backend.analyze(request, tmp_path / "analysis"), + request, + tmp_path / "unused", + seed=0, + ) + config_path = project / "scene_export" / "scene_config.json" + config = json.loads(config_path.read_text(encoding="utf-8")) + config["rigid_object"][0]["init_pos"] = [0.25, 0.0, 0.0] + config_path.write_text(json.dumps(config), encoding="utf-8") + + with pytest.raises(RuntimeError, match="changed before geometry inspection"): + backend.inspect(revision, tmp_path / "inspection.json") diff --git a/tests/gen_sim/task_engine/test_workflow.py b/tests/gen_sim/task_engine/test_workflow.py new file mode 100644 index 000000000..4a1564606 --- /dev/null +++ b/tests/gen_sim/task_engine/test_workflow.py @@ -0,0 +1,337 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from embodichain.gen_sim.task_engine.config import ( + TaskEngineExecutionCfg, + TaskEnginePlanningCfg, + TaskEngineWorkflowCfg, + load_task_engine_config, +) +from embodichain.gen_sim.task_engine.state_machine import ( + StageStatus, + WorkflowStage, + complete_stage, + fail_stage, + initial_state, + replay_events, + start_stage, + skip_stage, +) +from embodichain.gen_sim.task_engine.workflow_contracts import ( + TASK_RUN_REQUEST_SCHEMA, + scene_input_kind, + validate_scene_history_root, + validate_task_run_request, +) + + +def _request(tmp_path: Path, *, image: bool, edit: bool) -> dict[str, object]: + return { + "schema_version": TASK_RUN_REQUEST_SCHEMA, + "task_id": "pick-cup", + "task_instruction": "Pick up the red cup.", + "image_path": str(tmp_path / "input.png") if image else None, + "gym_project": None if image else str(tmp_path / "gym_project"), + "scene_edit_prompt": "Add a tray." if edit else None, + "output_dir": str(tmp_path / "output"), + } + + +@pytest.mark.parametrize("image", [False, True]) +@pytest.mark.parametrize("edit", [False, True]) +def test_run_request_accepts_all_four_input_combinations( + tmp_path: Path, + image: bool, + edit: bool, +) -> None: + request = validate_task_run_request(_request(tmp_path, image=image, edit=edit)) + assert scene_input_kind(request) == ("image" if image else "gym_project") + assert request["scene_edit_prompt"] == ("Add a tray." if edit else None) + + +def test_run_request_rejects_two_scene_inputs(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + request["gym_project"] = str(tmp_path / "gym_project") + with pytest.raises(ValueError, match="exactly one"): + validate_task_run_request(request) + + +def test_run_request_rejects_scene_generation_prompt(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + request["scene_generation_prompt"] = "Make a kitchen." + with pytest.raises(ValueError, match="fields differ"): + validate_task_run_request(request) + + +def test_run_request_rejects_output_inside_gym_project(tmp_path: Path) -> None: + request = _request(tmp_path, image=False, edit=False) + request["output_dir"] = str(tmp_path / "gym_project" / "task_run") + + with pytest.raises(ValueError, match="must not overlap"): + validate_task_run_request(request) + + +def test_run_request_rejects_output_containing_explicit_gym_config( + tmp_path: Path, +) -> None: + project = tmp_path / "gym_project" + project.mkdir() + config_path = project / "gym_config.json" + config_path.write_text("{}", encoding="utf-8") + request = _request(tmp_path, image=False, edit=False) + request["gym_project"] = str(config_path) + request["output_dir"] = str(project) + + with pytest.raises(ValueError, match="must not overlap"): + validate_task_run_request(request) + + +def test_scene_history_root_allows_a_source_from_a_prior_run( + tmp_path: Path, +) -> None: + history = tmp_path / "task_history" + source = history / "20260820_105939" / "attempts" / "scene_export" + source.mkdir(parents=True) + + validate_scene_history_root(source, history) + + request = _request(tmp_path, image=False, edit=False) + request["gym_project"] = str(source) + request["output_dir"] = str(history / "20260820_130000") + assert validate_task_run_request(request)["gym_project"] == source.as_posix() + + +@pytest.mark.parametrize("relative_output", [".", "new_runs", "new_runs/task"]) +def test_scene_history_root_rejects_writes_into_source_project( + tmp_path: Path, + relative_output: str, +) -> None: + source = tmp_path / "scene_export" + source.mkdir() + output_root = source / relative_output + + with pytest.raises(ValueError, match="read-only source"): + validate_scene_history_root(source, output_root) + + +def test_scene_history_root_resolves_symlinks_before_comparison( + tmp_path: Path, +) -> None: + source = tmp_path / "scene_export" + source.mkdir() + source_link = tmp_path / "scene_link" + source_link.symlink_to(source, target_is_directory=True) + + with pytest.raises(ValueError, match="read-only source"): + validate_scene_history_root(source_link, source / "new_runs") + + +def test_scene_history_root_protects_explicit_config_parent(tmp_path: Path) -> None: + source = tmp_path / "scene_export" + source.mkdir() + config = source / "scene_config.json" + config.write_text("{}\n", encoding="utf-8") + + with pytest.raises(ValueError, match="read-only source"): + validate_scene_history_root(config, source) + + +def test_task_and_scene_stages_can_run_concurrently(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=True, edit=False)) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + assert state.stages[WorkflowStage.TASK_CANDIDATES] == StageStatus.RUNNING + assert state.stages[WorkflowStage.SCENE_PREPARATION] == StageStatus.RUNNING + assert state.stages[WorkflowStage.SCENE_EDIT] == StageStatus.SKIPPED + + +def test_candidate_selection_waits_for_both_branches(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=False, edit=False)) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = complete_stage(state, WorkflowStage.TASK_CANDIDATES) + with pytest.raises(ValueError, match="incomplete dependencies"): + start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + + +def test_unbound_action_can_run_while_user_scene_edit_is_running( + tmp_path: Path, +) -> None: + state = initial_state(_request(tmp_path, image=True, edit=True)) + for stage in (WorkflowStage.TASK_CANDIDATES, WorkflowStage.SCENE_PREPARATION): + state = start_stage(state, stage) + state = complete_stage(state, stage) + state = start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + state = complete_stage(state, WorkflowStage.CANDIDATE_SELECTION) + state = start_stage(state, WorkflowStage.SCENE_EDIT) + state = start_stage(state, WorkflowStage.UNBOUND_ACTION) + + assert state.stages[WorkflowStage.SCENE_EDIT] == StageStatus.RUNNING + assert state.stages[WorkflowStage.UNBOUND_ACTION] == StageStatus.RUNNING + with pytest.raises(ValueError, match="incomplete dependencies"): + start_stage(state, WorkflowStage.SCENE_FINALIZATION) + + +def test_only_scene_edit_can_be_skipped(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=True, edit=True)) + + with pytest.raises(ValueError, match="Only the optional scene_edit stage"): + skip_stage(state, WorkflowStage.FINAL_BINDING) + + +def test_state_events_replay_to_the_same_snapshot(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + state = complete_stage(state, WorkflowStage.TASK_CANDIDATES) + state = complete_stage(state, WorkflowStage.SCENE_PREPARATION) + state = start_stage(state, WorkflowStage.CANDIDATE_SELECTION) + state = complete_stage(state, WorkflowStage.CANDIDATE_SELECTION) + + replayed = replay_events(request, state.events) + + assert replayed.to_dict() == state.to_dict() + + +def test_state_replay_rejects_tampered_transition(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + events = [dict(event) for event in state.events] + events[-1]["stage"] = WorkflowStage.FINAL_BINDING.value + + with pytest.raises(ValueError, match="event does not match"): + replay_events(request, events) + + +def test_state_replay_preserves_failure_reason(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + state = start_stage(state, WorkflowStage.TASK_CANDIDATES) + state = fail_stage(state, WorkflowStage.TASK_CANDIDATES, reason="model timeout") + + replayed = replay_events(request, state.events) + + assert replayed.terminal + assert replayed.to_dict() == state.to_dict() + + +def test_later_retry_can_fail_a_previously_successful_stage(tmp_path: Path) -> None: + request = _request(tmp_path, image=True, edit=False) + state = initial_state(request) + state = start_stage(state, WorkflowStage.SCENE_PREPARATION) + state = complete_stage(state, WorkflowStage.SCENE_PREPARATION) + state = fail_stage( + state, + WorkflowStage.SCENE_PREPARATION, + reason="later scene attempt failed", + ) + + assert state.terminal + assert replay_events(request, state.events).to_dict() == state.to_dict() + + +def test_state_snapshot_mappings_are_immutable(tmp_path: Path) -> None: + state = initial_state(_request(tmp_path, image=True, edit=False)) + + with pytest.raises(TypeError): + state.stages[WorkflowStage.TASK_CANDIDATES] = StageStatus.SUCCEEDED + with pytest.raises(TypeError): + state.request["task_id"] = "changed" + with pytest.raises(TypeError): + state.events[0]["to"] = StageStatus.FAILED.value + + +def test_workflow_configuration_rejects_non_positive_limits() -> None: + with pytest.raises(ValueError, match="max_scene_attempts"): + TaskEngineWorkflowCfg(max_scene_attempts=0) + + +def test_packaged_workflow_configuration_uses_recovery_defaults() -> None: + workflow, planning, execution = load_task_engine_config() + + assert workflow.max_scene_attempts == 2 + assert workflow.max_action_attempts == 3 + assert planning.candidate_count == 3 + assert planning.planning_mode == "offline" + assert planning.max_episodes == 1 + assert planning.max_episode_steps == 4000 + assert execution.num_envs == 1 + assert execution.required_successes == 1 + + +def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: + config = tmp_path / "task_engine.yaml" + config.write_text( + """\ +schema_version: embodichain.task-engine-defaults/v1 +workflow: + max_parallel_workers: 3 + max_scene_attempts: 4 + max_action_attempts: 5 +planning: + candidate_count: 7 + planning_mode: offline + max_episodes: 2 + max_episode_steps: 5000 +execution: + num_envs: 6 + success_policy: at_least + min_successful_envs: 2 +""", + encoding="utf-8", + ) + + workflow, planning, execution = load_task_engine_config(config) + + assert workflow.max_parallel_workers == 3 + assert workflow.max_scene_attempts == 4 + assert workflow.max_action_attempts == 5 + assert planning.candidate_count == 7 + assert planning.max_episodes == 2 + assert planning.max_episode_steps == 5000 + assert execution.num_envs == 6 + assert execution.required_successes == 2 + + +def test_execution_configuration_validates_success_policy() -> None: + assert TaskEngineExecutionCfg().num_envs == 1 + assert ( + TaskEngineExecutionCfg( + num_envs=4, + success_policy="at_least", + min_successful_envs=2, + ).required_successes + == 2 + ) + with pytest.raises(ValueError, match="success_policy=all"): + TaskEngineExecutionCfg( + num_envs=4, + success_policy="all", + min_successful_envs=1, + ) + + +def test_planning_configuration_rejects_invalid_values() -> None: + with pytest.raises(ValueError, match="candidate_count"): + TaskEnginePlanningCfg(candidate_count=0) + with pytest.raises(ValueError, match="planning_mode"): + TaskEnginePlanningCfg(planning_mode="unsupported") From 33f570d1295fc3eee0d95e7d1438245b1009587e Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Tue, 25 Aug 2026 17:59:30 +0800 Subject: [PATCH 65/85] feat(action-engine): add body-grasp planning and verified cleanup for E2 axis alignment --- .../action_engine/capabilities/atomic.py | 35 ++- .../action_engine/config/defaults.yaml | 26 +- .../action_engine/config/runtime_policy.py | 17 +- .../templates/dual_franka_robot.json | 8 +- .../generation/templates/dual_ur_robot.json | 8 +- .../gen_sim/action_engine/runtime/actions.py | 222 ++++++++++--- .../action_engine/runtime/body_grasp.py | 297 ++++++++++++++++++ .../gen_sim/action_engine/runtime/executor.py | 136 +++++++- .../action_engine/runtime/geometry_axes.py | 102 ++++++ .../action_engine/runtime/grounding.py | 52 ++- .../action_engine/runtime/predicates.py | 5 +- .../gen_sim/action_engine/tasks/recipes.py | 85 ++++- .../capabilities/test_atomic_v2.py | 70 +++++ .../config/test_runtime_policy.py | 68 +++- .../generation/test_generation.py | 16 +- .../action_engine/runtime/test_actions.py | 217 ++++++++++++- .../action_engine/runtime/test_body_grasp.py | 153 +++++++++ .../runtime/test_runtime_contracts.py | 99 +++++- tests/gen_sim/action_engine/task_fixtures.py | 220 ++++++++++++- .../action_engine/tasks/test_factory.py | 146 ++++++++- .../tasks/test_interpretation.py | 7 +- .../action_engine/test_motion_policy.py | 4 +- 22 files changed, 1854 insertions(+), 139 deletions(-) create mode 100644 embodichain/gen_sim/action_engine/runtime/body_grasp.py create mode 100644 embodichain/gen_sim/action_engine/runtime/geometry_axes.py create mode 100644 tests/gen_sim/action_engine/runtime/test_body_grasp.py diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 4731ed7e5..d76e83e05 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -346,11 +346,11 @@ def catalog_hash(self) -> str: def build_atomic_capability_registry() -> AtomicCapabilityRegistry: """Build the default catalog, including explicit planning-only skills.""" from embodichain.lab.sim.atomic_actions import ( + AxisAlign, CoordinatedPickment, CoordinatedPickmentOptions, CoordinatedPlacement, CoordinatedPlacementOptions, - AxisAlign, AxisAlignOptions, HandOver, HandOverOptions, @@ -385,8 +385,9 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: "single_arm_object", "preserve", "axis_align", - motion_base="PickUp", + motion_base="AxisAlign", verifier="postcondition", + verifier_hook=_verify_axis_alignment, failure_classifier="grasp", contract_resolver_hook=_resolve_axis_align_contract, allows_target_contact=True, @@ -787,7 +788,16 @@ def _resolve_end_effector_contract( raise ValueError("MoveEndEffector contract requires actor and target_binding.") arm = _required_arm(_actor_arms(actor)[0], "MoveEndEffector") if binding.get("operation") == "retreat" or node.get("role") == "cleanup": - requires = [StateAtom("arm_free", arm=arm)] + requires = [ + StateAtom( + ( + "arm_clear" + if binding.get("operation") == "retreat_after_lift" + else "arm_free" + ), + arm=arm, + ) + ] if binding.get("source") == "handover": requires.append(StateAtom("handover_complete", object_uid=object_uid)) return ResolvedActionContract( @@ -828,6 +838,20 @@ def _resolve_axis_align_contract(node: Mapping[str, Any]) -> ResolvedActionContr ) +def _verify_axis_alignment( + *, + executor: Any, + step: Any, + arm: str, + outcome: Any, + attempted: torch.Tensor, +) -> torch.Tensor: + """Reuse the E2 live predicate before committing AxisAlign completion.""" + del arm, outcome + verified_failed, success, _ = executor._verify_step(step, ~attempted) + return attempted & success & ~verified_failed + + def _resolve_pour_contract(node: Mapping[str, Any]) -> ResolvedActionContract: """Retain one verified holder until observable content transfer succeeds.""" object_uid = _required_string(node.get("object_uid"), "node.object_uid") @@ -924,6 +948,9 @@ def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: ), ) if node.get("role") == "cleanup": + required_home = ( + node.get("task_type") == "E2" and binding.get("operation") == "e2_home" + ) return ResolvedActionContract( requires=(StateAtom("arm_clear", arm=arm),), effects=( @@ -932,7 +959,7 @@ def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: ), claims=(ResourceClaim(f"arm:{arm}"),), completion="terminal_barrier", - failure_policy="best_effort", + failure_policy="safety_required" if required_home else "best_effort", ) return ResolvedActionContract( requires=(StateAtom("arm_free", arm=arm),), diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index 117c76bac..44b2f02cf 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -169,12 +169,19 @@ runtime: force_grasp_reannotate: false motion_defaults: + AxisAlign: + sample_interval: 180 + pre_grasp_distance: 0.15 + lift_height: 0.16 + lower_distance: 0.03 + hand_interp_steps: 12 PickUp: - pre_grasp_distance: 0.08 - lift_height: 0.30 - sample_interval: 45 + pre_grasp_distance: 0.15 + lift_height: 0.16 + sample_interval: 120 + hand_interp_steps: 12 MoveHeldObject: - sample_interval: 45 + sample_interval: 120 relation_distance: 0.18 robot_relative_distance: 0.10 relation_clearance: 0.02 @@ -198,10 +205,11 @@ runtime: line_perpendicular_tolerance: 0.06 preserve_orientation_tolerance: 0.2617993877991494 Place: - sample_interval: 15 - lift_height: 0.0 - post_hold_steps: 0 + sample_interval: 120 + lift_height: 0.14 + post_hold_steps: 60 cartesian_waypoint_count: 2 + hand_interp_steps: 12 MoveEndEffector: sample_interval: 20 retreat_height: 0.30 @@ -264,8 +272,8 @@ runtime: upright_xy_tolerance: 0.05 upright_max_tilt: 0.2617993877991494 Place: - sample_interval: 64 - post_hold_steps: 12 + sample_interval: 120 + post_hold_steps: 60 hand_interp_steps: 12 MoveEndEffector: sample_interval: 30 diff --git a/embodichain/gen_sim/action_engine/config/runtime_policy.py b/embodichain/gen_sim/action_engine/config/runtime_policy.py index 56198211b..8999efbae 100644 --- a/embodichain/gen_sim/action_engine/config/runtime_policy.py +++ b/embodichain/gen_sim/action_engine/config/runtime_policy.py @@ -42,7 +42,8 @@ ] ACTION_ENGINE_DEFAULTS_SCHEMA: Final = "action_engine_defaults_v1" -RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v6" +RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v7" +_PRE_AXIS_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v6" _PREVIOUS_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v5" _PRE_GRASP_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v4" _PRE_PLANNER_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v3" @@ -129,6 +130,7 @@ "collision_activation_distance", } _MOTION_DEFAULT_ACTIONS = { + "AxisAlign", "CoordinatedPickment", "CoordinatedPlacement", "HandOver", @@ -645,6 +647,19 @@ def resolve_agent_runtime_policy(agent_config: Mapping[str, Any]) -> RuntimePoli merged.update(snapshot["arm_selection"]) policy.arm_selection = ArmSelectionPolicyCfg.from_mapping(merged) return policy + if snapshot.get("schema_version") == _PRE_AXIS_RUNTIME_POLICY_SCHEMA: + defaults = default_runtime_policy( + str(agent_config.get("robot_profile", "dual_ur10")) + ) + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + migrated_motion = deepcopy(dict(migrated.get("motion_defaults", {}))) + migrated_motion.setdefault( + "AxisAlign", + deepcopy(defaults.motion_defaults["AxisAlign"]), + ) + migrated["motion_defaults"] = migrated_motion + return RuntimePolicyCfg.from_mapping(migrated) if snapshot.get("schema_version") == _PREVIOUS_RUNTIME_POLICY_SCHEMA: defaults = default_runtime_policy( str(agent_config.get("robot_profile", "dual_ur10")) 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 index b5709f40d..6334e75e1 100644 --- a/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json @@ -21,8 +21,8 @@ "component_type": "left_hand", "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", "transform": [ + [1.0, 0.0, 0.0, 0.0], [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] ] @@ -41,8 +41,8 @@ "component_type": "right_hand", "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", "transform": [ + [1.0, 0.0, 0.0, 0.0], [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] ] @@ -161,8 +161,8 @@ "end_link_name": "left_fr3_link8", "root_link_name": "left_base", "tcp": [ + [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.2], [0.0, 0.0, 0.0, 1.0] ], @@ -174,8 +174,8 @@ "end_link_name": "right_fr3_link8", "root_link_name": "right_base", "tcp": [ + [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.2], [0.0, 0.0, 0.0, 1.0] ], 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 index 8a7496547..8b96a2a59 100644 --- a/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json @@ -18,8 +18,8 @@ "component_type": "left_hand", "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", "transform": [ + [1.0, 0.0, 0.0, 0.0], [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] ] @@ -38,8 +38,8 @@ "component_type": "right_hand", "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", "transform": [ + [1.0, 0.0, 0.0, 0.0], [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] ] @@ -102,8 +102,8 @@ "root_link_name": "left_base_link", "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], "tcp": [ + [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.2], [0.0, 0.0, 0.0, 1.0] ] @@ -116,8 +116,8 @@ "root_link_name": "right_base_link", "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], "tcp": [ + [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.2], [0.0, 0.0, 0.0, 1.0] ] diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 3a1411a9b..7c2549900 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -18,7 +18,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Callable, Mapping from copy import deepcopy from dataclasses import replace import math @@ -37,6 +37,7 @@ ActionPlan, AntipodalAffordance, AtomicActionEngine, + AxisAlignGoal, ControlPartCommandProfile, CoordinatedPickGoal, DynamicCollisionMode, @@ -60,13 +61,16 @@ MotionGenerator, ToppraPlannerCfg, ) +from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg from embodichain.toolkits.graspkit.pg_grasp import ( - AntipodalSamplerCfg, - GraspGeneratorCfg, - GripperCollisionCfg, + AntipodalGraspPoseGenerator, + AntipodalGraspPoseGeneratorCfg, + GraspAnnotationCfg, + ParallelJawGraspCollisionCfg, ) from embodichain.utils.logger import log_info +from .body_grasp import AxisAlignBodyGraspAdapter from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache from .models import ActionOutcome, GroundedAction from .state import ExecutionState @@ -96,6 +100,8 @@ # Preserve cuRobo's fixed world shape while disabling intentional-contact objects. _COLLISION_PARKING_Z_OFFSET = -100.0 +_BODY_GRASP_CANDIDATE_LIMIT = 500 +_BODY_GRASP_SEED = 17_392 def _collision_cache_for_world( @@ -304,27 +310,7 @@ def semantics(self, uid: str) -> ObjectSemantics: raise ValueError(f"Object {uid!r} has invalid mesh triangles.") grasp_options = self.grasp_policy - sampler = AntipodalSamplerCfg( - n_sample=int(grasp_options["antipodal_n_sample"]), - max_angle=float(grasp_options["antipodal_max_angle"]), - max_length=float(grasp_options["max_open_length"]), - min_length=float(grasp_options["min_open_length"]), - ) - generator = GraspGeneratorCfg( - viser_port=int(grasp_options["viser_port"]), - antipodal_sampler_cfg=sampler, - max_deviation_angle=float(grasp_options["max_deviation_angle"]), - n_deviated_approach_directions=int( - grasp_options["n_deviated_approach_directions"] - ), - ) max_hulls = int(grasp_options["max_decomposition_hulls"]) - collision = GripperCollisionCfg( - max_open_length=float(grasp_options["max_open_length"]), - finger_length=float(grasp_options["finger_length"]), - point_sample_dense=float(grasp_options["point_sample_dense"]), - max_decomposition_hulls=max_hulls, - ) cache_result = ensure_vhacd_grasp_collision_cache( mesh_vertices=vertices, mesh_triangles=triangles, @@ -341,9 +327,6 @@ def semantics(self, uid: str) -> ObjectSemantics: object_label=uid, mesh_vertices=vertices, mesh_triangles=triangles, - generator_cfg=generator, - gripper_collision_cfg=collision, - force_reannotate=bool(grasp_options["force_grasp_reannotate"]), ), ) self._semantics[uid] = semantics @@ -359,8 +342,25 @@ def plan( state = state or self.initial_state() grounded = self._select_upright_transport_yaw(grounded, state) context = self._planning_context(state, grounded) - invocation = self._invocation(grounded, capability) - plan = self._engine().plan(invocation, context) + grounded_candidates = self._adapt_axis_align_body_grasps( + grounded, + context, + capability, + ) + selected: tuple[GroundedAction, ActionInvocation, ActionPlan] | None = None + best_failure_count = self.num_envs + 1 + for candidate in grounded_candidates: + candidate_invocation = self._invocation(candidate, capability) + candidate_plan = self._engine().plan(candidate_invocation, context) + failure_count = int((~candidate_plan.plan_success).sum().item()) + if selected is None or failure_count < best_failure_count: + selected = (candidate, candidate_invocation, candidate_plan) + best_failure_count = failure_count + if failure_count == 0: + break + if selected is None: + raise RuntimeError("Atomic action adaptation produced no plan candidate.") + grounded, invocation, plan = selected selected_positions = self._positions_with_agent_holds( plan, grounded, @@ -499,9 +499,90 @@ def plan( # Auditability takes precedence over compactness here: every # selected planner route retains its complete joint path. "planned_trajectory": selected_positions.detach().clone(), + "primary_action_diagnostics": deepcopy(dict(plan.diagnostics.metadata)), + "fallback_action_diagnostics": ( + None + if fallback_plan is None + else deepcopy(dict(fallback_plan.diagnostics.metadata)) + ), + "action_segments": { + segment.name: { + "start": int(segment.start), + "stop": int(segment.stop), + } + for segment in plan.segments + }, }, ) + def _adapt_axis_align_body_grasps( + self, + grounded: GroundedAction, + context: PlanningContext, + capability: AtomicCapability, + ) -> tuple[GroundedAction, ...]: + if capability.target_materializer != "axis_align": + return (grounded,) + goal = grounded.target + if not isinstance(goal, AxisAlignGoal) or goal.grasp_xpos is not None: + return (grounded,) + object_pose = grounded.object_pose + if ( + not isinstance(object_pose, torch.Tensor) + or object_pose.shape != (self.num_envs, 4, 4) + or not torch.isfinite(object_pose).all() + ): + raise ValueError( + "AxisAlign body grasp requires a finite grounded live object pose " + f"with shape ({self.num_envs}, 4, 4)." + ) + if goal.semantics.entity_id is not None: + goal = replace( + goal, + semantics=replace(goal.semantics, entity_id=None), + ) + _, hand_part, _ = self._parts(grounded.arm) + if hand_part is None: + raise ValueError("AxisAlign body grasp requires a configured hand part.") + options = self._build_config(grounded, capability) + selected_approach = options.approach_direction + adaptation = AxisAlignBodyGraspAdapter().adapt( + goal, + object_pose=object_pose, + grasp_generator=self._engine().grasp_pose_generators[hand_part], + approach_direction=selected_approach, + target_axis=options.target_axis, + seed=_BODY_GRASP_SEED, + ) + cfg = dict(grounded.cfg) + cfg["approach_direction"] = selected_approach + candidates: list[GroundedAction] = [] + for adaptation_index, candidate_goal in enumerate(adaptation.alternative_goals): + rank = adaptation.alternative_rank_indices[adaptation_index] + policy = dict(grounded.motion_policy) + policy["body_grasp"] = { + "long_axis_index": adaptation.axes.long_axis_index, + "short_axis_index": adaptation.axes.short_axis_index, + "elongation_ratio": adaptation.axes.elongation_ratio, + "candidate_indices": ( + adaptation.selection.ranked_candidate_indices[:, rank].tolist() + ), + "candidate_counts": ( + adaptation.selection.body_candidate_counts.tolist() + ), + "candidate_rank": rank, + "approach_direction": selected_approach.detach().cpu().tolist(), + } + candidates.append( + replace( + grounded, + target=candidate_goal, + cfg=cfg, + motion_policy=policy, + ) + ) + return tuple(candidates) + def _search_reachable_retreat( self, *, @@ -752,6 +833,9 @@ def _planner_trace( (direction / norm).detach().cpu().tolist() ) trace["grasp_policy"] = grasp_policy + body_grasp = grounded.motion_policy.get("body_grasp") + if isinstance(body_grasp, Mapping): + trace["body_grasp"] = deepcopy(dict(body_grasp)) return trace def _select_upright_transport_yaw( @@ -999,11 +1083,9 @@ def _invocation( ) else: dynamic_mode = DynamicCollisionMode.OFF - goal = ( - self._coordinated_pickment_goal(grounded) - if capability.config_materializer == "coordinated_pickment" - else grounded.target - ) + goal = grounded.target + if capability.config_materializer == "coordinated_pickment": + self._validate_coordinated_pickment_goal(grounded) return ActionInvocation( skill_id=str(capability.action_type.skill_id), goal=goal, @@ -1018,28 +1100,22 @@ def _invocation( ) @staticmethod - def _coordinated_pickment_goal(grounded: GroundedAction) -> CoordinatedPickGoal: - """Apply GenSim-only coordinated grasp filtering to an owned goal copy.""" + def _validate_coordinated_pickment_goal( + grounded: GroundedAction, + ) -> CoordinatedPickGoal: + """Validate the coordinated goal against the engine-scoped generator.""" target = grounded.target if not isinstance(target, CoordinatedPickGoal): raise TypeError("CoordinatedPickment requires a CoordinatedPickGoal.") requested = grounded.cfg.get("is_filter_ground_collision") - if requested is None: - return target - if not isinstance(requested, bool): + if requested is not None and not isinstance(requested, bool): raise TypeError("is_filter_ground_collision must be a boolean.") - semantics = target.semantics - affordance = semantics.affordance - if not isinstance(affordance, AntipodalAffordance): + if not isinstance(target.semantics.affordance, AntipodalAffordance): raise TypeError( "CoordinatedPickment requires an AntipodalAffordance for GenSim " "grasp filtering." ) - generator_cfg = deepcopy(affordance.generator_cfg or GraspGeneratorCfg()) - generator_cfg.is_filter_ground_collision = requested - scoped_affordance = replace(affordance, generator_cfg=generator_cfg) - scoped_semantics = replace(semantics, affordance=scoped_affordance) - return replace(target, semantics=scoped_semantics) + return target def _binding( self, @@ -1341,6 +1417,7 @@ def execute_trajectory( trajectory: torch.Tensor, *, active: torch.Tensor, + waypoint_observer: Callable[[int], None] | None = None, ) -> list[torch.Tensor]: """Advance the environment while holding inactive vectorized rows.""" if trajectory.ndim != 3 or trajectory.shape[0] != self.num_envs: @@ -1351,13 +1428,15 @@ def execute_trajectory( dtype=trajectory.dtype, ) commands: list[torch.Tensor] = [] - for waypoint in trajectory.unbind(dim=1): + for waypoint_index, waypoint in enumerate(trajectory.unbind(dim=1)): command = torch.where(active[:, None], waypoint, current) self.env.step(command) self._scene_time += self._scene_step_duration() update = getattr(self.env, "update_obj_info", None) if callable(update): update() + if waypoint_observer is not None: + waypoint_observer(waypoint_index) commands.append(command.detach()) current = command sync = getattr(self.env, "sync_agent_state_from_qpos", None) @@ -1432,6 +1511,7 @@ def _engine(self) -> AtomicActionEngine: engine = AtomicActionEngine( self._generator(), control_profiles=self._control_profiles(), + grasp_pose_generators=self._grasp_pose_generators(), ) engine.register(ExactTargetMoveHeldObject(), replace=True) self._atomic_engine = engine @@ -1512,6 +1592,52 @@ def _control_profiles(self) -> dict[str, ControlPartCommandProfile]: ) return profiles + def _grasp_pose_generators(self) -> dict[str, AntipodalGraspPoseGenerator]: + """Build one mainline grasp service for each runtime hand endpoint.""" + options = self.grasp_policy + model = ParallelJawGripperModelCfg( + model_id="gen_sim_parallel_jaw", + min_opening_width=float(options["min_open_length"]), + max_opening_width=float(options["max_open_length"]), + finger_length=float(options["finger_length"]), + finger_width=0.03, + finger_thickness=0.01, + palm_depth=0.08, + ) + algorithm = AntipodalGraspPoseGeneratorCfg( + sample_count=int(options["antipodal_n_sample"]), + ray_deviation_angle=float(options["antipodal_max_angle"]), + approach_deviation_angle=float(options["max_deviation_angle"]), + approach_direction_samples=int(options["n_deviated_approach_directions"]), + max_candidates=_BODY_GRASP_CANDIDATE_LIMIT, + ) + collision = ParallelJawGraspCollisionCfg( + point_sample_density=float(options["point_sample_dense"]), + max_decomposition_hulls=int(options["max_decomposition_hulls"]), + opening_margin=0.01, + filter_ground_collision=True, + ) + annotation = GraspAnnotationCfg( + selection_mode="whole_mesh", + viser_port=int(options["viser_port"]), + force_refresh=bool(options["force_grasp_reannotate"]), + ) + generators: dict[str, AntipodalGraspPoseGenerator] = {} + for arm in ("left_arm", "right_arm"): + try: + _, hand_part, _ = self._parts(arm) + except ValueError: + continue + if hand_part is None or hand_part in generators: + continue + generators[hand_part] = AntipodalGraspPoseGenerator( + model, + algorithm_cfg=algorithm, + collision_cfg=collision, + annotation_cfg=annotation, + ) + return generators + def _parts(self, arm: str) -> tuple[str, str | None, int]: if arm not in {"left_arm", "right_arm"}: raise ValueError(f"Expected a physical arm, got {arm!r}.") diff --git a/embodichain/gen_sim/action_engine/runtime/body_grasp.py b/embodichain/gen_sim/action_engine/runtime/body_grasp.py new file mode 100644 index 000000000..560e34b5b --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/body_grasp.py @@ -0,0 +1,297 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Pure body-grasp candidate filtering for elongated rigid objects.""" + +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import replace +from collections.abc import Callable + +import torch + +from embodichain.lab.sim.atomic_actions import ( + AxisAlignAffordance, + AxisAlignGoal, +) +from .geometry_axes import LocalGeometryAxes +from .geometry_axes import analyze_local_geometry_axes + +__all__ = [ + "AxisAlignBodyGraspAdapter", + "BodyGraspAdaptation", + "BodyGraspSelection", + "select_body_grasp_candidates", +] + + +@dataclass(frozen=True, slots=True) +class BodyGraspSelection: + """One selected body grasp per environment row.""" + + success: torch.Tensor + grasp_xpos: torch.Tensor + candidate_indices: torch.Tensor + body_candidate_counts: torch.Tensor + central_candidate_counts: torch.Tensor + radial_candidate_counts: torch.Tensor + minimum_normalized_axial_offset: torch.Tensor + minimum_long_axis_opening_cosine: torch.Tensor + reachable_candidate_counts: torch.Tensor + ranked_grasp_xpos: torch.Tensor + ranked_candidate_indices: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class BodyGraspAdaptation: + """AxisAlign goal plus auditable body-grasp selection metadata.""" + + goal: AxisAlignGoal + alternative_goals: tuple[AxisAlignGoal, ...] + alternative_rank_indices: tuple[int, ...] + axes: LocalGeometryAxes + selection: BodyGraspSelection + + +class AxisAlignBodyGraspAdapter: + """Lower elongated-object semantics into an explicit mainline grasp goal.""" + + def __init__( + self, + *, + body_band_fraction: float = 0.80, + maximum_long_axis_opening_cosine: float = 0.50, + ) -> None: + self.body_band_fraction = body_band_fraction + self.maximum_long_axis_opening_cosine = maximum_long_axis_opening_cosine + + def adapt( + self, + goal: AxisAlignGoal, + *, + object_pose: torch.Tensor, + grasp_generator: object, + approach_direction: torch.Tensor, + target_axis: torch.Tensor, + seed: int, + candidate_feasibility: Callable[[torch.Tensor], torch.Tensor] | None = None, + maximum_adaptations: int = 12, + ) -> BodyGraspAdaptation: + affordance = goal.semantics.affordance + if not isinstance(affordance, AxisAlignAffordance): + raise TypeError("AxisAlign body grasp requires AxisAlignAffordance.") + if affordance.mesh_vertices is None or affordance.mesh_triangles is None: + raise ValueError("AxisAlign body grasp requires indexed mesh geometry.") + axes = analyze_local_geometry_axes(affordance.mesh_vertices) + sampled = self._sample( + grasp_generator, + seed=seed, + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + obj_poses=object_pose, + approach_direction=approach_direction, + obj_longest_axis=None, + is_positive_part=True, + ) + candidates, costs = self._pack(sampled, object_pose) + feasible = ( + None if candidate_feasibility is None else candidate_feasibility(candidates) + ) + selection = select_body_grasp_candidates( + candidates, + costs, + object_pose, + axes, + body_band_fraction=self.body_band_fraction, + maximum_long_axis_opening_cosine=(self.maximum_long_axis_opening_cosine), + feasible=feasible, + ) + if not bool(selection.success.all().item()): + failed = ( + torch.nonzero(~selection.success, as_tuple=False).flatten().tolist() + ) + raise ValueError( + "No central radial body grasp is available for rows " + f"{failed}; central_counts=" + f"{selection.central_candidate_counts.tolist()}, radial_counts=" + f"{selection.radial_candidate_counts.tolist()}, min_axial=" + f"{selection.minimum_normalized_axial_offset.tolist()}, " + "min_opening_cos=" + f"{selection.minimum_long_axis_opening_cosine.tolist()}, " + "reachable_counts=" + f"{selection.reachable_candidate_counts.tolist()}." + ) + adaptation_count = min( + maximum_adaptations, + selection.ranked_grasp_xpos.shape[1], + ) + alternative_ranks = tuple( + int(value) + for value in torch.linspace( + 0, + selection.ranked_grasp_xpos.shape[1] - 1, + adaptation_count, + ) + .round() + .to(torch.int64) + .tolist() + ) + goals = tuple( + replace(goal, grasp_xpos=selection.ranked_grasp_xpos[:, rank]) + for rank in alternative_ranks + ) + del target_axis + return BodyGraspAdaptation( + goal=goals[0], + alternative_goals=goals, + alternative_rank_indices=alternative_ranks, + axes=axes, + selection=selection, + ) + + @staticmethod + def _sample( + generator: object, + *, + seed: int, + **kwargs: object, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + poses = kwargs.get("obj_poses") + if not isinstance(poses, torch.Tensor): + raise TypeError("obj_poses must be a torch.Tensor.") + devices: list[int] = [] + if poses.device.type == "cuda": + devices.append( + torch.cuda.current_device() + if poses.device.index is None + else poses.device.index + ) + with torch.random.fork_rng(devices=devices): + torch.manual_seed(seed) + return generator.get_valid_grasp_poses( # type: ignore[attr-defined] + **kwargs + ) + + @staticmethod + def _pack( + sampled: list[tuple[torch.Tensor, torch.Tensor]], + object_pose: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if len(sampled) != object_pose.shape[0]: + raise ValueError("Grasp generator must return one result per object row.") + count = max((poses.shape[0] for poses, _ in sampled), default=0) + if count == 0: + raise ValueError("Grasp generator returned no candidates.") + candidates = torch.eye( + 4, + dtype=torch.float32, + device=object_pose.device, + ).repeat(object_pose.shape[0], count, 1, 1) + costs = torch.full( + (object_pose.shape[0], count), + torch.inf, + dtype=torch.float32, + device=object_pose.device, + ) + for env_index, (poses, values) in enumerate(sampled): + row_count = poses.shape[0] + if row_count == 0: + continue + candidates[env_index, :row_count] = poses.to( + device=object_pose.device, + dtype=torch.float32, + ) + costs[env_index, :row_count] = values.to( + device=object_pose.device, + dtype=torch.float32, + ) + return candidates, costs + + +def select_body_grasp_candidates( + candidates: torch.Tensor, + costs: torch.Tensor, + object_pose: torch.Tensor, + axes: LocalGeometryAxes, + *, + body_band_fraction: float = 0.80, + maximum_long_axis_opening_cosine: float = 0.50, + feasible: torch.Tensor | None = None, +) -> BodyGraspSelection: + """Keep central radial grasps and reject cap/end grasps.""" + if candidates.ndim != 4 or candidates.shape[-2:] != (4, 4): + raise ValueError("candidates must have shape (B, N, 4, 4).") + if costs.shape != candidates.shape[:2]: + raise ValueError("costs must have shape (B, N).") + if object_pose.shape != (candidates.shape[0], 4, 4): + raise ValueError("object_pose must have shape (B, 4, 4).") + if not 0.0 < body_band_fraction <= 1.0: + raise ValueError("body_band_fraction must be in (0, 1].") + if not 0.0 <= maximum_long_axis_opening_cosine < 1.0: + raise ValueError("maximum_long_axis_opening_cosine must be in [0, 1).") + if feasible is None: + feasible = torch.ones_like(costs, dtype=torch.bool) + if feasible.dtype != torch.bool or feasible.shape != costs.shape: + raise ValueError("feasible must be a bool tensor shaped (B, N).") + + rotation = object_pose[:, :3, :3] + translation = object_pose[:, :3, 3] + local_centers = torch.matmul( + candidates[..., :3, 3] - translation[:, None], + rotation, + ) + center = axes.bounds_center.to( + device=candidates.device, + dtype=candidates.dtype, + ) + long_axis = axes.long_axis.to( + device=candidates.device, + dtype=candidates.dtype, + ) + axial_offset = torch.abs(torch.sum((local_centers - center) * long_axis, dim=-1)) + normalized_axial = axial_offset / max(axes.long_half_extent, 1.0e-8) + within_body = normalized_axial <= body_band_fraction + + world_opening = torch.nn.functional.normalize(candidates[..., :3, 0], dim=-1) + local_opening = torch.matmul(world_opening, rotation) + long_axis_opening = torch.abs(torch.sum(local_opening * long_axis, dim=-1)) + radial = long_axis_opening <= maximum_long_axis_opening_cosine + valid = within_body & radial & feasible & torch.isfinite(costs) + + ranked = torch.where(valid, costs, torch.inf) + best_cost, best_index = ranked.min(dim=1) + env_index = torch.arange(candidates.shape[0], device=candidates.device) + valid_counts = valid.sum(dim=1) + rank_count = int(valid_counts.min().item()) + ranked_indices = torch.argsort(ranked, dim=1)[:, :rank_count] + ranked_grasps = candidates[ + env_index[:, None], + ranked_indices, + ].clone() + return BodyGraspSelection( + success=torch.isfinite(best_cost), + grasp_xpos=candidates[env_index, best_index].clone(), + candidate_indices=best_index.clone(), + body_candidate_counts=valid.sum(dim=1), + central_candidate_counts=within_body.sum(dim=1), + radial_candidate_counts=radial.sum(dim=1), + minimum_normalized_axial_offset=normalized_axial.min(dim=1).values, + minimum_long_axis_opening_cosine=long_axis_opening.min(dim=1).values, + reachable_candidate_counts=feasible.sum(dim=1), + ranked_grasp_xpos=ranked_grasps, + ranked_candidate_indices=ranked_indices.clone(), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index 31d26a847..cfa7ccd99 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -606,6 +606,7 @@ def run( else None ), ) + self._release_candidate_plans(step.id) revalidation_failures = self._revalidate_support_relations() for step_id, lost in revalidation_failures.items(): step = self.steps[step_id] @@ -1610,6 +1611,16 @@ def _consume_transitions(self, count: int) -> None: if self._transition_count > self.max_transitions: raise RuntimeError("Execution exceeded max_transitions.") + def _release_candidate_plans(self, step_id: str) -> None: + """Drop completed-step speculative plans and their retained trajectories.""" + for mapping in (self._candidate_cache, self._candidate_failures): + for key in tuple(mapping): + if key[0] == step_id: + mapping.pop(key, None) + self._candidate_diagnostics.pop(step_id, None) + self._candidate_blockers.pop(step_id, None) + self._reported_candidates.discard(step_id) + def _pack_ready_edges( self, ready: Sequence[ExecutionEdge], @@ -2907,7 +2918,46 @@ def _execute_edge( ) trajectory, action_success = self.adapter.combine(outcomes, masks) active = assigned & action_success & ~failed & ~planning_failed - actions = self.adapter.execute_trajectory(trajectory, active=active) + observation_points: dict[int, list[str]] = {} + execution_observations: dict[str, dict[str, torch.Tensor]] = {} + for outcome in outcomes.values(): + if outcome is None: + continue + diagnostics = outcome.planner_trace.get( + "primary_action_diagnostics", {} + ) + observed_segments = set( + diagnostics.get("execution_observation_segments", ()) + ) + for name, segment in outcome.planner_trace.get( + "action_segments", {} + ).items(): + stop = int(segment["stop"]) + if name in observed_segments and stop > 0: + observation_points.setdefault(stop - 1, []).append(name) + break + + def observe_waypoint(waypoint_index: int) -> None: + for name in observation_points.get(waypoint_index, ()): + execution_observations[name] = self._action_execution_observation( + step.object_uid + ) + + actions = ( + self.adapter.execute_trajectory( + trajectory, + active=active, + waypoint_observer=observe_waypoint, + ) + if observation_points + else self.adapter.execute_trajectory(trajectory, active=active) + ) + if execution_observations: + for outcome in outcomes.values(): + if outcome is not None: + outcome.planner_trace["execution_observations"] = ( + execution_observations + ) physical_failed = torch.zeros_like(failed) for arm, outcome in outcomes.items(): if outcome is not None: @@ -2977,6 +3027,90 @@ def _execute_edge( active, ) + def _action_execution_observation( + self, + object_uid: str, + ) -> dict[str, torch.Tensor]: + """Capture live object and TCP state at one requested segment boundary.""" + entity = self.env.sim.get_rigid_object(object_uid) + if entity is None: + raise ValueError(f"Unknown action observation object {object_uid!r}.") + observation = { + "object_pose": self._entity_pose(object_uid).detach().clone(), + } + for name, attribute_names in ( + ( + "linear_velocity", + ("lin_vel", "linear_velocity", "get_linear_velocity"), + ), + ( + "angular_velocity", + ("ang_vel", "angular_velocity", "get_angular_velocity"), + ), + ): + for attribute_name in attribute_names: + value = getattr(entity, attribute_name, None) + if callable(value): + value = value() + if value is None: + continue + tensor = torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if tensor.shape[-1:] == (3,): + observation[name] = tensor.detach().clone() + break + if not {"linear_velocity", "angular_velocity"} <= set(observation): + body_state = getattr(entity, "body_state", None) + if callable(body_state): + body_state = body_state() + if body_state is not None: + state = torch.as_tensor( + body_state, + dtype=torch.float32, + device=self.env.device, + ) + if state.ndim == 1: + state = state.unsqueeze(0) + if state.shape == (int(self.env.num_envs), 13): + observation.setdefault("linear_velocity", state[:, 7:10].clone()) + observation.setdefault( + "angular_velocity", state[:, 10:13].clone() + ) + body_data = getattr(entity, "body_data", None) + if body_data is not None: + for name, attribute_name in ( + ("linear_velocity", "lin_vel"), + ("angular_velocity", "ang_vel"), + ): + if name in observation: + continue + value = getattr(body_data, attribute_name, None) + if callable(value): + value = value() + if value is not None: + observation[name] = torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ).detach().clone() + getter = getattr(self.env, "get_current_xpos_agent", None) + if callable(getter): + left, right = getter() + observation["left_tcp_pose"] = torch.as_tensor( + left, + dtype=torch.float32, + device=self.env.device, + ).detach().clone() + observation["right_tcp_pose"] = torch.as_tensor( + right, + dtype=torch.float32, + device=self.env.device, + ).detach().clone() + return observation + def _plan_live_hold( self, edge: ExecutionEdge, diff --git a/embodichain/gen_sim/action_engine/runtime/geometry_axes.py b/embodichain/gen_sim/action_engine/runtime/geometry_axes.py new file mode 100644 index 000000000..33564ab1e --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/geometry_axes.py @@ -0,0 +1,102 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Object-local axis analysis shared by GenSim grounding and grasp lowering.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +__all__ = ["LocalGeometryAxes", "analyze_local_geometry_axes"] + + +@dataclass(frozen=True, slots=True) +class LocalGeometryAxes: + """Validated local AABB axes with a PCA alignment cross-check.""" + + bounds_center: torch.Tensor + extents: torch.Tensor + ordered_axis_indices: tuple[int, int, int] + long_axis_index: int + short_axis_index: int + long_axis: torch.Tensor + long_half_extent: float + elongation_ratio: float + principal_alignment: float + + +def analyze_local_geometry_axes( + vertices: torch.Tensor, + *, + minimum_elongation_ratio: float = 1.10, + minimum_principal_alignment: float = 0.90, +) -> LocalGeometryAxes: + """Resolve stable local long/short axes or fail on ambiguous geometry.""" + if ( + not isinstance(vertices, torch.Tensor) + or not vertices.is_floating_point() + or vertices.ndim != 2 + or vertices.shape[1] != 3 + or vertices.shape[0] < 3 + or not torch.isfinite(vertices).all() + ): + raise ValueError("vertices must be a finite floating tensor shaped (N, 3).") + if minimum_elongation_ratio <= 1.0: + raise ValueError("minimum_elongation_ratio must be greater than one.") + if not 0.0 < minimum_principal_alignment <= 1.0: + raise ValueError("minimum_principal_alignment must be in (0, 1].") + + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + extents = upper - lower + if torch.any(extents <= 1.0e-8): + raise ValueError("Object geometry must have non-zero extent on every axis.") + ordered = torch.argsort(extents, descending=True) + long_index = int(ordered[0].item()) + short_index = int(ordered[-1].item()) + elongation = float((extents[long_index] / extents[int(ordered[1])]).item()) + if elongation < minimum_elongation_ratio: + raise ValueError( + "Object long axis is ambiguous; provide an explicit upright_local_axis." + ) + + centered = vertices - vertices.mean(dim=0, keepdim=True) + covariance = centered.transpose(0, 1) @ centered + _, eigenvectors = torch.linalg.eigh(covariance) + principal = eigenvectors[:, -1] + principal_index = int(torch.argmax(torch.abs(principal)).item()) + alignment = float(torch.abs(principal[principal_index]).item()) + if principal_index != long_index or alignment < minimum_principal_alignment: + raise ValueError( + "Object principal axis is not aligned with its local AABB; provide an " + "explicit local axis instead of inferring long_axis." + ) + + long_axis = torch.zeros(3, dtype=vertices.dtype, device=vertices.device) + long_axis[long_index] = 1.0 + return LocalGeometryAxes( + bounds_center=((lower + upper) * 0.5).clone(), + extents=extents.clone(), + ordered_axis_indices=tuple(int(index) for index in ordered.tolist()), + long_axis_index=long_index, + short_axis_index=short_index, + long_axis=long_axis, + long_half_extent=float((extents[long_index] * 0.5).item()), + elongation_ratio=elongation, + principal_alignment=alignment, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 8687a12b7..2f0b6285e 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -66,6 +66,7 @@ ) from embodichain.utils.logger import log_info from .frames import arm_base_poses, relation_offset, robot_frame_axes +from .geometry_axes import analyze_local_geometry_axes from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache from .models import ExecutionProgram, GroundedAction, SemanticStep from .motion_policy import resolve_motion_policy, with_motion_modifiers @@ -74,6 +75,8 @@ __all__ = ["ActionGrounder", "LiveArrangementPlan", "LivePlacementPlan"] +_E2_CLEARANCE_RETREAT_DISTANCE = 0.20 + def _batched_pose(value: Any, env: Any) -> torch.Tensor: pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) @@ -959,13 +962,16 @@ def ground( ) elif kind == "policy_pose": source = binding.get("source") + operation = binding.get("operation") retreat_reference = self._retreat_reference_pose( arm, reference_eef_pose, ) - if binding.get("operation") == "retreat": + if operation == "retreat": policy["retreat_reachability_search"] = True policy["retreat_reference_pose"] = retreat_reference.clone() + if operation == "retreat_after_lift": + policy["retreat_distance"] = _E2_CLEARANCE_RETREAT_DISTANCE if source in {"release", "handover"}: policy["clearance_object_uid"] = step.object_uid policy["collision_safety"] = "required" @@ -988,6 +994,7 @@ def ground( policy, retreat_reference, clear_exchange=source == "handover", + retreat_after_lift=operation == "retreat_after_lift", ) ) elif kind == "visual_constraint": @@ -2537,9 +2544,8 @@ def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: axis = self._upright_local_axis(step) entity = _object(self.env, step.object_uid) vertices = _local_vertices(entity, self.env, 0) - extents = vertices.max(dim=0).values - vertices.min(dim=0).values if axis == "long_axis": - axis_index = int(torch.argmax(extents).item()) + axis_index = analyze_local_geometry_axes(vertices).long_axis_index else: axis_index = {"x": 0, "y": 1, "z": 2}[axis] direction = torch.zeros(3, dtype=torch.float32, device=self.env.device) @@ -2567,8 +2573,7 @@ def _uses_upright_yaw_search( return False entity = _object(self.env, step.object_uid) vertices = _local_vertices(entity, self.env, 0) - extents = vertices.max(dim=0).values - vertices.min(dim=0).values - axis_index = int(torch.argmax(extents).item()) + axis_index = analyze_local_geometry_axes(vertices).long_axis_index pose = _live_pose(self.env, step.object_uid) cosine = pose[:, 2, axis_index].abs().clamp(0.0, 1.0) tolerance = float(self.runtime_policy.predicate_fallbacks["upright_max_tilt"]) @@ -2623,11 +2628,9 @@ def _target_rotation( rotations = [] for env_id in range(int(self.env.num_envs)): vertices = _local_vertices(entity, self.env, env_id) - extents = vertices.max(dim=0).values - vertices.min(dim=0).values - longest_to_shortest = torch.argsort( - extents, - descending=True, - ).tolist() + longest_to_shortest = list( + analyze_local_geometry_axes(vertices).ordered_axis_indices + ) if goal == "upright": upright_axis = ( align_term.local_axis @@ -2699,8 +2702,7 @@ def _horizontal_orientation( if isinstance(align_to, str) and align_to: reference = _object(self.env, align_to) vertices = _local_vertices(reference, self.env, env_id) - extents = vertices.max(dim=0).values - vertices.min(dim=0).values - ordered = torch.argsort(extents, descending=True) + ordered = analyze_local_geometry_axes(vertices).ordered_axis_indices requested = str(step.goal.get("orientation_axis", "long_axis")) reference_axis = int( ordered[-1] if requested == "short_axis" else ordered[0] @@ -2997,8 +2999,34 @@ def _retreat_pose( reference: torch.Tensor | None, *, clear_exchange: bool = False, + retreat_after_lift: bool = False, ) -> torch.Tensor: target = self._retreat_reference_pose(arm, reference).clone() + if retreat_after_lift: + direction: torch.Tensor | None = None + clearance_uid = policy.get("clearance_object_uid") + if isinstance(clearance_uid, str) and clearance_uid: + entity = _object(self.env, clearance_uid) + object_pose = _batched_pose( + entity.get_local_pose(to_matrix=True), + self.env, + ) + direction = target[:, :2, 3] - object_pose[:, :2, 3] + if direction is None: + direction = torch.zeros_like(target[:, :2, 3]) + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + unresolved = norm <= 1.0e-6 + if unresolved.any(): + left_base, right_base = arm_base_poses(self.env) + base = left_base if arm == "left_arm" else right_base + baseward = base[:, :2, 3] - target[:, :2, 3] + baseward_norm = torch.linalg.vector_norm(baseward, dim=1, keepdim=True) + baseward = baseward / torch.clamp(baseward_norm, min=1.0e-6) + direction = torch.where(unresolved, baseward, direction) + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + direction = direction / torch.clamp(norm, min=1.0e-6) + target[:, :2, 3] += direction * float(policy.get("retreat_distance", 0.10)) + return target desired = float(self._policy_value(policy, "retreat_height")) if clear_exchange: _, lateral = robot_frame_axes(self.env) diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index 57fe8a774..4062204ba 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -23,6 +23,8 @@ import torch +from .geometry_axes import analyze_local_geometry_axes + from embodichain.gen_sim.action_engine.config import default_runtime_policy from .frames import relation_axes @@ -245,8 +247,7 @@ def _local_axis_index(env: Any, uid: str, axis: Any) -> int: vertices = vertices[0] if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") - extents = vertices.max(dim=0).values - vertices.min(dim=0).values - return int(torch.argmax(extents).item()) + return analyze_local_geometry_axes(vertices).long_axis_index def _arm_values( diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index 9beda20b6..3df54d5e3 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -409,23 +409,76 @@ def _recipe( "local_axis": goal["upright_local_axis"], } if incoming_held_arm is None and terminal_behavior == "place": + alignment = _node( + group_id, + 1, + "AxisAlign", + task_type, + object_uid, + actor, + "arm", + {"kind": "object", "object": object_uid}, + dependencies, + role, + success, + motion_policy(), + ) + lift_clear = _node( + group_id, + 2, + "MoveEndEffector", + task_type, + object_uid, + actor, + "arm", + { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + }, + [alignment["id"]], + "cleanup", + {}, + motion_policy(), + ) + retreat = _node( + group_id, + 3, + "MoveEndEffector", + task_type, + object_uid, + actor, + "arm", + { + "kind": "policy_pose", + "source": "release", + "operation": "retreat_after_lift", + }, + [lift_clear["id"]], + "cleanup", + {}, + motion_policy(("orientation", "upright")), + ) + home = _node( + group_id, + 4, + "MoveJoints", + task_type, + object_uid, + actor, + "arm", + { + "kind": "joint_state", + "source": "initial", + "operation": "e2_home", + }, + [retreat["id"]], + "cleanup", + {}, + motion_policy(), + ) return ( - [ - _node( - group_id, - 1, - "AxisAlign", - task_type, - object_uid, - actor, - "arm", - {"kind": "object", "object": object_uid}, - dependencies, - role, - success, - motion_policy(), - ) - ], + [alignment, lift_clear, retreat, home], "orient_object", goal, success, diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py index 07245b274..9e065b1a9 100644 --- a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -119,6 +119,76 @@ def get_rigid_object(self, _uid: str): return _Entity() +def test_axis_align_uses_its_tutorial_motion_policy_base() -> None: + capability = build_atomic_capability_registry().get("AxisAlign") + + assert capability.motion_base == "AxisAlign" + + +def test_axis_align_verifies_the_live_semantic_postcondition() -> None: + capability = build_atomic_capability_registry().get("AxisAlign") + attempted = torch.tensor([True, False]) + calls = [] + + def verify_step(step, failed): + calls.append((step, failed.clone())) + return failed.clone(), torch.tensor([True, False]), torch.zeros(2, 3) + + executor = SimpleNamespace( + _verify_step=verify_step, + _action_execution_observation=lambda _uid: { + "linear_velocity": torch.zeros(2, 3), + "angular_velocity": torch.zeros(2, 3), + }, + runtime_policy=SimpleNamespace( + execution={ + "support_linear_velocity_tolerance": 0.02, + "support_angular_velocity_tolerance": 0.2, + } + ), + ) + step = SimpleNamespace(id="orient", object_uid="can") + + verified = capability.verifier_hook( + executor=executor, + step=step, + arm="left_arm", + outcome=SimpleNamespace(), + attempted=attempted, + ) + + assert verified.tolist() == [True, False] + assert calls[0][0] is step + assert calls[0][1].tolist() == [False, True] + + +def test_e2_home_is_required_while_generic_cleanup_home_is_best_effort() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + base = { + "atomic_action": "MoveJoints", + "object_uid": "can", + "actor": {"mode": "required", "arm": "right_arm"}, + "control": "arm", + "role": "cleanup", + "target_binding": {"kind": "joint_state", "source": "initial"}, + } + + generic = capability.resolve_contract(base) + e2_home = capability.resolve_contract( + { + **base, + "task_type": "E2", + "target_binding": { + **base["target_binding"], + "operation": "e2_home", + }, + } + ) + + assert generic.failure_policy == "best_effort" + assert e2_home.failure_policy == "safety_required" + + def test_new_descriptor_reuses_loader_and_adapter_without_dispatch_changes() -> None: registry = build_atomic_capability_registry() calls = [] diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py index 3e0b603d5..c74a964b5 100644 --- a/tests/gen_sim/action_engine/config/test_runtime_policy.py +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -120,6 +120,43 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: ] +def test_e1_motion_defaults_match_atomic_action_tutorial_cadence() -> None: + policy = default_runtime_policy("dual_franka") + motion = policy.motion_defaults + + assert motion["PickUp"] == { + "pre_grasp_distance": pytest.approx(0.15), + "lift_height": pytest.approx(0.16), + "sample_interval": 120, + "hand_interp_steps": 12, + } + assert motion["MoveHeldObject"]["sample_interval"] == 120 + assert motion["Place"] == { + "sample_interval": 120, + "lift_height": pytest.approx(0.14), + "post_hold_steps": 60, + "cartesian_waypoint_count": 2, + "hand_interp_steps": 12, + } + assert policy.motion_modifiers["orientation"]["upright"]["Place"] == { + "sample_interval": 120, + "post_hold_steps": 60, + "hand_interp_steps": 12, + } + + +def test_axis_align_defaults_match_atomic_action_tutorial_cadence() -> None: + axis_align = default_runtime_policy("dual_franka").motion_defaults["AxisAlign"] + + assert axis_align == { + "sample_interval": 180, + "pre_grasp_distance": pytest.approx(0.15), + "lift_height": pytest.approx(0.16), + "lower_distance": pytest.approx(0.03), + "hand_interp_steps": 12, + } + + def test_place_defaults_fit_the_mainline_motion_sample_budget() -> None: place = default_runtime_policy("dual_ur10").motion_defaults["Place"] sample_count = int(place["sample_interval"]) @@ -142,7 +179,7 @@ def test_default_runtime_policy_returns_detached_profile_snapshots() -> None: first.motion_defaults["PickUp"]["lift_height"] = 9.0 assert second.arm_selection.pickup_crossing_weight == 1.0 - assert second.motion_defaults["PickUp"]["lift_height"] == 0.30 + assert second.motion_defaults["PickUp"]["lift_height"] == 0.16 assert franka.arm_selection.pickup_crossing_weight == 1.0 assert franka.motion_defaults["MoveEndEffector"]["retreat_height"] == 0.10 @@ -228,6 +265,31 @@ def test_agent_policy_snapshot_is_hash_verified_and_legacy_config_falls_back() - assert legacy.as_mapping() == snapshot +def test_v6_policy_snapshot_adds_axis_align_defaults_without_rewriting_e1() -> None: + snapshot = default_runtime_policy("dual_franka").as_mapping() + snapshot["schema_version"] = "action_engine_runtime_policy_v6" + snapshot["motion_defaults"].pop("AxisAlign") + snapshot["motion_defaults"]["PickUp"]["lift_height"] = 0.11 + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + resolved = resolve_agent_runtime_policy( + { + "robot_profile": "dual_franka", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert resolved.schema_version == "action_engine_runtime_policy_v7" + assert resolved.motion_defaults["AxisAlign"]["sample_interval"] == 180 + assert resolved.motion_defaults["PickUp"]["lift_height"] == pytest.approx(0.11) + + def test_narrow_v1_policy_snapshot_is_migrated_to_complete_runtime_policy() -> None: snapshot = { "schema_version": "action_engine_runtime_policy_v1", @@ -255,7 +317,7 @@ def test_narrow_v1_policy_snapshot_is_migrated_to_complete_runtime_policy() -> N ) assert resolved.arm_selection.pickup_crossing_weight == 2.0 - assert resolved.motion_defaults["PickUp"]["lift_height"] == 0.30 + assert resolved.motion_defaults["PickUp"]["lift_height"] == 0.16 def test_v3_policy_snapshot_is_migrated_with_default_planner_policy() -> None: @@ -278,7 +340,7 @@ def test_v3_policy_snapshot_is_migrated_with_default_planner_policy() -> None: } ) - assert resolved.schema_version == "action_engine_runtime_policy_v6" + assert resolved.schema_version == "action_engine_runtime_policy_v7" assert resolved.planner == expected.planner diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py index 13f7d7f7d..4f5116272 100644 --- a/tests/gen_sim/action_engine/generation/test_generation.py +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -669,15 +669,15 @@ def test_fast_gym_config_supports_all_robot_profiles( robot_uid: str, solver_type: str | None, ) -> None: - expected_tcp = [ + rotated_tcp = [ + [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.2], [0.0, 0.0, 0.0, 1.0], ] - expected_hand_mount = [ + identity_hand_mount = [ + [1.0, 0.0, 0.0, 0.0], [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], ] @@ -695,13 +695,13 @@ def test_fast_gym_config_supports_all_robot_profiles( 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 + assert config["robot"]["solver_cfg"][arm]["tcp"] == rotated_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 + assert components[hand]["transform"] == identity_hand_mount if solver_type is not None: assert config["robot"]["solver_cfg"]["left_arm"]["ur_type"] == solver_type @@ -998,7 +998,7 @@ def capture_writer(*args, **kwargs): 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" + "action_engine_runtime_policy_v7" ) assert agent_config["runtime_policy"]["planner"]["dynamic_collision"] is True assert agent_config["runtime_policy"]["planner"]["static_obstacle_uids"] == [ @@ -1414,7 +1414,7 @@ def test_agent_config_uses_relative_program_paths(gym_export: Path) -> None: assert config["runtime_policy"]["arm_selection"]["pickup_crossing_weight"] == 1.0 assert config["runtime_policy"]["motion_defaults"]["PickUp"][ "lift_height" - ] == pytest.approx(0.30) + ] == pytest.approx(0.16) assert config["runtime_policy"]["grasp"]["max_open_length"] == pytest.approx(0.15) assert len(config["runtime_policy_hash"]) == 64 diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index ba2d35951..e2ee2cb83 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -25,6 +25,9 @@ import torch from embodichain.gen_sim.action_engine.runtime import actions +from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ExactTargetMoveHeldObject, +) from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter from embodichain.gen_sim.action_engine.runtime.models import ( ActionOutcome, @@ -36,8 +39,11 @@ ActionBinding, ActionPlan, AntipodalAffordance, + AxisAlignAffordance, + AxisAlignGoal, CoordinatedPickGoal, EndEffectorPoseGoal, + EntityState, GraspGoal, HeldObjectState, JointPositionGoal, @@ -49,9 +55,10 @@ StateDelta, TimedCommandSequence, TimedTrajectory, + TrackingPolicy, ) from embodichain.lab.sim.planners import CuroboPlannerCfg -from embodichain.toolkits.graspkit.pg_grasp import GraspGeneratorCfg +from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator class _MeshEntity: @@ -102,6 +109,16 @@ def get_control_part_base_pose(self, *, name: str, to_matrix: bool) -> torch.Ten pose[:, 1, 3] = 0.3 if name == "physical_left_arm" else -0.3 return pose + def compute_batch_ik( + self, + *, + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + del pose, name + return torch.ones(joint_seed.shape[:2], dtype=torch.bool), joint_seed.clone() + def _commands_for(trajectory: TimedTrajectory) -> TimedCommandSequence: """Build timing-only frames for retained test trajectories.""" @@ -139,6 +156,157 @@ def plan(self, invocation, context) -> ActionPlan: return self._plan(invocation, context) +def test_adapter_registers_only_move_held_object_compat_action( + monkeypatch: Any, +) -> None: + registered: list[tuple[type, bool]] = [] + + class Engine: + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + pass + + def register(self, action: Any, *, replace: bool = False) -> None: + registered.append((type(action), replace)) + + adapter = AtomicActionAdapter(_planner_env()) + monkeypatch.setattr(actions, "AtomicActionEngine", Engine) + monkeypatch.setattr(adapter, "_generator", lambda: object()) + monkeypatch.setattr(adapter, "_control_profiles", lambda: {}) + monkeypatch.setattr(adapter, "_grasp_pose_generators", lambda: {}) + + engine = adapter._engine() + + assert isinstance(engine, Engine) + assert registered == [ + (ExactTargetMoveHeldObject, True), + ] + + +def test_adapter_lowers_axis_align_from_live_pose_with_a_stable_seed() -> None: + vertices = torch.tensor( + [[x, y, z] for x in (-0.03, 0.03) for y in (-0.06, 0.06) for z in (-0.03, 0.03)] + ) + semantics = ObjectSemantics( + label="can", + entity_id="can", + geometry={}, + affordance=AxisAlignAffordance( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2]]), + internal_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + ) + + live_pose = torch.eye(4).repeat(2, 1, 1) + live_pose[:, 2, 3] = 1.10 + parked_pose = live_pose.clone() + parked_pose[:, 2, 3] -= 100.0 + sampled_poses: list[torch.Tensor] = [] + sampled_random_values: list[float] = [] + + class Generator: + def get_valid_grasp_poses(self, **kwargs: Any): + poses = kwargs["obj_poses"].clone() + sampled_poses.append(poses) + sampled_random_values.append(float(torch.rand(()))) + return [ + (pose.unsqueeze(0), torch.tensor([0.1])) for pose in poses.unbind(dim=0) + ] + + adapter = AtomicActionAdapter(_planner_env()) + adapter._atomic_engine = SimpleNamespace( + grasp_pose_generators={"physical_left_eef": Generator()} + ) + grounded = GroundedAction( + "AxisAlign", + "left_arm", + "arm", + AxisAlignGoal(semantics=semantics), + {}, + object_pose=live_pose, + object_uid="can", + ) + contexts = [ + SimpleNamespace( + robot=SimpleNamespace(qpos=torch.zeros(2, 8)), + scene=SceneSnapshot( + timestamp=0.0, + version=version, + entities={"can": EntityState(parked_pose)}, + ), + ) + for version in (3, 97) + ] + + adaptations = [ + adapter._adapt_axis_align_body_grasps( + grounded, + context, + adapter.capabilities.get("AxisAlign"), + ) + for context in contexts + ] + adapted_items = adaptations[0] + adapted = adapted_items[0] + + assert len(adapted_items) == 1 + assert isinstance(adapted.target, AxisAlignGoal) + assert adapted.target.semantics.entity_id is None + assert adapted.target.grasp_xpos is not None + assert adapted.motion_policy["body_grasp"]["long_axis_index"] == 1 + assert adapted.motion_policy["body_grasp"]["candidate_counts"] == [1, 1] + assert len(sampled_poses) == 2 + torch.testing.assert_close(sampled_poses[0], live_pose) + torch.testing.assert_close(sampled_poses[1], live_pose) + assert not torch.equal(sampled_poses[0], parked_pose) + assert sampled_random_values[0] == sampled_random_values[1] + torch.testing.assert_close( + adaptations[0][0].target.grasp_xpos, + adaptations[1][0].target.grasp_xpos, + ) + + +def test_axis_align_body_grasp_does_not_fall_back_to_scene_snapshot_pose() -> None: + vertices = torch.tensor( + [[x, y, z] for x in (-0.03, 0.03) for y in (-0.06, 0.06) for z in (-0.03, 0.03)] + ) + grounded = GroundedAction( + "AxisAlign", + "left_arm", + "arm", + AxisAlignGoal( + semantics=ObjectSemantics( + label="can", + geometry={}, + affordance=AxisAlignAffordance( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2]]), + internal_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + ) + ), + {}, + object_uid="can", + ) + parked_pose = torch.eye(4).repeat(2, 1, 1) + parked_pose[:, 2, 3] = -98.9 + context = SimpleNamespace( + scene=SceneSnapshot( + timestamp=0.0, + version=3, + entities={"can": EntityState(parked_pose)}, + ) + ) + adapter = AtomicActionAdapter(_planner_env()) + + with pytest.raises(ValueError, match="grounded live object pose"): + adapter._adapt_axis_align_body_grasps( + grounded, + context, + adapter.capabilities.get("AxisAlign"), + ) + + def _planner_env( *, table: Any | None = None, @@ -190,8 +358,7 @@ def fake_prepare(**kwargs: Any) -> SimpleNamespace: def fake_affordance(**kwargs: Any) -> Affordance: events.append("affordance") - observed["generator_cfg"] = kwargs["generator_cfg"] - observed["gripper_collision_cfg"] = kwargs["gripper_collision_cfg"] + observed["affordance_kwargs"] = kwargs return Affordance() monkeypatch.setattr( @@ -206,12 +373,16 @@ def fake_affordance(**kwargs: Any) -> Affordance: second = adapter.semantics("cube") assert first is second + assert first.entity_id is None assert events == ["cache", "affordance"] assert observed["max_decomposition_hulls"] == 8 assert observed["mesh_vertices"].dtype == torch.float32 assert observed["mesh_triangles"].dtype == torch.int64 - assert observed["generator_cfg"].n_deviated_approach_directions == 4 - assert observed["gripper_collision_cfg"] is not None + assert set(observed["affordance_kwargs"]) == { + "object_label", + "mesh_vertices", + "mesh_triangles", + } def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None: @@ -257,11 +428,10 @@ def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None assert hand.motion_policy.strategy == "ik_interp" -def test_coordinated_pickment_scopes_ground_filter_to_gensim_goal_copy() -> None: +def test_coordinated_pickment_uses_engine_scoped_grasp_generator() -> None: adapter = AtomicActionAdapter(_planner_env()) adapter._atomic_engine = _FakeEngine() - original_cfg = GraspGeneratorCfg(is_filter_ground_collision=True) - affordance = AntipodalAffordance(generator_cfg=original_cfg) + affordance = AntipodalAffordance() goal = CoordinatedPickGoal( semantics=ObjectSemantics( label="tray", @@ -289,14 +459,25 @@ def test_coordinated_pickment_scopes_ground_filter_to_gensim_goal_copy() -> None scoped_affordance = invocation.goal.semantics.affordance assert isinstance(scoped_affordance, AntipodalAffordance) - assert scoped_affordance is not affordance - assert affordance.generator_cfg is original_cfg - assert original_cfg.is_filter_ground_collision is True - assert scoped_affordance.generator_cfg is not original_cfg - assert scoped_affordance.generator_cfg.is_filter_ground_collision is False + assert scoped_affordance is affordance assert invocation.skill_options.middle_empty_ratio == pytest.approx(0.7) +def test_grasp_generators_follow_mainline_service_contract() -> None: + adapter = AtomicActionAdapter(_planner_env()) + + generators = adapter._grasp_pose_generators() + + assert set(generators) == {"physical_left_eef", "physical_right_eef"} + generator = generators["physical_left_eef"] + assert isinstance(generator, AntipodalGraspPoseGenerator) + assert generator.algorithm_cfg.sample_count == 10000 + assert generator.algorithm_cfg.approach_direction_samples == 4 + assert generator.algorithm_cfg.max_candidates == 500 + assert generator.collision_cfg.max_decomposition_hulls == 16 + assert generator.collision_cfg.filter_ground_collision is True + + def test_retreat_uses_row_local_motion_planner_reachability_search( monkeypatch: Any, ) -> None: @@ -328,6 +509,7 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: commands=_commands_for(trajectory), joint_trajectory=trajectory, recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), planned_scene_version=0, planned_collision_world_revision=(0, 0), diagnostics=PlannerDiagnostics(backend="fake"), @@ -717,9 +899,13 @@ def action_plan( commands=_commands_for(trajectory), joint_trajectory=trajectory, recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), planned_scene_version=0, planned_collision_world_revision=(0, 0), - diagnostics=PlannerDiagnostics(backend="fake"), + diagnostics=PlannerDiagnostics( + backend="fake", + metadata={"marker": terminal}, + ), expected_effects=StateDelta( held_object_updates={"physical_left_arm": held} ), @@ -772,6 +958,8 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: assert torch.equal( outcome.planner_trace["fallback_used"], torch.tensor([False, True]) ) + assert outcome.planner_trace["primary_action_diagnostics"]["marker"] == 1.0 + assert outcome.planner_trace["fallback_action_diagnostics"]["marker"] == 2.0 def test_collision_required_cleanup_does_not_use_unsafe_fallback( @@ -796,6 +984,7 @@ def test_collision_required_cleanup_does_not_use_unsafe_fallback( commands=_commands_for(failed_trajectory), joint_trajectory=failed_trajectory, recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), planned_scene_version=1, planned_collision_world_revision=(1, 1), diagnostics=PlannerDiagnostics(backend="fake"), diff --git a/tests/gen_sim/action_engine/runtime/test_body_grasp.py b/tests/gen_sim/action_engine/runtime/test_body_grasp.py new file mode 100644 index 000000000..3eb9a8cb0 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_body_grasp.py @@ -0,0 +1,153 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for elongated-object axis analysis and body-grasp filtering.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime.body_grasp import ( + AxisAlignBodyGraspAdapter, + select_body_grasp_candidates, +) +from embodichain.gen_sim.action_engine.runtime.geometry_axes import ( + analyze_local_geometry_axes, +) +from embodichain.lab.sim.atomic_actions import ( + AxisAlignAffordance, + AxisAlignGoal, + ObjectSemantics, +) + + +def _box_vertices(extents: tuple[float, float, float]) -> torch.Tensor: + half = torch.tensor(extents) * 0.5 + return torch.tensor( + [ + [sx * half[0], sy * half[1], sz * half[2]] + for sx in (-1.0, 1.0) + for sy in (-1.0, 1.0) + for sz in (-1.0, 1.0) + ] + ) + + +def test_can_geometry_resolves_local_y_as_the_long_axis() -> None: + axes = analyze_local_geometry_axes(_box_vertices((0.0611, 0.1143, 0.0632))) + + assert axes.long_axis_index == 1 + assert axes.short_axis_index == 0 + torch.testing.assert_close(axes.long_axis, torch.tensor([0.0, 1.0, 0.0])) + assert axes.elongation_ratio == pytest.approx(0.1143 / 0.0632) + + +def test_axis_analysis_rejects_ambiguous_or_rotated_local_geometry() -> None: + with pytest.raises(ValueError, match="ambiguous"): + analyze_local_geometry_axes(_box_vertices((0.06, 0.06, 0.06))) + + vertices = _box_vertices((0.04, 0.12, 0.05)) + angle = math.radians(35.0) + rotation = torch.tensor( + [ + [math.cos(angle), -math.sin(angle), 0.0], + [math.sin(angle), math.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ] + ) + with pytest.raises(ValueError, match="not aligned"): + analyze_local_geometry_axes(vertices @ rotation.T) + + +def test_body_grasp_rejects_caps_and_longitudinal_closing() -> None: + axes = analyze_local_geometry_axes(_box_vertices((0.06, 0.12, 0.06))) + candidates = torch.eye(4).repeat(1, 3, 1, 1) + candidates[0, 0, 1, 3] = 0.055 # End-cap candidate with the best raw cost. + candidates[0, 1, 1, 3] = 0.0 # Central radial body grasp. + candidates[0, 2, 1, 3] = 0.0 + candidates[0, 2, :3, 0] = torch.tensor([0.0, 1.0, 0.0]) + candidates[0, 2, :3, 1] = torch.tensor([1.0, 0.0, 0.0]) + costs = torch.tensor([[0.0, 0.2, 0.1]]) + + selected = select_body_grasp_candidates( + candidates, + costs, + torch.eye(4).unsqueeze(0), + axes, + ) + + assert selected.success.tolist() == [True] + assert selected.candidate_indices.tolist() == [1] + assert selected.body_candidate_counts.tolist() == [1] + assert selected.ranked_candidate_indices.tolist() == [[1]] + torch.testing.assert_close(selected.grasp_xpos[0], candidates[0, 1]) + + +def test_body_grasp_chooses_a_reachable_body_candidate() -> None: + axes = analyze_local_geometry_axes(_box_vertices((0.06, 0.12, 0.06))) + candidates = torch.eye(4).repeat(1, 2, 1, 1) + candidates[0, 1, 0, 3] = 0.01 + costs = torch.tensor([[0.0, 0.2]]) + + selected = select_body_grasp_candidates( + candidates, + costs, + torch.eye(4).unsqueeze(0), + axes, + feasible=torch.tensor([[False, True]]), + ) + + assert selected.candidate_indices.tolist() == [1] + assert selected.reachable_candidate_counts.tolist() == [1] + + +def test_axis_align_adapter_injects_the_selected_body_grasp_unchanged() -> None: + vertices = _box_vertices((0.06, 0.12, 0.06)) + triangles = torch.tensor([[0, 1, 2], [0, 2, 3]]) + goal = AxisAlignGoal( + semantics=ObjectSemantics( + label="can", + geometry={}, + affordance=AxisAlignAffordance( + mesh_vertices=vertices, + mesh_triangles=triangles, + internal_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + ) + ) + candidate = torch.eye(4).unsqueeze(0) + + class Generator: + def get_valid_grasp_poses(self, **_kwargs): + return [(candidate, torch.tensor([0.1]))] + + adapted = AxisAlignBodyGraspAdapter().adapt( + goal, + object_pose=torch.eye(4).unsqueeze(0), + grasp_generator=Generator(), + approach_direction=torch.tensor([0.0, 0.0, -1.0]), + target_axis=torch.tensor([0.0, 0.0, 1.0]), + seed=7, + ) + + assert adapted.goal.grasp_xpos is not None + assert len(adapted.alternative_goals) == 1 + assert adapted.alternative_rank_indices == (0,) + explicit = adapted.goal.grasp_xpos + torch.testing.assert_close(explicit, candidate) diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index 81db74238..ede1101fe 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -94,6 +94,7 @@ CoordinatedPickGoal, CoordinatedPlacementGoal, CoordinatedPlacementOptions, + EndEffectorPoseGoal, HandOverOptions, HeldObjectPoseGoal, HeldObjectState, @@ -553,7 +554,7 @@ def test_dual_ur5_policy_uses_short_reach_upright_lifts() -> None: assert ur5_pickup["lift_height"] == pytest.approx(0.12) assert ur5_transport["staging_lift_height"] == pytest.approx(0.12) - assert ur10_pickup["lift_height"] == pytest.approx(0.30) + assert ur10_pickup["lift_height"] == pytest.approx(0.16) assert ur10_transport["staging_lift_height"] == pytest.approx(0.25) @@ -619,7 +620,7 @@ def test_runtime_policy_v4_migrates_grasp_direction_count() -> None: } ) - assert policy.schema_version == "action_engine_runtime_policy_v6" + assert policy.schema_version == "action_engine_runtime_policy_v7" assert policy.grasp["n_deviated_approach_directions"] == 4 @@ -665,7 +666,7 @@ def test_runtime_policy_v5_migrates_support_geometry_thresholds() -> None: } ) - assert policy.schema_version == "action_engine_runtime_policy_v6" + assert policy.schema_version == "action_engine_runtime_policy_v7" assert policy.predicate_fallbacks["support_min_overlap_ratio"] == 0.25 assert policy.grounding["placement"]["clearance"] == 0.019 assert policy.grounding["placement"]["candidate_count"] == 5 @@ -1659,7 +1660,9 @@ def _handover_then_place_task() -> dict[str, Any]: def test_handover_continuation_uses_stable_upright_policies() -> None: entities = { - "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "can": _FakeEntity( + "can", _pose(0.0, 0.2, 0.75), _rect_vertices(0.03, 0.03, 0.10) + ), "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), } @@ -3452,7 +3455,7 @@ def test_handover_defers_clearance_verification_to_retreat_action() -> None: def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() -> None: - entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) + entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _rect_vertices(0.10, 0.03, 0.03)) env = _FakeEnv( { "can": entity, @@ -3507,13 +3510,26 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - if candidate.id in orient_step.edge_ids and candidate.actions[0]["atomic_action_class"] == "AxisAlign" ) + orient_lift_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["target_binding"].get("operation") == "lift_clear" + ) + orient_retreat_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["target_binding"].get("operation") + == "retreat_after_lift" + ) handover_edge = next( candidate for candidate in program.edges if candidate.id in handover_step.edge_ids if candidate.actions[0]["atomic_action_class"] == "PickUp" ) - vertices = _box_vertices(0.03) + vertices = _rect_vertices(0.10, 0.03, 0.03) semantics = ObjectSemantics( affordance=AntipodalAffordance( object_label="can", @@ -3531,6 +3547,21 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - arm="right_arm", state=ExecutionState(last_qpos=env.robot.get_qpos()), ) + release_pose = _pose(0.05, 0.2, 0.78) + orient_lift = grounder.ground( + orient_lift_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + reference_eef_pose=release_pose, + ) + orient_retreat = grounder.ground( + orient_retreat_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + reference_eef_pose=orient_lift.target.xpos, + ) handover_pickup = grounder.ground( handover_edge.actions[0], handover_step, @@ -3542,6 +3573,33 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - assert "approach_direction_mode" not in handover_pickup.cfg assert handover_pickup.cfg["pick_object_part"] == "top" assert isinstance(orient_alignment.target, AxisAlignGoal) + assert isinstance(orient_lift.target, EndEffectorPoseGoal) + assert torch.equal( + orient_lift.target.xpos[:, :2, 3], + release_pose[:, :2, 3], + ) + assert bool((orient_lift.target.xpos[:, 2, 3] > release_pose[:, 2, 3]).all()) + assert "retreat_reachability_search" not in orient_lift.motion_policy + assert isinstance(orient_retreat.target, EndEffectorPoseGoal) + retreat_distance = torch.linalg.vector_norm( + orient_retreat.target.xpos[:, :2, 3] - orient_lift.target.xpos[:, :2, 3], + dim=1, + ) + torch.testing.assert_close( + retreat_distance, torch.full_like(retreat_distance, 0.20) + ) + assert bool( + (orient_retreat.target.xpos[:, 0, 3] > orient_lift.target.xpos[:, 0, 3]).all() + ) + assert torch.equal( + orient_retreat.target.xpos[:, 2, 3], + orient_lift.target.xpos[:, 2, 3], + ) + assert bool( + ( + orient_retreat.target.xpos[:, :2, 3] != orient_lift.target.xpos[:, :2, 3] + ).any() + ) assert orient_alignment.target.grasp_xpos is None assert torch.equal( orient_alignment.target.object_target_pose, @@ -4046,6 +4104,35 @@ def plan(grounded: GroundedAction, state: ExecutionState) -> ActionOutcome: assert executor._object_owners["can"] == [None] +def test_completed_step_releases_speculative_candidate_plans() -> None: + executor = object.__new__(ProgramExecutor) + executor._candidate_cache = { + ("completed", "left_arm"): object(), + ("next", "right_arm"): object(), + } + executor._candidate_failures = { + ("completed", "left_arm"): "failed", + ("next", "right_arm"): "pending", + } + executor._candidate_diagnostics = { + "completed": {"large": "trace"}, + "next": {"small": "trace"}, + } + executor._candidate_blockers = { + "completed": ({"reason": "old"},), + "next": ({"reason": "new"},), + } + executor._reported_candidates = {"completed", "next"} + + executor._release_candidate_plans("completed") + + assert set(executor._candidate_cache) == {("next", "right_arm")} + assert set(executor._candidate_failures) == {("next", "right_arm")} + assert set(executor._candidate_diagnostics) == {"next"} + assert set(executor._candidate_blockers) == {"next"} + assert executor._reported_candidates == {"next"} + + def test_live_pickup_planning_exception_is_a_retryable_edge_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/gen_sim/action_engine/task_fixtures.py b/tests/gen_sim/action_engine/task_fixtures.py index c845cdfab..f7308dc1a 100644 --- a/tests/gen_sim/action_engine/task_fixtures.py +++ b/tests/gen_sim/action_engine/task_fixtures.py @@ -30,7 +30,38 @@ TASK_SPEC_SCHEMA, ) -__all__ = ["make_task_level", "make_task_spec"] +__all__ = [ + "TASK2_1_HISTORICAL_ROLE_BINDINGS", + "TASK2_1_HISTORICAL_SCENE_FINGERPRINT", + "make_task2_1_historical_spec", + "make_task_level", + "make_task_spec", +] + + +TASK2_1_HISTORICAL_ROLE_BINDINGS = { + "object_01": "interact_purple_soda_can", + "object_02": "interact_orange_soda_can", + "object_03": "interact_spiral_notebook", +} +"""Runtime role bindings from clean v14 run ``20260820_233507``.""" + +TASK2_1_HISTORICAL_SCENE_FINGERPRINT = { + "config_sha256": "1042967c9b7021f518e82ace62aa824015a3ad50639fa8a326b7dc0474277481", + "asset_sha256": { + "interact_orange_soda_can": ( + "ae6c8b9922d4a20746241daf0a607d29f33eb5d96f3dfa2244ffa6ba1d89f5ce" + ), + "interact_purple_soda_can": ( + "ee3c0d53d298f2be2db778d8a1227122b57b9492651e1825cd7335a8bf4cec42" + ), + "interact_spiral_notebook": ( + "25cbbf49e89da935c32ea939e523997ca538f09c4a29cff334dcbf785380afc8" + ), + "table": "99caec34e31d43e34f9326fb16c6e8660288d24e0482f463ac6d90c99368b76a", + }, +} +"""Path-independent source fingerprint for the historical Task 2-1 scene.""" _OBJECT_FIXTURES = { "E1": ("can", ["graspable", "placeable"], {}), @@ -45,6 +76,193 @@ } +def make_task2_1_historical_spec() -> dict[str, Any]: + """Build the deterministic ten-step Task 2-1 behavior from clean v14. + + The fixture preserves task semantics and ownership transitions from commit + ``f70138c6`` run ``20260820_233507``. It intentionally does not freeze the + Atomic Action node count or the v14 E2 lowering topology. + """ + instances = [ + { + "id": "task_01", + "task_type": "E2", + "params": { + "object_role": "object_01", + "required_arm": "right_arm", + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E2", + "params": { + "object_role": "object_02", + "required_arm": "left_arm", + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_03", + "task_type": "E4", + "params": { + "object_role": "object_02", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "none", + }, + "depends_on": ["task_02"], + "role": "primary", + }, + { + "id": "task_04", + "task_type": "E1", + "params": { + "object_role": "object_02", + "target_role": "object_01", + "relation": "behind", + "relation_frame": "robot", + "required_arm": "right_arm", + "orientation_goal": "none", + "orientation_axis": "none", + }, + "depends_on": ["task_03"], + "role": "primary", + }, + { + "id": "task_05", + "task_type": "E4", + "params": { + "object_role": "object_01", + "transfer_arm": "right_arm", + "receive_arm": "left_arm", + "orientation_goal": "none", + }, + "depends_on": ["task_04", "task_01"], + "role": "primary", + }, + { + "id": "task_06", + "task_type": "E1", + "params": { + "object_role": "object_01", + "target_role": "object_03", + "relation": "left_of", + "relation_frame": "robot", + "required_arm": "left_arm", + "orientation_goal": "none", + "orientation_axis": "none", + }, + "depends_on": ["task_05"], + "role": "primary", + }, + { + "id": "task_07", + "task_type": "E1", + "params": { + "object_role": "object_02", + "target_role": "object_03", + "relation": "front_of", + "relation_frame": "robot", + "required_arm": "right_arm", + "orientation_goal": "none", + "orientation_axis": "none", + }, + "depends_on": ["task_04"], + "role": "primary", + }, + { + "id": "task_08", + "task_type": "E1", + "params": { + "object_role": "object_02", + "target_role": "object_03", + "relation": "on", + "relation_frame": "robot", + "required_arm": "left_arm", + "orientation_goal": "none", + "orientation_axis": "none", + }, + "depends_on": ["task_07"], + "role": "primary", + }, + { + "id": "task_09", + "task_type": "E4", + "params": { + "object_role": "object_01", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "orientation_goal": "none", + }, + "depends_on": ["task_06"], + "role": "primary", + }, + { + "id": "task_10", + "task_type": "E1", + "params": { + "object_role": "object_01", + "target_role": "object_02", + "relation": "above", + "relation_frame": "robot", + "required_arm": "right_arm", + "orientation_goal": "none", + "orientation_axis": "none", + }, + "depends_on": ["task_09", "task_08"], + "role": "primary", + }, + ] + success_types = ( + "object_upright", + "object_upright", + "handover_complete", + "semantic_goal", + "handover_complete", + "semantic_goal", + "semantic_goal", + "semantic_goal", + "handover_complete", + "semantic_goal", + ) + return validate_task_spec( + { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "task2_1", + "level": "L3", + "instruction": "historical-task2_1-ten-step-regression", + "reasoning_type": "none", + "task_instances": instances, + "success": { + "op": "all", + "terms": [ + {"type": success_type, "task_instance_id": instance["id"]} + for instance, success_type in zip(instances, success_types) + ], + }, + "oracle": { + "task_order": [instance["id"] for instance in instances], + "role_bindings": dict(TASK2_1_HISTORICAL_ROLE_BINDINGS), + }, + "metadata": { + "fixture": True, + "historical_commit": "f70138c626daf84918b15b954765493000cb40a5", + "historical_run": "20260820_233507", + "role_bindings": dict(TASK2_1_HISTORICAL_ROLE_BINDINGS), + }, + } + ) + + def make_task_spec( task_type: str = "E1", *, diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index b9af243ee..45718bbe1 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -24,6 +24,11 @@ ground_instruction_draft, instantiate_seed_graph, ) +from tests.gen_sim.action_engine.task_fixtures import ( + TASK2_1_HISTORICAL_ROLE_BINDINGS, + TASK2_1_HISTORICAL_SCENE_FINGERPRINT, + make_task2_1_historical_spec, +) def _selector( @@ -87,6 +92,102 @@ def _ground_draft( ) +def _historical_task2_1_graph() -> dict: + return instantiate_seed_graph( + make_task2_1_historical_spec(), + TASK2_1_HISTORICAL_ROLE_BINDINGS, + ) + + +def _actions_by_task_group(graph: dict) -> dict[str, list[str]]: + nodes = {node["id"]: node for node in graph["nodes"]} + return { + group["id"]: [nodes[node_id]["atomic_action"] for node_id in group["node_ids"]] + for group in graph["task_groups"] + } + + +def test_historical_task2_1_fixture_preserves_ten_step_semantics() -> None: + task = make_task2_1_historical_spec() + + assert [instance["task_type"] for instance in task["task_instances"]] == [ + "E2", + "E2", + "E4", + "E1", + "E4", + "E1", + "E1", + "E1", + "E4", + "E1", + ] + assert [term["task_instance_id"] for term in task["success"]["terms"]] == [ + instance["id"] for instance in task["task_instances"] + ] + assert TASK2_1_HISTORICAL_SCENE_FINGERPRINT["config_sha256"] == ( + "1042967c9b7021f518e82ace62aa824015a3ad50639fa8a326b7dc0474277481" + ) + + +def test_historical_task2_1_uses_axis_align_and_explicit_handover_arms() -> None: + task = make_task2_1_historical_spec() + graph = _historical_task2_1_graph() + actions = _actions_by_task_group(graph) + + expected_orient_actions = [ + "AxisAlign", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + ] + assert actions["task_01"] == expected_orient_actions + assert actions["task_02"] == expected_orient_actions + assert [ + ( + instance["params"]["transfer_arm"], + instance["params"]["receive_arm"], + ) + for instance in task["task_instances"] + if instance["task_type"] == "E4" + ] == [ + ("left_arm", "right_arm"), + ("right_arm", "left_arm"), + ("left_arm", "right_arm"), + ] + + +def test_historical_task2_1_handover_continuations_preserve_receiver_hold() -> None: + graph = _historical_task2_1_graph() + actions = _actions_by_task_group(graph) + groups = {group["id"]: group for group in graph["task_groups"]} + + for group_id, expected_arm in ( + ("task_04", "right_arm"), + ("task_06", "left_arm"), + ("task_10", "right_arm"), + ): + assert actions[group_id][0] == "MoveHeldObject" + assert "PickUp" not in actions[group_id] + assert groups[group_id]["contract"]["entry_requires"] == [ + { + "predicate": "object_held", + "object_uid": groups[group_id]["object_uid"], + "arm": expected_arm, + } + ] + + +def test_historical_task2_1_reacquires_objects_after_release() -> None: + graph = _historical_task2_1_graph() + actions = _actions_by_task_group(graph) + + assert actions["task_05"][0] == "PickUp" + assert actions["task_07"][0] == "PickUp" + assert actions["task_08"][0] == "PickUp" + assert actions["task_09"][0] == "PickUp" + + def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() -> None: task = { "schema_version": "action_engine_task_spec_v2", @@ -131,8 +232,38 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - ] orient = next(group for group in graph["task_groups"] if group["id"] == "orient") - assert [node["atomic_action"] for node in orient_nodes] == ["AxisAlign"] + assert [node["atomic_action"] for node in orient_nodes] == [ + "AxisAlign", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + ] assert orient_nodes[0]["motion_policy"] == {"modifiers": []} + assert [node["role"] for node in orient_nodes] == [ + "primary", + "cleanup", + "cleanup", + "cleanup", + ] + assert orient_nodes[1]["target_binding"] == { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + } + assert orient_nodes[1]["motion_policy"] == {"modifiers": []} + assert orient_nodes[2]["target_binding"] == { + "kind": "policy_pose", + "source": "release", + "operation": "retreat_after_lift", + } + assert orient_nodes[3]["target_binding"] == { + "kind": "joint_state", + "source": "initial", + "operation": "e2_home", + } + assert orient_nodes[1]["depends_on"] == [orient_nodes[0]["id"]] + assert orient_nodes[2]["depends_on"] == [orient_nodes[1]["id"]] + assert orient_nodes[3]["depends_on"] == [orient_nodes[2]["id"]] assert [node["atomic_action"] for node in handover_nodes] == [ "PickUp", "MoveHeldObject", @@ -147,8 +278,17 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - 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[-1]["contract"]["failure_policy"] == "task_required" - assert not any( + assert orient_nodes[0]["contract"]["failure_policy"] == "task_required" + assert orient_nodes[1]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[2]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[-1]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[1]["contract"]["requires"] == [ + {"predicate": "arm_free", "arm": "left_arm"} + ] + assert orient_nodes[2]["contract"]["requires"] == [ + {"predicate": "arm_clear", "arm": "left_arm"} + ] + assert any( effect["atom"]["predicate"] == "arm_home" for effect in orient["contract"]["exit_effects"] ) diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index 3cc64182b..9c7541ed9 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -317,7 +317,12 @@ def caller(**kwargs): 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 == ["AxisAlign"] + assert orient_actions == [ + "AxisAlign", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + ] assert [node["atomic_action"] for node in handover_nodes] == [ "PickUp", "MoveHeldObject", diff --git a/tests/gen_sim/action_engine/test_motion_policy.py b/tests/gen_sim/action_engine/test_motion_policy.py index aa253c6fb..5d39dd66a 100644 --- a/tests/gen_sim/action_engine/test_motion_policy.py +++ b/tests/gen_sim/action_engine/test_motion_policy.py @@ -29,8 +29,8 @@ def test_upright_policy_matches_mature_runtime_across_robot_profiles() -> None: franka = resolve_motion_policy("dual_franka", "PickUp", upright) ur10 = resolve_motion_policy("dual_ur10", "PickUp", upright) - assert franka["lift_height"] == pytest.approx(0.30) - assert ur10["lift_height"] == pytest.approx(0.30) + assert franka["lift_height"] == pytest.approx(0.16) + assert ur10["lift_height"] == pytest.approx(0.16) assert ur10["rotate_upright"] == pytest.approx(0.7853981633974483) From e8a724af6f9161cd5721cae7568523d6407b1419 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Tue, 25 Aug 2026 20:32:07 +0800 Subject: [PATCH 66/85] feat(action-engine): add body-grasp planning and verified cleanup for E2 axis alignment --- .../action_engine/capabilities/__init__.py | 3 + .../action_engine/capabilities/atomic.py | 7 +- .../action_engine/capabilities/builtins.py | 12 +- .../capabilities/held_hand_over.py | 660 ++++++++++++++++++ .../action_engine/config/defaults.yaml | 2 - .../gen_sim/action_engine/orientation.py | 29 +- .../gen_sim/action_engine/runtime/actions.py | 51 +- .../action_engine/runtime/atomic_compat.py | 41 +- .../action_engine/runtime/grounding.py | 51 +- .../gen_sim/action_engine/runtime/models.py | 2 + .../action_engine/runtime/predicates.py | 23 +- .../capabilities/test_held_hand_over.py | 311 +++++++++ .../config/test_runtime_policy.py | 8 + .../action_engine/runtime/test_actions.py | 78 ++- .../runtime/test_atomic_compat.py | 28 +- .../runtime/test_runtime_contracts.py | 122 +++- .../gen_sim/action_engine/test_orientation.py | 62 +- 17 files changed, 1355 insertions(+), 135 deletions(-) create mode 100644 embodichain/gen_sim/action_engine/capabilities/held_hand_over.py create mode 100644 tests/gen_sim/action_engine/capabilities/test_held_hand_over.py diff --git a/embodichain/gen_sim/action_engine/capabilities/__init__.py b/embodichain/gen_sim/action_engine/capabilities/__init__.py index 5e7b6935e..b23dc8867 100644 --- a/embodichain/gen_sim/action_engine/capabilities/__init__.py +++ b/embodichain/gen_sim/action_engine/capabilities/__init__.py @@ -30,6 +30,7 @@ capability_precondition, ) from .builtins import build_default_registry +from .held_hand_over import HeldObjectHandOver, HeldObjectHandOverOptions from .registry import ( ActionCapability, ActionTemplate, @@ -45,6 +46,8 @@ "AtomicCapability", "AtomicCapabilityRegistry", "CapabilityRegistry", + "HeldObjectHandOver", + "HeldObjectHandOverOptions", "OperatorCapability", "PhaseTemplate", "ResolvedActionContract", diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index d76e83e05..8d3d58657 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -352,8 +352,6 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: CoordinatedPlacement, CoordinatedPlacementOptions, AxisAlignOptions, - HandOver, - HandOverOptions, MoveEndEffector, MoveEndEffectorOptions, MoveHeldObject, @@ -373,6 +371,7 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: Twist, TwistOptions, ) + from .held_hand_over import HeldObjectHandOver, HeldObjectHandOverOptions registry = AtomicCapabilityRegistry() definitions = ( @@ -545,8 +544,8 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: ), AtomicCapability( "HandOver", - HandOver, - HandOverOptions, + HeldObjectHandOver, + HeldObjectHandOverOptions, frozenset({"handover_goal"}), frozenset({"coordinated"}), "coordinated_object", diff --git a/embodichain/gen_sim/action_engine/capabilities/builtins.py b/embodichain/gen_sim/action_engine/capabilities/builtins.py index 839797e37..1af8b9950 100644 --- a/embodichain/gen_sim/action_engine/capabilities/builtins.py +++ b/embodichain/gen_sim/action_engine/capabilities/builtins.py @@ -346,7 +346,6 @@ def _expand_hold_hover(step: Mapping[str, Any]) -> list[dict[str, Any]]: orientation_goal, orientation_axis = _orientation( goal, "hold_hover", - allow_change=False, ) reference = str(goal.get("reference_object", object_uid)) return [ @@ -984,17 +983,10 @@ def _required_string( def _orientation( goal: Mapping[str, Any], operator: str, - *, - allow_change: bool = True, ) -> tuple[str, str]: - default_goal = "none" if allow_change else "preserve" - orientation_goal = str(goal.get("orientation_goal", default_goal)) + orientation_goal = str(goal.get("orientation_goal", "none")) orientation_axis = str(goal.get("orientation_axis", "none")) - allowed_goals = ( - {"none", "preserve", "upright", "lay_flat", "axis_align"} - if allow_change - else {"preserve"} - ) + allowed_goals = {"none", "preserve", "upright", "lay_flat", "axis_align"} if orientation_goal not in allowed_goals: raise ValueError( f"{operator} orientation_goal {orientation_goal!r} is unsupported." diff --git a/embodichain/gen_sim/action_engine/capabilities/held_hand_over.py b/embodichain/gen_sim/action_engine/capabilities/held_hand_over.py new file mode 100644 index 000000000..d269dcb4f --- /dev/null +++ b/embodichain/gen_sim/action_engine/capabilities/held_hand_over.py @@ -0,0 +1,660 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""GenSim-local handover for transferring an already-held object.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionOptions, + ActionPlan, + AntipodalAffordance, + AtomicAction, + CARTESIAN_POSE_CAPABILITY, + DisjointResourceSlots, + FORWARD_KINEMATICS_CAPABILITY, + GRASP_COMMAND, + GraspGoal, + HeldObjectState, + JointPositionCommand, + JointPositionTarget, + ObjectSemantics, + OPEN_COMMAND, + PlanningContext, + ResolvedActionRequest, + SkillBindingContract, + StateDelta, + TimedTrajectory, +) +from embodichain.lab.sim.atomic_actions.goals import ( + PoseGoalValue, + collect_scene_dependencies, + resolve_pose_goal, + validate_pose_goal, +) +from embodichain.lab.sim.atomic_actions.primitives._binding_contracts import ( + make_manipulation_slot, +) +from embodichain.lab.sim.atomic_actions.primitives._helpers import ( + assemble_full_robot_trajectory, + plan_named_arm_trajectory, + repeat_qpos, + require_shared_task_state_key, + resolve_batched_pose, +) +from embodichain.lab.sim.atomic_actions.trajectory_ops import ( + interpolate_hand_qpos, + translate_pose_world, +) +from embodichain.lab.sim.planners.utils import normalize_success_mask +from embodichain.utils.math import pose_inv + + +@dataclass(frozen=True, slots=True, eq=False) +class HeldObjectHandOverOptions(ActionOptions): + """Per-invocation behavior for transferring an already-held object.""" + + receive_pick_object_part: str = "bottom" + middle_object_pose: PoseGoalValue | None = None + final_object_pose: PoseGoalValue | None = None + receive_approach_direction: torch.Tensor = torch.tensor([0.0, 0.0, -1.0]) + pre_grasp_distance: float = 0.10 + lift_height: float = 0.08 + hand_interp_steps: int = 10 + hold_steps: int = 4 + retreat_steps: int = 24 + + def __post_init__(self) -> None: + if self.receive_pick_object_part not in frozenset({"center", "top", "bottom"}): + raise ValueError( + "receive_pick_object_part must be 'center', 'top', or 'bottom'." + ) + direction = self.receive_approach_direction + if ( + not isinstance(direction, torch.Tensor) + or direction.shape != (3,) + or not torch.isfinite(direction).all() + or torch.linalg.vector_norm(direction) <= 1.0e-6 + ): + raise ValueError( + "receive_approach_direction must be a finite non-zero (3,) tensor." + ) + if self.pre_grasp_distance < 0.0 or self.lift_height < 0.0: + raise ValueError("Handover distances must be non-negative.") + for name in ("hand_interp_steps", "hold_steps", "retreat_steps"): + value = getattr(self, name) + if type(value) is not int or value < 0: + raise ValueError(f"{name} must be a non-negative integer.") + + object.__setattr__(self, "receive_approach_direction", direction.clone()) + for name in ("middle_object_pose", "final_object_pose"): + value = getattr(self, name) + if value is None: + continue + validate_pose_goal(value, name, allow_waypoints=False) + object.__setattr__( + self, + name, + value.clone() if isinstance(value, torch.Tensor) else value.snapshot(), + ) + + +@dataclass(frozen=True, slots=True) +class _HandoverResources: + source_state_key: str + destination_state_key: str + source_arm: JointPositionTarget + destination_arm: JointPositionTarget + source_hand: JointPositionTarget + destination_hand: JointPositionTarget + source_hand_open_qpos: torch.Tensor + source_hand_grasp_qpos: torch.Tensor + destination_hand_open_qpos: torch.Tensor + destination_hand_grasp_qpos: torch.Tensor + + +class HeldObjectHandOver(AtomicAction[GraspGoal, HeldObjectHandOverOptions]): + """Transfer an existing attachment while leaving the receiver holding it.""" + + skill_id: ClassVar[str] = "hand_over" + GoalType: ClassVar[type] = GraspGoal + OptionsType: ClassVar[type] = HeldObjectHandOverOptions + manipulator_roles: ClassVar[tuple[str, ...]] = ("source", "destination") + end_effector_roles: ClassVar[tuple[str, ...]] = ("source", "destination") + open_loop: ClassVar[bool] = True + binding_contract: ClassVar[SkillBindingContract] = SkillBindingContract( + slots=( + make_manipulation_slot( + "source", + motion_capabilities=frozenset( + {CARTESIAN_POSE_CAPABILITY, FORWARD_KINEMATICS_CAPABILITY} + ), + grasp_commands={ + OPEN_COMMAND: JointPositionCommand, + GRASP_COMMAND: JointPositionCommand, + }, + ), + make_manipulation_slot( + "destination", + motion_capabilities=frozenset({CARTESIAN_POSE_CAPABILITY}), + grasp_commands={ + OPEN_COMMAND: JointPositionCommand, + GRASP_COMMAND: JointPositionCommand, + }, + ), + ), + constraints=(DisjointResourceSlots(("source", "destination")),), + ) + + def __init__( + self, default_options: HeldObjectHandOverOptions | None = None + ) -> None: + super().__init__(default_options) + + def _scene_dependencies( + self, + request: ResolvedActionRequest[GraspGoal, HeldObjectHandOverOptions], + ) -> tuple[str, ...]: + """Return scene entities whose poses materially affect this plan.""" + return collect_scene_dependencies( + tuple( + value + for value in ( + request.skill_options.middle_object_pose, + request.skill_options.final_object_pose, + ) + if value is not None + ) + ) + + def _plan( + self, + request: ResolvedActionRequest[GraspGoal, HeldObjectHandOverOptions], + context: PlanningContext, + ) -> ActionPlan: + goal = self.require_goal(request) + options = request.skill_options + self._require_exchange_pose(options) + resources = self._resolve_resources(request) + + if ( + request.motion_policy.strategy == "motion_gen" + and self.motion_generator.planner.cfg.planner_type == "curobo" + ): + raise ValueError( + "Coordinated dual-arm planning is not supported by cuRobo." + ) + + held = context.get_held_object(resources.source_state_key) + if held is None: + raise ValueError( + "HeldObjectHandOver requires the source participant to hold an object." + ) + self._require_same_object(goal.semantics, held.semantics) + eligible = context.task.exclusive_held_object_mask(resources.source_state_key) + if not eligible.any(): + return self.failed_plan( + request, + context, + message="Source object must be held exclusively.", + ) + + source_start, destination_start = self._start_qpos(context, resources) + object_to_source = self._pose(held.object_to_eef, "held.object_to_eef") + source_eef = self.robot.compute_fk( + qpos=source_start, + name=resources.source_arm.control_part, + to_matrix=True, + ) + current_object_pose = torch.bmm(source_eef, pose_inv(object_to_source)) + assert options.middle_object_pose is not None + assert options.final_object_pose is not None + middle_object_pose = self._pose( + resolve_pose_goal( + options.middle_object_pose, + context, + name="middle_object_pose", + ), + "middle_object_pose", + ) + final_object_pose = self._pose( + resolve_pose_goal( + options.final_object_pose, + context, + name="final_object_pose", + ), + "final_object_pose", + ) + middle_object_pose[:, :3, :3] = current_object_pose[:, :3, :3] + final_object_pose[:, :3, :3] = current_object_pose[:, :3, :3] + if not torch.allclose( + middle_object_pose, + final_object_pose, + atol=1.0e-5, + rtol=1.0e-5, + ): + raise ValueError( + "HeldObjectHandOver requires final_object_pose to match the exchange " + "pose so the receiver remains stationary." + ) + + source_middle_eef = torch.bmm(middle_object_pose, object_to_source) + destination_grasp, grasp_success = self._destination_grasp( + held.semantics, + middle_object_pose, + resources.destination_hand.control_part, + options, + ) + success = normalize_success_mask( + grasp_success, + num_envs=self.num_envs, + device=self.device, + name="Receiving-grasp success", + ) + success &= eligible + if not success.any(): + return self.failed_plan( + request, + context, + message="No receiving grasp was available.", + ) + + object_to_destination = torch.bmm( + pose_inv(middle_object_pose), destination_grasp + ) + destination_pre_grasp = translate_pose_world( + destination_grasp, + -destination_grasp[:, :3, 2] * options.pre_grasp_distance, + ) + source_retreat_eef = translate_pose_world( + source_middle_eef, + source_middle_eef.new_tensor([0.0, 0.0, options.lift_height]), + ) + lengths = self._segment_lengths(request.motion_policy.sample_count, options) + + segment_success, source_transfer = plan_named_arm_trajectory( + self.motion_generator, + resources.source_arm.control_part, + source_start, + source_middle_eef.unsqueeze(1), + lengths["transfer"], + request.motion_policy, + context.control_dt, + ) + success &= self._success(segment_success, "Source transfer") + segment_success, destination_approach = plan_named_arm_trajectory( + self.motion_generator, + resources.destination_arm.control_part, + destination_start, + torch.stack((destination_pre_grasp, destination_grasp), dim=1), + lengths["approach"], + request.motion_policy, + context.control_dt, + ) + success &= self._success(segment_success, "Destination approach") + source_hold = source_transfer[:, -1] + destination_hold = destination_approach[:, -1] + segment_success, source_retreat = plan_named_arm_trajectory( + self.motion_generator, + resources.source_arm.control_part, + source_hold, + source_retreat_eef.unsqueeze(1), + lengths["retreat"], + request.motion_policy, + context.control_dt, + ) + success &= self._success(segment_success, "Source retreat") + if not success.any(): + return self.failed_plan( + request, + context, + message="Handover arm planning failed.", + ) + + segments = [ + ( + "transfer", + self._segment( + context, + resources, + source_transfer, + repeat_qpos(destination_start, lengths["transfer"]), + repeat_qpos(resources.source_hand_grasp_qpos, lengths["transfer"]), + repeat_qpos( + resources.destination_hand_open_qpos, lengths["transfer"] + ), + ), + ), + ( + "approach", + self._segment( + context, + resources, + repeat_qpos(source_hold, lengths["approach"]), + destination_approach, + repeat_qpos(resources.source_hand_grasp_qpos, lengths["approach"]), + repeat_qpos( + resources.destination_hand_open_qpos, lengths["approach"] + ), + ), + ), + ( + "close", + self._segment( + context, + resources, + repeat_qpos(source_hold, lengths["close"]), + repeat_qpos(destination_hold, lengths["close"]), + repeat_qpos(resources.source_hand_grasp_qpos, lengths["close"]), + interpolate_hand_qpos( + resources.destination_hand_open_qpos, + resources.destination_hand_grasp_qpos, + n_waypoints=lengths["close"], + ), + ), + ), + ] + if lengths["hold"]: + segments.append( + ( + "hold", + self._segment( + context, + resources, + repeat_qpos(source_hold, lengths["hold"]), + repeat_qpos(destination_hold, lengths["hold"]), + repeat_qpos(resources.source_hand_grasp_qpos, lengths["hold"]), + repeat_qpos( + resources.destination_hand_grasp_qpos, lengths["hold"] + ), + ), + ) + ) + segments.extend( + ( + ( + "release", + self._segment( + context, + resources, + repeat_qpos(source_hold, lengths["release"]), + repeat_qpos(destination_hold, lengths["release"]), + interpolate_hand_qpos( + resources.source_hand_grasp_qpos, + resources.source_hand_open_qpos, + n_waypoints=lengths["release"], + ), + repeat_qpos( + resources.destination_hand_grasp_qpos, + lengths["release"], + ), + ), + ), + ( + "retreat", + self._segment( + context, + resources, + source_retreat, + repeat_qpos(destination_hold, lengths["retreat"]), + repeat_qpos( + resources.source_hand_open_qpos, lengths["retreat"] + ), + repeat_qpos( + resources.destination_hand_grasp_qpos, + lengths["retreat"], + ), + ), + ), + ) + ) + + trajectory = torch.cat([value for _, value in segments], dim=1) + received = HeldObjectState( + semantics=held.semantics, + object_to_eef=object_to_destination, + grasp_xpos=destination_grasp, + ) + return self.build_plan( + request, + context, + success=success, + trajectory=TimedTrajectory.from_uniform_step( + trajectory, + env_ids=context.env_ids, + step_dt=context.require_control_dt(), + ), + expected_effects=StateDelta( + held_object_updates={ + resources.source_state_key: None, + resources.destination_state_key: received, + } + ), + segment_lengths={name: value.shape[1] for name, value in segments}, + ) + + def _resolve_resources( + self, + request: ResolvedActionRequest[GraspGoal, HeldObjectHandOverOptions], + ) -> _HandoverResources: + binding = request.binding + source_motion = binding.endpoint("source", "motion") + source_grasp = binding.endpoint("source", "grasp") + destination_motion = binding.endpoint("destination", "motion") + destination_grasp = binding.endpoint("destination", "grasp") + source_arm = source_motion.require_target(JointPositionTarget) + source_hand = source_grasp.require_target(JointPositionTarget) + destination_arm = destination_motion.require_target(JointPositionTarget) + destination_hand = destination_grasp.require_target(JointPositionTarget) + source_key = require_shared_task_state_key( + source_motion, + source_grasp, + participant="HeldObjectHandOver source", + ) + destination_key = require_shared_task_state_key( + destination_motion, + destination_grasp, + participant="HeldObjectHandOver destination", + ) + if source_key == destination_key: + raise ValueError("Handover participants require different state keys.") + return _HandoverResources( + source_state_key=source_key, + destination_state_key=destination_key, + source_arm=source_arm, + destination_arm=destination_arm, + source_hand=source_hand, + destination_hand=destination_hand, + source_hand_open_qpos=source_grasp.joint_positions( + OPEN_COMMAND, + num_envs=self.num_envs, + device=self.device, + dtype=torch.float32, + ), + source_hand_grasp_qpos=source_grasp.joint_positions( + GRASP_COMMAND, + num_envs=self.num_envs, + device=self.device, + dtype=torch.float32, + ), + destination_hand_open_qpos=destination_grasp.joint_positions( + OPEN_COMMAND, + num_envs=self.num_envs, + device=self.device, + dtype=torch.float32, + ), + destination_hand_grasp_qpos=destination_grasp.joint_positions( + GRASP_COMMAND, + num_envs=self.num_envs, + device=self.device, + dtype=torch.float32, + ), + ) + + def _destination_grasp( + self, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + grasp_target_id: str, + options: HeldObjectHandOverOptions, + ) -> tuple[torch.Tensor, torch.Tensor]: + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise ValueError("HeldObjectHandOver requires AntipodalAffordance.") + direction = options.receive_approach_direction.to( + device=self.device, dtype=torch.float32 + ) + direction = direction / torch.linalg.vector_norm(direction) + direction = direction.expand(self.num_envs, -1) + axis = None + positive: bool | torch.Tensor = True + if options.receive_pick_object_part != "center": + local_axis = object_pose.new_tensor([0.0, 0.0, 1.0]) + axis = torch.matmul(object_pose[:, :3, :3], local_axis) + positive = torch.full( + (self.num_envs,), + options.receive_pick_object_part == "top", + dtype=torch.bool, + device=self.device, + ) + + generator = self.planning_services.grasp_pose_generator(grasp_target_id) + sampled = generator.get_valid_grasp_poses( + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + obj_poses=object_pose, + approach_direction=direction, + obj_longest_axis=axis, + is_positive_part=positive, + ) + poses = torch.eye(4, dtype=torch.float32, device=self.device).repeat( + self.num_envs, 1, 1 + ) + success = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + for env_index, (candidates, costs) in enumerate(sampled): + candidates = candidates.to(device=self.device, dtype=torch.float32) + costs = costs.to(device=self.device, dtype=torch.float32) + finite = torch.isfinite(costs) + if candidates.shape[0] == 0 or not finite.any(): + continue + ranked = torch.where(finite, costs, torch.inf) + poses[env_index] = candidates[torch.argmin(ranked)] + success[env_index] = True + return poses, success + + @staticmethod + def _segment( + context: PlanningContext, + resources: _HandoverResources, + source_arm: torch.Tensor, + destination_arm: torch.Tensor, + source_hand: torch.Tensor, + destination_hand: torch.Tensor, + ) -> torch.Tensor: + return assemble_full_robot_trajectory( + context.robot.qpos, + ( + (resources.source_arm.joint_ids, source_arm), + (resources.destination_arm.joint_ids, destination_arm), + (resources.source_hand.joint_ids, source_hand), + (resources.destination_hand.joint_ids, destination_hand), + ), + ) + + def _success(self, value: torch.Tensor, name: str) -> torch.Tensor: + return normalize_success_mask( + value, + num_envs=self.num_envs, + device=self.device, + name=name, + ) + + def _pose(self, value: torch.Tensor, name: str) -> torch.Tensor: + return resolve_batched_pose( + value, + num_envs=self.num_envs, + device=self.device, + name=name, + ) + + @staticmethod + def _require_exchange_pose(options: HeldObjectHandOverOptions) -> None: + if options.middle_object_pose is None or options.final_object_pose is None: + raise ValueError( + "middle_object_pose and final_object_pose are required for " + "HeldObjectHandOver." + ) + + @staticmethod + def _require_same_object( + requested: ObjectSemantics, + held: ObjectSemantics, + ) -> None: + if requested.entity_id is not None and held.entity_id is not None: + matches = requested.entity_id == held.entity_id + elif requested.entity is not None and held.entity is not None: + matches = requested.entity is held.entity + else: + matches = bool(requested.label) and requested.label == held.label + if not matches: + raise ValueError( + "Handover goal must identify the object held by the source." + ) + + @staticmethod + def _start_qpos( + context: PlanningContext, + resources: _HandoverResources, + ) -> tuple[torch.Tensor, torch.Tensor]: + qpos = context.robot.qpos.to(dtype=torch.float32) + return ( + qpos[:, list(resources.source_arm.joint_ids)], + qpos[:, list(resources.destination_arm.joint_ids)], + ) + + @staticmethod + def _segment_lengths( + sample_count: int, + options: HeldObjectHandOverOptions, + ) -> dict[str, int]: + close = max(2, options.hand_interp_steps) + release = max(2, options.hand_interp_steps) + retreat = max(2, options.retreat_steps) + hold = options.hold_steps + reserved = close + release + retreat + hold + transfer = max(2, (sample_count - reserved) // 2) + approach = sample_count - reserved - transfer + if approach < 2: + raise ValueError( + "Not enough handover waypoints; increase sample_count or reduce " + "handover segment lengths." + ) + return { + "transfer": transfer, + "approach": approach, + "close": close, + "hold": hold, + "release": release, + "retreat": retreat, + } + + +__all__ = ["HeldObjectHandOver", "HeldObjectHandOverOptions"] diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index 44b2f02cf..f9c2a7122 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -264,11 +264,9 @@ runtime: upright: PickUp: rotate_upright: 0.7853981633974483 - upright_yaw_samples: 8 MoveHeldObject: staging_lift_height: 0.25 surface_clearance: 0.05 - upright_yaw_samples: 8 upright_xy_tolerance: 0.05 upright_max_tilt: 0.2617993877991494 Place: diff --git a/embodichain/gen_sim/action_engine/orientation.py b/embodichain/gen_sim/action_engine/orientation.py index 3cac92e27..465384dcb 100644 --- a/embodichain/gen_sim/action_engine/orientation.py +++ b/embodichain/gen_sim/action_engine/orientation.py @@ -74,8 +74,16 @@ def requires_reference(self) -> bool: ) @property - def allows_upright_yaw_search(self) -> bool: - """Return whether all hard terms leave world-up yaw unconstrained.""" + def allows_yaw_search(self) -> bool: + """Return whether hard constraints leave world-Z yaw unconstrained.""" + return all( + isinstance(term, AlignAxisConstraint) and term.target_axis == "world_up" + for term in self.terms + ) + + @property + def requires_upright_axis_alignment(self) -> bool: + """Return whether a hard term requires alignment with world up.""" return bool(self.terms) and all( isinstance(term, AlignAxisConstraint) and term.target_axis == "world_up" for term in self.terms @@ -116,9 +124,18 @@ def compile_orientation_constraint( directed=directed, ), ) - elif orientation_goal in {"lay_flat", "axis_align"}: - # These established modes materialize a full target rotation. Keep that - # contract until semantic face/axis metadata can express narrower terms. + elif orientation_goal == "lay_flat": + terms = ( + AlignAxisConstraint( + local_axis="short_axis", + target_axis="world_up", + directed=False, + ), + ) + elif orientation_goal == "axis_align": + # axis_align explicitly requests a horizontal heading. Its established + # target-pose contract remains strict until it is replaced by a typed + # non-world-up axis term. terms = (MatchRotationConstraint(reference="target_pose"),) else: raise ValueError(f"Unsupported orientation_goal {orientation_goal!r}.") @@ -172,7 +189,7 @@ def _compile_term(value: Any, index: int) -> OrientationTerm: f"{context} contains unsupported fields: {sorted(unknown)}." ) local_axis = str(value.get("local_axis", "")) - if local_axis not in {"x", "y", "z", "long_axis"}: + if local_axis not in {"x", "y", "z", "long_axis", "short_axis"}: raise ValueError(f"{context}.local_axis {local_axis!r} is unsupported.") target_axis = str(value.get("target_axis", "world_up")) if target_axis != "world_up": diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 7c2549900..922e904a6 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -102,6 +102,7 @@ _COLLISION_PARKING_Z_OFFSET = -100.0 _BODY_GRASP_CANDIDATE_LIMIT = 500 _BODY_GRASP_SEED = 17_392 +_FREE_YAW_SAMPLE_COUNT = 8 def _collision_cache_for_world( @@ -226,7 +227,7 @@ def start_session( """ capability = self.capabilities.require_executable(grounded.action_class) state = state or self.initial_state() - grounded = self._select_upright_transport_yaw(grounded, state) + grounded = self._select_transport_yaw(grounded, state) context = self._planning_context(state, grounded) invocation = self._invocation(grounded, capability) return self._engine().start((invocation,), context) @@ -340,7 +341,7 @@ def plan( """Plan one grounded primitive through the mainline typed contract.""" capability = self.capabilities.require_executable(grounded.action_class) state = state or self.initial_state() - grounded = self._select_upright_transport_yaw(grounded, state) + grounded = self._select_transport_yaw(grounded, state) context = self._planning_context(state, grounded) grounded_candidates = self._adapt_axis_align_body_grasps( grounded, @@ -838,13 +839,13 @@ def _planner_trace( trace["body_grasp"] = deepcopy(dict(body_grasp)) return trace - def _select_upright_transport_yaw( + def _select_transport_yaw( self, grounded: GroundedAction, state: ExecutionState, ) -> GroundedAction: - """Choose the closest IK-feasible yaw for an upright object target.""" - sample_count = int(grounded.cfg.get("upright_yaw_samples", 1)) + """Choose the closest IK-feasible yaw when task semantics leave it free.""" + sample_count = _FREE_YAW_SAMPLE_COUNT if grounded.allow_yaw_search else 1 capability = self.capabilities.get(grounded.action_class) if ( capability.target_materializer != "semantic_held_object" @@ -858,9 +859,7 @@ def _select_upright_transport_yaw( if target_pose.shape == (4, 4): target_pose = target_pose.unsqueeze(0).repeat(self.num_envs, 1, 1) if target_pose.shape != (self.num_envs, 4, 4): - raise ValueError( - "Upright transport target must have shape (4, 4) or (N, 4, 4)." - ) + raise ValueError("Transport target must have shape (4, 4) or (N, 4, 4).") arm_part, _, _ = self._parts(grounded.arm) held = state.get_held_object(arm_part) @@ -872,7 +871,7 @@ def _select_upright_transport_yaw( ) if object_to_eef.shape == (4, 4): object_to_eef = object_to_eef.unsqueeze(0).repeat(self.num_envs, 1, 1) - variants = self._upright_yaw_variants(target_pose, sample_count) + variants = self._yaw_variants(target_pose, sample_count) eef_variants = torch.matmul(variants, object_to_eef[:, None]) joint_ids = list(self.env.robot.get_joint_ids(name=arm_part)) start_qpos = state.last_qpos[:, joint_ids] @@ -895,7 +894,30 @@ def _select_upright_transport_yaw( distance, torch.full_like(distance, torch.inf), ) - best = distance.argmin(dim=1) + yaw_offsets = torch.matmul( + variants[:, :, :3, :3], + target_pose[:, None, :3, :3].transpose(-1, -2), + ) + yaw_distance = torch.atan2( + yaw_offsets[:, :, 1, 0], + yaw_offsets[:, :, 0, 0], + ).abs() + minimum_yaw = torch.where( + success, + yaw_distance, + torch.full_like(yaw_distance, torch.inf), + ).amin(dim=1) + minimum_rotation = success & torch.isclose( + yaw_distance, + minimum_yaw[:, None], + atol=1.0e-6, + rtol=0.0, + ) + best = torch.where( + minimum_rotation, + distance, + torch.full_like(distance, torch.inf), + ).argmin(dim=1) env_ids = torch.arange(self.num_envs, device=self.device) selected = variants[env_ids, best] selected = torch.where( @@ -910,7 +932,7 @@ def _select_upright_transport_yaw( ) @staticmethod - def _upright_yaw_variants( + def _yaw_variants( target_pose: torch.Tensor, sample_count: int, ) -> torch.Tensor: @@ -1254,8 +1276,6 @@ def _build_single_arm_config( from .atomic_compat import ExactTargetMoveHeldObjectOptions config_type = ExactTargetMoveHeldObjectOptions - if int(action.cfg.get("upright_yaw_samples", 1)) > 1: - policy["allow_automatic_transport_rotation"] = False if capability.target_materializer == "press": press_depth = policy.pop("press_depth", None) if press_depth is not None and "press_distance" not in policy: @@ -1506,6 +1526,10 @@ def joint_ids(self, arm: str, *, include_hand: bool) -> list[int]: def _engine(self) -> AtomicActionEngine: if self._atomic_engine is None: + from embodichain.gen_sim.action_engine.capabilities import ( + HeldObjectHandOver, + ) + from .atomic_compat import ExactTargetMoveHeldObject engine = AtomicActionEngine( @@ -1514,6 +1538,7 @@ def _engine(self) -> AtomicActionEngine: grasp_pose_generators=self._grasp_pose_generators(), ) engine.register(ExactTargetMoveHeldObject(), replace=True) + engine.register(HeldObjectHandOver(), replace=True) self._atomic_engine = engine return self._atomic_engine diff --git a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py index 8c6a0d63b..1953360de 100644 --- a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py +++ b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py @@ -23,12 +23,8 @@ import torch from embodichain.lab.sim.atomic_actions import ( - ActionPlan, - HeldObjectPoseGoal, MoveHeldObject, MoveHeldObjectOptions, - PlanningContext, - ResolvedActionRequest, ) __all__ = ["ExactTargetMoveHeldObject", "ExactTargetMoveHeldObjectOptions"] @@ -36,10 +32,7 @@ @dataclass(frozen=True, slots=True, eq=False) class ExactTargetMoveHeldObjectOptions(MoveHeldObjectOptions): - """Action Engine transport options with an exact-orientation switch.""" - - allow_automatic_transport_rotation: bool = True - """Whether the mainline transport heuristic may replace target rotation.""" + """Action Engine transport options for a grounded object target.""" class ExactTargetMoveHeldObject(MoveHeldObject): @@ -48,38 +41,10 @@ class ExactTargetMoveHeldObject(MoveHeldObject): OptionsType = ExactTargetMoveHeldObjectOptions binding_contract = MoveHeldObject.binding_contract - def __init__( - self, - default_options: ExactTargetMoveHeldObjectOptions | None = None, - ) -> None: - super().__init__(default_options) - self._allow_automatic_transport_rotation = True - - def _plan( - self, - request: ResolvedActionRequest[ - HeldObjectPoseGoal, - ExactTargetMoveHeldObjectOptions, - ], - context: PlanningContext, - ) -> ActionPlan: - previous = self._allow_automatic_transport_rotation - self._allow_automatic_transport_rotation = ( - request.skill_options.allow_automatic_transport_rotation - ) - try: - return super()._plan(request, context) - finally: - self._allow_automatic_transport_rotation = previous - def _apply_automatic_transport_rotation( self, move_eef_xpos: torch.Tensor, end_arm_xpos: torch.Tensor, ) -> None: - """Apply the heuristic unless semantic grounding selected exact yaw.""" - if self._allow_automatic_transport_rotation: - super()._apply_automatic_transport_rotation( - move_eef_xpos, - end_arm_xpos, - ) + """Keep target shaping in GenSim grounding and free-yaw search.""" + del move_eef_xpos, end_arm_xpos diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 2f0b6285e..788dbed15 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -698,16 +698,16 @@ def ground( kind == "handover_staging" and capability.target_materializer == "semantic_held_object" ) - use_upright_yaw_search = ( + use_upright_transport_policy = ( is_handover_continuation or uses_handover_staging - ) and self._uses_upright_yaw_search( + ) and self._uses_upright_transport_policy( step, orientation, ) extra_modifiers: tuple[tuple[str, str], ...] = () if ( is_handover_continuation - and use_upright_yaw_search + and use_upright_transport_policy and capability.target_materializer in { "semantic_held_object", @@ -733,15 +733,6 @@ def ground( # collision-aware planner cannot find a route, do not silently # replace it with collision-unaware joint interpolation. policy["collision_safety"] = "required" - if uses_handover_staging and use_upright_yaw_search: - # Handover consumes the live payload pose immediately after this - # move. Use the existing upright-yaw feasibility search instead - # of the generic transport orientation heuristic, which can tilt - # a payload while moving it to the exchange point. - policy["upright_yaw_samples"] = max( - int(policy.get("upright_yaw_samples", 1)), - 8, - ) object_pose = _live_pose(self.env, step.object_uid) if step.operator == "orient_object": policy["upright_local_axis"] = self._upright_local_axis(step) @@ -769,7 +760,10 @@ def ground( f"AtomicAction {action_class!r} target materializer must " "return GroundedAction." ) - return grounded + return replace( + grounded, + allow_yaw_search=orientation.allows_yaw_search, + ) if kind == "object": semantics = self.semantics_factory( @@ -1063,6 +1057,7 @@ def ground( target_object_pose=target_object_pose, motion_policy=policy, object_uid=step.object_uid, + allow_yaw_search=orientation.allows_yaw_search, ) def _handover_role_axis( @@ -2544,27 +2539,28 @@ def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: axis = self._upright_local_axis(step) entity = _object(self.env, step.object_uid) vertices = _local_vertices(entity, self.env, 0) - if axis == "long_axis": - axis_index = analyze_local_geometry_axes(vertices).long_axis_index + if axis in {"long_axis", "short_axis"}: + axes = analyze_local_geometry_axes(vertices) + axis_index = ( + axes.long_axis_index if axis == "long_axis" else axes.short_axis_index + ) else: axis_index = {"x": 0, "y": 1, "z": 2}[axis] direction = torch.zeros(3, dtype=torch.float32, device=self.env.device) direction[axis_index] = 1.0 return direction - def _uses_upright_yaw_search( + def _uses_upright_transport_policy( self, step: SemanticStep, constraint: OrientationConstraint, ) -> bool: - """Preserve a live upright state as a planning preference. + """Return whether transport should retain upright-specific tuning. - Explicit full-frame matching cannot admit yaw search. With no hard - orientation terms, yaw search is enabled only when the live object's - long axis is already upright, so a preceding upright operation remains - stable without turning that state into a sticky acceptance constraint. + A preceding upright operation may use higher-clearance motion settings + without turning that live posture into a hard terminal constraint. """ - if constraint.allows_upright_yaw_search: + if constraint.requires_upright_axis_alignment: return True if ( constraint.terms @@ -2573,7 +2569,10 @@ def _uses_upright_yaw_search( return False entity = _object(self.env, step.object_uid) vertices = _local_vertices(entity, self.env, 0) - axis_index = analyze_local_geometry_axes(vertices).long_axis_index + try: + axis_index = analyze_local_geometry_axes(vertices).long_axis_index + except ValueError: + return False pose = _live_pose(self.env, step.object_uid) cosine = pose[:, 2, axis_index].abs().clamp(0.0, 1.0) tolerance = float(self.runtime_policy.predicate_fallbacks["upright_max_tilt"]) @@ -2640,7 +2639,11 @@ def _target_rotation( vertical_axis = ( int(longest_to_shortest[0]) if upright_axis == "long_axis" - else {"x": 0, "y": 1, "z": 2}[upright_axis] + else ( + int(longest_to_shortest[-1]) + if upright_axis == "short_axis" + else {"x": 0, "y": 1, "z": 2}[upright_axis] + ) ) horizontal_axis = next( int(axis) diff --git a/embodichain/gen_sim/action_engine/runtime/models.py b/embodichain/gen_sim/action_engine/runtime/models.py index b4596a241..103613a83 100644 --- a/embodichain/gen_sim/action_engine/runtime/models.py +++ b/embodichain/gen_sim/action_engine/runtime/models.py @@ -144,6 +144,8 @@ class GroundedAction: 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.""" + allow_yaw_search: bool = False + """Whether planning may vary world-Z yaw without changing task semantics.""" @dataclass diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index 4062204ba..68e924cf6 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -234,7 +234,15 @@ def _local_axis_index(env: Any, uid: str, axis: Any) -> int: name = str(axis).lower() if name in {"x", "y", "z"}: return {"x": 0, "y": 1, "z": 2}[name] - if name not in {"long", "long_axis", "longest"}: + geometry_axis = { + "long": "long", + "long_axis": "long", + "longest": "long", + "short": "short", + "short_axis": "short", + "shortest": "short", + }.get(name) + if geometry_axis is None: raise ValueError(f"Unsupported upright local axis {axis!r}.") entity = env.sim.get_rigid_object(uid) if entity is None: @@ -247,7 +255,8 @@ def _local_axis_index(env: Any, uid: str, axis: Any) -> int: vertices = vertices[0] if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") - return analyze_local_geometry_axes(vertices).long_axis_index + axes = analyze_local_geometry_axes(vertices) + return axes.long_axis_index if geometry_axis == "long" else axes.short_axis_index def _arm_values( @@ -592,7 +601,15 @@ def evaluate_predicate( cosine = axis[:, 2].clamp(-1.0, 1.0) directed = spec.get( "directed", - str(local_axis).lower() not in {"long", "long_axis", "longest"}, + str(local_axis).lower() + not in { + "long", + "long_axis", + "longest", + "short", + "short_axis", + "shortest", + }, ) if not isinstance(directed, bool): raise ValueError("object_upright directed must be a boolean.") diff --git a/tests/gen_sim/action_engine/capabilities/test_held_hand_over.py b/tests/gen_sim/action_engine/capabilities/test_held_hand_over.py new file mode 100644 index 000000000..113bcb631 --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/test_held_hand_over.py @@ -0,0 +1,311 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Focused contracts for GenSim's receiver-hold handover action.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + HeldObjectHandOver, + HeldObjectHandOverOptions, + build_atomic_capability_registry, +) +from embodichain.lab.sim.atomic_actions import ( + ActionInvocation, + AntipodalAffordance, + AtomicActionEngine, + ControlPartCommandProfile, + GraspGoal, + HeldObjectState, + MotionPolicy, + ObjectSemantics, + PlanningContext, + RobotObservation, + SceneSnapshot, + TaskState, +) +from embodichain.lab.sim.planners import MotionGenerator +from embodichain.toolkits.graspkit import ( + ParallelJawGraspPoseGenerator, + ParallelJawGripperModelCfg, +) + +_HAND_DOF = 1 +_ROBOT_DOF = 6 +_CONTROL_DT = 1.0 / 60.0 + + +class _GraspGenerator(ParallelJawGraspPoseGenerator): + def __init__(self) -> None: + super().__init__(ParallelJawGripperModelCfg(model_id="test_gripper")) + + def get_valid_grasp_poses( + self, + *, + mesh_vertices: torch.Tensor, + mesh_triangles: torch.Tensor, + obj_poses: torch.Tensor, + approach_direction: torch.Tensor, + obj_longest_axis: torch.Tensor | None = None, + is_positive_part: bool | torch.Tensor = True, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + del ( + mesh_vertices, + mesh_triangles, + approach_direction, + obj_longest_axis, + is_positive_part, + ) + return [ + (torch.eye(4).unsqueeze(0), torch.zeros(1)) + for _ in range(obj_poses.shape[0]) + ] + + def get_best_grasp_poses(self, **kwargs: object): + poses = kwargs["obj_poses"] + assert isinstance(poses, torch.Tensor) + return ( + torch.ones(poses.shape[0], dtype=torch.bool), + poses, + torch.zeros(poses.shape[0]), + ) + + def get_dual_arm_valid_grasp_poses(self, **kwargs: object): + del kwargs + raise AssertionError("Receiver-hold handover uses one destination grasp.") + + +def _motion_generator() -> MotionGenerator: + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = _ROBOT_DOF + robot.control_parts = { + "left_arm": object(), + "left_hand": object(), + "right_arm": object(), + "right_hand": object(), + } + joint_ids = { + "left_arm": [0, 1], + "left_hand": [2], + "right_arm": [3, 4], + "right_hand": [5], + } + robot.get_joint_ids.side_effect = lambda name: list(joint_ids[name]) + robot.get_qpos.return_value = torch.zeros(1, _ROBOT_DOF) + robot.compute_fk.side_effect = lambda qpos, **_kwargs: torch.eye(4).repeat( + qpos.shape[0], 1, 1 + ) + + generator = object.__new__(MotionGenerator) + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner = Mock() + generator.planner.cfg.planner_type = "stub" + generator.planner.collision_world_info = None + generator.planner.preserve_plan_samples = False + return generator + + +def _engine() -> AtomicActionEngine: + generator = _motion_generator() + profiles = { + hand: ControlPartCommandProfile.joint_positions( + open=torch.zeros(_HAND_DOF), + grasp=torch.ones(_HAND_DOF), + ) + for hand in ("left_hand", "right_hand") + } + grasp = _GraspGenerator() + engine = AtomicActionEngine( + generator, + control_profiles=profiles, + grasp_pose_generators={"left_hand": grasp, "right_hand": grasp}, + load_builtins=False, + ) + engine.register(HeldObjectHandOver()) + return engine + + +def _semantics(label: str = "can") -> ObjectSemantics: + return ObjectSemantics( + label=label, + entity_id=label, + geometry={}, + affordance=AntipodalAffordance( + object_label=label, + mesh_vertices=torch.tensor( + [[0.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.0, 0.1, 0.0]] + ), + mesh_triangles=torch.tensor([[0, 1, 2]]), + ), + ) + + +def _context(semantics: ObjectSemantics) -> PlanningContext: + relation = torch.eye(4).unsqueeze(0) + task = TaskState( + batch_size=1, + device="cpu", + held_objects={ + "left_arm": HeldObjectState( + semantics=semantics, + object_to_eef=relation, + grasp_xpos=relation, + ) + }, + ) + qpos = torch.zeros(1, _ROBOT_DOF) + return PlanningContext( + robot=RobotObservation(0.0, qpos, torch.zeros_like(qpos)), + task=task, + scene=SceneSnapshot.empty(), + env_ids=torch.tensor([0]), + control_dt=_CONTROL_DT, + ) + + +def _invocation( + engine: AtomicActionEngine, + semantics: ObjectSemantics, + *, + final_x: float = 0.0, +) -> ActionInvocation: + middle = torch.eye(4) + final = middle.clone() + final[0, 3] = final_x + binding = engine.bind_control_parts( + "hand_over", + { + "source": {"motion": "left_arm", "grasp": "left_hand"}, + "destination": {"motion": "right_arm", "grasp": "right_hand"}, + }, + ) + return ActionInvocation( + skill_id="hand_over", + goal=GraspGoal(semantics=semantics), + binding=binding, + motion_policy=MotionPolicy(sample_count=24), + skill_options=HeldObjectHandOverOptions( + middle_object_pose=middle, + final_object_pose=final, + hand_interp_steps=2, + hold_steps=2, + retreat_steps=4, + ), + ) + + +def _install_planner(monkeypatch: pytest.MonkeyPatch) -> None: + def plan( + _generator: MotionGenerator, + _control_part: str, + start_qpos: torch.Tensor, + _target_poses: torch.Tensor, + n_waypoints: int, + _motion_policy: MotionPolicy, + _control_dt: float | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.ones(start_qpos.shape[0], dtype=torch.bool), + start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1), + ) + + monkeypatch.setattr( + "embodichain.gen_sim.action_engine.capabilities.held_hand_over." + "plan_named_arm_trajectory", + plan, + ) + + +def test_handover_capability_matches_installed_receiver_hold_action() -> None: + capability = build_atomic_capability_registry().get("HandOver") + + assert capability.action_type is HeldObjectHandOver + assert capability.config_type is HeldObjectHandOverOptions + assert capability.action_type.GoalType is GraspGoal + + +def test_receiver_hold_plan_transfers_ownership_and_preserves_phase_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_planner(monkeypatch) + engine = _engine() + semantics = _semantics() + context = _context(semantics) + + plan = engine.plan(_invocation(engine, semantics), context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) + + assert plan.plan_success.tolist() == [True] + assert projected.get_held_object("left_arm") is None + received = projected.get_held_object("right_arm") + assert received is not None + assert received.semantics.entity_id == "can" + assert [segment.name for segment in plan.segments] == [ + "transfer", + "approach", + "close", + "hold", + "release", + "retreat", + ] + assert plan.joint_trajectory is not None + positions = plan.joint_trajectory.positions + close = plan.segment("close") + release = plan.segment("release") + retreat = plan.segment("retreat") + torch.testing.assert_close(positions[:, close.stop - 1, 5], torch.ones(1)) + torch.testing.assert_close(positions[:, release.stop - 1, 2], torch.zeros(1)) + torch.testing.assert_close( + positions[:, release.start :, 5], + torch.ones_like(positions[:, release.start :, 5]), + ) + torch.testing.assert_close( + positions[:, retreat.start : retreat.stop, 3:5], + positions[:, retreat.start : retreat.start + 1, 3:5].expand_as( + positions[:, retreat.start : retreat.stop, 3:5] + ), + ) + + +def test_receiver_hold_rejects_delivery_away_from_exchange( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_planner(monkeypatch) + engine = _engine() + semantics = _semantics() + + with pytest.raises(ValueError, match="receiver remains stationary"): + engine.plan(_invocation(engine, semantics, final_x=0.1), _context(semantics)) + + +def test_receiver_hold_rejects_a_different_requested_object( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_planner(monkeypatch) + engine = _engine() + + with pytest.raises(ValueError, match="object held by the source"): + engine.plan( + _invocation(engine, _semantics("other")), + _context(_semantics("held")), + ) diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py index c74a964b5..2c254773f 100644 --- a/tests/gen_sim/action_engine/config/test_runtime_policy.py +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -86,6 +86,14 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: assert runtime.motion_modifiers["orientation"]["upright"]["MoveHeldObject"][ "surface_clearance" ] == pytest.approx(0.05) + assert ( + "upright_yaw_samples" + not in runtime.motion_modifiers["orientation"]["upright"]["PickUp"] + ) + assert ( + "upright_yaw_samples" + not in runtime.motion_modifiers["orientation"]["upright"]["MoveHeldObject"] + ) assert runtime.motion_modifiers["handover_role"]["transfer"]["PickUp"] == { "sample_interval": 80, "hand_interp_steps": 5, diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index e2ee2cb83..2c36ce651 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -24,6 +24,7 @@ import pytest import torch +from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver from embodichain.gen_sim.action_engine.runtime import actions from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( ExactTargetMoveHeldObject, @@ -45,6 +46,7 @@ EndEffectorPoseGoal, EntityState, GraspGoal, + HeldObjectPoseGoal, HeldObjectState, JointPositionGoal, ObjectSemantics, @@ -156,7 +158,7 @@ def plan(self, invocation, context) -> ActionPlan: return self._plan(invocation, context) -def test_adapter_registers_only_move_held_object_compat_action( +def test_adapter_registers_gen_sim_compat_actions( monkeypatch: Any, ) -> None: registered: list[tuple[type, bool]] = [] @@ -179,9 +181,83 @@ def register(self, action: Any, *, replace: bool = False) -> None: assert isinstance(engine, Engine) assert registered == [ (ExactTargetMoveHeldObject, True), + (HeldObjectHandOver, True), ] +def test_free_yaw_search_uses_an_internal_reachability_sample_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + target = torch.eye(4).repeat(2, 1, 1) + target[:, :3, 3] = torch.tensor([0.1, -0.2, 0.9]) + semantics = ObjectSemantics(label="can", geometry={}, affordance=Affordance()) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + state = ExecutionState( + last_qpos=torch.zeros(2, 8), + held_objects={"physical_left_arm": held}, + ) + grounded = GroundedAction( + "MoveHeldObject", + "left_arm", + "arm", + HeldObjectPoseGoal(object_target_pose=target), + {}, + target_object_pose=target, + allow_yaw_search=True, + ) + + def compute_batch_ik( + *, + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert name == "physical_left_arm" + assert pose.shape[:2] == (2, 8) + success = torch.zeros(2, 8, dtype=torch.bool) + success[:, 3] = True + return success, joint_seed.clone() + + monkeypatch.setattr(env.robot, "compute_batch_ik", compute_batch_ik) + + selected = adapter._select_transport_yaw(grounded, state) + + assert selected.target_object_pose is not None + torch.testing.assert_close(selected.target_object_pose[:, :3, 3], target[:, :3, 3]) + torch.testing.assert_close( + selected.target_object_pose[:, :3, :3], + torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]).repeat( + 2, 1, 1 + ), + atol=1.0e-6, + rtol=1.0e-6, + ) + + def all_yaws_reachable( + *, + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + del pose, name + qpos = joint_seed.clone() + qpos[:, 0] += 1.0 + return torch.ones(2, 8, dtype=torch.bool), qpos + + monkeypatch.setattr(env.robot, "compute_batch_ik", all_yaws_reachable) + + minimum_rotation = adapter._select_transport_yaw(grounded, state) + + assert minimum_rotation.target_object_pose is not None + torch.testing.assert_close(minimum_rotation.target_object_pose, target) + + def test_adapter_lowers_axis_align_from_live_pose_with_a_stable_seed() -> None: vertices = torch.tensor( [[x, y, z] for x in (-0.03, 0.03) for y in (-0.06, 0.06) for z in (-0.03, 0.03)] diff --git a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py index f04c13b6e..2859095c2 100644 --- a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py +++ b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py @@ -29,7 +29,7 @@ from embodichain.lab.sim.atomic_actions import MoveHeldObject, MoveHeldObjectOptions -def test_exact_target_transport_only_disables_rotation_when_requested( +def test_grounded_target_transport_never_applies_an_implicit_rotation( monkeypatch: pytest.MonkeyPatch, ) -> None: applied = [] @@ -52,31 +52,15 @@ def fake_plan(self, request, context): monkeypatch.setattr(MoveHeldObject, "_plan", fake_plan) action = ExactTargetMoveHeldObject() assert type(action).__dict__["binding_contract"] is MoveHeldObject.binding_contract - disabled_request = SimpleNamespace( - skill_options=ExactTargetMoveHeldObjectOptions( - allow_automatic_transport_rotation=False, - ) - ) - enabled_request = SimpleNamespace( - skill_options=ExactTargetMoveHeldObjectOptions(), - ) + request = SimpleNamespace(skill_options=ExactTargetMoveHeldObjectOptions()) - assert action._plan(disabled_request, object()) is result + assert action._plan(request, object()) is result assert not applied - assert action._plan(enabled_request, object()) is result - assert applied == [True] -@pytest.mark.parametrize( - ("yaw_samples", "expected"), - [(1, True), (8, False)], -) -def test_semantic_transport_config_scopes_rotation_override( - yaw_samples: int, - expected: bool, -) -> None: +def test_semantic_transport_config_has_no_task_facing_rotation_switch() -> None: adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) - action = SimpleNamespace(cfg={"upright_yaw_samples": yaw_samples}) + action = SimpleNamespace(cfg={}) capability = SimpleNamespace( config_type=MoveHeldObjectOptions, target_materializer="semantic_held_object", @@ -85,4 +69,4 @@ def test_semantic_transport_config_scopes_rotation_override( options = adapter._build_single_arm_config(action, capability) assert isinstance(options, ExactTargetMoveHeldObjectOptions) - assert options.allow_automatic_transport_rotation is expected + assert not hasattr(options, "allow_automatic_transport_rotation") diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index ede1101fe..9ccf9e6fc 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -35,6 +35,7 @@ resolve_agent_runtime_policy, runtime_policy_hash, ) +from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOverOptions from embodichain.gen_sim.action_engine.compiler import ( compile_task_agent, compile_task_agent_v2, @@ -95,7 +96,6 @@ CoordinatedPlacementGoal, CoordinatedPlacementOptions, EndEffectorPoseGoal, - HandOverOptions, HeldObjectPoseGoal, HeldObjectState, ObjectSemantics, @@ -1559,7 +1559,7 @@ def test_handover_grounding_uses_bottom_region_and_diagonal_receive() -> None: ) middle = grounded.cfg["middle_object_pose"] final = grounded.cfg["final_object_pose"] - cfg = AtomicActionAdapter(env)._build_config(grounded, HandOverOptions) + cfg = AtomicActionAdapter(env)._build_config(grounded, HeldObjectHandOverOptions) staging_edge = next( edge @@ -1579,7 +1579,7 @@ def test_handover_grounding_uses_bottom_region_and_diagonal_receive() -> None: assert cfg.receive_pick_object_part == "bottom" assert cfg.receive_approach_direction[1] < 0.0 assert cfg.receive_approach_direction[2] < 0.0 - assert staging.motion_policy["upright_yaw_samples"] == 8 + assert staging.allow_yaw_search def test_handover_rejects_receiver_motion_during_internal_final_phase() -> None: @@ -1591,7 +1591,7 @@ def test_handover_rejects_receiver_motion_during_internal_final_phase() -> None: target=SimpleNamespace(), cfg={"transfer_arm": "left_arm"}, ) - options = HandOverOptions(retreat_steps=4) + options = HeldObjectHandOverOptions(retreat_steps=4) trajectory = torch.zeros(1, 12, adapter.env.robot.dof) assert bool( @@ -1732,8 +1732,8 @@ def test_handover_continuation_uses_stable_upright_policies() -> None: release_defaults = resolve_motion_policy("dual_ur10", "Place", upright) retreat_defaults = resolve_motion_policy("dual_ur10", "MoveEndEffector", upright) - assert grounded_staging.cfg["upright_yaw_samples"] == 8 - assert grounded_final.cfg["upright_yaw_samples"] == 8 + assert grounded_staging.allow_yaw_search + assert grounded_final.allow_yaw_search assert grounded_final_with_reference.target_object_pose is not None assert grounded_final_with_reference.target_object_pose[0, 2, 3] == pytest.approx( 0.90 @@ -1797,7 +1797,7 @@ def test_preserve_handover_continuation_does_not_enable_yaw_search() -> None: orientation_reference_pose=reference, ) - assert "upright_yaw_samples" not in grounded.cfg + assert not grounded.allow_yaw_search assert grounded.target_object_pose is not None torch.testing.assert_close( grounded.target_object_pose[:, :3, :3], @@ -1805,6 +1805,52 @@ def test_preserve_handover_continuation_does_not_enable_yaw_search() -> None: ) +def test_unconstrained_handover_continuation_has_free_yaw() -> None: + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "target": _FakeEntity("target", _pose(0.2, 0.0, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + task = deepcopy(_handover_then_place_task()) + task["task_instances"][0]["params"]["orientation_goal"] = "none" + program = load_execution_program( + instantiate_seed_graph(task, {"can": "can", "target": "target"}) + ) + step = next( + candidate + for candidate in program.semantic_steps + if candidate.operator == "place_relative" + ) + state = _held_state(env, entities["can"], arm="right_arm") + held = state.get_held_object("physical_right_arm") + assert held is not None + grounder = ActionGrounder(program, env, lambda _uid: held.semantics) + grounded = [ + grounder.ground(edge.actions[0], step, arm="right_arm", state=state) + for edge in program.edges + if edge.id in step.edge_ids + and edge.actions[0]["atomic_action_class"] in {"MoveHeldObject", "Place"} + ] + + assert [item.action_class for item in grounded] == [ + "MoveHeldObject", + "MoveHeldObject", + "Place", + ] + assert all(item.allow_yaw_search for item in grounded) + + yawed = entities["can"]._pose.clone() + yawed[:, :3, :3] = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] + ) + entities["can"]._pose = yawed + executor = ProgramExecutor(program, env, settle_steps=0, record_runtime=False) + executor._target_poses[step.id] = _pose(0.0, 0.2, 0.75) + + assert bool(executor._placement_orientation_satisfied(step, yawed)[0]) + + def test_dual_franka_handover_uses_explicit_exchange_clearance() -> None: entities = { "can": _FakeEntity("can", _pose(0.0, 0.2, 1.20), _box_vertices(0.03)), @@ -3944,7 +3990,7 @@ def test_handover_receiver_uses_the_mirrored_diagonal_approach( }, ) - cfg = AtomicActionAdapter(env)._build_config(action, HandOverOptions) + cfg = AtomicActionAdapter(env)._build_config(action, HeldObjectHandOverOptions) diagonal = 2.0**-0.5 assert cfg.receive_approach_direction[0] == pytest.approx(0.0) @@ -3989,7 +4035,7 @@ def get_link_pose(*, link_name: str, to_matrix: bool) -> torch.Tensor: }, ) - cfg = AtomicActionAdapter(env)._build_config(action, HandOverOptions) + cfg = AtomicActionAdapter(env)._build_config(action, HeldObjectHandOverOptions) diagonal = 2.0**-0.5 assert cfg.receive_approach_direction[0] == pytest.approx(expected_x * diagonal) @@ -5808,10 +5854,10 @@ def test_orient_grounding_uses_mature_robot_profile_policy() -> None: torch.tensor([0.0, 0.0, 1.0]), ) assert pickup.motion_policy["rotate_upright"] == pytest.approx(torch.pi / 4) - assert pickup.motion_policy["upright_yaw_samples"] == 8 + assert "upright_yaw_samples" not in pickup.motion_policy assert final.target_object_pose[0, 2, 3] == pytest.approx(expected_z) assert final.motion_policy["upright_local_axis"] == "long_axis" - assert final.motion_policy["upright_yaw_samples"] == 8 + assert final.allow_yaw_search def test_orient_verification_requires_upright_pose_near_initial_xy() -> None: @@ -5863,6 +5909,60 @@ def test_orient_verification_requires_upright_pose_near_initial_xy() -> None: assert bool(failed[0]) +@pytest.mark.parametrize("yaw", [0.0, torch.pi / 3, torch.pi, -torch.pi / 2]) +def test_orient_verification_accepts_any_upright_yaw(yaw: float) -> None: + pose = _pose(0.10, 0.20, 0.823) + pose[:, :3, :3] = torch.tensor( + [ + [torch.cos(torch.tensor(yaw)), -torch.sin(torch.tensor(yaw)), 0.0], + [torch.sin(torch.tensor(yaw)), torch.cos(torch.tensor(yaw)), 0.0], + [0.0, 0.0, 1.0], + ] + ) + bottle = _FakeEntity( + "bottle", + pose, + _rect_vertices(0.02, 0.03, 0.10), + ) + env = _FakeEnv({"bottle": bottle}) + env.agent_initial_object_poses = {"bottle": _pose(0.10, 0.20, 0.78)} + executor = ProgramExecutor( + load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "orient", + "operator": "orient_object", + "object": "bottle", + "actor": {"mode": "auto"}, + "goal": { + "orientation_goal": "upright", + "orientation_axis": "none", + }, + "depends_on": [], + } + ) + ) + ), + env, + settle_steps=0, + record_runtime=False, + ) + step = executor.program.semantic_steps[0] + executor._target_poses[step.id] = _pose(0.10, 0.20, 0.823) + executor._policies[step.id] = { + "upright_max_tilt": torch.pi / 12, + "upright_xy_tolerance": 0.05, + "upright_local_axis": "z", + } + + failed, success, _ = executor._verify_step(step, torch.tensor([False])) + + assert bool(success[0]) + assert not bool(failed[0]) + assert executor.retry_count == 0 + + def test_orient_verification_accepts_grounded_live_xy_anchor() -> None: bottle = _FakeEntity( "bottle", diff --git a/tests/gen_sim/action_engine/test_orientation.py b/tests/gen_sim/action_engine/test_orientation.py index 57aa7ccd0..e30e9567d 100644 --- a/tests/gen_sim/action_engine/test_orientation.py +++ b/tests/gen_sim/action_engine/test_orientation.py @@ -23,7 +23,11 @@ MatchRotationConstraint, compile_orientation_constraint, ) -from embodichain.gen_sim.action_engine.protocol import TASK_SPEC_SCHEMA +from embodichain.gen_sim.action_engine.compiler import compile_task_agent +from embodichain.gen_sim.action_engine.protocol import ( + TASK_AGENT_SCHEMA, + TASK_SPEC_SCHEMA, +) from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph @@ -42,6 +46,24 @@ def test_explicit_preserve_compiles_to_rotation_match() -> None: assert constraint.requires_reference +def test_match_rotation_requires_an_explicit_serialized_term() -> None: + constraint = compile_orientation_constraint( + { + "orientation_constraint": { + "terms": [ + { + "type": "match_rotation", + "reference": "target_pose", + } + ] + } + } + ) + + assert constraint.terms == (MatchRotationConstraint(reference="target_pose"),) + assert not constraint.allows_yaw_search + + def test_upright_compiles_to_directed_axis_when_requested() -> None: constraint = compile_orientation_constraint( { @@ -85,6 +107,44 @@ def test_legacy_long_axis_upright_remains_undirected() -> None: directed=False, ), ) + assert constraint.allows_yaw_search + + +def test_lay_flat_compiles_to_short_axis_alignment_with_free_yaw() -> None: + constraint = compile_orientation_constraint({"orientation_goal": "lay_flat"}) + + assert constraint.terms == ( + AlignAxisConstraint( + local_axis="short_axis", + target_axis="world_up", + directed=False, + ), + ) + assert constraint.allows_yaw_search + + +def test_hold_hover_without_orientation_request_has_no_rotation_match() -> None: + compiled = compile_task_agent( + { + "schema_version": TASK_AGENT_SCHEMA, + "task": "hold", + "goal": "Hold the can above its initial position.", + "semantic_steps": [ + { + "id": "hold", + "operator": "hold_hover", + "object": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": {}, + "depends_on": [], + } + ], + } + ) + + goal = compiled["semantic_steps"][0]["goal"] + assert goal["orientation_goal"] == "none" + assert compile_orientation_constraint(goal).terms == () def test_serialized_constraint_keeps_term_local_tolerance() -> None: From 063cd5e4476bb6451b7bbb78851a4d3a95f6e4cc Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:42:40 +0800 Subject: [PATCH 67/85] feat(gen-sim): add failure continuation and task video archival --- .../gen_sim/action_engine/cli/run_agent.py | 22 ++ .../action_engine/environment/agent_env.py | 4 + .../gen_sim/action_engine/runtime/executor.py | 29 ++- .../action_engine/runtime/recording.py | 5 + embodichain/gen_sim/task_engine/cli.py | 16 ++ embodichain/gen_sim/task_engine/config.py | 2 +- embodichain/gen_sim/task_engine/defaults.yaml | 2 +- embodichain/gen_sim/task_engine/workflow.py | 17 +- embodichain/gen_sim/video_archive.py | 182 +++++++++++++++++ .../action_engine/cli/test_run_agent.py | 84 ++++++++ .../runtime/test_runtime_contracts.py | 190 ++++++++++++++++++ .../orchestration/test_coordinator_cli.py | 5 + .../task_engine/test_parallel_workflow.py | 10 +- tests/gen_sim/task_engine/test_workflow.py | 2 +- tests/gen_sim/test_video_archive.py | 145 +++++++++++++ 15 files changed, 704 insertions(+), 11 deletions(-) create mode 100644 embodichain/gen_sim/video_archive.py create mode 100644 tests/gen_sim/test_video_archive.py diff --git a/embodichain/gen_sim/action_engine/cli/run_agent.py b/embodichain/gen_sim/action_engine/cli/run_agent.py index 13cffe627..cb60fedf7 100644 --- a/embodichain/gen_sim/action_engine/cli/run_agent.py +++ b/embodichain/gen_sim/action_engine/cli/run_agent.py @@ -44,6 +44,7 @@ load_agent_execution_program, write_execution_report, ) +from embodichain.gen_sim.video_archive import _archive_task_recording from embodichain.lab.gym.utils.gym_utils import ( add_env_launcher_args_to_parser, build_env_cfg_from_args, @@ -89,6 +90,15 @@ def build_parser() -> argparse.ArgumentParser: default="independent", help="Execution backend. Action Engine owns the production runtime.", ) + parser.add_argument( + "--failure-policy", + choices=("stop", "continue"), + default="stop", + help=( + "Whether dependency failures stop affected downstream execution or " + "allow diagnostic continuation." + ), + ) parser.add_argument( "--vlm-model", default=None, @@ -181,6 +191,7 @@ def cli() -> int | None: runtime_arguments = { "agent_config": str(Path(args.agent_config).expanduser().resolve()), "base_seed": args.seed, + "failure_policy": str(args.failure_policy), "gym_config": str(Path(args.gym_config).expanduser().resolve()), "max_episodes": episodes, "planning_mode": planning_mode, @@ -211,6 +222,7 @@ def cli() -> int | None: execute = env.get_wrapper_attr("create_demo_action_list") result = execute( regenerate=bool(args.regenerate), + failure_policy=str(args.failure_policy), runtime_run_id=run_id, episode_index=episode_index, ) @@ -257,6 +269,9 @@ def cli() -> int | None: # the final episode as well; otherwise only episodes followed by a next # iteration reach the configured dataset recorder. env.reset(options={"final": True}) + archived_video = _archive_task_recording(env, str(args.task_name)) + if archived_video is not None: + log_info(f"Archived task video: {archived_video}", color="green") except KeyboardInterrupt: log_warning("Action Engine run interrupted by user.") return 130 if args.task_engine_report else None @@ -344,10 +359,12 @@ def __init__( env: gymnasium.Env, *, record_root: Path, + failure_policy: str = "stop", ) -> None: self.graph = graph self.env = env self.record_root = record_root + self.failure_policy = failure_policy def preflight(self) -> bool: """Compile and capability-check the branch without sending motion.""" @@ -390,6 +407,7 @@ def run(self, *, run_id: str, episode_index: int) -> Any: runtime_run_id=run_id, episode_index=episode_index, record_root=self.record_root.as_posix(), + failure_policy=self.failure_policy, ) @@ -408,6 +426,7 @@ class _ABWorkerConfig: seed: int camera_uids: tuple[str, ...] staging_dir: str + failure_policy: str = "stop" class _ABBranchWorker: @@ -911,6 +930,7 @@ def _ab_worker_main(connection: Any, config: _ABWorkerConfig) -> None: graph, env, record_root=Path(config.staging_dir).parent / "runtime", + failure_policy=config.failure_policy, ).preflight() elif operation == "run": graph = _worker_graph(request.get("graph")) @@ -928,6 +948,7 @@ def _ab_worker_main(connection: Any, config: _ABWorkerConfig) -> None: graph, env, record_root=Path(record_root).expanduser().resolve(), + failure_policy=config.failure_policy, ).run( run_id=run_id, episode_index=int(request.get("episode_index", 0)), @@ -1172,6 +1193,7 @@ def _run_ab( seed=episode_seed, camera_uids=tuple(str(uid) for uid in camera_uids), staging_dir=(episode_root / ".work" / route / "video").as_posix(), + failure_policy=str(args.failure_policy), ) for route in ("offline", "online") } diff --git a/embodichain/gen_sim/action_engine/environment/agent_env.py b/embodichain/gen_sim/action_engine/environment/agent_env.py index ed927cfce..7ed4477d3 100644 --- a/embodichain/gen_sim/action_engine/environment/agent_env.py +++ b/embodichain/gen_sim/action_engine/environment/agent_env.py @@ -384,6 +384,7 @@ def update_obj_info(self) -> None: def create_demo_action_list( self, regenerate: bool = False, + failure_policy: str = "stop", **kwargs: Any, ) -> Any: """Compile in memory when requested, then execute the program online.""" @@ -408,6 +409,7 @@ def create_demo_action_list( record_runtime=bool(getattr(self, "action_engine_record_runtime", True)), record_root=getattr(self, "action_engine_record_root", None), runtime_policy=self.runtime_policy, + failure_policy=failure_policy, ) self.last_execution = executor.run( run_id=kwargs.get("runtime_run_id"), @@ -422,6 +424,7 @@ def execute_seed_graph( runtime_run_id: str, episode_index: int, record_root: str | None = None, + failure_policy: str = "stop", ) -> Any: """Execute one already validated branch graph without rewriting config.""" program = self.preflight_seed_graph(seed_graph) @@ -448,6 +451,7 @@ def execute_seed_graph( record_runtime=bool(getattr(self, "action_engine_record_runtime", True)), record_root=record_root, runtime_policy=self.runtime_policy, + failure_policy=failure_policy, ) self.last_execution = executor.run( run_id=runtime_run_id, diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index cfa7ccd99..5e0bf3248 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -184,11 +184,17 @@ def __init__( runtime_policy: RuntimePolicyCfg | None = None, capability_registry: Any | None = None, scene_provider: SceneProvider | None = None, + failure_policy: str = "stop", ) -> None: self.program = program self.env = env self.record_runtime = bool(record_runtime) self.record_root = record_root + if failure_policy not in {"stop", "continue"}: + raise ValueError( + "ProgramExecutor failure_policy must be 'stop' or 'continue'." + ) + self.failure_policy = str(failure_policy) if runtime_policy is None: profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) runtime_policy = default_runtime_policy(profile) @@ -359,6 +365,7 @@ def run( enabled=self.record_runtime, runtime_policy=self.runtime_policy.as_mapping(), runtime_policy_hash=runtime_policy_hash(self.runtime_policy), + failure_policy=self.failure_policy, ) aggregate_failed = torch.zeros( int(self.env.num_envs), @@ -366,6 +373,7 @@ def run( device=self.env.device, ) edge_failures: dict[str, torch.Tensor] = {} + step_failures: dict[str, torch.Tensor] = {} semantic_success: dict[str, torch.Tensor] = {} failure_events: list[dict[str, Any]] = [] completed: set[str] = set() @@ -383,16 +391,24 @@ def run( raise RuntimeError( "Execution program is deadlocked: no remaining edge is ready." ) - ready_blocked = { + dependency_failed = { edge.id: self._dependency_failures(edge, edge_failures) for edge in ready } + scheduling_blocked = { + edge_id: ( + failed + if self.failure_policy == "stop" + else torch.zeros_like(failed) + ) + for edge_id, failed in dependency_failed.items() + } batch = self._pack_ready_edges( ready, - inactive=ready_blocked, + inactive=scheduling_blocked, completed=completed, ) - blocked = {edge.id: ready_blocked[edge.id] for edge in batch} + blocked = {edge.id: scheduling_blocked[edge.id] for edge in batch} # A synchronized pair needs the same active rows. Execute a # healthy independent branch separately when its peer is blocked. if len(batch) == 2 and not torch.equal( @@ -530,9 +546,14 @@ def run( completed.add(edge.id) remaining.remove(edge.id) step = self.step_by_edge[edge.id] + historical_step_failure = step_failures.setdefault( + step.id, + torch.zeros_like(edge_failures[edge.id]), + ) + historical_step_failure |= edge_failures[edge.id] if edge.id != step.edge_ids[-1]: continue - prior_failed = edge_failures[edge.id] + prior_failed = historical_step_failure verified_failed, step_success, observed = self._verify_step( step, prior_failed ) diff --git a/embodichain/gen_sim/action_engine/runtime/recording.py b/embodichain/gen_sim/action_engine/runtime/recording.py index 4c38a64a2..e8f7e6e66 100644 --- a/embodichain/gen_sim/action_engine/runtime/recording.py +++ b/embodichain/gen_sim/action_engine/runtime/recording.py @@ -86,7 +86,10 @@ def __init__( enabled: bool = True, runtime_policy: Mapping[str, Any] | None = None, runtime_policy_hash: str | None = None, + failure_policy: str = "stop", ) -> None: + if failure_policy not in {"stop", "continue"}: + raise ValueError("failure_policy must be 'stop' or 'continue'.") self.enabled = enabled self.num_envs = int(num_envs) self.run_id = _safe_name( @@ -125,6 +128,7 @@ def __init__( "episode_index": int(episode_index), "program_schema_version": self.seed_topology.get("schema_version"), "seed_graph_hash": self.program_hash, + "failure_policy": failure_policy, } if runtime_policy is not None: if not isinstance(runtime_policy_hash, str) or not runtime_policy_hash: @@ -291,6 +295,7 @@ def _write_step_checkpoint( "run_id": self.run_id, "episode_index": self.program_metadata["episode_index"], "env_id": env_id, + "failure_policy": self.program_metadata["failure_policy"], "semantic_step": deepcopy(self.step_specs[step.id]), "status": event["status"], "events": related_events, diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py index fec6de86a..b1fd60550 100644 --- a/embodichain/gen_sim/task_engine/cli.py +++ b/embodichain/gen_sim/task_engine/cli.py @@ -72,6 +72,7 @@ def build_parser() -> argparse.ArgumentParser: run_parser.add_argument("--seed", type=int, default=0) run_parser.add_argument("--num-envs", type=int, default=None) run_parser.add_argument("--dataset-saving", action="store_true") + _add_failure_policy_argument(run_parser) return parser @@ -99,6 +100,19 @@ def _add_workflow_arguments(parser: argparse.ArgumentParser) -> None: choices=_ROBOT_PROFILES, default="franka", ) + _add_failure_policy_argument(parser) + + +def _add_failure_policy_argument(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--failure-policy", + choices=("stop", "continue"), + default="stop", + help=( + "Whether dependency failures stop affected downstream execution or " + "allow diagnostic continuation." + ), + ) def main(argv: Sequence[str] | None = None) -> int: @@ -156,6 +170,7 @@ def _run_workflow( vlm_model=args.vlm_model, base_seed=args.base_seed, dataset_saving=args.dataset_saving, + failure_policy=args.failure_policy, run_id=allocation.run_id, created_at=allocation.created_at, execute=execute, @@ -188,6 +203,7 @@ def _run_prepared_bundle(args: argparse.Namespace) -> int: seed=int(args.seed), num_envs=num_envs, dataset_saving=bool(args.dataset_saving), + failure_policy=args.failure_policy, ) environments = report.get("environments", ()) successes = [ diff --git a/embodichain/gen_sim/task_engine/config.py b/embodichain/gen_sim/task_engine/config.py index 60cda86ef..1a6df29bd 100644 --- a/embodichain/gen_sim/task_engine/config.py +++ b/embodichain/gen_sim/task_engine/config.py @@ -108,7 +108,7 @@ class TaskEnginePlanningCfg: candidate_count: int = 3 planning_mode: str = "offline" max_episodes: int = 1 - max_episode_steps: int = 4000 + max_episode_steps: int = 6000 def __post_init__(self) -> None: for field_name in ( diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml index 33169e461..c880ba50f 100644 --- a/embodichain/gen_sim/task_engine/defaults.yaml +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -25,7 +25,7 @@ planning: candidate_count: 3 planning_mode: offline max_episodes: 1 - max_episode_steps: 4000 + max_episode_steps: 6000 execution: num_envs: 1 diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py index fca9778c4..88968d021 100644 --- a/embodichain/gen_sim/task_engine/workflow.py +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -106,6 +106,7 @@ def __call__( seed: int, num_envs: int, dataset_saving: bool = False, + failure_policy: str = "stop", ) -> Mapping[str, Any]: """Run one simulator attempt and preserve its report and trajectory. @@ -115,12 +116,16 @@ def __call__( seed: Action Engine random seed. num_envs: Number of vectorized scene replicas. dataset_saving: Whether to enable the Gym project's dataset recorder. + failure_policy: Whether failed dependencies stop or permit downstream + diagnostic execution. Returns: Validated Action Engine execution report. """ bundle_root = Path(bundle).expanduser().resolve() attempt_root = Path(output_root).expanduser().resolve() + if failure_policy not in {"stop", "continue"}: + raise ValueError("failure_policy must be 'stop' or 'continue'.") attempt_root.mkdir(parents=True, exist_ok=False) command = [ sys.executable, @@ -137,11 +142,12 @@ def __call__( ] if not dataset_saving: command.append("--filter_dataset_saving") + command.extend(["--failure-policy", failure_policy]) log_path = attempt_root / "action.log" print( "[Task Engine] Starting " f"{attempt_root.name}: seed={seed}, num_envs={num_envs}, " - f"dataset_saving={dataset_saving}", + f"dataset_saving={dataset_saving}, failure_policy={failure_policy}", flush=True, ) completed = _run_streaming_process(command, log_path) @@ -177,6 +183,7 @@ def __call__( "seed": seed, "num_envs": num_envs, "dataset_saving": dataset_saving, + "failure_policy": failure_policy, "returncode": completed.returncode, "trajectory_copy": trajectory_copy, "report": report, @@ -288,6 +295,7 @@ def run( vlm_model: str | None = None, base_seed: int = 0, dataset_saving: bool = False, + failure_policy: str = "stop", run_id: str | None = None, created_at: datetime | None = None, overwrite: bool = False, @@ -305,6 +313,8 @@ def run( vlm_model: Optional Action Engine VLM override. base_seed: First audited scene and action attempt seed. dataset_saving: Whether Action attempts may initialize dataset recording. + failure_policy: Whether failed dependencies stop or permit downstream + diagnostic execution. run_id: Optional externally allocated run identifier. created_at: Optional timezone-aware run creation timestamp. overwrite: Whether to atomically replace an existing run directory. @@ -316,6 +326,8 @@ def run( normalized = validate_task_run_request(request) if not isinstance(dataset_saving, bool): raise TypeError("dataset_saving must be a boolean.") + if failure_policy not in {"stop", "continue"}: + raise ValueError("failure_policy must be 'stop' or 'continue'.") if workflow_cfg is None or planning_cfg is None or execution_cfg is None: loaded_workflow, loaded_planning, loaded_execution = ( load_task_engine_config(config_path) @@ -337,6 +349,7 @@ def run( "run_id": effective_run_id, "created_at": effective_created_at.isoformat(), "dataset_saving": bool(dataset_saving), + "failure_policy": str(failure_policy), } state = initial_state(normalized) attempts: list[dict[str, Any]] = [] @@ -851,6 +864,7 @@ def run( seed=action_seed, num_envs=execution_cfg.num_envs, dataset_saving=bool(dataset_saving), + failure_policy=failure_policy, ) successes = _environment_successes( report, @@ -1029,6 +1043,7 @@ def _publish( "success_policy": execution_cfg.success_policy, "min_successful_envs": execution_cfg.min_successful_envs, "dataset_saving": bool(run_metadata["dataset_saving"]), + "failure_policy": str(run_metadata["failure_policy"]), }, }, "attempts": deepcopy(list(attempts)), diff --git a/embodichain/gen_sim/video_archive.py b/embodichain/gen_sim/video_archive.py new file mode 100644 index 000000000..9782f5348 --- /dev/null +++ b/embodichain/gen_sim/video_archive.py @@ -0,0 +1,182 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Rename one completed GenSim recording to its task ID.""" + +from __future__ import annotations + +import argparse +from pathlib import Path +import sys +from typing import Any, Sequence + +__all__: list[str] = [] + + +def _archive_task_recording(env: Any, task_id: str) -> Path | None: + """Archive the generated audience video from a completed GenSim task. + + Args: + env: Completed GenSim environment whose final reset flushed recording. + task_id: ID of the task that produced the recording. + + Returns: + Archived video path, or ``None`` when video recording is disabled. + + Raises: + RuntimeError: If configured recorders do not identify one task video. + ValueError: If the task ID can escape the video directory. + FileNotFoundError: If the expected source recording does not exist. + FileExistsError: If the task archive already exists. + """ + manager = getattr(env.unwrapped, "event_manager", None) + mode_cfgs = getattr(manager, "_mode_functor_cfgs", {}) + recorders: list[Any] = [] + for configured in mode_cfgs.values(): + for functor_cfg in configured: + functor = getattr(functor_cfg, "func", None) + if getattr(type(functor), "__name__", "") in { + "record_camera_data", + "record_camera_data_async", + }: + recorders.append(functor) + if not recorders: + return None + + audience = [ + recorder + for recorder in recorders + if getattr(recorder, "_name", None) == "record_cam_audience_view" + ] + if len(audience) == 1: + recorder = audience[0] + elif len(recorders) == 1: + recorder = recorders[0] + else: + raise RuntimeError( + "GenSim task video archival found multiple camera recorders without " + "one audience recorder." + ) + recorder_name = str(getattr(recorder, "_name", "")).strip() + save_path = getattr(recorder, "_save_path", None) + if not recorder_name or not isinstance(save_path, (str, Path)): + raise RuntimeError( + f"Cannot archive video for task {task_id!r}: camera recorder does not " + "expose its output path and name." + ) + return _archive_task_video( + save_path, + source_stem=f"episode_0_{recorder_name}", + task_id=task_id, + ) + + +def _archive_task_video( + video_directory: str | Path, + *, + source_stem: str, + task_id: str, +) -> Path: + """Move a completed recording to ``.``. + + Args: + video_directory: Directory containing the completed recording. + source_stem: Source file name without its video extension. + task_id: ID of the task that produced the recording. + + Returns: + Path to the archived recording. + + Raises: + ValueError: If the task ID can escape the video directory. + FileNotFoundError: If the expected source recording does not exist. + FileExistsError: If the task archive already exists. + RuntimeError: If more than one source extension matches. + """ + _validate_task_id(task_id) + directory = Path(video_directory).expanduser().resolve() + source_prefix = f"{source_stem}." + candidates = ( + sorted( + path + for path in directory.iterdir() + if path.is_file() and path.name.startswith(source_prefix) + ) + if directory.is_dir() + else [] + ) + expected = directory / f"{source_stem}." + if not candidates: + raise FileNotFoundError( + f"Cannot archive video for task {task_id!r}: " + f"expected source video at {expected}." + ) + if len(candidates) != 1: + matches = ", ".join(path.as_posix() for path in candidates) + raise RuntimeError( + f"Cannot archive video for task {task_id!r}: expected exactly one " + f"source video at {expected}, found {matches}." + ) + + source = candidates[0] + extension = source.name[len(source_stem) :] + destination = directory / f"{task_id}{extension}" + if destination.exists(): + raise FileExistsError( + f"Cannot archive video for task {task_id!r}: destination already " + f"exists at {destination}." + ) + source.rename(destination) + return destination + + +def _validate_task_id(task_id: str) -> None: + if ( + not isinstance(task_id, str) + or not task_id + or task_id in {".", ".."} + or "/" in task_id + or "\\" in task_id + or "\x00" in task_id + ): + raise ValueError( + f"Invalid task ID {task_id!r}: task IDs must be non-empty file names " + "without path separators." + ) + + +def _main(argv: Sequence[str] | None = None) -> int: + """Run task-video archival as a standalone GenSim command.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--video-directory", required=True) + parser.add_argument("--source-stem", required=True) + parser.add_argument("--task-id", required=True) + args = parser.parse_args(argv) + try: + destination = _archive_task_video( + args.video_directory, + source_stem=args.source_stem, + task_id=args.task_id, + ) + except (OSError, RuntimeError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + print(destination) + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/tests/gen_sim/action_engine/cli/test_run_agent.py b/tests/gen_sim/action_engine/cli/test_run_agent.py index 53b2dc99b..3a83339db 100644 --- a/tests/gen_sim/action_engine/cli/test_run_agent.py +++ b/tests/gen_sim/action_engine/cli/test_run_agent.py @@ -18,10 +18,12 @@ import json from pathlib import Path +import sys from types import SimpleNamespace import pytest +from embodichain.gen_sim.action_engine.cli import run_agent as run_agent_module from embodichain.gen_sim.action_engine.cli.run_agent import ( _ABWorkerConfig, _SerializedABBranch, @@ -80,6 +82,88 @@ def test_capture_ab_initial_frame_requires_audience_recorder() -> None: _capture_ab_initial_frame(_FakeEnv()) +def test_cli_archives_task_video_after_final_reset_and_before_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + events = [] + + class Env: + def __init__(self) -> None: + self.unwrapped = self + self.final_reset = False + + def reset(self, *, seed=None, options=None) -> None: + if options == {"final": True}: + self.final_reset = True + events.append("final_reset") + + def get_wrapper_attr(self, _name): + return lambda **_kwargs: SimpleNamespace( + already_executed=True, + runtime_success=[True], + runtime_graph_output_dir=None, + ) + + def close(self) -> None: + events.append("close") + + env = Env() + monkeypatch.setattr( + run_agent_module, + "build_env_cfg_from_args", + lambda _args: ( + SimpleNamespace(seed=None), + {"id": "ActionEngine-v1", "max_episodes": 1}, + None, + ), + ) + monkeypatch.setattr( + run_agent_module, + "load_config", + lambda _path: {"planning_mode": "offline"}, + ) + monkeypatch.setattr(run_agent_module, "_validate_gym_id", lambda _cfg: None) + monkeypatch.setattr( + run_agent_module, + "_validate_run_contract", + lambda *_args: None, + ) + monkeypatch.setattr( + run_agent_module, + "load_agent_execution_program", + lambda *_args, **_kwargs: SimpleNamespace(seed_graph=None), + ) + monkeypatch.setattr( + run_agent_module, + "_load_grounded_task_plan", + lambda _path: None, + ) + monkeypatch.setattr(run_agent_module.gymnasium, "make", lambda **_kwargs: env) + + def archive(completed_env, task_id): + assert completed_env is env + assert env.final_reset is True + events.append(f"archive:{task_id}") + + monkeypatch.setattr(run_agent_module, "_archive_task_recording", archive) + monkeypatch.setattr( + sys, + "argv", + [ + "run_agent", + "--task_name", + "task2_1", + "--gym_config", + "gym.json", + "--agent_config", + "agent.json", + ], + ) + + assert run_agent_module.cli() is None + assert events == ["final_reset", "archive:task2_1", "close"] + + def _worker_config(route: str) -> _ABWorkerConfig: return _ABWorkerConfig( route=route, diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index 9ccf9e6fc..3decac78e 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -543,6 +543,24 @@ def test_documented_run_command_arguments_remain_compatible() -> None: assert args.headless is True assert args.seed == 17 assert args.runtime_backend == "independent" + assert args.failure_policy == "stop" + + +def test_run_command_accepts_continue_failure_policy() -> None: + args = build_run_parser().parse_args( + [ + "--task_name", + "task4_2", + "--gym_config", + "/tmp/fast_gym_config.json", + "--agent_config", + "/tmp/agent_config.json", + "--failure-policy", + "continue", + ] + ) + + assert args.failure_policy == "continue" def test_dual_ur5_policy_uses_short_reach_upright_lifts() -> None: @@ -4958,6 +4976,139 @@ def execute(edge, step, *, failed): assert not bool(result.success[0]) +def _run_dependent_failure_policy_chain( + monkeypatch: pytest.MonkeyPatch, + *, + failure_policy: str, + record_root: Path | None = None, +) -> tuple[ExecutionResult, dict[str, list[list[bool]]]]: + first = _hold_step("first", "can_a", "left_arm") + second = _hold_step("second", "can_b", "right_arm") + second["depends_on"] = ["first"] + env = _FakeEnv() + env.num_envs = 2 + env.robot = _FakeRobot(env.num_envs) + executor = ProgramExecutor( + load_execution_program(compile_task_agent(_task_agent(first, second))), + env, + settle_steps=0, + record_runtime=record_root is not None, + record_root=record_root, + failure_policy=failure_policy, + ) + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, failed: executor._assignments.setdefault( + step.id, + [ + None if bool(failed[env_id]) else str(step.actor["arm"]) + for env_id in range(env.num_envs) + ], + ), + ) + active_by_step: dict[str, list[list[bool]]] = {"first": [], "second": []} + + def execute( + edge: ExecutionEdge, + step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + active = ~failed + active_by_step[step.id].append(active.tolist()) + result_failed = failed.clone() + if step.id == "first" and edge.id == step.edge_ids[0]: + result_failed[1] = True + return _EdgeResult( + actions=[], + failed=result_failed, + grounded=[], + executed=active, + ) + + monkeypatch.setattr(executor, "_execute_edge", execute) + monkeypatch.setattr( + executor, + "_step_runtime_metadata", + lambda _step: [{} for _ in range(env.num_envs)], + ) + monkeypatch.setattr( + executor, + "_verify_step", + lambda _step, failed: ( + failed, + ~failed, + torch.zeros(env.num_envs, 3), + ), + ) + return executor.run(run_id=failure_policy), active_by_step + + +def test_stop_failure_policy_blocks_only_failed_environment_downstream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + result, active_by_step = _run_dependent_failure_policy_chain( + monkeypatch, + failure_policy="stop", + ) + + assert all(active == [True, False] for active in active_by_step["second"]) + assert result.semantic_success["second"].tolist() == [True, False] + assert result.success.tolist() == [True, False] + + +def test_continue_failure_policy_executes_downstream_without_clearing_history( + monkeypatch: pytest.MonkeyPatch, +) -> None: + result, active_by_step = _run_dependent_failure_policy_chain( + monkeypatch, + failure_policy="continue", + ) + + assert all(active == [True, True] for active in active_by_step["second"]) + assert result.semantic_success["first"].tolist() == [True, False] + assert result.semantic_success["second"].tolist() == [True, True] + assert result.success.tolist() == [True, False] + + +def test_continue_failure_policy_records_failed_then_executed_checkpoints( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + visualization = ModuleType("embodichain.gen_sim.action_engine.graph_visualization") + visualization.render_task_graph_png = lambda _document: b"\x89PNG\r\n\x1a\n" + monkeypatch.setitem(sys.modules, visualization.__name__, visualization) + + result, _ = _run_dependent_failure_policy_chain( + monkeypatch, + failure_policy="continue", + record_root=tmp_path, + ) + + env_dir = Path(result.record_dir) / "env_0001" + checkpoints = { + document["semantic_step"]["id"]: document + for path in env_dir.joinpath("checkpoints").glob("*.json") + for document in [json.loads(path.read_text(encoding="utf-8"))] + } + task_graph = json.loads( + env_dir.joinpath("task_graph.json").read_text(encoding="utf-8") + ) + + assert checkpoints["first"]["status"] == "failed" + assert checkpoints["second"]["status"] == "success" + assert checkpoints["first"]["failure_policy"] == "continue" + assert any( + event["event"] == "edge" + and event["semantic_step_id"] == "second" + and event["status"] == "executed" + for event in task_graph["runtime"]["events"] + ) + assert task_graph["runtime"]["failure_policy"] == "continue" + assert task_graph["runtime"]["status"] == "failed" + + def test_resource_ordering_waits_without_propagating_semantic_failure() -> None: task = { "schema_version": TASK_SPEC_SCHEMA, @@ -6325,6 +6476,45 @@ def fake_install(robot: Any) -> int: assert env._normalize_demo_action_list(result) is result +def test_environment_passes_failure_policy_to_program_executor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + expected = object() + + class FakeExecutor: + def __init__(self, _program: Any, _env: Any, **kwargs: Any) -> None: + captured.update(kwargs) + + def run(self, **kwargs: Any) -> Any: + captured["run"] = kwargs + return expected + + monkeypatch.setattr( + env_module, + "load_agent_execution_program", + lambda *_args, **_kwargs: object(), + ) + monkeypatch.setattr(env_module, "ProgramExecutor", FakeExecutor) + env = SimpleNamespace( + agent_config={}, + agent_config_path="/tmp/agent_config.json", + runtime_policy=object(), + last_execution=None, + ) + + result = env_module.ActionEngineEnv.create_demo_action_list.__wrapped__( + env, + failure_policy="continue", + runtime_run_id="run", + episode_index=3, + ) + + assert result is expected + assert captured["failure_policy"] == "continue" + assert captured["run"] == {"run_id": "run", "episode_index": 3} + + def test_solver_compat_repairs_only_stale_action_engine_ur_dh_defaults() -> None: stale_ur5 = URSolverCfg() stale_ur5.ur_type = "ur5" diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py index 053b3e93f..8096eab98 100644 --- a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -833,6 +833,7 @@ def run(self, request, **kwargs): assert request["scene_edit_prompt"] == edit assert captured["kwargs"]["base_seed"] == 9 assert captured["kwargs"]["dataset_saving"] is (mode == "image") + assert captured["kwargs"]["failure_policy"] == "stop" assert captured["kwargs"]["execute"] is True payload = json.loads(capsys.readouterr().out) assert payload["status"] == "succeeded" @@ -975,6 +976,7 @@ def test_public_cli_exposes_prepare_run_and_run_all_modes() -> None: ) assert arguments.command == "prepare" assert arguments.dataset_saving is True + assert arguments.failure_policy == "stop" def test_prepare_cli_stops_before_simulator_execution( @@ -1033,6 +1035,7 @@ class Executor: def __call__(self, _bundle, output, **kwargs): Path(output).mkdir() assert kwargs["num_envs"] == 2 + assert kwargs["failure_policy"] == "continue" return { "status": "failed", "environments": [ @@ -1052,6 +1055,8 @@ def __call__(self, _bundle, output, **kwargs): str(tmp_path / "history"), "--num-envs", "2", + "--failure-policy", + "continue", ] ) diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py index 23aaab655..f135c5ce0 100644 --- a/tests/gen_sim/task_engine/test_parallel_workflow.py +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -286,11 +286,13 @@ def __call__( seed: int, num_envs: int, dataset_saving: bool = False, + failure_policy: str = "stop", ): values = self.successes[min(self.calls, len(self.successes) - 1)] self.calls += 1 assert len(values) == num_envs assert dataset_saving is self.expected_dataset_saving + assert failure_policy == "stop" return { "status": "succeeded" if all(values) else "failed", "seed": seed, @@ -384,7 +386,7 @@ def test_parallel_workflow_accepts_one_success_and_publishes_all_graphs( candidate_count=3, planning_mode="offline", max_episodes=1, - max_episode_steps=4000, + max_episode_steps=6000, ), execution_cfg=TaskEngineExecutionCfg(num_envs=4), base_seed=11, @@ -403,10 +405,10 @@ def test_parallel_workflow_accepts_one_success_and_publishes_all_graphs( "candidate_count": 3, "planning_mode": "offline", "max_episodes": 1, - "max_episode_steps": 4000, + "max_episode_steps": 6000, } assert manifest["configuration"]["execution"]["dataset_saving"] is False - assert coordinator.kwargs[0]["max_episode_steps"] == 4000 + assert coordinator.kwargs[0]["max_episode_steps"] == 6000 assert coordinator.kwargs[0]["final_inspection"]["scene_revision_id"] == "0" * 64 assert ( coordinator.kwargs[0]["unbound_action_plan"]["candidate_id"] == "candidate_01" @@ -546,6 +548,7 @@ def fake_run(command, log_path): seed=7, num_envs=4, dataset_saving=dataset_saving, + failure_policy="continue", ) assert report["status"] == "succeeded" @@ -558,6 +561,7 @@ def fake_run(command, log_path): assert " prepare" not in " ".join(captured["command"]) assert " workflow" not in " ".join(captured["command"]) assert ("--filter_dataset_saving" in captured["command"]) is expects_filter + assert captured["command"][-2:] == ["--failure-policy", "continue"] assert captured["log_path"] == attempt / "action.log" assert (attempt / "action.log").read_text(encoding="utf-8") == "child output\n" assert (attempt / "trajectory" / "episode.json").is_file() diff --git a/tests/gen_sim/task_engine/test_workflow.py b/tests/gen_sim/task_engine/test_workflow.py index 4a1564606..e9e40fa35 100644 --- a/tests/gen_sim/task_engine/test_workflow.py +++ b/tests/gen_sim/task_engine/test_workflow.py @@ -273,7 +273,7 @@ def test_packaged_workflow_configuration_uses_recovery_defaults() -> None: assert planning.candidate_count == 3 assert planning.planning_mode == "offline" assert planning.max_episodes == 1 - assert planning.max_episode_steps == 4000 + assert planning.max_episode_steps == 6000 assert execution.num_envs == 1 assert execution.required_successes == 1 diff --git a/tests/gen_sim/test_video_archive.py b/tests/gen_sim/test_video_archive.py new file mode 100644 index 000000000..4b479ed9e --- /dev/null +++ b/tests/gen_sim/test_video_archive.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 + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from embodichain.gen_sim.video_archive import ( + _archive_task_recording, + _archive_task_video, +) + +SOURCE_STEM = "episode_0_record_cam_audience_view" + + +class record_camera_data: + def __init__(self, save_path: Path) -> None: + self._name = "record_cam_audience_view" + self._save_path = save_path + + +def _env(recorder: record_camera_data | None = None) -> SimpleNamespace: + event_manager = SimpleNamespace( + _mode_functor_cfgs={ + "interval": ( + [SimpleNamespace(func=recorder)] if recorder is not None else [] + ) + } + ) + env = SimpleNamespace(event_manager=event_manager) + env.unwrapped = env + return env + + +def _write_source(directory: Path, extension: str, content: bytes = b"video") -> Path: + source = directory / f"{SOURCE_STEM}{extension}" + source.write_bytes(content) + return source + + +def test_archive_task_video_renames_source_and_preserves_extension( + tmp_path: Path, +) -> None: + source = _write_source(tmp_path, ".webm") + + destination = _archive_task_video( + tmp_path, + source_stem=SOURCE_STEM, + task_id="task2_1", + ) + + assert destination == tmp_path / "task2_1.webm" + assert destination.read_bytes() == b"video" + assert not source.exists() + + +@pytest.mark.parametrize("task_id", ["../task2_1", "task2/1", r"task2\1", ".."]) +def test_archive_task_video_rejects_path_characters( + tmp_path: Path, + task_id: str, +) -> None: + source = _write_source(tmp_path, ".mp4") + + with pytest.raises(ValueError, match="Invalid task ID"): + _archive_task_video(tmp_path, source_stem=SOURCE_STEM, task_id=task_id) + + assert source.is_file() + + +def test_archive_task_video_reports_missing_source_with_task_and_path( + tmp_path: Path, +) -> None: + with pytest.raises(FileNotFoundError) as error: + _archive_task_video( + tmp_path, + source_stem=SOURCE_STEM, + task_id="task2_1", + ) + + message = str(error.value) + assert "task2_1" in message + assert str(tmp_path / f"{SOURCE_STEM}.") in message + + +def test_archive_task_video_does_not_overwrite_existing_target( + tmp_path: Path, +) -> None: + source = _write_source(tmp_path, ".mp4", b"new") + destination = tmp_path / "task2_1.mp4" + destination.write_bytes(b"existing") + + with pytest.raises(FileExistsError, match="destination already exists"): + _archive_task_video( + tmp_path, + source_stem=SOURCE_STEM, + task_id="task2_1", + ) + + assert source.read_bytes() == b"new" + assert destination.read_bytes() == b"existing" + + +def test_consecutive_tasks_keep_independent_videos(tmp_path: Path) -> None: + for task_id, content in (("task2_1", b"first"), ("task2_2", b"second")): + _write_source(tmp_path, ".mp4", content) + _archive_task_video( + tmp_path, + source_stem=SOURCE_STEM, + task_id=task_id, + ) + + assert (tmp_path / "task2_1.mp4").read_bytes() == b"first" + assert (tmp_path / "task2_2.mp4").read_bytes() == b"second" + assert not (tmp_path / f"{SOURCE_STEM}.mp4").exists() + + +def test_task_recording_uses_runtime_recorder_path(tmp_path: Path) -> None: + recorder = record_camera_data(tmp_path) + source = _write_source(tmp_path, ".mkv") + + destination = _archive_task_recording(_env(recorder), "task2_1") + + assert destination == tmp_path / "task2_1.mkv" + assert destination.read_bytes() == b"video" + assert not source.exists() + + +def test_task_recording_is_noop_when_recording_is_disabled() -> None: + assert _archive_task_recording(_env(), "task2_1") is None From 0219d5ae43486bb64f01513359cc36bddb57029b Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:33:57 +0800 Subject: [PATCH 68/85] feat(action-engine): implement single-arm E3 pour workflow --- .../action_engine/capabilities/atomic.py | 4 +- .../action_engine/domain/task_contracts.py | 2 +- .../gen_sim/action_engine/domain/v2.py | 11 ++ .../action_engine/runtime/grounding.py | 28 ++-- .../action_engine/runtime/predicates.py | 5 + .../gen_sim/action_engine/tasks/recipes.py | 73 +++++++--- embodichain/gen_sim/task_engine/ontology.py | 4 +- .../gen_sim/task_engine/scene/feasibility.py | 25 ---- .../runtime/test_runtime_contracts.py | 39 ++---- .../action_engine/tasks/test_e3_pour.py | 132 ++++++++++++++++++ .../action_engine/tasks/test_factory.py | 28 ++-- .../task_engine/scene/test_scene_boundary.py | 13 +- 12 files changed, 254 insertions(+), 110 deletions(-) create mode 100644 tests/gen_sim/action_engine/tasks/test_e3_pour.py diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 8d3d58657..504bb2a15 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -852,7 +852,7 @@ def _verify_axis_alignment( def _resolve_pour_contract(node: Mapping[str, Any]) -> ResolvedActionContract: - """Retain one verified holder until observable content transfer succeeds.""" + """Retain one verified holder until the E3 action chain completes.""" object_uid = _required_string(node.get("object_uid"), "node.object_uid") actor = node.get("actor", {}) binding = node.get("target_binding", {}) @@ -949,7 +949,7 @@ def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: if node.get("role") == "cleanup": required_home = ( node.get("task_type") == "E2" and binding.get("operation") == "e2_home" - ) + ) or (node.get("task_type") == "E3" and binding.get("operation") == "e3_home") return ResolvedActionContract( requires=(StateAtom("arm_clear", arm=arm),), effects=( diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py index 6d7a66cad..d160a5d6c 100644 --- a/embodichain/gen_sim/action_engine/domain/task_contracts.py +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -71,7 +71,7 @@ def normalize_placement_relation(value: Any) -> str: { "E1": ("PickUp", "MoveHeldObject", "Place"), "E2": ("AxisAlign",), - "E3": ("Pour",), + "E3": ("PickUp", "MoveHeldObject", "Pour", "Place"), "E4": ("PickUp", "MoveHeldObject", "HandOver"), "E5": ("CoordinatedPickment",), "E6": ("PullArticulatedPart",), diff --git a/embodichain/gen_sim/action_engine/domain/v2.py b/embodichain/gen_sim/action_engine/domain/v2.py index bae0cbbbe..7496d1717 100644 --- a/embodichain/gen_sim/action_engine/domain/v2.py +++ b/embodichain/gen_sim/action_engine/domain/v2.py @@ -637,6 +637,17 @@ def _validate_task_group_semantics( } for group in groups: task_type = str(group["task_type"]) + if task_type == "E3": + goal = group.get("goal", {}) + unsupported = sorted( + {"pour_mode", "pouring_arm", "holding_arm"} & set(goal) + ) + if unsupported: + raise ValueError( + f"SeedGraph TaskGroup {group['id']!r} uses unsupported " + f"dual-arm E3 fields {unsupported}; regenerate it as a " + "single-arm pour over a fixed target container." + ) group_nodes = [node_by_id[node_id] for node_id in group["node_ids"]] actions = {str(node["atomic_action"]) for node in group_nodes} missing = required_actions[task_type] - actions diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 788dbed15..46366837b 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -848,25 +848,6 @@ def ground( f"Target materializer {capability.target_materializer!r} " "cannot resolve a pour_goal." ) - contents = step.goal.get("contents", ()) - if not isinstance(contents, Sequence) or isinstance( - contents, (str, bytes, bytearray) - ): - raise ValueError("Pour contents must be a list of object bindings.") - content_uids = [ - item.get("object") if isinstance(item, Mapping) else item - for item in contents - ] - if not content_uids or any( - not isinstance(uid, str) or not uid for uid in content_uids - ): - raise ValueError( - "Pour requires independently observable content objects; " - "a texture or contents baked into the source mesh is not " - "physical transfer evidence." - ) - for uid in content_uids: - _object(self.env, str(uid)) policy.setdefault("rotate_angle", math.pi / 2.0) target = PourGoal() elif kind == "articulation_goal": @@ -1797,6 +1778,15 @@ def _semantic_target( target[:, 2, 3] += float(policy["staging_lift_height"]) return target if step.operator == "pour": + if phase == "return": + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + if initial is None: + raise ValueError( + f"Pour return requires an initial pose for {step.object_uid!r}." + ) + return _batched_pose(initial, self.env).clone() if reference_pose is None: raise ValueError("Pour requires a live target-container pose.") target = object_pose.clone() diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index 68e924cf6..210aa7225 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -738,6 +738,11 @@ def evaluate_predicate( ) return _constant(env, False) if kind == "poured": + if spec.get("verification") == "action_completion": + # Reaching semantic-step verification means every required E3 edge + # already completed without a fatal planning or execution failure. + return _constant(env, True) + raw_contents = spec.get("contents", ()) if not isinstance(raw_contents, Sequence) or isinstance( raw_contents, (str, bytes, bytearray) diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index 3df54d5e3..29cefacc0 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -499,30 +499,31 @@ def _recipe( success, ) if task_type == "E3": + unsupported = sorted({"pour_mode", "pouring_arm", "holding_arm"} & set(params)) + if unsupported: + raise ValueError( + "Dual-arm E3 is not supported; remove fields " + f"{unsupported} and use required_arm with a fixed target container." + ) target = str(params["target_role"]) - contents = [{"object": str(uid)} for uid in params.get("content_roles", [])] goal = { "reference_object": target, "relation": "above", "amount": "task_defined", - "contents": deepcopy(contents), } success = { "type": "poured", + "verification": "action_completion", "object": object_uid, "reference_object": target, - "contents": deepcopy(contents), } - specs = [] + specs: list[tuple[str, Mapping[str, Any], str]] = [] if incoming_held_arm is None: specs.append( ( "PickUp", - { - "kind": "object", - "object": object_uid, - "payloads": deepcopy(contents), - }, + {"kind": "object", "object": object_uid}, + role, ) ) specs.extend( @@ -533,8 +534,8 @@ def _recipe( "kind": "semantic_goal", "semantic_step": group_id, "phase": "final", - "payloads": deepcopy(contents), }, + role, ), ( "Pour", @@ -542,14 +543,46 @@ def _recipe( "kind": "pour_goal", "object": object_uid, "reference_object": target, - "payloads": deepcopy(contents), }, + role, + ), + ( + "MoveHeldObject", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "return", + }, + role, + ), + ( + "Place", + {"kind": "current_held_pose"}, + role, + ), + ( + "MoveEndEffector", + { + "kind": "policy_pose", + "source": "release", + "operation": "retreat", + }, + "cleanup", + ), + ( + "MoveJoints", + { + "kind": "joint_state", + "source": "initial", + "operation": "e3_home", + }, + "cleanup", ), ) ) nodes = [] previous = list(dependencies) - for index, (action, binding) in enumerate(specs, start=1): + for index, (action, binding, node_role) in enumerate(specs, start=1): node = _node( group_id, index, @@ -560,8 +593,8 @@ def _recipe( "arm", binding, previous, - role, - success if action == "Pour" else {}, + node_role, + success if index == len(specs) else {}, motion_policy(), ) nodes.append(node) @@ -1054,10 +1087,6 @@ def _terminal_hold( ) -> tuple[str, str] | None: if task_type == "E4": return object_uid, str(params.get("receive_arm", "right_arm")) - if task_type == "E3": - arm = str(params.get("required_arm", "")) - if arm in {"left_arm", "right_arm"}: - return object_uid, 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"}: @@ -1077,7 +1106,13 @@ def _validate_bindings( 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"])) + references = _role_references(instance["params"]) + if instance["task_type"] == "E3": + references -= _role_references( + instance["params"].get("content_roles", []), + "content_roles", + ) + required.update(references) required.discard("table") missing = sorted(required - set(bindings)) if missing: diff --git a/embodichain/gen_sim/task_engine/ontology.py b/embodichain/gen_sim/task_engine/ontology.py index eb3d090a1..574d1362d 100644 --- a/embodichain/gen_sim/task_engine/ontology.py +++ b/embodichain/gen_sim/task_engine/ontology.py @@ -135,8 +135,8 @@ def _contract( ), "E3": _contract( "E3", - "Pick up when needed and pour contents from a source container " - "into a target container.", + "Pick up a source container, execute a tilt-and-restore pour over " + "a fixed target container, then place and home.", frozenset({"target", "relation", "required_arm"}), "rigid_object", frozenset({"graspable", "pourable"}), diff --git a/embodichain/gen_sim/task_engine/scene/feasibility.py b/embodichain/gen_sim/task_engine/scene/feasibility.py index 363eb4be9..e0d6f00b8 100644 --- a/embodichain/gen_sim/task_engine/scene/feasibility.py +++ b/embodichain/gen_sim/task_engine/scene/feasibility.py @@ -116,31 +116,6 @@ def assess( evidence={"action": str(action_name)}, ) ) - if task_type == "E3": - runtime_observation = bool( - manifest.get("adapter_capabilities", {}).get( - "runtime_scene_observation", False - ) - ) - checks.append( - _check( - "content_observation", - step_id, - "runtime_probe" if runtime_observation else "contradicted", - ( - "Runtime scene observation can verify independently " - "modeled contents after pouring." - if runtime_observation - else "E3 requires independently observable content " - "bodies or fluid state; contents baked into one source " - "mesh cannot prove physical transfer." - ), - evidence={ - "runtime_scene_observation": runtime_observation, - "required_evidence": "content_inside_target_container", - }, - ) - ) if task_type == "E8": reference_id = f"{step_id}.object" raw_uids = bindings.get(reference_id, ()) diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index 3decac78e..939e75117 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -3686,7 +3686,7 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - assert handover_pickup.target.semantics.affordance is semantics.affordance -def test_pour_grounding_targets_receiver_and_requires_physical_contents() -> None: +def test_pour_grounding_targets_receiver_without_physical_contents() -> None: task = { "schema_version": TASK_SPEC_SCHEMA, "task_id": "pour_grounding", @@ -3700,7 +3700,6 @@ def test_pour_grounding_targets_receiver_and_requires_physical_contents() -> Non "params": { "source_role": "cup", "target_role": "bin", - "content_roles": ["ball"], "required_arm": "right_arm", }, "depends_on": [], @@ -3711,14 +3710,23 @@ def test_pour_grounding_targets_receiver_and_requires_physical_contents() -> Non "oracle": {}, "metadata": {}, } - bindings = {"cup": "source", "bin": "target", "ball": "content"} + bindings = {"cup": "source", "bin": "target"} program = load_execution_program(instantiate_seed_graph(task, bindings)) step = program.semantic_steps[0] edges = {edge.actions[0]["atomic_action_class"]: edge for edge in program.edges} + staging_edge = next( + edge + for edge in program.edges + if edge.actions[0]["atomic_action_class"] == "MoveHeldObject" + and edge.actions[0]["target_binding"].get("phase") == "final" + ) source = _FakeEntity("source", _pose(-0.2, 0.0, 0.7), _box_vertices(0.05)) target = _FakeEntity("target", _pose(0.2, 0.0, 0.7), _box_vertices(0.10)) - content = _FakeEntity("content", _pose(-0.2, 0.0, 0.7), _box_vertices(0.01)) - env = _FakeEnv({"source": source, "target": target, "content": content}) + env = _FakeEnv({"source": source, "target": target}) + env.agent_initial_object_poses = { + "source": source.get_local_pose(to_matrix=True), + "target": target.get_local_pose(to_matrix=True), + } semantics = ObjectSemantics( affordance=AntipodalAffordance( object_label="source", @@ -3735,7 +3743,7 @@ def test_pour_grounding_targets_receiver_and_requires_physical_contents() -> Non edges["PickUp"].actions[0], step, arm="right_arm", state=state ) staging = grounder.ground( - edges["MoveHeldObject"].actions[0], + staging_edge.actions[0], step, arm="right_arm", state=state, @@ -3755,24 +3763,7 @@ def test_pour_grounding_targets_receiver_and_requires_physical_contents() -> Non assert isinstance(pouring.target, PourGoal) assert pouring.cfg["rotate_angle"] == pytest.approx(torch.pi / 2.0) - task["task_instances"][0]["params"]["content_roles"] = [] - blocked = load_execution_program( - instantiate_seed_graph(task, {"cup": "source", "bin": "target"}) - ) - blocked_step = blocked.semantic_steps[0] - blocked_pour = next( - edge - for edge in blocked.edges - if edge.actions[0]["atomic_action_class"] == "Pour" - ) - blocked_grounder = ActionGrounder(blocked, env, lambda _uid: semantics) - with pytest.raises(ValueError, match="independently observable"): - blocked_grounder.ground( - blocked_pour.actions[0], - blocked_step, - arm="right_arm", - state=state, - ) + assert step.postcondition["verification"] == "action_completion" @pytest.mark.parametrize( diff --git a/tests/gen_sim/action_engine/tasks/test_e3_pour.py b/tests/gen_sim/action_engine/tasks/test_e3_pour.py new file mode 100644 index 000000000..54c61d8b9 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_e3_pour.py @@ -0,0 +1,132 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 contracts for the E3 approximate pouring workflow.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.domain import validate_seed_graph +from embodichain.gen_sim.action_engine.runtime.predicates import evaluate_predicate +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + + +def _task() -> dict: + params = { + "source_role": "source", + "target_role": "target", + "required_arm": "right_arm", + } + return { + "schema_version": "action_engine_task_spec_v2", + "task_id": "e3_single", + "level": "L1", + "instruction": "Pour from the source container into the target container.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "pour", + "task_type": "E3", + "params": params, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "poured"}, + "oracle": {}, + "metadata": {}, + } + + +def _graph() -> dict: + return instantiate_seed_graph( + _task(), + {"source": "source_container", "target": "target_container"}, + ) + + +def test_single_arm_mode_preserves_the_historical_auto_actor_contract() -> None: + task = _task() + task["task_instances"][0]["params"].pop("required_arm") + + graph = instantiate_seed_graph( + task, + {"source": "source_container", "target": "target_container"}, + ) + + assert graph["task_groups"][0]["actor"] == {"mode": "auto"} + assert all(node["actor"] == {"mode": "auto"} for node in graph["nodes"]) + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "PickUp", + "MoveHeldObject", + "Pour", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + + +def test_single_arm_recipe_binds_only_the_source_and_requested_arm() -> None: + graph = _graph() + + assert all(node["object_uid"] == "source_container" for node in graph["nodes"]) + assert all( + node["actor"] == {"mode": "required", "arm": "right_arm"} + for node in graph["nodes"] + ) + assert graph["task_groups"][0]["success"]["verification"] == "action_completion" + + +@pytest.mark.parametrize("field", ["pour_mode", "pouring_arm", "holding_arm"]) +def test_single_arm_recipe_rejects_legacy_dual_arm_fields(field: str) -> None: + task = _task() + task["task_instances"][0]["params"][field] = "dual_arm" + + with pytest.raises(ValueError, match="Dual-arm E3 is not supported"): + instantiate_seed_graph( + task, + {"source": "source_container", "target": "target_container"}, + ) + + +def test_seed_graph_rejects_legacy_dual_arm_e3_goal() -> None: + graph = _graph() + graph["task_groups"][0]["goal"]["pour_mode"] = "dual_arm" + + with pytest.raises(ValueError, match="unsupported dual-arm E3 fields"): + validate_seed_graph(graph) + + +def test_approximate_poured_predicate_needs_no_scene_or_content_state() -> None: + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + ) + + result = evaluate_predicate( + env, + { + "type": "poured", + "verification": "action_completion", + }, + ) + + assert result.tolist() == [True] diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index 45718bbe1..3609bca94 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -298,7 +298,7 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - ) -def test_pour_recipe_establishes_hold_and_requires_observable_contents() -> None: +def test_pour_recipe_completes_release_and_home_without_observable_contents() -> None: task = { "schema_version": "action_engine_task_spec_v2", "task_id": "pour_contents", @@ -312,7 +312,7 @@ def test_pour_recipe_establishes_hold_and_requires_observable_contents() -> None "params": { "source_role": "cup", "target_role": "bin", - "content_roles": ["ball"], + "content_roles": ["unmodeled_liquid"], "required_arm": "right_arm", }, "depends_on": [], @@ -326,7 +326,7 @@ def test_pour_recipe_establishes_hold_and_requires_observable_contents() -> None graph = instantiate_seed_graph( task, - {"cup": "source_cup", "bin": "target_bin", "ball": "content_ball"}, + {"cup": "source_cup", "bin": "target_bin"}, ) nodes = graph["nodes"] group = graph["task_groups"][0] @@ -335,13 +335,25 @@ def test_pour_recipe_establishes_hold_and_requires_observable_contents() -> None "PickUp", "MoveHeldObject", "Pour", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", ] - assert nodes[1]["depends_on"] == [nodes[0]["id"]] - assert nodes[2]["depends_on"] == [nodes[1]["id"]] - assert nodes[2]["contract"]["completion"] == "terminal_barrier" + assert all( + node["depends_on"] == [nodes[index - 1]["id"]] + for index, node in enumerate(nodes[1:], start=1) + ) assert nodes[2]["contract"]["failure_policy"] == "task_required" - assert group["success"]["contents"] == [{"object": "content_ball"}] - assert nodes[2]["target_binding"]["payloads"] == [{"object": "content_ball"}] + assert nodes[3]["target_binding"]["phase"] == "return" + assert nodes[4]["contract"]["effects"][-1]["atom"] == { + "predicate": "object_free", + "object_uid": "source_cup", + } + assert nodes[-1]["contract"]["completion"] == "terminal_barrier" + assert group["success"]["verification"] == "action_completion" + assert "contents" not in group["goal"] + assert all("payloads" not in node["target_binding"] for node in nodes) def test_handover_to_place_uses_receiver_hold_without_repickup() -> None: diff --git a/tests/gen_sim/task_engine/scene/test_scene_boundary.py b/tests/gen_sim/task_engine/scene/test_scene_boundary.py index 69b81af89..5321f83dc 100644 --- a/tests/gen_sim/task_engine/scene/test_scene_boundary.py +++ b/tests/gen_sim/task_engine/scene/test_scene_boundary.py @@ -263,7 +263,7 @@ def test_planning_only_action_is_reported_as_contradicted(tmp_path: Path) -> Non assert any("planning-only" in blocker for blocker in report["blockers"]) -def test_e3_requires_runtime_content_observation_before_execution( +def test_e3_does_not_require_runtime_content_observation_before_execution( tmp_path: Path, ) -> None: manifest = SceneEngineV1Adapter().adapt_prepared_scene( @@ -279,15 +279,8 @@ def test_e3_requires_runtime_content_observation_before_execution( task_actions={"E3": ("PickUp", "MoveHeldObject", "Pour")}, ) - assert report["status"] == "contradicted" - assert report["remediation_class"] == "terminal" - assert any( - check["kind"] == "content_observation" and check["status"] == "contradicted" - for check in report["checks"] - ) - assert any( - "baked into one source mesh" in blocker for blocker in report["blockers"] - ) + assert report["status"] != "contradicted" + assert all(check["kind"] != "content_observation" for check in report["checks"]) def test_e8_requires_explicit_setting_to_angle_mapping(tmp_path: Path) -> None: From ebdb66a5a5b49ab6a06111444fe1f5c93cea0d29 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Wed, 26 Aug 2026 04:32:01 +0800 Subject: [PATCH 69/85] feat(action-engine): add contract-driven payload-aware coordinated manipulation --- .../action_engine/capabilities/atomic.py | 11 +- .../gen_sim/action_engine/compiler/v2.py | 7 +- .../action_engine/config/defaults.yaml | 2 + .../action_engine/domain/task_contracts.py | 14 + .../gen_sim/action_engine/domain/v2.py | 200 +++------ .../action_engine/environment/agent_env.py | 6 +- .../gen_sim/action_engine/planning/linker.py | 43 +- .../gen_sim/action_engine/runtime/actions.py | 420 ++++++++++++++---- .../gen_sim/action_engine/runtime/executor.py | 185 +++++++- .../action_engine/runtime/grounding.py | 21 +- .../action_engine/runtime/predicates.py | 15 +- .../gen_sim/action_engine/tasks/assembly.py | 22 +- .../gen_sim/action_engine/tasks/recipes.py | 112 +++-- embodichain/gen_sim/task_engine/defaults.yaml | 2 +- embodichain/gen_sim/task_engine/ontology.py | 69 ++- .../gen_sim/task_engine/scene/feasibility.py | 15 +- embodichain/gen_sim/video_archive.py | 9 +- .../capabilities/test_atomic_v2.py | 36 +- .../config/test_runtime_policy.py | 6 + .../domain/test_task_contracts.py | 14 + .../action_engine/runtime/test_actions.py | 141 +++++- .../runtime/test_runtime_contracts.py | 188 +++++++- .../action_engine/tasks/test_factory.py | 1 + .../tasks/test_payload_contracts.py | 253 +++++++++++ tests/gen_sim/test_video_archive.py | 18 +- 25 files changed, 1434 insertions(+), 376 deletions(-) create mode 100644 tests/gen_sim/action_engine/tasks/test_payload_contracts.py diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 504bb2a15..86404ea9e 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -903,13 +903,12 @@ def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: release_role = binding.get("coordinated_release_role") if release_role is not None: if ( - node.get("task_type") != "E5" - or node.get("control") != "hand" + node.get("control") != "hand" or binding.get("source") != "gripper_open" or not node.get("sync_group") ): raise ValueError( - "Coordinated MoveJoints release requires an E5 synchronized " + "Coordinated MoveJoints release requires a synchronized " "hand action targeting gripper_open." ) if release_role not in {"participant", "commit"}: @@ -947,9 +946,9 @@ def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: ), ) if node.get("role") == "cleanup": - required_home = ( - node.get("task_type") == "E2" and binding.get("operation") == "e2_home" - ) or (node.get("task_type") == "E3" and binding.get("operation") == "e3_home") + required_home = binding.get("required_home", False) + if not isinstance(required_home, bool): + raise TypeError("joint_state required_home must be a boolean.") return ResolvedActionContract( requires=(StateAtom("arm_clear", arm=arm),), effects=( diff --git a/embodichain/gen_sim/action_engine/compiler/v2.py b/embodichain/gen_sim/action_engine/compiler/v2.py index 13942afe7..9677ce780 100644 --- a/embodichain/gen_sim/action_engine/compiler/v2.py +++ b/embodichain/gen_sim/action_engine/compiler/v2.py @@ -397,7 +397,12 @@ def _v2_actor(value: Mapping[str, Any]) -> dict[str, Any]: def _task_type(operator: str) -> str: - return _OPERATOR_TASK_TYPES.get(operator, "E1") + try: + return _OPERATOR_TASK_TYPES[operator] + except KeyError as exc: + raise ValueError( + f"Semantic operator {operator!r} has no registered task contract." + ) from exc def _level(groups: Sequence[Mapping[str, Any]]) -> str: diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index f9c2a7122..8592b7151 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -230,6 +230,8 @@ runtime: lift_height: 0.08 middle_empty_ratio: 0.4 is_filter_ground_collision: false + release_sample_interval: 60 + release_gripper_tolerance: 0.08 postcondition_tolerance: 0.06 HandOver: sample_interval: 140 diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py index d160a5d6c..5ee1228d3 100644 --- a/embodichain/gen_sim/action_engine/domain/task_contracts.py +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -94,6 +94,13 @@ class TaskContract: required_affordances: frozenset[str] success_type: str scene_affordances: frozenset[str] + primary_role_field: str + resource_mode: str + moves_primary_object: bool + accepts_direct_payloads: bool + direct_payload_relations: frozenset[str] + accepts_incoming_hold: bool + terminal_success_types: tuple[tuple[str, str], ...] def _action_contract(value: SemanticTaskContract) -> TaskContract: @@ -106,6 +113,13 @@ def _action_contract(value: SemanticTaskContract) -> TaskContract: required_affordances=value.required_affordances, success_type=value.success_type, scene_affordances=value.scene_affordances, + primary_role_field=value.primary_role_field, + resource_mode=value.resource_mode, + moves_primary_object=value.moves_primary_object, + accepts_direct_payloads=value.accepts_direct_payloads, + direct_payload_relations=value.direct_payload_relations, + accepts_incoming_hold=value.accepts_incoming_hold, + terminal_success_types=value.terminal_success_types, ) diff --git a/embodichain/gen_sim/action_engine/domain/v2.py b/embodichain/gen_sim/action_engine/domain/v2.py index 7496d1717..d8f581a59 100644 --- a/embodichain/gen_sim/action_engine/domain/v2.py +++ b/embodichain/gen_sim/action_engine/domain/v2.py @@ -691,156 +691,78 @@ def _validate_ownership_transitions( nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]], ) -> None: - """Check release/reacquire and explicit single-arm hold transitions. - - An ordinary E2 -> E4 transition persists the supported upright state, ends - the predecessor resource lease, and lets E4 acquire a fresh transfer grasp. - E4 -> E1 keeps receiver ownership because the exchanged object is not yet - supported. Recovery groups may preserve an explicitly requested hold. - """ + """Validate object ownership by folding persisted action contracts.""" node_by_id = {str(node["id"]): node for node in nodes} - group_by_id = {str(group["id"]): group for group in groups} - nodes_by_group = { - group_id: [node_by_id[node_id] for node_id in group["node_ids"]] - for group_id, group in group_by_id.items() - } + ownership: dict[str, str] = {} - def direct_predecessor( - group: Mapping[str, Any], task_type: str - ) -> Mapping[str, Any] | None: - for dependency in group.get("depends_on", []): - candidate = group_by_id.get(str(dependency)) - if candidate is not None and candidate.get("task_type") == task_type: - return candidate - return None - - def held_arm(node: Mapping[str, Any]) -> str | None: - precondition = node.get("precondition", {}) - if ( - isinstance(precondition, Mapping) - and precondition.get("type") == "object_held" - ): - arm = str(precondition.get("arm", "")) - if arm in {"left_arm", "right_arm"}: - return arm - actor = node.get("actor", {}) - if isinstance(actor, Mapping) and actor.get("mode") == "required": - arm = str(actor.get("arm", "")) - if arm in {"left_arm", "right_arm"}: - return arm - return None - - for group_id, group in group_by_id.items(): - task_type = str(group.get("task_type")) - group_nodes = nodes_by_group[group_id] - actions = [str(node.get("atomic_action")) for node in group_nodes] + for group in groups: + group_id = str(group["id"]) object_uid = str(group.get("object_uid")) - - if task_type == "E2": - handover = next( - ( - candidate - for candidate in groups - if candidate.get("task_type") == "E4" - and group_id - in {str(item) for item in candidate.get("depends_on", [])} - and str(candidate.get("object_uid")) == object_uid - ), - None, - ) - if handover is None: - continue - if ( - group.get("goal", {}).get("terminal_behavior") == "hold" - and group.get("role") != "recovery" - ): - raise ValueError( - f"SeedGraph E2 group {group_id!r} may not preserve a holder " - "across an ordinary E2->E4 TaskGroup boundary." - ) - if ( - group.get("role") != "recovery" - and "Place" not in actions - and "AxisAlign" not in actions - ): + current = ownership.get(object_uid, "free") + may_rebase_recovery_entry = group.get("role") == "recovery" + for node_id in group["node_ids"]: + node = node_by_id[str(node_id)] + contract = node.get("contract", {}) + if not isinstance(contract, Mapping): raise ValueError( - f"SeedGraph E2 group {group_id!r} must release its supported " - "object before E4 reacquires it." - ) - - if task_type == "E4": - predecessor = direct_predecessor(group, "E2") - if ( - predecessor is not None - and str(predecessor.get("object_uid")) == object_uid - ): - predecessor_nodes = nodes_by_group[str(predecessor["id"])] - preserves_hold = ( - predecessor.get("role") == "recovery" - and predecessor.get("goal", {}).get("terminal_behavior") == "hold" + f"SeedGraph node {node_id!r} requires an Action Contract." ) - if preserves_hold: - if "PickUp" in actions or not group_nodes: - raise ValueError( - f"SeedGraph E4 group {group_id!r} must consume the " - "recovery-held object without PickUp." - ) - first = group_nodes[0] - holder_arm = next( - ( - held_arm(node) - for node in reversed(predecessor_nodes) - if held_arm(node) is not None - ), - None, + for requirement in contract.get("requires", []): + if not isinstance(requirement, Mapping): + continue + if requirement.get("object_uid") != object_uid: + continue + predicate = str(requirement.get("predicate", "")) + expected = ( + "free" + if predicate == "object_free" + else ( + "coordinated" + if predicate == "object_coordinated_held" + else str(requirement.get("arm", "")) ) - if ( - str(first.get("atomic_action")) != "MoveHeldObject" - or held_arm(first) is None - or held_arm(first) != holder_arm - ): - raise ValueError( - f"SeedGraph E2->E4 recovery holder mismatch for object " - f"{object_uid!r}." - ) - else: - predecessor_actions = { - str(node.get("atomic_action")) for node in predecessor_nodes + ) + if ( + predicate + in { + "object_free", + "object_held", + "object_coordinated_held", } - if not predecessor_actions.intersection({"Place", "AxisAlign"}): - raise ValueError( - f"SeedGraph E2 predecessor {predecessor['id']!r} must " - "release its object before E4." - ) - if ( - not group_nodes - or str(group_nodes[0].get("atomic_action")) != "PickUp" - ): + and current != expected + ): + if may_rebase_recovery_entry: + current = expected + else: raise ValueError( - f"SeedGraph E4 group {group_id!r} must reacquire the " - "supported E2 object with PickUp." + f"SeedGraph group {group_id!r} requires {object_uid!r} " + f"ownership {expected!r}, but the preceding contract " + f"flow provides {current!r}." ) - - if task_type == "E1": - predecessor = direct_predecessor(group, "E4") - if ( - predecessor is not None - and str(predecessor.get("object_uid")) == object_uid - ): - if "PickUp" in actions or not group_nodes: - raise ValueError( - f"SeedGraph E1 group {group_id!r} must preserve the E4 receiver hold " - "without PickUp." - ) - first = group_nodes[0] + may_rebase_recovery_entry = False + for effect in contract.get("effects", []): + if not isinstance(effect, Mapping) or effect.get("op") != "add": + continue + atom = effect.get("atom", {}) if ( - str(first.get("atomic_action")) != "MoveHeldObject" - or held_arm(first) is None + not isinstance(atom, Mapping) + or atom.get("object_uid") != object_uid ): - raise ValueError( - f"SeedGraph E1 group {group_id!r} must start with MoveHeldObject " - "from the receiver hold." - ) + continue + predicate = str(atom.get("predicate", "")) + if predicate == "object_free": + current = "free" + elif predicate == "object_coordinated_held": + current = "coordinated" + elif predicate == "object_held": + arm = str(atom.get("arm", "")) + if not arm: + raise ValueError( + f"SeedGraph node {node_id!r} adds object_held without " + "an ownership resource." + ) + current = arm + ownership[object_uid] = current def _validate_group_dependency_alignment( diff --git a/embodichain/gen_sim/action_engine/environment/agent_env.py b/embodichain/gen_sim/action_engine/environment/agent_env.py index 7ed4477d3..93694dba1 100644 --- a/embodichain/gen_sim/action_engine/environment/agent_env.py +++ b/embodichain/gen_sim/action_engine/environment/agent_env.py @@ -125,8 +125,10 @@ def _capture_runtime_state(self) -> None: dtype=self.init_qpos.dtype, device=self.init_qpos.device, ).flatten() - self.left_arm_current_gripper_state = self._hand_qpos("left") - self.right_arm_current_gripper_state = self._hand_qpos("right") + for side in ("left", "right"): + hand_qpos = self._hand_qpos(side) + setattr(self, f"{side}_arm_init_gripper_state", hand_qpos.clone()) + setattr(self, f"{side}_arm_current_gripper_state", hand_qpos) self.update_obj_info() self.agent_initial_object_poses = { uid: item["pose"].clone() for uid, item in self.obj_info.items() diff --git a/embodichain/gen_sim/action_engine/planning/linker.py b/embodichain/gen_sim/action_engine/planning/linker.py index c1db21e84..6a391ae0f 100644 --- a/embodichain/gen_sim/action_engine/planning/linker.py +++ b/embodichain/gen_sim/action_engine/planning/linker.py @@ -29,6 +29,7 @@ build_atomic_capability_registry, ) from embodichain.gen_sim.action_engine.domain import ( + task_contract, validate_seed_graph, validate_task_spec, ) @@ -98,31 +99,28 @@ def link_task_dependencies( distinct_arm_pairs = _distinct_arm_pairs(task.get("metadata", {})) linked: list[dict[str, str]] = [] - latest_by_object: dict[str, tuple[str, str]] = {} + latest_by_object: dict[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) + previous is not None + and previous not in dependencies[instance_id] + and not _reaches(dependencies, previous, instance_id) ): - dependencies[instance_id].add(previous[0]) - dependency_order[instance_id].append(previous[0]) + dependencies[instance_id].add(previous) + dependency_order[instance_id].append(previous) linked.append( { - "from": previous[0], + "from": previous, "to": instance_id, "reason": "causal", - "detail": f"object_held:{primary}", + "detail": f"object_flow:{primary}", } ) _assert_acyclic(dependencies, "TaskSpec causal linking") - latest_by_object[primary] = (instance_id, task_type) + latest_by_object[primary] = instance_id for later_index, later_id in enumerate(order): for earlier_id in order[:later_index]: @@ -392,8 +390,9 @@ def _task_claims( instance: Mapping[str, Any], bindings: Mapping[str, str] ) -> list[dict[str, str]]: task_type = str(instance["task_type"]) + contract = task_contract(task_type) params = _resolve_roles(instance.get("params", {}), bindings) - primary_key = "source_role" if task_type == "E3" else "object_role" + primary_key = contract.primary_role_field primary = params.get(primary_key) claims: list[dict[str, str]] = [] if isinstance(primary, str) and primary: @@ -408,7 +407,7 @@ def _task_claims( for payload in payloads: if isinstance(payload, str) and payload and payload != primary: claims.append(_claim(f"object:{payload}", "exclusive")) - if task_type == "E4": + if contract.resource_mode == "handover": transfer = str(params.get("transfer_arm", "")) receive = str(params.get("receive_arm", "")) if transfer not in {"left_arm", "right_arm"} or receive not in { @@ -416,19 +415,24 @@ def _task_claims( "right_arm", }: raise ValueError( - "E4 contract linking requires explicit transfer/receive arms." + "Handover resource mode requires explicit transfer/receive arms." ) if transfer == receive: - raise ValueError("E4 transfer_arm and receive_arm must be distinct.") + raise ValueError("Handover transfer_arm and receive_arm must be distinct.") claims.extend((_claim(f"arm:{transfer}"), _claim(f"arm:{receive}"))) - elif task_type == "E5": + elif contract.resource_mode == "coordinated": claims.extend((_claim("arm:left_arm"), _claim("arm:right_arm"))) - else: + elif contract.resource_mode == "single_arm": 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")) + else: + raise ValueError( + f"TaskGroup {instance.get('id')!r} has unsupported resource mode " + f"{contract.resource_mode!r}." + ) return _merge_claims(claims) @@ -436,8 +440,9 @@ def _task_primary_object( instance: Mapping[str, Any], bindings: Mapping[str, str] ) -> str: task_type = str(instance["task_type"]) + contract = task_contract(task_type) params = _resolve_roles(instance.get("params", {}), bindings) - key = "source_role" if task_type == "E3" else "object_role" + key = contract.primary_role_field value = params.get(key) if not isinstance(value, str) or not value: raise ValueError( diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 922e904a6..b0f88e1d8 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -192,6 +192,7 @@ def __init__( self.capabilities = capability_registry or build_atomic_capability_registry() self._motion_generator: MotionGenerator | None = None self._atomic_engine: AtomicActionEngine | None = None + self._coordinated_engines: dict[bool, AtomicActionEngine] = {} self._semantics: dict[str, ObjectSemantics] = {} self._scene_time = 0.0 if scene_provider is not None and not isinstance(scene_provider, SceneProvider): @@ -228,9 +229,14 @@ def start_session( capability = self.capabilities.require_executable(grounded.action_class) state = state or self.initial_state() grounded = self._select_transport_yaw(grounded, state) + grounded = self._adapt_coordinated_pickment_grasps( + grounded, + capability, + )[0] context = self._planning_context(state, grounded) - invocation = self._invocation(grounded, capability) - return self._engine().start((invocation,), context) + engine = self._engine_for(grounded, capability) + invocation = self._invocation(grounded, capability, engine=engine) + return engine.start((invocation,), context) def _build_scene_provider(self) -> SceneProvider | None: """Create the shared live rigid-object provider when entities are available.""" @@ -343,25 +349,51 @@ def plan( state = state or self.initial_state() grounded = self._select_transport_yaw(grounded, state) context = self._planning_context(state, grounded) - grounded_candidates = self._adapt_axis_align_body_grasps( + coordinated_candidates = self._adapt_coordinated_pickment_grasps( grounded, - context, capability, ) - selected: tuple[GroundedAction, ActionInvocation, ActionPlan] | None = None + grounded_candidates = tuple( + candidate + for coordinated in coordinated_candidates + for candidate in self._adapt_axis_align_body_grasps( + coordinated, + context, + capability, + ) + ) + selected: ( + tuple[ + GroundedAction, + ActionInvocation, + ActionPlan, + AtomicActionEngine, + ] + | None + ) = None best_failure_count = self.num_envs + 1 for candidate in grounded_candidates: - candidate_invocation = self._invocation(candidate, capability) - candidate_plan = self._engine().plan(candidate_invocation, context) + candidate_engine = self._engine_for(candidate, capability) + candidate_invocation = self._invocation( + candidate, + capability, + engine=candidate_engine, + ) + candidate_plan = candidate_engine.plan(candidate_invocation, context) failure_count = int((~candidate_plan.plan_success).sum().item()) if selected is None or failure_count < best_failure_count: - selected = (candidate, candidate_invocation, candidate_plan) + selected = ( + candidate, + candidate_invocation, + candidate_plan, + candidate_engine, + ) best_failure_count = failure_count if failure_count == 0: break if selected is None: raise RuntimeError("Atomic action adaptation produced no plan candidate.") - grounded, invocation, plan = selected + grounded, invocation, plan, selected_engine = selected selected_positions = self._positions_with_agent_holds( plan, grounded, @@ -409,7 +441,7 @@ def plan( dynamic_collision_mode=DynamicCollisionMode.OFF, plan_opts=None, ) - fallback_plan = self._engine().plan( + fallback_plan = selected_engine.plan( replace(invocation, motion_policy=fallback_policy), context, ) @@ -584,6 +616,163 @@ def _adapt_axis_align_body_grasps( ) return tuple(candidates) + def _adapt_coordinated_pickment_grasps( + self, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> tuple[GroundedAction, ...]: + """Build deterministic geometry-ranked E5 partition candidates. + + ``left_to_right_arm_direction`` remains the live base-to-base direction: + it labels the two participant regions and is not the transport direction. + Object geometry only adjusts how much of the projected middle is excluded. + """ + if capability.config_materializer != "coordinated_pickment": + return (grounded,) + target = self._validate_coordinated_pickment_goal(grounded) + live_pose = grounded.object_pose + expected_shape = (self.num_envs, 4, 4) + if ( + not isinstance(live_pose, torch.Tensor) + or live_pose.shape != expected_shape + or not bool(torch.isfinite(live_pose).all()) + ): + raise ValueError( + "CoordinatedPickment requires a finite grounded live object pose " + f"with shape {expected_shape}." + ) + live_pose = live_pose.to(device=self.device, dtype=torch.float32).clone() + affordance = target.semantics.affordance + assert isinstance(affordance, AntipodalAffordance) + vertices = torch.as_tensor( + affordance.mesh_vertices, + dtype=torch.float32, + device=self.device, + ) + if ( + vertices.ndim != 2 + or vertices.shape[1] != 3 + or vertices.shape[0] < 3 + or not bool(torch.isfinite(vertices).all()) + ): + raise ValueError( + "CoordinatedPickment mesh vertices must be finite with shape (N, 3)." + ) + + left_base, right_base = self._coordinated_arm_bases() + arm_directions = right_base[:, :3, 3] - left_base[:, :3, 3] + arm_norms = torch.linalg.vector_norm(arm_directions, dim=1, keepdim=True) + if not bool(torch.isfinite(arm_directions).all()) or bool( + (arm_norms <= 1.0e-6).any() + ): + raise ValueError( + "Coordinated pickup requires distinct finite left/right arm bases." + ) + arm_directions = arm_directions / arm_norms + shared_direction = arm_directions[0] + if bool( + (torch.matmul(arm_directions, shared_direction).abs() < 1.0 - 1.0e-4).any() + ): + raise ValueError( + "CoordinatedPickment requires one shared base-to-base direction " + "across vectorized environments." + ) + + centered = vertices - vertices.mean(dim=0, keepdim=True) + covariance = centered.transpose(0, 1) @ centered / float(vertices.shape[0]) + eigenvalues, eigenvectors = torch.linalg.eigh(covariance) + principal_local = eigenvectors[:, -1] + principal_world = torch.matmul( + live_pose[:, :3, :3], + principal_local, + ) + principal_world = principal_world / torch.linalg.vector_norm( + principal_world, + dim=1, + keepdim=True, + ).clamp_min(1.0e-6) + arm_alignment = torch.abs((principal_world * arm_directions).sum(dim=1)) + elongation_ratio = torch.sqrt( + eigenvalues[-1].clamp_min(1.0e-12) / eigenvalues[-2].clamp_min(1.0e-12) + ) + elongation_confidence = torch.clamp( + (elongation_ratio - 1.0) / 1.5, + min=0.0, + max=1.0, + ) + base_ratio = float(grounded.cfg.get("middle_empty_ratio", 0.4)) + if not math.isfinite(base_ratio) or not 0.0 <= base_ratio < 1.0: + raise ValueError("middle_empty_ratio must be finite and in [0, 1).") + geometric_ratio = 0.25 + 0.45 * float(arm_alignment.mean()) + confidence = float(elongation_confidence) + preferred_ratio = (1.0 - confidence) * base_ratio + confidence * geometric_ratio + raw_ratios = ( + preferred_ratio, + base_ratio, + preferred_ratio - 0.15, + preferred_ratio + 0.15, + ) + ratios: list[float] = [] + for raw_ratio in raw_ratios: + ratio = min(0.90, max(0.05, float(raw_ratio))) + if not any(abs(ratio - existing) <= 1.0e-6 for existing in ratios): + ratios.append(ratio) + + approach = grounded.cfg.get("approach_direction", (0.0, 0.0, -1.0)) + approach = torch.as_tensor( + approach, + dtype=torch.float32, + device=self.device, + ) + if approach.shape != (3,) or not bool(torch.isfinite(approach).all()): + raise ValueError("approach_direction must be a finite vector shaped (3,).") + approach_norm = torch.linalg.vector_norm(approach) + if float(approach_norm) <= 1.0e-6: + raise ValueError("approach_direction must be non-zero.") + approach = approach / approach_norm + trace = { + "strategy": "live_geometry_partition_search", + "local_principal_axis": principal_local.detach().cpu().tolist(), + "world_principal_axes": principal_world.detach().cpu().tolist(), + "elongation_ratio": float(elongation_ratio), + "elongation_confidence": confidence, + "arm_axis_alignment": arm_alignment.detach().cpu().tolist(), + "left_to_right_arm_direction": shared_direction.detach().cpu().tolist(), + "approach_direction": approach.detach().cpu().tolist(), + "candidate_middle_empty_ratios": list(ratios), + } + candidates: list[GroundedAction] = [] + for candidate_index, ratio in enumerate(ratios): + cfg = { + **grounded.cfg, + "left_to_right_arm_direction": shared_direction.clone(), + "approach_direction": approach.clone(), + "middle_empty_ratio": ratio, + } + motion_policy = { + **grounded.motion_policy, + "coordinated_grasp": { + **trace, + "candidate_index": candidate_index, + "selected_middle_empty_ratio": ratio, + }, + } + candidates.append( + replace( + grounded, + target=replace(target, object_initial_pose=live_pose.clone()), + cfg=cfg, + object_pose=live_pose.clone(), + motion_policy=motion_policy, + ) + ) + return tuple(candidates) + + def _coordinated_arm_bases(self) -> tuple[torch.Tensor, torch.Tensor]: + from .frames import arm_base_poses + + return arm_base_poses(self.env) + def _search_reachable_retreat( self, *, @@ -837,6 +1026,9 @@ def _planner_trace( body_grasp = grounded.motion_policy.get("body_grasp") if isinstance(body_grasp, Mapping): trace["body_grasp"] = deepcopy(dict(body_grasp)) + coordinated_grasp = grounded.motion_policy.get("coordinated_grasp") + if isinstance(coordinated_grasp, Mapping): + trace["coordinated_grasp"] = deepcopy(dict(coordinated_grasp)) return trace def _select_transport_yaw( @@ -1085,6 +1277,8 @@ def _invocation( self, grounded: GroundedAction, capability: AtomicCapability, + *, + engine: AtomicActionEngine | None = None, ) -> ActionInvocation: if capability.resource_mode == "coordinated_object": strategy = str(self.planner_policy["coordinated_strategy"]) @@ -1111,7 +1305,7 @@ def _invocation( return ActionInvocation( skill_id=str(capability.action_type.skill_id), goal=goal, - binding=self._binding(grounded, capability), + binding=self._binding(grounded, capability, engine=engine), motion_policy=MotionPolicy( strategy=strategy, sample_count=sample_count, @@ -1143,8 +1337,10 @@ def _binding( self, action: GroundedAction, capability: AtomicCapability, + *, + engine: AtomicActionEngine | None = None, ) -> ActionBinding: - engine = self._engine() + engine = self._engine() if engine is None else engine contract = getattr(capability.action_type, "binding_contract", None) if contract is None: return ActionBinding(owner_id=engine.binding_owner_id) @@ -1303,17 +1499,20 @@ def _build_coordinated_pickment_config( action: GroundedAction, capability: AtomicCapability, ) -> Any: - from .frames import arm_base_poses - policy = self._config_policy(action) - left_base, right_base = arm_base_poses(self.env) + left_base, right_base = self._coordinated_arm_bases() direction = right_base[0, :3, 3] - left_base[0, :3, 3] norm = torch.linalg.vector_norm(direction) if not torch.isfinite(direction).all() or norm <= 1.0e-6: raise ValueError( "Coordinated pickup requires distinct finite left/right arm bases." ) - policy["left_to_right_arm_direction"] = direction / norm + policy.setdefault("left_to_right_arm_direction", direction / norm) + for name in ("approach_direction", "left_to_right_arm_direction"): + if name in policy and not isinstance(policy[name], torch.Tensor): + policy[name] = torch.as_tensor( + policy[name], dtype=torch.float32, device=self.device + ) return capability.config_type( **_supported_kwargs(capability.config_type, policy) ) @@ -1526,81 +1725,113 @@ def joint_ids(self, arm: str, *, include_hand: bool) -> list[int]: def _engine(self) -> AtomicActionEngine: if self._atomic_engine is None: - from embodichain.gen_sim.action_engine.capabilities import ( - HeldObjectHandOver, - ) - - from .atomic_compat import ExactTargetMoveHeldObject - - engine = AtomicActionEngine( + self._atomic_engine = self._new_engine( self._generator(), - control_profiles=self._control_profiles(), - grasp_pose_generators=self._grasp_pose_generators(), + filter_ground_collision=True, ) - engine.register(ExactTargetMoveHeldObject(), replace=True) - engine.register(HeldObjectHandOver(), replace=True) - self._atomic_engine = engine return self._atomic_engine + def _engine_for( + self, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> AtomicActionEngine: + if capability.config_materializer != "coordinated_pickment": + return self._engine() + filter_ground_collision = grounded.cfg.get( + "is_filter_ground_collision", + True, + ) + if not isinstance(filter_ground_collision, bool): + raise TypeError("is_filter_ground_collision must be a boolean.") + if filter_ground_collision: + return self._engine() + cached = self._coordinated_engines.get(filter_ground_collision) + if cached is None: + cached = self._new_engine( + MotionGenerator(cfg=self._motion_generator_cfg()), + filter_ground_collision=filter_ground_collision, + ) + self._coordinated_engines[filter_ground_collision] = cached + return cached + + def _new_engine( + self, + motion_generator: MotionGenerator, + *, + filter_ground_collision: bool, + ) -> AtomicActionEngine: + from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver + + from .atomic_compat import ExactTargetMoveHeldObject + + engine = AtomicActionEngine( + motion_generator, + control_profiles=self._control_profiles(), + grasp_pose_generators=self._grasp_pose_generators( + filter_ground_collision=filter_ground_collision, + ), + ) + engine.register(ExactTargetMoveHeldObject(), replace=True) + engine.register(HeldObjectHandOver(), replace=True) + return engine + def _generator(self) -> MotionGenerator: if self._motion_generator is None: - backend = str(self.planner_policy.get("backend", "curobo")) - if backend == "curobo": - options = dict(self.planner_policy.get("curobo", {})) - obstacle_uids = tuple( - dict.fromkeys( - [ - *self.planner_policy.get("static_obstacle_uids", ()), - *self.planner_policy.get("dynamic_obstacle_uids", ()), - ] - ) - ) - rigid_objects: dict[str, Any] = {} - for uid in obstacle_uids: - obstacle_uid = str(uid) - entity = self.env.sim.get_rigid_object(obstacle_uid) - if entity is None: - raise ValueError(f"Unknown cuRobo obstacle {uid!r}.") - rigid_objects[obstacle_uid] = entity - obstacle_representation = str( - options.get("obstacle_representation", "cuboid") - ) - world = CuroboWorldCfg( - rigid_objects=rigid_objects or None, - obstacle_representation=obstacle_representation, - collision_cache=_collision_cache_for_world( - obstacle_representation, - len(rigid_objects), - ), - dynamic_obstacle_names=[ - str(uid) - for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) - ], - multi_env=bool(options.get("multi_env", False)), - ) - planner_cfg = CuroboPlannerCfg( - robot_uid=self.env.robot.uid, - log_level=str(options.get("log_level", "error")), - world=world, - use_cuda_graph=bool(options.get("use_cuda_graph", True)), - preserve_plan_samples=bool( - options.get("preserve_plan_samples", False) - ), - max_attempts=int(options.get("max_attempts", 5)), - collision_activation_distance=float( - options.get("collision_activation_distance", 0.01) - ), - ) - elif backend == "toppra": - planner_cfg = ToppraPlannerCfg(robot_uid=self.env.robot.uid) - else: - raise ValueError( - f"Unsupported Action Engine planner backend {backend!r}." + self._motion_generator = MotionGenerator(cfg=self._motion_generator_cfg()) + return self._motion_generator + + def _motion_generator_cfg(self) -> MotionGenCfg: + backend = str(self.planner_policy.get("backend", "curobo")) + if backend == "curobo": + options = dict(self.planner_policy.get("curobo", {})) + obstacle_uids = tuple( + dict.fromkeys( + [ + *self.planner_policy.get("static_obstacle_uids", ()), + *self.planner_policy.get("dynamic_obstacle_uids", ()), + ] ) - self._motion_generator = MotionGenerator( - cfg=MotionGenCfg(planner_cfg=planner_cfg) ) - return self._motion_generator + rigid_objects: dict[str, Any] = {} + for uid in obstacle_uids: + obstacle_uid = str(uid) + entity = self.env.sim.get_rigid_object(obstacle_uid) + if entity is None: + raise ValueError(f"Unknown cuRobo obstacle {uid!r}.") + rigid_objects[obstacle_uid] = entity + obstacle_representation = str( + options.get("obstacle_representation", "cuboid") + ) + world = CuroboWorldCfg( + rigid_objects=rigid_objects or None, + obstacle_representation=obstacle_representation, + collision_cache=_collision_cache_for_world( + obstacle_representation, + len(rigid_objects), + ), + dynamic_obstacle_names=[ + str(uid) + for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ], + multi_env=bool(options.get("multi_env", False)), + ) + planner_cfg = CuroboPlannerCfg( + robot_uid=self.env.robot.uid, + log_level=str(options.get("log_level", "error")), + world=world, + use_cuda_graph=bool(options.get("use_cuda_graph", True)), + preserve_plan_samples=bool(options.get("preserve_plan_samples", False)), + max_attempts=int(options.get("max_attempts", 5)), + collision_activation_distance=float( + options.get("collision_activation_distance", 0.01) + ), + ) + elif backend == "toppra": + planner_cfg = ToppraPlannerCfg(robot_uid=self.env.robot.uid) + else: + raise ValueError(f"Unsupported Action Engine planner backend {backend!r}.") + return MotionGenCfg(planner_cfg=planner_cfg) def _control_profiles(self) -> dict[str, ControlPartCommandProfile]: profiles: dict[str, ControlPartCommandProfile] = {} @@ -1617,8 +1848,14 @@ def _control_profiles(self) -> dict[str, ControlPartCommandProfile]: ) return profiles - def _grasp_pose_generators(self) -> dict[str, AntipodalGraspPoseGenerator]: + def _grasp_pose_generators( + self, + *, + filter_ground_collision: bool = True, + ) -> dict[str, AntipodalGraspPoseGenerator]: """Build one mainline grasp service for each runtime hand endpoint.""" + if not isinstance(filter_ground_collision, bool): + raise TypeError("filter_ground_collision must be a boolean.") options = self.grasp_policy model = ParallelJawGripperModelCfg( model_id="gen_sim_parallel_jaw", @@ -1640,13 +1877,19 @@ def _grasp_pose_generators(self) -> dict[str, AntipodalGraspPoseGenerator]: point_sample_density=float(options["point_sample_dense"]), max_decomposition_hulls=int(options["max_decomposition_hulls"]), opening_margin=0.01, - filter_ground_collision=True, + filter_ground_collision=filter_ground_collision, ) annotation = GraspAnnotationCfg( selection_mode="whole_mesh", viser_port=int(options["viser_port"]), force_refresh=bool(options["force_grasp_reannotate"]), ) + shared_generator = AntipodalGraspPoseGenerator( + model, + algorithm_cfg=algorithm, + collision_cfg=collision, + annotation_cfg=annotation, + ) generators: dict[str, AntipodalGraspPoseGenerator] = {} for arm in ("left_arm", "right_arm"): try: @@ -1655,12 +1898,7 @@ def _grasp_pose_generators(self) -> dict[str, AntipodalGraspPoseGenerator]: continue if hand_part is None or hand_part in generators: continue - generators[hand_part] = AntipodalGraspPoseGenerator( - model, - algorithm_cfg=algorithm, - collision_cfg=collision, - annotation_cfg=annotation, - ) + generators[hand_part] = shared_generator return generators def _parts(self, arm: str) -> tuple[str, str | None, int]: diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index 5e0bf3248..6c263045d 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -2944,9 +2944,7 @@ def _execute_edge( for outcome in outcomes.values(): if outcome is None: continue - diagnostics = outcome.planner_trace.get( - "primary_action_diagnostics", {} - ) + diagnostics = outcome.planner_trace.get("primary_action_diagnostics", {}) observed_segments = set( diagnostics.get("execution_observation_segments", ()) ) @@ -3097,9 +3095,7 @@ def _action_execution_observation( state = state.unsqueeze(0) if state.shape == (int(self.env.num_envs), 13): observation.setdefault("linear_velocity", state[:, 7:10].clone()) - observation.setdefault( - "angular_velocity", state[:, 10:13].clone() - ) + observation.setdefault("angular_velocity", state[:, 10:13].clone()) body_data = getattr(entity, "body_data", None) if body_data is not None: for name, attribute_name in ( @@ -3112,24 +3108,36 @@ def _action_execution_observation( if callable(value): value = value() if value is not None: - observation[name] = torch.as_tensor( - value, - dtype=torch.float32, - device=self.env.device, - ).detach().clone() + observation[name] = ( + torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + .detach() + .clone() + ) getter = getattr(self.env, "get_current_xpos_agent", None) if callable(getter): left, right = getter() - observation["left_tcp_pose"] = torch.as_tensor( - left, - dtype=torch.float32, - device=self.env.device, - ).detach().clone() - observation["right_tcp_pose"] = torch.as_tensor( - right, - dtype=torch.float32, - device=self.env.device, - ).detach().clone() + observation["left_tcp_pose"] = ( + torch.as_tensor( + left, + dtype=torch.float32, + device=self.env.device, + ) + .detach() + .clone() + ) + observation["right_tcp_pose"] = ( + torch.as_tensor( + right, + dtype=torch.float32, + device=self.env.device, + ) + .detach() + .clone() + ) return observation def _plan_live_hold( @@ -3267,6 +3275,98 @@ def _release_ownership( if arm not in owners: self._object_states.pop((uid, arm), None) + def _physical_coordinated_hold( + self, + uid: str, + state: ExecutionState, + grounded: GroundedAction, + attempted: torch.Tensor, + ) -> torch.Tensor: + """Require two live closed grasps and the measured transport target.""" + held = evaluate_predicate( + self.env, + {"type": "held_by_both_grippers", "object": uid}, + coordinated_state=state, + ) + target = grounded.target_object_pose + if not isinstance(target, torch.Tensor): + return torch.zeros_like(attempted) + target = torch.as_tensor( + target, + dtype=torch.float32, + device=self.env.device, + ) + if target.shape == (4, 4): + target = target.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + if target.shape != (int(self.env.num_envs), 4, 4): + return torch.zeros_like(attempted) + actual = self._entity_pose(uid).to(dtype=target.dtype, device=target.device) + tolerance = float( + grounded.cfg.get( + "postcondition_tolerance", + self.runtime_policy.predicate_fallbacks["position_tolerance"], + ) + ) + reached = ( + torch.linalg.vector_norm( + actual[:, :3, 3] - target[:, :3, 3], + dim=1, + ) + <= tolerance + ) + return attempted & held & reached + + def _commit_coordinated_ownership( + self, + uid: str, + held: torch.Tensor, + ) -> None: + owners = self._object_owners.setdefault( + uid, + [None] * int(self.env.num_envs), + ) + for env_id in torch.nonzero(held, as_tuple=False).flatten().tolist(): + owners[env_id] = "coordinated" + self._arm_owners["left_arm"][env_id] = uid + self._arm_owners["right_arm"][env_id] = uid + + def _release_coordinated_ownership( + self, + uid: str, + released: torch.Tensor, + ) -> None: + owners = self._object_owners.setdefault( + uid, + [None] * int(self.env.num_envs), + ) + released_rows = torch.nonzero(released, as_tuple=False).flatten().tolist() + for env_id in released_rows: + if owners[env_id] == "coordinated": + owners[env_id] = None + for arm in ("left_arm", "right_arm"): + if self._arm_owners[arm][env_id] == uid: + self._arm_owners[arm][env_id] = None + + for key in tuple(self._object_states): + if key[0] != uid: + continue + state = self._object_states[key] + updates = {name: None for name in state.held_objects} + if not updates: + continue + task = StateDelta(held_object_updates=updates).apply( + state.to_task_state(), + released, + ) + updated = ExecutionState.from_task_state( + task, + last_qpos=self.env.robot.get_qpos().clone(), + ) + if updated.held_objects: + self._object_states[key] = updated + else: + self._object_states.pop(key, None) + def _execute_coordinated( self, edge: ExecutionEdge, @@ -3379,7 +3479,25 @@ def _execute_coordinated( physical_failed = torch.zeros_like(failed) committed_state = outcome.state_after(successful) if capability.state_effect == "coordinated_hold": + physical = self._physical_coordinated_hold( + step.object_uid, + committed_state, + grounded, + successful, + ) + physical_failed |= successful & ~physical + successful = physical + committed_state = outcome.state_after(successful) + for participant_arm in ("left_arm", "right_arm"): + committed_state = self._rebase_held_state( + step.object_uid, + participant_arm, + committed_state, + successful, + from_planned_qpos=False, + ) self._clear_support_relation(step.object_uid, successful) + self._commit_coordinated_ownership(step.object_uid, successful) if capability.state_effect == "transfer_hold": if bool(successful.any()): current_owners = list( @@ -3466,7 +3584,7 @@ def _execute_coordinated( | physical_failed, [grounded], [outcome.planner_trace], - active & outcome.success, + successful, ) def _rebase_held_state( @@ -3592,7 +3710,27 @@ def _execute_explicit_dual( ) physical_failed = torch.zeros_like(failed) if is_coordinated_release: - opened = evaluate_predicate(self.env, {"type": "both_grippers_open"}) + coordinated_name = next( + name + for name in self.adapter.capabilities.executable_names() + if self.adapter.capabilities.get(name).config_materializer + == "coordinated_pickment" + ) + coordinated_policy = self.runtime_policy.motion_defaults[coordinated_name] + opened = evaluate_predicate( + self.env, + { + "type": "both_grippers_open", + "tolerance": float( + coordinated_policy.get( + "release_gripper_tolerance", + self.runtime_policy.predicate_fallbacks[ + "gripper_state_tolerance" + ], + ) + ), + }, + ) released = active & opened physical_failed = active & ~opened control_parts = ( @@ -3608,6 +3746,7 @@ def _execute_explicit_dual( ) for key in ("coordinated", "left_arm", "right_arm"): self._step_states[(step.id, key)] = released_state + self._release_coordinated_ownership(step.object_uid, released) else: for arm, outcome in outcomes.items(): if outcome is not None: diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 46366837b..b1855c1f9 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -725,9 +725,24 @@ def ground( joint_defaults["hand_close_sample_interval"] ) elif source == "gripper_open": - policy["sample_interval"] = int( - joint_defaults["hand_open_sample_interval"] - ) + if binding.get("coordinated_release_role") is not None: + coordinated_name = next( + name + for name in self.capabilities.executable_names() + if self.capabilities.get(name).config_materializer + == "coordinated_pickment" + ) + coordinated = self.runtime_policy.motion_defaults[coordinated_name] + policy["sample_interval"] = int( + coordinated.get( + "release_sample_interval", + joint_defaults["hand_open_sample_interval"], + ) + ) + else: + policy["sample_interval"] = int( + joint_defaults["hand_open_sample_interval"] + ) elif source == "initial" and control == "arm": # Returning home after release is a safety motion. If the # collision-aware planner cannot find a route, do not silently diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index 210aa7225..c7ef5c0e3 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -688,16 +688,19 @@ def evaluate_predicate( if not hasattr(env, "get_current_gripper_state_agent"): return _constant(env, False) left, right = env.get_current_gripper_state_agent() - expected = torch.as_tensor( - env.open_state, - dtype=torch.float32, - device=env.device, - ) results = [] - for value in (left, right): + for side, value in zip(("left", "right"), (left, right)): value = torch.as_tensor(value, dtype=torch.float32, device=env.device) if value.ndim == 1: value = value.unsqueeze(0).repeat(int(env.num_envs), 1) + expected = getattr(env, f"{side}_arm_init_gripper_state", env.open_state) + expected = torch.as_tensor( + expected, + dtype=torch.float32, + device=env.device, + ) + if expected.ndim == 1: + expected = expected.unsqueeze(0).repeat(int(env.num_envs), 1) results.append( torch.linalg.vector_norm(value - expected, dim=-1) <= float(spec.get("tolerance", defaults["gripper_state_tolerance"])) diff --git a/embodichain/gen_sim/action_engine/tasks/assembly.py b/embodichain/gen_sim/action_engine/tasks/assembly.py index 7fdfe7246..c4afa5f65 100644 --- a/embodichain/gen_sim/action_engine/tasks/assembly.py +++ b/embodichain/gen_sim/action_engine/tasks/assembly.py @@ -150,7 +150,7 @@ def __init__( self.requirements: dict[str, dict[str, Any]] = {} self.previous_object_uid: str | None = None self.previous_arm: str | None = None - self.last_task_by_object_uid: dict[str, tuple[str, str]] = {} + self.last_task_by_object_uid: dict[str, str] = {} def add( self, @@ -162,6 +162,7 @@ def add( depends_on: Sequence[str] | None = None, ) -> str: values = deepcopy(dict(params or {})) + contract = task_contract(task_type) relation = str(values.get("relation", "none")) validate_source_compatibility(task_type, (object_entity,)) validate_target_compatibility(task_type, target, relation=relation) @@ -169,12 +170,10 @@ def add( instance_id = f"task_{len(self.instances) + 1:02d}" object_role = self._role( object_entity, - required_affordances=task_contract(task_type).required_affordances, + required_affordances=contract.required_affordances, initial_state={"orientation": "fallen"} if task_type == "E2" else {}, ) - values = {"object_role": object_role, **values} - if task_type == "E3": - values["source_role"] = values.pop("object_role") + values = {contract.primary_role_field: object_role, **values} if target is not None: values["target_role"] = self._role( target, @@ -185,13 +184,8 @@ def add( else: dependencies = list(depends_on) previous_for_object = self.last_task_by_object_uid.get(object_entity.uid) - if ( - task_type == "E4" - and previous_for_object is not None - and previous_for_object[1] == "E2" - and previous_for_object[0] not in dependencies - ): - dependencies.append(previous_for_object[0]) + if previous_for_object is not None and previous_for_object not in dependencies: + dependencies.append(previous_for_object) self.instances.append( { "id": instance_id, @@ -201,9 +195,9 @@ def add( "role": "primary", } ) - self.last_task_by_object_uid[object_entity.uid] = (instance_id, task_type) + self.last_task_by_object_uid[object_entity.uid] = instance_id self.previous_object_uid = object_entity.uid - if task_type == "E4": + if contract.resource_mode == "handover": receive_arm = str(values.get("receive_arm", "")) self.previous_arm = ( receive_arm if receive_arm in {"left_arm", "right_arm"} else None diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index 29cefacc0..8966eb82c 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -18,7 +18,7 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from typing import Any @@ -29,8 +29,10 @@ ) from embodichain.gen_sim.action_engine.domain import ( TERMINAL_BEHAVIORS, + TaskContract, TRANSPORT_DIRECTIONS, motion_policy, + task_contract, task_success_type, validate_task_spec, ) @@ -97,10 +99,10 @@ def instantiate_seed_graph( ) nodes.extend(recipe_nodes) terminal_by_group[group_id] = _terminal_nodes(recipe_nodes) - held_after_group[group_id] = _terminal_hold( - task_type, + held_after_group[group_id] = _terminal_hold_from_contracts( object_uid, - params, + recipe_nodes, + capabilities, ) groups.append( { @@ -192,8 +194,10 @@ def _topological_instances( def _propagate_direct_payloads( task: Mapping[str, Any], bindings: Mapping[str, str], + *, + contract_resolver: Callable[[str], TaskContract] = task_contract, ) -> tuple[dict[str, Any], list[dict[str, str]]]: - """Carry direct E1 support relations into a later single-arm E1 move. + """Propagate direct support through declared carrier-flow contracts. 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 @@ -208,25 +212,29 @@ def _propagate_direct_payloads( for instance in _topological_instances(result["task_instances"]): task_type = str(instance["task_type"]) + contract = contract_resolver(task_type) params = instance["params"] - primary_key = "source_role" if task_type == "E3" else "object_role" + primary_key = contract.primary_role_field 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": + if not contract.accepts_direct_payloads: 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." + f"TaskGroup {instance['id']!r} consumes carrier " + f"{primary_uid!r}, but its {task_type!r} contract does not " + "accept direct payloads." ) 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: + if producer_id not in instance["depends_on"]: + instance["depends_on"].append(producer_id) + changed = True links.append( { "producer": producer_id, @@ -237,7 +245,7 @@ def _propagate_direct_payloads( } ) - if task_type in {"E1", "E2", "E3", "E4", "E5"}: + if contract.moves_primary_object: old_carrier = carrier_by_payload.pop(primary_uid, None) if old_carrier is not None: direct_by_carrier[old_carrier] = [ @@ -246,7 +254,7 @@ def _propagate_direct_payloads( if item[0] != primary_uid ] - if task_type != "E1" or str(params.get("relation")) not in {"on", "inside"}: + if str(params.get("relation")) not in contract.direct_payload_relations: continue target_role = params.get("target_role") if not isinstance(target_role, str) or not target_role: @@ -272,14 +280,14 @@ def _payload_goal(params: Mapping[str, Any], object_uid: str) -> list[dict[str, if not isinstance(raw_payloads, Sequence) or isinstance( raw_payloads, (str, bytes, bytearray) ): - raise ValueError("E1 payload_roles must be a list.") + raise ValueError("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.") + raise ValueError("payload_roles must contain non-empty object IDs.") if object_uid in payloads: - raise ValueError("An E1 carrier cannot be its own payload.") + raise ValueError("A carrier cannot be its own payload.") if len(payloads) != len(set(payloads)): - raise ValueError("E1 direct payload objects must be unique.") + raise ValueError("Direct payload objects must be unique.") return [{"object": value, "slot": "center"} for value in payloads] @@ -471,6 +479,7 @@ def _recipe( "kind": "joint_state", "source": "initial", "operation": "e2_home", + "required_home": True, }, [retreat["id"]], "cleanup", @@ -575,6 +584,7 @@ def _recipe( "kind": "joint_state", "source": "initial", "operation": "e3_home", + "required_home": True, }, "cleanup", ), @@ -745,6 +755,9 @@ def _recipe( **_orientation_extensions(params), "relation_frame": str(params.get("relation_frame", "robot")), } + payloads = _payload_goal(params, object_uid) + if payloads: + goal["payloads"] = payloads target = params.get("target_role") relation = str(params.get("relation", "none")) if isinstance(target, str) and target: @@ -768,7 +781,11 @@ def _recipe( object_uid, actor, "coordinated", - {"kind": "coordinated_goal", "object": object_uid}, + { + "kind": "coordinated_goal", + "object": object_uid, + **({"payloads": deepcopy(payloads)} if payloads else {}), + }, dependencies, role, {"type": "held_by_both_grippers", "object": object_uid}, @@ -1024,7 +1041,7 @@ def _terminal_nodes(nodes: list[Mapping[str, Any]]) -> list[str]: def _primary_object(task_type: str, params: Mapping[str, Any]) -> str: - key = "source_role" if task_type == "E3" else "object_role" + key = task_contract(task_type).primary_role_field value = params.get(key) if not isinstance(value, str) or not value: raise ValueError(f"{task_type} requires resolved parameter {key!r}.") @@ -1037,6 +1054,7 @@ def _actor( *, incoming_held_arm: str | None = None, ) -> dict[str, Any]: + contract = task_contract(task_type) required_arm = params.get("required_arm") if ( incoming_held_arm is not None @@ -1051,9 +1069,9 @@ def _actor( 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": + if contract.resource_mode == "coordinated": return {"mode": "coordinated", "arms": ["left_arm", "right_arm"]} - if task_type == "E4": + if contract.resource_mode == "handover": return {"mode": "required", "arm": str(params.get("transfer_arm", "left_arm"))} return {"mode": "auto"} @@ -1065,7 +1083,7 @@ def _incoming_held_arm( 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", "E3", "E4"}: + if not task_contract(task_type).accepts_incoming_hold: return None candidates = { held[1] @@ -1077,23 +1095,51 @@ def _incoming_held_arm( raise ValueError( f"Task instance has conflicting predecessor holders for {object_uid!r}." ) - return next(iter(candidates), None) + holder = next(iter(candidates), None) + if holder is not None and holder not in {"left_arm", "right_arm"}: + raise ValueError( + f"Task instance cannot consume holder kind {holder!r} for " + f"{object_uid!r}; its contract accepts only single-arm ownership." + ) + return holder -def _terminal_hold( - task_type: str, +def _terminal_hold_from_contracts( object_uid: str, - params: Mapping[str, Any], + nodes: Sequence[Mapping[str, Any]], + capabilities: AtomicCapabilityRegistry, ) -> 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": + """Fold action effects into the terminal holder of one recipe.""" + holders: set[str] = set() + coordinated = False + for node in nodes: + contract = capabilities.get(str(node["atomic_action"])).resolve_contract(node) + for effect in contract.effects: + atom = effect.atom + if atom.object_uid != object_uid: + continue + if atom.predicate == "object_held" and atom.arm is not None: + if effect.op == "add": + holders.add(atom.arm) + else: + holders.discard(atom.arm) + elif atom.predicate == "object_coordinated_held": + coordinated = effect.op == "add" + elif atom.predicate == "object_free" and effect.op == "add": + holders.clear() + coordinated = False + if coordinated and holders: + raise ValueError( + f"Recipe for {object_uid!r} ends with conflicting single and " + "coordinated ownership effects." + ) + if coordinated: return object_uid, "coordinated" - return None + if len(holders) > 1: + raise ValueError( + f"Recipe for {object_uid!r} ends with multiple single-arm holders." + ) + return (object_uid, next(iter(holders))) if holders else None def _validate_bindings( diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml index c880ba50f..e2ef9e5ad 100644 --- a/embodichain/gen_sim/task_engine/defaults.yaml +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -19,7 +19,7 @@ schema_version: embodichain.task-engine-defaults/v1 workflow: max_parallel_workers: 2 max_scene_attempts: 2 - max_action_attempts: 3 + max_action_attempts: 1 planning: candidate_count: 3 diff --git a/embodichain/gen_sim/task_engine/ontology.py b/embodichain/gen_sim/task_engine/ontology.py index 574d1362d..4ae707ca7 100644 --- a/embodichain/gen_sim/task_engine/ontology.py +++ b/embodichain/gen_sim/task_engine/ontology.py @@ -70,6 +70,7 @@ } ) TERMINAL_BEHAVIORS = frozenset({"none", "hold", "place"}) +_RESOURCE_MODES = frozenset({"single_arm", "handover", "coordinated"}) @dataclass(frozen=True, slots=True) @@ -83,6 +84,26 @@ class TaskContract: required_affordances: frozenset[str] success_type: str scene_affordances: frozenset[str] + primary_role_field: str + resource_mode: str + moves_primary_object: bool + accepts_direct_payloads: bool + direct_payload_relations: frozenset[str] + accepts_incoming_hold: bool + terminal_success_types: tuple[tuple[str, str], ...] + + def __post_init__(self) -> None: + if not self.primary_role_field.endswith("_role"): + raise ValueError("primary_role_field must name one role parameter.") + if self.resource_mode not in _RESOURCE_MODES: + raise ValueError(f"Unknown task resource_mode {self.resource_mode!r}.") + if self.direct_payload_relations - RELATIONS: + raise ValueError("direct_payload_relations contain unknown relations.") + terminal_behaviors = [item[0] for item in self.terminal_success_types] + if len(terminal_behaviors) != len(set(terminal_behaviors)): + raise ValueError("terminal_success_types must use unique behaviors.") + if set(terminal_behaviors) - TERMINAL_BEHAVIORS: + raise ValueError("terminal_success_types contain unknown behaviors.") def _contract( @@ -94,6 +115,13 @@ def _contract( success_type: str, *, scene_affordances: frozenset[str] | None = None, + primary_role_field: str = "object_role", + resource_mode: str = "single_arm", + moves_primary_object: bool = False, + accepts_direct_payloads: bool = False, + direct_payload_relations: frozenset[str] = frozenset(), + accepts_incoming_hold: bool = False, + terminal_success_types: tuple[tuple[str, str], ...] = (), ) -> TaskContract: return TaskContract( task_type=task_type, @@ -103,6 +131,13 @@ def _contract( required_affordances=required_affordances, success_type=success_type, scene_affordances=scene_affordances or required_affordances, + primary_role_field=primary_role_field, + resource_mode=resource_mode, + moves_primary_object=moves_primary_object, + accepts_direct_payloads=accepts_direct_payloads, + direct_payload_relations=direct_payload_relations, + accepts_incoming_hold=accepts_incoming_hold, + terminal_success_types=terminal_success_types, ) @@ -124,6 +159,10 @@ def _contract( "rigid_object", frozenset({"graspable", "placeable"}), "semantic_goal", + moves_primary_object=True, + accepts_direct_payloads=True, + direct_payload_relations=frozenset({"on", "inside"}), + accepts_incoming_hold=True, ), "E2": _contract( "E2", @@ -132,6 +171,8 @@ def _contract( "rigid_object", frozenset({"graspable", "orientable"}), "object_upright", + moves_primary_object=True, + accepts_incoming_hold=True, ), "E3": _contract( "E3", @@ -141,6 +182,9 @@ def _contract( "rigid_object", frozenset({"graspable", "pourable"}), "poured", + primary_role_field="source_role", + moves_primary_object=True, + accepts_incoming_hold=True, ), "E4": _contract( "E4", @@ -149,6 +193,9 @@ def _contract( "rigid_object", frozenset({"graspable", "handover"}), "handover_complete", + resource_mode="handover", + moves_primary_object=True, + accepts_incoming_hold=True, ), "E5": _contract( "E5", @@ -158,6 +205,13 @@ def _contract( frozenset({"dual_graspable"}), "held_by_both_grippers", scene_affordances=frozenset({"dual_graspable", "rigid"}), + resource_mode="coordinated", + moves_primary_object=True, + accepts_direct_payloads=True, + terminal_success_types=( + ("hold", "held_by_both_grippers"), + ("place", "semantic_goal"), + ), ), "E6": _contract( "E6", @@ -211,11 +265,14 @@ def task_success_type( ) -> str: """Resolve a TaskSpec success type, including E5's terminal behavior.""" contract = task_contract(task_type) - if contract.task_type != "E5": + if not contract.terminal_success_types: return contract.success_type terminal_behavior = str((params or {}).get("terminal_behavior", "hold")) - if terminal_behavior == "hold": - return "held_by_both_grippers" - if terminal_behavior == "place": - return "semantic_goal" - raise ValueError("E5 terminal_behavior must be 'hold' or 'place'.") + success_by_behavior = dict(contract.terminal_success_types) + try: + return success_by_behavior[terminal_behavior] + except KeyError as exc: + raise ValueError( + f"{contract.task_type} terminal_behavior must be one of " + f"{sorted(success_by_behavior)}." + ) from exc diff --git a/embodichain/gen_sim/task_engine/scene/feasibility.py b/embodichain/gen_sim/task_engine/scene/feasibility.py index e0d6f00b8..db2a7a69d 100644 --- a/embodichain/gen_sim/task_engine/scene/feasibility.py +++ b/embodichain/gen_sim/task_engine/scene/feasibility.py @@ -23,6 +23,8 @@ import math from typing import Any +from ..ontology import task_contract + from .contracts import ( ASSESSMENT_STATUSES, FEASIBILITY_REPORT_SCHEMA, @@ -334,8 +336,9 @@ def _workspace_checks( object_uids_by_step, ) task_type = str(step.get("task_type", "")) + contract = task_contract(task_type) required_arm = str(step.get("required_arm", "auto")) - if task_type == "E4": + if contract.resource_mode == "handover": required_arm = str(step.get("transfer_arm", "none")) if required_arm in {"left_arm", "right_arm"}: for uid in object_uids: @@ -375,9 +378,10 @@ def _workspace_checks( phases.extend( _workflow_phases( step_id, - task_type, object_uids, target_uids, + resource_mode=contract.resource_mode, + moves_primary_object=contract.moves_primary_object, transfer_arm=str(step.get("transfer_arm", "none")), receive_arm=str(step.get("receive_arm", "none")), ) @@ -626,10 +630,11 @@ def _step_selector_uids( def _workflow_phases( step_id: str, - task_type: str, object_uids: Sequence[str], target_uids: Sequence[str], *, + resource_mode: str, + moves_primary_object: bool, transfer_arm: str, receive_arm: str, ) -> list[dict[str, Any]]: @@ -643,7 +648,7 @@ def _workflow_phases( "object_uids": list(object_uids), } ) - if task_type == "E4": + if resource_mode == "handover": phases.append( { "step_id": step_id, @@ -662,7 +667,7 @@ def _workflow_phases( "target_uids": list(target_uids), } ) - if task_type in {"E1", "E2", "E3", "E4", "E5"}: + if moves_primary_object: phases.append( { "step_id": step_id, diff --git a/embodichain/gen_sim/video_archive.py b/embodichain/gen_sim/video_archive.py index 9782f5348..9d49ee19a 100644 --- a/embodichain/gen_sim/video_archive.py +++ b/embodichain/gen_sim/video_archive.py @@ -40,7 +40,6 @@ def _archive_task_recording(env: Any, task_id: str) -> Path | None: RuntimeError: If configured recorders do not identify one task video. ValueError: If the task ID can escape the video directory. FileNotFoundError: If the expected source recording does not exist. - FileExistsError: If the task archive already exists. """ manager = getattr(env.unwrapped, "event_manager", None) mode_cfgs = getattr(manager, "_mode_functor_cfgs", {}) @@ -103,7 +102,6 @@ def _archive_task_video( Raises: ValueError: If the task ID can escape the video directory. FileNotFoundError: If the expected source recording does not exist. - FileExistsError: If the task archive already exists. RuntimeError: If more than one source extension matches. """ _validate_task_id(task_id) @@ -134,12 +132,7 @@ def _archive_task_video( source = candidates[0] extension = source.name[len(source_stem) :] destination = directory / f"{task_id}{extension}" - if destination.exists(): - raise FileExistsError( - f"Cannot archive video for task {task_id!r}: destination already " - f"exists at {destination}." - ) - source.rename(destination) + source.replace(destination) return destination diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py index 9e065b1a9..09e00b6d4 100644 --- a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -24,6 +24,7 @@ from embodichain.gen_sim.action_engine.capabilities import ( AtomicCapability, + StateAtom, build_atomic_capability_registry, ) from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter @@ -162,7 +163,7 @@ def verify_step(step, failed): assert calls[0][1].tolist() == [False, True] -def test_e2_home_is_required_while_generic_cleanup_home_is_best_effort() -> None: +def test_explicit_required_home_is_safety_required_for_any_task_type() -> None: capability = build_atomic_capability_registry().get("MoveJoints") base = { "atomic_action": "MoveJoints", @@ -174,19 +175,44 @@ def test_e2_home_is_required_while_generic_cleanup_home_is_best_effort() -> None } generic = capability.resolve_contract(base) - e2_home = capability.resolve_contract( + required_home = capability.resolve_contract( { **base, - "task_type": "E2", + "task_type": "test_carrier_consumer", "target_binding": { **base["target_binding"], - "operation": "e2_home", + "operation": "custom_home", + "required_home": True, }, } ) assert generic.failure_policy == "best_effort" - assert e2_home.failure_policy == "safety_required" + assert required_home.failure_policy == "safety_required" + + +def test_coordinated_release_contract_uses_binding_not_task_number() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + node = { + "atomic_action": "MoveJoints", + "task_type": "test_carrier_consumer", + "object_uid": "tray", + "actor": {"mode": "required", "arm": "left_arm"}, + "control": "hand", + "role": "primary", + "sync_group": "release_pair", + "target_binding": { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": "participant", + }, + } + + contract = capability.resolve_contract(node) + + assert contract.requires == ( + StateAtom("object_coordinated_held", object_uid="tray"), + ) def test_new_descriptor_reuses_loader_and_adapter_without_dispatch_changes() -> None: diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py index 2c254773f..6581c2707 100644 --- a/tests/gen_sim/action_engine/config/test_runtime_policy.py +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -107,6 +107,12 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: runtime.motion_defaults["CoordinatedPickment"]["is_filter_ground_collision"] is False ) + assert ( + runtime.motion_defaults["CoordinatedPickment"]["release_sample_interval"] == 60 + ) + assert runtime.motion_defaults["CoordinatedPickment"][ + "release_gripper_tolerance" + ] == pytest.approx(0.08) assert runtime.predicate_fallbacks["upright_max_tilt"] == pytest.approx( 0.2617993877991494 ) diff --git a/tests/gen_sim/action_engine/domain/test_task_contracts.py b/tests/gen_sim/action_engine/domain/test_task_contracts.py index 765e05aa8..a1b76fb88 100644 --- a/tests/gen_sim/action_engine/domain/test_task_contracts.py +++ b/tests/gen_sim/action_engine/domain/test_task_contracts.py @@ -44,6 +44,20 @@ def test_task_contract_catalog_covers_the_canonical_protocol() -> None: } +def test_task_contracts_declare_carrier_flow_and_resource_semantics() -> None: + e1 = task_contract("E1") + e3 = task_contract("E3") + e5 = task_contract("E5") + + assert e1.direct_payload_relations == {"on", "inside"} + assert e1.accepts_direct_payloads + assert e3.primary_role_field == "source_role" + assert e5.accepts_direct_payloads + assert e5.moves_primary_object + assert e5.resource_mode == "coordinated" + assert not task_contract("E2").accepts_direct_payloads + + def test_e5_success_depends_only_on_terminal_behavior() -> None: assert task_success_type("E5", {"terminal_behavior": "hold"}) == ( "held_by_both_grippers" diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 2c36ce651..5a70dc228 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -81,6 +81,38 @@ def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: return torch.tensor([[0, 1, 2]], dtype=torch.int64) +def _cuboid_vertices(x: float, y: float, z: float) -> torch.Tensor: + return torch.tensor( + [ + [sx * x, sy * y, sz * z] + for sx in (-1.0, 1.0) + for sy in (-1.0, 1.0) + for sz in (-1.0, 1.0) + ], + dtype=torch.float32, + ) + + +def _rotation_x(degrees: float) -> torch.Tensor: + angle = torch.deg2rad(torch.tensor(degrees, dtype=torch.float32)) + rotation = torch.eye(3) + rotation[1, 1] = torch.cos(angle) + rotation[1, 2] = -torch.sin(angle) + rotation[2, 1] = torch.sin(angle) + rotation[2, 2] = torch.cos(angle) + return rotation + + +def _rotation_z(degrees: float) -> torch.Tensor: + angle = torch.deg2rad(torch.tensor(degrees, dtype=torch.float32)) + rotation = torch.eye(3) + rotation[0, 0] = torch.cos(angle) + rotation[0, 1] = -torch.sin(angle) + rotation[1, 0] = torch.sin(angle) + rotation[1, 1] = torch.cos(angle) + return rotation + + class _PoseEntity: def __init__(self, pose: torch.Tensor) -> None: self.pose = pose @@ -174,7 +206,7 @@ def register(self, action: Any, *, replace: bool = False) -> None: monkeypatch.setattr(actions, "AtomicActionEngine", Engine) monkeypatch.setattr(adapter, "_generator", lambda: object()) monkeypatch.setattr(adapter, "_control_profiles", lambda: {}) - monkeypatch.setattr(adapter, "_grasp_pose_generators", lambda: {}) + monkeypatch.setattr(adapter, "_grasp_pose_generators", lambda **_kwargs: {}) engine = adapter._engine() @@ -539,6 +571,95 @@ def test_coordinated_pickment_uses_engine_scoped_grasp_generator() -> None: assert invocation.skill_options.middle_empty_ratio == pytest.approx(0.7) +def _coordinated_grounded( + rotation: torch.Tensor, + *, + vertices: torch.Tensor | None = None, +) -> GroundedAction: + object_pose = torch.eye(4).repeat(2, 1, 1) + object_pose[:, :3, :3] = rotation + object_pose[:, :3, 3] = torch.tensor([0.05, 0.0, 0.75]) + mesh_vertices = _cuboid_vertices(0.03, 0.04, 0.20) if vertices is None else vertices + goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="test_object", + geometry={}, + affordance=AntipodalAffordance( + object_label="test_object", + mesh_vertices=mesh_vertices, + mesh_triangles=torch.tensor([[0, 1, 2]], dtype=torch.int64), + ), + ), + object_target_pose=object_pose.clone(), + object_initial_pose=object_pose.clone(), + ) + return GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + goal, + {"middle_empty_ratio": 0.4}, + object_pose=object_pose, + object_uid="test_object", + ) + + +def test_coordinated_pickment_geometry_candidates_are_live_and_continuous() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + + vertical = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(torch.eye(3)), capability + ) + tilted = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(_rotation_x(45.0)), capability + ) + horizontal = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(_rotation_x(90.0)), capability + ) + yawed = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(_rotation_z(90.0) @ _rotation_x(90.0)), + capability, + ) + + preferred = [ + candidates[0].cfg["middle_empty_ratio"] + for candidates in (vertical, tilted, horizontal) + ] + assert preferred[0] < preferred[1] < preferred[2] + assert yawed[0].cfg["middle_empty_ratio"] == pytest.approx(preferred[0]) + for candidates in (vertical, tilted, horizontal, yawed): + assert candidates + assert torch.allclose( + candidates[0].cfg["left_to_right_arm_direction"], + torch.tensor([0.0, -1.0, 0.0]), + ) + assert torch.allclose( + candidates[0].cfg["approach_direction"], + torch.tensor([0.0, 0.0, -1.0]), + ) + + +def test_coordinated_pickment_geometry_candidates_are_deterministic_for_tray() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + grounded = _coordinated_grounded( + _rotation_z(31.0), + vertices=_cuboid_vertices(0.20, 0.14, 0.02), + ) + + first = adapter._adapt_coordinated_pickment_grasps(grounded, capability) + second = adapter._adapt_coordinated_pickment_grasps(grounded, capability) + + assert [item.cfg["middle_empty_ratio"] for item in first] == pytest.approx( + [item.cfg["middle_empty_ratio"] for item in second] + ) + assert ( + first[0].motion_policy["coordinated_grasp"] + == second[0].motion_policy["coordinated_grasp"] + ) + + def test_grasp_generators_follow_mainline_service_contract() -> None: adapter = AtomicActionAdapter(_planner_env()) @@ -546,6 +667,7 @@ def test_grasp_generators_follow_mainline_service_contract() -> None: assert set(generators) == {"physical_left_eef", "physical_right_eef"} generator = generators["physical_left_eef"] + assert generators["physical_right_eef"] is generator assert isinstance(generator, AntipodalGraspPoseGenerator) assert generator.algorithm_cfg.sample_count == 10000 assert generator.algorithm_cfg.approach_direction_samples == 4 @@ -554,6 +676,17 @@ def test_grasp_generators_follow_mainline_service_contract() -> None: assert generator.collision_cfg.filter_ground_collision is True +def test_coordinated_grasp_generator_honors_ground_filter_policy() -> None: + adapter = AtomicActionAdapter(_planner_env()) + + generators = adapter._grasp_pose_generators(filter_ground_collision=False) + + assert generators["physical_left_eef"] is generators["physical_right_eef"] + assert ( + generators["physical_left_eef"].collision_cfg.filter_ground_collision is False + ) + + def test_retreat_uses_row_local_motion_planner_reachability_search( monkeypatch: Any, ) -> None: @@ -841,7 +974,11 @@ def test_start_session_delegates_to_shared_atomic_engine(monkeypatch: Any) -> No captured: dict[str, Any] = {} monkeypatch.setattr(adapter, "_planning_context", lambda *_args: "context") - monkeypatch.setattr(adapter, "_invocation", lambda *_args: "invocation") + monkeypatch.setattr( + adapter, + "_invocation", + lambda *_args, **_kwargs: "invocation", + ) class _Engine: def start(self, invocations: tuple[Any, ...], context: Any) -> object: diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index 939e75117..bc48cb88f 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -35,7 +35,10 @@ resolve_agent_runtime_policy, runtime_policy_hash, ) -from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOverOptions +from embodichain.gen_sim.action_engine.capabilities import ( + HeldObjectHandOverOptions, + build_atomic_capability_registry, +) from embodichain.gen_sim.action_engine.compiler import ( compile_task_agent, compile_task_agent_v2, @@ -106,6 +109,7 @@ PressOptions, SlideAffordance, SlideGoal, + StateDelta, TwistAffordance, TwistGoal, ) @@ -602,6 +606,46 @@ def test_joint_state_binding_selects_hand_timing_without_a_named_policy() -> Non assert grounded.cfg["sample_interval"] == 10 +def test_e5_synchronized_release_uses_physics_verified_opening_time() -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.20)) + env = _FakeEnv({"tray": entity}) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "transport", + "operator": "coordinated_transport", + "object": "tray", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": {"direction": "front", "terminal_behavior": "place"}, + "depends_on": [], + } + ) + ) + ) + step = program.semantic_steps[0] + release = next( + action + for edge in program.edges + for action in edge.actions + if action.get("target_binding", {}).get("coordinated_release_role") + == "participant" + ) + grounder = ActionGrounder(program, env, lambda _uid: None) + + grounded = grounder.ground( + release, + step, + arm="left_arm", + state=_coordinated_held_state(env, entity), + ) + + assert grounded.cfg["sample_interval"] == 60 + + def test_runtime_policy_discards_legacy_support_z_fallbacks() -> None: snapshot = default_runtime_policy("dual_franka").as_mapping() snapshot["predicate_fallbacks"].update( @@ -1419,9 +1463,15 @@ def test_explicit_dual_gripper_release_commits_only_after_both_hands_open( state = _coordinated_held_state(env, entity) executor = object.__new__(ProgramExecutor) executor.env = env + executor.runtime_policy = default_runtime_policy("dual_ur10") executor._assignments = {"task_01": ["coordinated"]} executor._step_states = {("task_01", "coordinated"): state} executor._object_states = {} + executor._object_owners = {"tray": ["coordinated"]} + executor._arm_owners = { + "left_arm": ["tray"], + "right_arm": ["tray"], + } executor._orientation_references = {} def ground( @@ -1461,6 +1511,7 @@ def execute_trajectory( executor.grounder = SimpleNamespace(ground=ground) executor.adapter = SimpleNamespace( + capabilities=build_atomic_capability_registry(), plan=plan, combine=lambda _outcomes, _masks: ( torch.zeros(1, 2, env.robot.dof), @@ -1487,7 +1538,7 @@ def execute_trajectory( result = executor._execute_explicit_dual( SimpleNamespace(id="release", actions=actions), - SimpleNamespace(id="task_01"), + SimpleNamespace(id="task_01", object_uid="tray"), torch.zeros(1, dtype=torch.bool), ) @@ -1496,6 +1547,95 @@ def execute_trajectory( right_held = released_state.get_held_object("physical_right_arm") assert result.failed.tolist() == [expected_failed] assert (left_held is not None and right_held is not None) is expect_held + expected_owner = "coordinated" if expect_held else None + assert executor._object_owners["tray"] == [expected_owner] + assert executor._arm_owners["left_arm"] == (["tray"] if expect_held else [None]) + assert executor._arm_owners["right_arm"] == (["tray"] if expect_held else [None]) + + +@pytest.mark.parametrize( + ("closes_both", "expected_failed", "expect_held"), + ((True, False, True), (False, True, False)), +) +def test_coordinated_pickment_commits_only_after_physical_dual_hold( + monkeypatch: pytest.MonkeyPatch, + closes_both: bool, + expected_failed: bool, + expect_held: bool, +) -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.2)) + env = _FakeEnv({"tray": entity}) + program = load_execution_program( + compile_task_agent( + _task_agent( + { + "id": "task_01", + "operator": "coordinated_transport", + "object": "tray", + "actor": { + "mode": "coordinated", + "arms": ["left_arm", "right_arm"], + }, + "goal": {"direction": "up", "terminal_behavior": "hold"}, + "depends_on": [], + } + ) + ) + ) + executor = ProgramExecutor(program, env, record_runtime=False) + step = program.semantic_steps[0] + edge = program.edges[0] + executor._assignments[step.id] = ["coordinated"] + prior_state = ExecutionState(last_qpos=env.robot.get_qpos().clone()) + planned_state = _coordinated_held_state(env, entity) + grounded = GroundedAction( + action_class="CoordinatedPickment", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={"postcondition_tolerance": 0.06}, + object_pose=entity.get_local_pose(to_matrix=True), + target_object_pose=entity.get_local_pose(to_matrix=True), + object_uid="tray", + ) + outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.tensor([True]), + next_state=planned_state, + grounded=grounded, + prior_state=prior_state, + expected_effects=StateDelta( + held_object_updates=dict(planned_state.held_objects) + ), + ) + monkeypatch.setattr( + executor.grounder, + "ground_candidates", + lambda *_args, **_kwargs: (grounded,), + ) + monkeypatch.setattr(executor.adapter, "plan", lambda *_args, **_kwargs: outcome) + + def execute_trajectory(*_args: Any, **_kwargs: Any) -> list[torch.Tensor]: + env.robot._qpos[:, env.left_eef_joints] = env.close_state + if closes_both: + env.robot._qpos[:, env.right_eef_joints] = env.close_state + return [] + + monkeypatch.setattr(executor.adapter, "execute_trajectory", execute_trajectory) + + result = executor._execute_coordinated(edge, step, torch.tensor([False])) + + committed = executor._step_states[(step.id, "coordinated")] + held = tuple( + committed.get_held_object(f"physical_{arm}") + for arm in ("left_arm", "right_arm") + ) + assert result.failed.tolist() == [expected_failed] + assert all(item is not None for item in held) is expect_held + expected_owner = "coordinated" if expect_held else None + assert executor._object_owners["tray"] == [expected_owner] + assert executor._arm_owners["left_arm"] == (["tray"] if expect_held else [None]) + assert executor._arm_owners["right_arm"] == (["tray"] if expect_held else [None]) def _handover_held_state( @@ -4368,6 +4508,28 @@ def test_coordinated_held_predicate_uses_per_arm_held_relations() -> None: assert not bool(evaluate_predicate(env, predicate, coordinated_state=state)[0]) +def test_both_grippers_open_uses_the_live_reset_posture() -> None: + env = _FakeEnv() + physical_open = torch.tensor([[0.20, 0.45]]) + env.left_arm_init_gripper_state = physical_open.clone() + env.right_arm_init_gripper_state = physical_open.clone() + env.robot._qpos[:, env.left_eef_joints] = physical_open + env.robot._qpos[:, env.right_eef_joints] = physical_open + + opened = evaluate_predicate( + env, + {"type": "both_grippers_open", "tolerance": 0.08}, + ) + assert opened.tolist() == [True] + + env.robot._qpos[:, env.right_eef_joints] += 0.20 + opened = evaluate_predicate( + env, + {"type": "both_grippers_open", "tolerance": 0.08}, + ) + assert opened.tolist() == [False] + + def test_object_supported_by_requires_overlap_and_vertical_contact() -> None: support_z = 0.75 payload_z = support_z + 0.05 + 0.02 + 0.005 @@ -5677,6 +5839,7 @@ def test_coordinated_transport_direction_is_grounded_from_live_pose( ) } env = _FakeEnv(entities) + env.agent_initial_object_poses = {"shared_box": _pose(9.0, 9.0, 9.0)} program = load_execution_program( compile_task_agent( _task_agent( @@ -5720,6 +5883,14 @@ def semantics(uid: str) -> ObjectSemantics: ) assert isinstance(grounded.target, CoordinatedPickGoal) + assert torch.allclose( + grounded.target.object_initial_pose, + entities["shared_box"].get_local_pose(to_matrix=True), + ) + assert not torch.allclose( + grounded.target.object_initial_pose, + env.agent_initial_object_poses["shared_box"], + ) assert torch.allclose( grounded.target.object_target_pose[0, :3, 3], torch.tensor(expected_position), @@ -5747,6 +5918,11 @@ def get_vertices( _pose(0.0, 0.0, 0.80), _rect_vertices(0.03, 0.03, 0.08), ), + "cup": _FakeEntity( + "cup", + _pose(0.06, 0.0, 0.79), + _rect_vertices(0.025, 0.025, 0.06), + ), } program = compile_task_agent( _task_agent( @@ -5760,7 +5936,10 @@ def get_vertices( }, "goal": { "terminal_behavior": "place", - "payloads": [{"object": "bottle", "slot": "center"}], + "payloads": [ + {"object": "bottle", "slot": "center"}, + {"object": "cup", "slot": "center"}, + ], }, "depends_on": [], } @@ -5778,6 +5957,9 @@ def get_vertices( entities["bottle"]._pose[:, 0, 3] += 0.20 assert not bool(executor._verify_payloads(step)[0]) entities["bottle"]._pose = _pose(0.0, 0.0, 0.80) + entities["cup"]._pose[:, 1, 3] += 0.20 + assert not bool(executor._verify_payloads(step)[0]) + entities["cup"]._pose = _pose(0.06, 0.0, 0.79) entities["tray"]._pose[:, :3, :3] = torch.tensor( [[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]] ) diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index 3609bca94..07c84fbf3 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -260,6 +260,7 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "kind": "joint_state", "source": "initial", "operation": "e2_home", + "required_home": True, } assert orient_nodes[1]["depends_on"] == [orient_nodes[0]["id"]] assert orient_nodes[2]["depends_on"] == [orient_nodes[1]["id"]] diff --git a/tests/gen_sim/action_engine/tasks/test_payload_contracts.py b/tests/gen_sim/action_engine/tasks/test_payload_contracts.py new file mode 100644 index 000000000..11e537bd3 --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_payload_contracts.py @@ -0,0 +1,253 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Contract-driven direct-payload propagation tests.""" + +from __future__ import annotations + +import ast +from dataclasses import replace +import inspect +from textwrap import dedent + +import pytest + +import embodichain.gen_sim.action_engine.domain.v2 as domain_v2_module +import embodichain.gen_sim.action_engine.tasks.recipes as recipes_module +import embodichain.gen_sim.action_engine.planning.linker as linker_module +import embodichain.gen_sim.action_engine.runtime.executor as executor_module +from embodichain.gen_sim.action_engine.domain import task_contract +from embodichain.gen_sim.action_engine.protocol import TASK_SPEC_SCHEMA +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_spec + + +def _loaded_carrier_task( + *, + consumer_type: str = "E5", + terminal_behavior: str = "hold", +) -> tuple[dict, dict[str, str]]: + consumer_params = { + "object_role": "tray", + "direction": "up", + "terminal_behavior": terminal_behavior, + } + if consumer_type == "E2": + consumer_params = { + "object_role": "tray", + "orientation_goal": "upright", + "support_role": "table", + "upright_local_axis": "long_axis", + } + elif consumer_type == "E6": + consumer_params = { + "object_role": "tray", + "target_state": "open", + } + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "loaded-carrier", + "level": "L3", + "instruction": "Place two objects in a tray, then move the loaded tray.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "task_01", + "task_type": "E1", + "params": { + "object_role": "cube", + "target_role": "tray", + "relation": "inside", + "required_arm": "left_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_02", + "task_type": "E1", + "params": { + "object_role": "apple", + "target_role": "tray", + "relation": "inside", + "required_arm": "right_arm", + }, + "depends_on": [], + "role": "primary", + }, + { + "id": "task_03", + "task_type": consumer_type, + "params": consumer_params, + "depends_on": [], + "role": "primary", + }, + ], + "success": {"op": "all", "terms": []}, + "oracle": {}, + "metadata": {}, + } + return task, { + "cube": "interact_cube", + "apple": "interact_apple", + "tray": "interact_tray", + } + + +def test_loaded_e5_propagates_payloads_through_goal_binding_and_claims() -> None: + task, bindings = _loaded_carrier_task() + + graph = instantiate_seed_graph(task, bindings) + + carrier_group = next( + group for group in graph["task_groups"] if group["id"] == "task_03" + ) + payloads = [ + {"object": "interact_cube", "slot": "center"}, + {"object": "interact_apple", "slot": "center"}, + ] + assert carrier_group["goal"]["payloads"] == payloads + assert carrier_group["depends_on"] == ["task_01", "task_02"] + pick = next( + node + for node in graph["nodes"] + if node["task_instance_id"] == "task_03" + and node["atomic_action"] == "CoordinatedPickment" + ) + assert pick["target_binding"]["payloads"] == payloads + payload_claims = { + claim["resource"] + for claim in pick["contract"]["claims"] + if claim["resource"].startswith("object:") + } + assert payload_claims == { + "object:interact_tray", + "object:interact_cube", + "object:interact_apple", + } + assert graph["metadata"]["direct_payload_links"] == [ + { + "producer": "task_01", + "consumer": "task_03", + "carrier": "interact_tray", + "payload": "interact_cube", + "relation": "direct_support", + }, + { + "producer": "task_02", + "consumer": "task_03", + "carrier": "interact_tray", + "payload": "interact_apple", + "relation": "direct_support", + }, + ] + + +def test_loaded_carrier_graph_and_payload_order_are_deterministic() -> None: + task, bindings = _loaded_carrier_task() + + first = instantiate_seed_graph(task, bindings) + second = instantiate_seed_graph(task, bindings) + + assert first == second + + +def test_loaded_e5_place_keeps_payload_contract_through_synchronized_release() -> None: + task, bindings = _loaded_carrier_task(terminal_behavior="place") + + graph = instantiate_seed_graph(task, bindings) + + group = next(item for item in graph["task_groups"] if item["id"] == "task_03") + nodes = [node for node in graph["nodes"] if node["task_instance_id"] == "task_03"] + assert len(group["goal"]["payloads"]) == 2 + assert [node["atomic_action"] for node in nodes] == [ + "CoordinatedPickment", + "MoveJoints", + "MoveJoints", + ] + assert len({node.get("sync_group") for node in nodes[1:]}) == 1 + + +def test_consumer_without_direct_payload_capability_is_rejected() -> None: + task, bindings = _loaded_carrier_task(consumer_type="E2") + + with pytest.raises(ValueError, match="does not accept direct payloads"): + instantiate_seed_graph(task, bindings) + + +def test_test_consumer_contract_needs_no_propagation_type_change() -> None: + task, bindings = _loaded_carrier_task(consumer_type="E6") + custom_consumer = replace( + task_contract("E6"), + accepts_direct_payloads=True, + moves_primary_object=True, + ) + + propagated, links = recipes_module._propagate_direct_payloads( + task, + bindings, + contract_resolver=lambda task_type: ( + custom_consumer if task_type == "E6" else task_contract(task_type) + ), + ) + + assert propagated["task_instances"][2]["params"]["payload_roles"] == [ + "cube", + "apple", + ] + assert propagated["task_instances"][2]["depends_on"] == ["task_01", "task_02"] + assert [link["producer"] for link in links] == ["task_01", "task_02"] + + +def test_payload_infrastructure_does_not_branch_on_task_numbers() -> None: + functions = ( + recipes_module._propagate_direct_payloads, + linker_module._task_claims, + linker_module._task_primary_object, + domain_v2_module._validate_ownership_transitions, + executor_module.ProgramExecutor._capture_payloads, + executor_module.ProgramExecutor._verify_payloads, + ) + task_numbers = {f"E{index}" for index in range(1, 10)} + + for function in functions: + tree = ast.parse(dedent(inspect.getsource(function))) + literals = { + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + assert literals.isdisjoint(task_numbers), function.__qualname__ + + +@pytest.mark.parametrize("task_type", [f"E{index}" for index in range(1, 10)]) +def test_standalone_tasks_gain_no_payload_dependency(task_type: str) -> None: + task, _requirements = make_task_spec(task_type) + params = task["task_instances"][0]["params"] + role_names = { + value + for key, value in params.items() + if (key.endswith("_role") or key.endswith("_roles")) + and isinstance(value, str) + and value != "table" + } + bindings = {role: f"runtime_{role}" for role in role_names} + + graph = instantiate_seed_graph(task, bindings) + + assert graph["task_groups"][0]["depends_on"] == [] + assert graph["metadata"]["direct_payload_links"] == [] diff --git a/tests/gen_sim/test_video_archive.py b/tests/gen_sim/test_video_archive.py index 4b479ed9e..658ea9120 100644 --- a/tests/gen_sim/test_video_archive.py +++ b/tests/gen_sim/test_video_archive.py @@ -98,22 +98,22 @@ def test_archive_task_video_reports_missing_source_with_task_and_path( assert str(tmp_path / f"{SOURCE_STEM}.") in message -def test_archive_task_video_does_not_overwrite_existing_target( +def test_archive_task_video_overwrites_existing_target( tmp_path: Path, ) -> None: source = _write_source(tmp_path, ".mp4", b"new") destination = tmp_path / "task2_1.mp4" destination.write_bytes(b"existing") - with pytest.raises(FileExistsError, match="destination already exists"): - _archive_task_video( - tmp_path, - source_stem=SOURCE_STEM, - task_id="task2_1", - ) + result = _archive_task_video( + tmp_path, + source_stem=SOURCE_STEM, + task_id="task2_1", + ) - assert source.read_bytes() == b"new" - assert destination.read_bytes() == b"existing" + assert result == destination + assert destination.read_bytes() == b"new" + assert not source.exists() def test_consecutive_tasks_keep_independent_videos(tmp_path: Path) -> None: From a371e8578fe162cea71641eeb1a3f77831edbc40 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Wed, 26 Aug 2026 06:03:34 +0800 Subject: [PATCH 70/85] feat(gen-sim): add configurable PGI and Robotiq gripper profiles --- .../cli/generate_action_agent_config.py | 8 + .../gen_sim/action_engine/cli/run_agent.py | 11 + .../action_engine/config/defaults.yaml | 4 +- .../action_engine/config/runtime_policy.py | 25 +- .../action_engine/environment/agent_env.py | 15 + .../action_engine/generation/__init__.py | 7 +- .../generation/config_builder.py | 167 ++++++++- .../action_engine/generation/generator.py | 9 + .../templates/dual_franka_robot.json | 52 +-- .../generation/templates/dual_ur_robot.json | 42 +-- .../generation/templates/robot_profiles.json | 16 +- .../gen_sim/action_engine/gripper_profiles.py | 326 ++++++++++++++++++ .../gen_sim/action_engine/runtime/actions.py | 22 +- embodichain/gen_sim/task_engine/cli.py | 13 + embodichain/gen_sim/task_engine/config.py | 7 + embodichain/gen_sim/task_engine/defaults.yaml | 1 + .../task_engine/orchestration/coordinator.py | 4 + embodichain/gen_sim/task_engine/workflow.py | 2 + .../config/test_runtime_policy.py | 39 ++- .../generation/test_generation.py | 96 +++++- .../action_engine/runtime/test_actions.py | 77 ++++- .../runtime/test_runtime_contracts.py | 4 +- .../action_engine/test_gripper_profiles.py | 102 ++++++ .../orchestration/test_coordinator_cli.py | 1 + .../task_engine/test_parallel_workflow.py | 2 + tests/gen_sim/task_engine/test_workflow.py | 5 + 26 files changed, 929 insertions(+), 128 deletions(-) create mode 100644 embodichain/gen_sim/action_engine/gripper_profiles.py create mode 100644 tests/gen_sim/action_engine/test_gripper_profiles.py 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 index 542dc1c70..65b4665b8 100644 --- a/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py +++ b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py @@ -97,6 +97,13 @@ def build_parser() -> argparse.ArgumentParser: default=str(_TASK_DEFAULTS["default_robot_profile"]), help="Robot template used in fast_gym_config.json.", ) + parser.add_argument( + "--gripper-model", + "--gripper_model", + choices=("pgi", "robotiq"), + default=str(_TASK_DEFAULTS["default_gripper_model"]), + help="Gripper asset, control, TCP, and grasp profile used by both arms.", + ) parser.add_argument( "--llm_model", "--llm-model", @@ -183,6 +190,7 @@ def cli() -> None: task_description=task_description, task_spec=args.task_spec, robot_profile=args.robot_profile, + gripper_model=args.gripper_model, llm_model=args.llm_model, source_scene_z_rotation_degrees=args.source_scene_z_rotation_degrees, body_scale_policy=args.body_scale_policy, diff --git a/embodichain/gen_sim/action_engine/cli/run_agent.py b/embodichain/gen_sim/action_engine/cli/run_agent.py index cb60fedf7..06bd72e8d 100644 --- a/embodichain/gen_sim/action_engine/cli/run_agent.py +++ b/embodichain/gen_sim/action_engine/cli/run_agent.py @@ -126,6 +126,8 @@ def _validate_run_contract( task_name: str, ) -> None: """Validate the small cross-artifact contract before simulator startup.""" + from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + configured_task = agent_config.get("task_name") if configured_task != task_name: raise ValueError( @@ -146,6 +148,15 @@ def _validate_run_contract( f"Gym and agent configs have different planning modes: " f"gym={gym_mode!r}, agent={agent_mode!r}." ) + gym_gripper = extension.get("gripper_model") + agent_gripper = agent_config.get("gripper_model") + get_gripper_profile(gym_gripper) + get_gripper_profile(agent_gripper) + if gym_gripper != agent_gripper: + raise ValueError( + "Gym and agent configs have different gripper models: " + f"gym={gym_gripper!r}, agent={agent_gripper!r}." + ) def cli() -> int | None: diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index 8592b7151..c5442dd04 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -21,6 +21,7 @@ schema_version: action_engine_defaults_v1 generation: task: default_robot_profile: ur10 + default_gripper_model: pgi max_episodes: 1 max_episode_steps: 2000 environment: @@ -158,9 +159,6 @@ runtime: grasp: antipodal_n_sample: 10000 antipodal_max_angle: 0.2617993877991494 - max_open_length: 0.15 - min_open_length: 0.01 - finger_length: 0.13 point_sample_dense: 0.012 max_deviation_angle: 0.3490658503988659 n_deviated_approach_directions: 4 diff --git a/embodichain/gen_sim/action_engine/config/runtime_policy.py b/embodichain/gen_sim/action_engine/config/runtime_policy.py index 8999efbae..cfcba3c22 100644 --- a/embodichain/gen_sim/action_engine/config/runtime_policy.py +++ b/embodichain/gen_sim/action_engine/config/runtime_policy.py @@ -27,6 +27,7 @@ from typing import Any, Final from embodichain.gen_sim.action_engine.domain.motion import MOTION_MODIFIER_MODES +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.utils import configclass from embodichain.utils.utility import load_config @@ -42,7 +43,8 @@ ] ACTION_ENGINE_DEFAULTS_SCHEMA: Final = "action_engine_defaults_v1" -RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v7" +RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v8" +_PRE_GRIPPER_PROFILE_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v7" _PRE_AXIS_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v6" _PREVIOUS_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v5" _PRE_GRASP_RUNTIME_POLICY_SCHEMA: Final = "action_engine_runtime_policy_v4" @@ -99,9 +101,6 @@ _GRASP_KEYS = { "antipodal_n_sample", "antipodal_max_angle", - "max_open_length", - "min_open_length", - "finger_length", "point_sample_dense", "max_deviation_angle", "n_deviated_approach_directions", @@ -340,12 +339,6 @@ def __post_init__(self) -> None: _PREDICATE_KEYS, "predicate_fallbacks", ) - if float(self.grasp.get("min_open_length", -1.0)) < 0.0: - raise ValueError("grasp.min_open_length must be non-negative.") - if float(self.grasp.get("max_open_length", 0.0)) <= float( - self.grasp.get("min_open_length", 0.0) - ): - raise ValueError("grasp.max_open_length must exceed min_open_length.") direction_count = self.grasp.get("n_deviated_approach_directions") if ( isinstance(direction_count, bool) @@ -445,6 +438,10 @@ def generation_defaults() -> dict[str, Any]: } if set(value) != required: raise ValueError("Generation defaults do not match the expected sections.") + task = value.get("task") + if not isinstance(task, Mapping): + raise ValueError("Generation task defaults must be a mapping.") + get_gripper_profile(task.get("default_gripper_model")) return deepcopy(dict(value)) @@ -635,6 +632,14 @@ def resolve_agent_runtime_policy(agent_config: Mapping[str, Any]) -> RuntimePoli raise ValueError( "agent_config runtime policy hash does not match its snapshot." ) + if snapshot.get("schema_version") == _PRE_GRIPPER_PROFILE_RUNTIME_POLICY_SCHEMA: + migrated = deepcopy(dict(snapshot)) + migrated["schema_version"] = RUNTIME_POLICY_SCHEMA + grasp = deepcopy(dict(migrated.get("grasp", {}))) + for key in ("min_open_length", "max_open_length", "finger_length"): + grasp.pop(key, None) + migrated["grasp"] = grasp + return RuntimePolicyCfg.from_mapping(migrated) if snapshot.get("schema_version") == _LEGACY_RUNTIME_POLICY_SCHEMA: if set(snapshot) != {"schema_version", "arm_selection"} or not isinstance( snapshot.get("arm_selection"), Mapping diff --git a/embodichain/gen_sim/action_engine/environment/agent_env.py b/embodichain/gen_sim/action_engine/environment/agent_env.py index 93694dba1..d970b38e0 100644 --- a/embodichain/gen_sim/action_engine/environment/agent_env.py +++ b/embodichain/gen_sim/action_engine/environment/agent_env.py @@ -33,6 +33,7 @@ ) from embodichain.gen_sim.action_engine.domain import validate_seed_graph from embodichain.gen_sim.action_engine.protocol import ACTION_ENGINE_ENV_ID +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.gen_sim.action_engine.runtime import ( ProgramExecutor, evaluate_predicate, @@ -88,6 +89,20 @@ def __init__(self, cfg: EmbodiedEnvCfg | None = None, **kwargs: Any) -> None: self._runtime_state_ready = False repair_action_engine_ur5_solver_cfg(getattr(cfg, "robot", None)) super().__init__(cfg, **kwargs) + selected_gripper = get_gripper_profile( + getattr( + self, + "agent_gripper_model", + self.agent_config.get("gripper_model", "pgi"), + ) + ) + agent_gripper = self.agent_config.get( + "gripper_model", selected_gripper.model.value + ) + if agent_gripper != selected_gripper.model.value: + raise ValueError( + "ActionEngineEnv gym and agent gripper selections do not match." + ) install_action_engine_solver_compat(self.robot) if bool(getattr(self, "ignore_terminations_during_agent", False)): # Atomic trajectories execute online through env.step(). Prevent a diff --git a/embodichain/gen_sim/action_engine/generation/__init__.py b/embodichain/gen_sim/action_engine/generation/__init__.py index 4446010d6..733b84406 100644 --- a/embodichain/gen_sim/action_engine/generation/__init__.py +++ b/embodichain/gen_sim/action_engine/generation/__init__.py @@ -18,7 +18,11 @@ from __future__ import annotations -from .config_builder import VLM_CAMERA_UIDS, canonical_robot_profile +from .config_builder import ( + VLM_CAMERA_UIDS, + canonical_gripper_model, + canonical_robot_profile, +) from .assets import normalize_scene_assets from .generator import generate_action_engine_config from .models import GeneratedConfigPaths, PreparedScene @@ -27,6 +31,7 @@ "GeneratedConfigPaths", "PreparedScene", "VLM_CAMERA_UIDS", + "canonical_gripper_model", "canonical_robot_profile", "generate_action_engine_config", "normalize_scene_assets", diff --git a/embodichain/gen_sim/action_engine/generation/config_builder.py b/embodichain/gen_sim/action_engine/generation/config_builder.py index f1193494a..5ca8501b9 100644 --- a/embodichain/gen_sim/action_engine/generation/config_builder.py +++ b/embodichain/gen_sim/action_engine/generation/config_builder.py @@ -33,6 +33,10 @@ generation_defaults, runtime_policy_hash, ) +from embodichain.gen_sim.action_engine.gripper_profiles import ( + GripperProfile, + get_gripper_profile, +) from embodichain.gen_sim.action_engine.protocol import ( ACTION_ENGINE_CONFIG_SCHEMA, ACTION_ENGINE_ENV_ID, @@ -46,6 +50,7 @@ __all__ = [ "build_agent_config", "build_fast_gym_config", + "canonical_gripper_model", "canonical_robot_profile", "VLM_CAMERA_UIDS", "validate_fast_gym_config", @@ -54,6 +59,7 @@ _TEMPLATE_DIR = Path(__file__).resolve().parent / "templates" _GENERATION_DEFAULTS = generation_defaults() _DEFAULT_TABLETOP_Z = float(_GENERATION_DEFAULTS["scene"]["default_tabletop_z"]) +_DEFAULT_GRIPPER_MODEL = str(_GENERATION_DEFAULTS["task"]["default_gripper_model"]) _ARM_SLOTS = { "left": {"arm": "left_arm", "eef": "left_eef"}, @@ -85,6 +91,11 @@ def canonical_robot_profile(profile: str) -> str: ) +def canonical_gripper_model(model: str) -> str: + """Validate and return one exact GenSim gripper model ID.""" + return get_gripper_profile(model).model.value + + def build_agent_config( *, task_name: str, @@ -92,6 +103,7 @@ def build_agent_config( execution_program_hash: str, source_config_path: Path, uid_map: dict[str, str], + gripper_model: str = _DEFAULT_GRIPPER_MODEL, static_obstacle_uids: Sequence[str] | None = None, dynamic_obstacle_uids: Sequence[str] | None = None, table_top_z: float | None = None, @@ -103,6 +115,7 @@ def build_agent_config( ) -> dict[str, Any]: """Build the small manifest consumed by ``run_agent``.""" profile = canonical_robot_profile(robot_profile) + selected_gripper = canonical_gripper_model(gripper_model) runtime_policy = default_runtime_policy(profile) if ( static_obstacle_uids is not None @@ -143,6 +156,7 @@ def build_agent_config( "schema_version": ACTION_ENGINE_CONFIG_SCHEMA, "task_name": task_name, "robot_profile": profile, + "gripper_model": selected_gripper, "planning_mode": planning_mode, "task_spec": TASK_SPEC_FILENAME, "scene_requirements": SCENE_REQUIREMENTS_FILENAME, @@ -213,6 +227,7 @@ def build_fast_gym_config( execution_program_hash: str, max_episodes: int, max_episode_steps: int, + gripper_model: str = _DEFAULT_GRIPPER_MODEL, randomize_scene: bool = False, randomize_table_material: bool = False, planning_mode: str = "offline", @@ -228,10 +243,16 @@ def build_fast_gym_config( if planning_mode == "ab" and graph_path == EXECUTION_PROGRAM_FILENAME: graph_path = f"offline/{EXECUTION_PROGRAM_FILENAME}" profile = canonical_robot_profile(robot_profile) + gripper_profile = get_gripper_profile(gripper_model) profile_config = _profile(profile) - robot = _make_robot(profile, profile_config, scene.table_top_z) - observations = _make_observations(robot) + robot = _make_robot( + profile, + profile_config, + scene.table_top_z, + gripper_profile=gripper_profile, + ) + observations = _make_observations(robot, gripper_profile) # These two template fields describe serialization order to generation, not # RobotCfg. Remove them after deriving observation IDs to avoid parser noise. robot.pop("observation_joint_parts", None) @@ -260,6 +281,7 @@ def build_fast_gym_config( "defaults_schema_version": ACTION_ENGINE_DEFAULTS_SCHEMA, "task_name": task_name, "robot_profile": profile, + "gripper_model": gripper_profile.model.value, "planning_mode": planning_mode, "task_spec": TASK_SPEC_FILENAME, "scene_requirements": SCENE_REQUIREMENTS_FILENAME, @@ -277,11 +299,18 @@ def build_fast_gym_config( extensions = { "action_engine": engine_extension, "agent_robot_profile": profile, + "agent_gripper_model": gripper_profile.model.value, "agent_arm_slots": deepcopy(_ARM_SLOTS), "agent_static_obstacle_uids": background_uids, "agent_dynamic_obstacle_uids": rigid_uids, - "gripper_open_state": list(profile_config["gripper_open_state"]), - "gripper_close_state": list(profile_config["gripper_close_state"]), + "gripper_open_state": list(gripper_profile.open_positions), + "gripper_close_state": list(gripper_profile.close_positions), + "gripper_profile": gripper_profile.runtime_manifest( + tcp_parent_frames={ + "left": str(robot["solver_cfg"]["left_arm"]["end_link_name"]), + "right": str(robot["solver_cfg"]["right_arm"]["end_link_name"]), + } + ), "arm_aim_yaw_offset": deepcopy(environment_policy["arm_aim_yaw_offset"]), "ignore_terminations_during_agent": bool( environment_policy["ignore_terminations_during_agent"] @@ -375,6 +404,11 @@ def validate_fast_gym_config(config: dict[str, Any]) -> None: raise ValueError("Gym config points to an unexpected TaskSpec artifact.") if action_engine.get("scene_requirements") != SCENE_REQUIREMENTS_FILENAME: raise ValueError("Gym config points to unexpected SceneRequirements.") + gripper_profile = get_gripper_profile(action_engine.get("gripper_model")) + extensions = config["env"]["extensions"] + if extensions.get("agent_gripper_model") != gripper_profile.model.value: + raise ValueError("Gym config gripper model fields do not match.") + _validate_robot_gripper_contract(config["robot"], gripper_profile) graph_path = action_engine.get("seed_task_graph") if ( not isinstance(graph_path, str) @@ -413,6 +447,8 @@ def _make_robot( profile_id: str, profile: dict[str, Any], table_top_z: float | None, + *, + gripper_profile: GripperProfile, ) -> dict[str, Any]: robot = _load_template(str(profile["template"])) tabletop_z = _DEFAULT_TABLETOP_Z if table_top_z is None else float(table_top_z) @@ -427,7 +463,9 @@ def _make_robot( display = family.upper() urdf_dir = display robot["uid"] = f"Dual{display}" - robot["urdf_cfg"]["fname"] = f"dual_{family}_robotiq_arg2f_140_basket" + robot["urdf_cfg"][ + "fname" + ] = f"dual_{family}_{gripper_profile.assembly_name}_basket" for component in robot["urdf_cfg"]["components"]: if str(component.get("component_type", "")).endswith("_arm"): component["urdf_path"] = f"UniversalRobots/{urdf_dir}/{urdf_dir}.urdf" @@ -443,11 +481,103 @@ def _make_robot( "right_eef", ] robot["observation_joint_parts"] = ["left_eef", "right_eef"] + else: + robot["urdf_cfg"][ + "fname" + ] = f"dual_{family}_{gripper_profile.assembly_name}_basket" + _apply_gripper_profile(robot, gripper_profile) if profile_id != canonical_robot_profile(profile_id): raise ValueError(f"Invalid canonical robot profile {profile_id!r}.") return robot +def _apply_gripper_profile( + robot: dict[str, Any], + profile: GripperProfile, +) -> None: + """Apply one profile atomically to simulator, controller, and solver config.""" + control_parts = robot.get("control_parts") + init_qpos = robot.get("init_qpos") + if not isinstance(control_parts, dict) or not isinstance(init_qpos, list): + raise ValueError("Robot template requires control_parts and init_qpos.") + arm_dof = sum( + len(control_parts.get(f"{side}_arm", ())) for side in ("left", "right") + ) + if arm_dof <= 0 or len(init_qpos) < arm_dof: + raise ValueError("Robot template has an invalid initial arm posture.") + arm_init_qpos = list(init_qpos[:arm_dof]) + + components = robot.get("urdf_cfg", {}).get("components") + if not isinstance(components, list): + raise ValueError("Robot template requires a URDF component list.") + hands = { + str(component.get("component_type")): component + for component in components + if str(component.get("component_type", "")).endswith("_hand") + } + if set(hands) != {"left_hand", "right_hand"}: + raise ValueError("Robot template requires exactly one left and right hand.") + for component in hands.values(): + component["urdf_path"] = profile.asset_path + + for side in ("left", "right"): + control_parts[f"{side}_eef"] = list(profile.control_joint_names(side)) + robot["init_qpos"] = ( + arm_init_qpos + list(profile.simulated_joint_initial_positions) * 2 + ) + + drive = robot.get("drive_pros") + if not isinstance(drive, dict): + raise ValueError("Robot template requires drive_pros.") + for section, value in ( + ("stiffness", profile.drive_stiffness), + ("damping", profile.drive_damping), + ("max_effort", profile.drive_max_effort), + ): + values = drive.get(section) + if not isinstance(values, dict): + raise ValueError(f"Robot drive_pros.{section} must be a mapping.") + for side in ("left", "right"): + values[f"{side}_eef"] = value + + solvers = robot.get("solver_cfg") + if not isinstance(solvers, dict): + raise ValueError("Robot template requires solver_cfg.") + tcp = [list(row) for row in profile.tcp_transform] + for arm in ("left_arm", "right_arm"): + if not isinstance(solvers.get(arm), dict): + raise ValueError(f"Robot template requires solver_cfg.{arm}.") + solvers[arm]["tcp"] = deepcopy(tcp) + + +def _validate_robot_gripper_contract( + robot: Mapping[str, Any], + profile: GripperProfile, +) -> None: + """Reject generated artifacts whose physical and planning profiles drift.""" + components = robot.get("urdf_cfg", {}).get("components", []) + hand_assets = { + str(component.get("urdf_path")) + for component in components + if isinstance(component, Mapping) + and str(component.get("component_type", "")).endswith("_hand") + } + if hand_assets != {profile.asset_path}: + raise ValueError("Robot hand assets do not match the selected gripper profile.") + control_parts = robot.get("control_parts", {}) + for side in ("left", "right"): + if control_parts.get(f"{side}_eef") != list(profile.control_joint_names(side)): + raise ValueError( + f"Robot {side} gripper controls do not match the selected profile." + ) + expected_tcp = [list(row) for row in profile.tcp_transform] + for arm in ("left_arm", "right_arm"): + if robot.get("solver_cfg", {}).get(arm, {}).get("tcp") != expected_tcp: + raise ValueError( + f"Robot {arm} TCP does not match the selected gripper profile." + ) + + @lru_cache(maxsize=1) def _robot_profiles() -> dict[str, dict[str, Any]]: value = _read_template("robot_profiles.json") @@ -464,8 +594,6 @@ def _profile(profile_id: str) -> dict[str, Any]: "robot_family", "tabletop_clearance", "arm_component_z", - "gripper_open_state", - "gripper_close_state", } missing = sorted(required - set(profile)) if missing: @@ -665,17 +793,22 @@ def _recording_policy(planning_mode: str) -> tuple[bool, tuple[int, int], int]: ) -def _make_observations(robot: dict[str, Any]) -> dict[str, Any]: - control_parts = robot["control_parts"] - qpos_order = robot["qpos_control_part_order"] - observed_parts = set(robot["observation_joint_parts"]) - offset = 0 +def _make_observations( + robot: dict[str, Any], + gripper_profile: GripperProfile, +) -> dict[str, Any]: + per_hand_dof = len(gripper_profile.simulated_joint_initial_positions) + arm_dof = len(robot["init_qpos"]) - 2 * per_hand_dof + if arm_dof <= 0: + raise ValueError("Robot initial posture does not contain arm joints.") joint_ids: list[int] = [] - for part in qpos_order: - count = len(control_parts[part]) - if part in observed_parts: - joint_ids.extend(range(offset, offset + count)) - offset += count + for side_index, side in enumerate(("left", "right")): + simulated = gripper_profile.simulated_joint_names(side) + base = arm_dof + side_index * per_hand_dof + joint_ids.extend( + base + simulated.index(name) + for name in gripper_profile.control_joint_names(side) + ) return { "norm_robot_eef_joint": { "func": "normalize_robot_joint_data", diff --git a/embodichain/gen_sim/action_engine/generation/generator.py b/embodichain/gen_sim/action_engine/generation/generator.py index ac9da7be2..b9f07fda7 100644 --- a/embodichain/gen_sim/action_engine/generation/generator.py +++ b/embodichain/gen_sim/action_engine/generation/generator.py @@ -63,6 +63,7 @@ def generate_action_engine_config( task_description: str | None = None, task_spec: Mapping[str, Any] | str | Path | None = None, robot_profile: str = str(_TASK_DEFAULTS["default_robot_profile"]), + gripper_model: str = str(_TASK_DEFAULTS["default_gripper_model"]), llm_model: str | None = None, source_scene_z_rotation_degrees: float | None = None, source_scene_xy_translation: Sequence[float] | None = None, @@ -92,6 +93,9 @@ def generate_action_engine_config( 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'.") + from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + + gripper_model = get_gripper_profile(gripper_model).model.value _raise_if_outputs_exist( output_dir, overwrite=overwrite, @@ -226,6 +230,7 @@ def generate_action_engine_config( agent_config = build_agent_config( task_name=task_name, robot_profile=robot_profile, + gripper_model=gripper_model, execution_program_hash=program_hash, source_config_path=scene.source_config_path, uid_map=scene.uid_map, @@ -251,6 +256,7 @@ def generate_action_engine_config( task_name=task_name, task_description=task_description, robot_profile=robot_profile, + gripper_model=gripper_model, execution_program_hash=program_hash, max_episodes=max_episodes, max_episode_steps=max_episode_steps, @@ -735,6 +741,9 @@ def _validate_agent_config(config: Mapping[str, Any]) -> None: 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.") + from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + + get_gripper_profile(config.get("gripper_model")) graph_path = config.get("seed_task_graph") if ( not isinstance(graph_path, str) 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 index 6334e75e1..d5daac633 100644 --- a/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_franka_robot.json @@ -1,7 +1,7 @@ { "uid": "DualFrankaPanda", "urdf_cfg": { - "fname": "dual_franka_panda_basket", + "fname": "dual_franka_dh_pgi_140_80_basket", "name_case": { "joint": "original", "link": "original" @@ -19,7 +19,7 @@ }, { "component_type": "left_hand", - "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "urdf_path": "DH_PGI_140_80/DH_PGI_140_80.urdf", "transform": [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], @@ -39,7 +39,7 @@ }, { "component_type": "right_hand", - "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "urdf_path": "DH_PGI_140_80/DH_PGI_140_80.urdf", "transform": [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], @@ -67,14 +67,6 @@ 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, @@ -84,20 +76,20 @@ "stiffness": { "left_arm": 10000.0, "right_arm": 10000.0, - "left_eef": 50.0, - "right_eef": 50.0 + "left_eef": 1000.0, + "right_eef": 1000.0 }, "damping": { "left_arm": 1000.0, "right_arm": 1000.0, - "left_eef": 5.0, - "right_eef": 5.0 + "left_eef": 100.0, + "right_eef": 100.0 }, "max_effort": { "left_arm": 10000.0, "right_arm": 10000.0, - "left_eef": 500.0, - "right_eef": 500.0 + "left_eef": 10000.0, + "right_eef": 10000.0 } }, "control_parts": { @@ -110,14 +102,7 @@ "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" - ], + "left_eef": ["left_gripper_finger1_joint_1"], "right_arm": [ "right_fr3_joint1", "right_fr3_joint2", @@ -127,14 +112,7 @@ "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" - ], + "right_eef": ["right_gripper_finger1_joint_1"], "dual_arm": [ "left_fr3_joint1", "left_fr3_joint2", @@ -161,9 +139,9 @@ "end_link_name": "left_fr3_link8", "root_link_name": "left_base", "tcp": [ - [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.2], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.121], [0.0, 0.0, 0.0, 1.0] ], "num_samples": 15 @@ -174,9 +152,9 @@ "end_link_name": "right_fr3_link8", "root_link_name": "right_base", "tcp": [ - [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.2], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.121], [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 index 8b96a2a59..a8472d25e 100644 --- a/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json +++ b/embodichain/gen_sim/action_engine/generation/templates/dual_ur_robot.json @@ -1,7 +1,7 @@ { "uid": "DualUR5", "urdf_cfg": { - "fname": "dual_ur5_robotiq_arg2f_140_basket", + "fname": "dual_ur5_dh_pgi_140_80_basket", "name_case": {"joint": "lower", "link": "lower"}, "components": [ { @@ -16,7 +16,7 @@ }, { "component_type": "left_hand", - "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "urdf_path": "DH_PGI_140_80/DH_PGI_140_80.urdf", "transform": [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], @@ -36,7 +36,7 @@ }, { "component_type": "right_hand", - "urdf_path": "Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + "urdf_path": "DH_PGI_140_80/DH_PGI_140_80.urdf", "transform": [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], @@ -50,27 +50,27 @@ "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 + -1.57, -1.57, 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 + "left_eef": 1000.0, + "right_eef": 1000.0 }, "damping": { "left_arm": 1000.0, "right_arm": 1000.0, - "left_eef": 5.0, - "right_eef": 5.0 + "left_eef": 100.0, + "right_eef": 100.0 }, "max_effort": { "left_arm": 10000.0, "right_arm": 10000.0, - "left_eef": 500.0, - "right_eef": 500.0 + "left_eef": 10000.0, + "right_eef": 10000.0 } }, "control_parts": { @@ -78,20 +78,12 @@ "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" - ], + "left_eef": ["left_gripper_finger1_joint_1"], "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" - ] + "right_eef": ["right_gripper_finger1_joint_1"] }, "solver_cfg": { "left_arm": { @@ -102,9 +94,9 @@ "root_link_name": "left_base_link", "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], "tcp": [ - [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.2], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.121], [0.0, 0.0, 0.0, 1.0] ] }, @@ -116,9 +108,9 @@ "root_link_name": "right_base_link", "ik_nearest_weight": [1.0, 4.0, 1.0, 1.0, 1.0, 1.0], "tcp": [ - [0.0, -1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.2], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.121], [0.0, 0.0, 0.0, 1.0] ] } diff --git a/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json b/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json index 31084a497..b8759243e 100644 --- a/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json +++ b/embodichain/gen_sim/action_engine/generation/templates/robot_profiles.json @@ -5,9 +5,7 @@ "robot_family": "franka", "tabletop_clearance": 0.05, "arm_component_z": 0.4, - "arm_base_x": -1.45, - "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + "arm_base_x": -1.45 }, "dual_ur3": { "aliases": ["ur3", "dual_ur3_dh_pgi", "dual_ur3_robotiq", "dual_ur3_robotiq_arg2f_140"], @@ -16,9 +14,7 @@ "tabletop_clearance": 0.05, "arm_component_z": 0.4, "arm_base_x": -1.45, - "max_effort": 56.0, - "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + "max_effort": 56.0 }, "dual_ur5": { "aliases": ["ur5", "dual_ur5_dh_pgi", "dual_ur5_robotiq", "dual_ur5_robotiq_arg2f_140"], @@ -27,9 +23,7 @@ "tabletop_clearance": 0.05, "arm_component_z": 0.4, "arm_base_x": -1.45, - "max_effort": 10000.0, - "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + "max_effort": 10000.0 }, "dual_ur10": { "aliases": ["ur10", "dual_ur10_dh_pgi", "dual_ur10_robotiq", "dual_ur10_robotiq_arg2f_140"], @@ -38,8 +32,6 @@ "tabletop_clearance": 0.05, "arm_component_z": 0.3, "arm_base_x": -1.1, - "max_effort": 330.0, - "gripper_open_state": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - "gripper_close_state": [0.7, -0.7, 0.7, -0.7, -0.7, 0.7] + "max_effort": 330.0 } } diff --git a/embodichain/gen_sim/action_engine/gripper_profiles.py b/embodichain/gen_sim/action_engine/gripper_profiles.py new file mode 100644 index 000000000..69e921f69 --- /dev/null +++ b/embodichain/gen_sim/action_engine/gripper_profiles.py @@ -0,0 +1,326 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Validated GenSim gripper assets, controls, TCPs, and grasp geometry.""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any + +__all__ = [ + "GraspModelSpec", + "GripperModel", + "GripperProfile", + "get_gripper_profile", +] + +_Side = str +_Transform = tuple[ + tuple[float, float, float, float], + tuple[float, float, float, float], + tuple[float, float, float, float], + tuple[float, float, float, float], +] + + +class GripperModel(str, Enum): + """Gripper models supported by the GenSim composition root.""" + + PGI = "pgi" + ROBOTIQ = "robotiq" + + +@dataclass(frozen=True, slots=True) +class GraspModelSpec: + """Collision geometry used to interpret sampled poses as gripper TCP poses.""" + + model_id: str + min_opening_width: float + max_opening_width: float + finger_length: float + finger_width: float + finger_thickness: float + palm_depth: float + opening_margin: float + + def as_mapping(self) -> dict[str, float | str]: + """Return a detached JSON-compatible geometry description.""" + return { + "model_id": self.model_id, + "min_opening_width": self.min_opening_width, + "max_opening_width": self.max_opening_width, + "finger_length": self.finger_length, + "finger_width": self.finger_width, + "finger_thickness": self.finger_thickness, + "palm_depth": self.palm_depth, + "opening_margin": self.opening_margin, + } + + +@dataclass(frozen=True, slots=True) +class GripperProfile: + """One indivisible simulator, controller, kinematics, and grasp contract. + + ``tcp_transform`` is the row-major homogeneous transform from each solver's + configured ``end_link_name`` frame to the tool center point. No quaternion + conversion is involved in this contract. + """ + + model: GripperModel + asset_path: str + assembly_name: str + tcp_transform: _Transform + left_control_joints: tuple[str, ...] + right_control_joints: tuple[str, ...] + left_mimic_joints: tuple[str, ...] + right_mimic_joints: tuple[str, ...] + mimic_multipliers: tuple[float, ...] + mimic_offsets: tuple[float, ...] + simulated_joint_initial_positions: tuple[float, ...] + open_positions: tuple[float, ...] + close_positions: tuple[float, ...] + control_limits: tuple[tuple[float, float], ...] + drive_stiffness: float + drive_damping: float + drive_max_effort: float + grasp_model: GraspModelSpec + + def __post_init__(self) -> None: + control_count = len(self.left_control_joints) + if not control_count or len(self.right_control_joints) != control_count: + raise ValueError( + "Gripper profiles require matching non-empty hand controls." + ) + if not ( + len(self.open_positions) + == len(self.close_positions) + == len(self.control_limits) + == control_count + ): + raise ValueError( + "Gripper control states and limits must match control joints." + ) + mimic_count = len(self.left_mimic_joints) + if not ( + len(self.right_mimic_joints) + == len(self.mimic_multipliers) + == len(self.mimic_offsets) + == mimic_count + ): + raise ValueError("Gripper mimic metadata must have matching lengths.") + if len(self.simulated_joint_initial_positions) != len( + self.simulated_joint_names("left") + ): + raise ValueError( + "Gripper simulated initial positions must match physical joints." + ) + + def control_joint_names(self, side: _Side) -> tuple[str, ...]: + """Return the exact assembled control-joint names for one hand.""" + self._validate_side(side) + return self.left_control_joints if side == "left" else self.right_control_joints + + def mimic_joint_names(self, side: _Side) -> tuple[str, ...]: + """Return exact assembled mimic-joint names for one hand.""" + self._validate_side(side) + return self.left_mimic_joints if side == "left" else self.right_mimic_joints + + def simulated_joint_names(self, side: _Side) -> tuple[str, ...]: + """Return physical movable joints in assembled qpos order for one hand.""" + return tuple( + dict.fromkeys( + (*self.control_joint_names(side), *self.mimic_joint_names(side)) + ) + ) + + def runtime_manifest( + self, + *, + tcp_parent_frames: dict[str, str], + ) -> dict[str, Any]: + """Describe the selected physical and planning contract for diagnostics.""" + if set(tcp_parent_frames) != {"left", "right"} or not all( + isinstance(value, str) and value for value in tcp_parent_frames.values() + ): + raise ValueError( + "TCP parent frames require non-empty left and right links." + ) + return { + "model": self.model.value, + "asset_path": self.asset_path, + "control_joints": { + side: list(self.control_joint_names(side)) for side in ("left", "right") + }, + "mimic_joints": { + side: [ + { + "name": name, + "source": self.control_joint_names(side)[0], + "multiplier": self.mimic_multipliers[index], + "offset": self.mimic_offsets[index], + } + for index, name in enumerate(self.mimic_joint_names(side)) + ] + for side in ("left", "right") + }, + "open_positions": list(self.open_positions), + "close_positions": list(self.close_positions), + "control_limits": [list(limit) for limit in self.control_limits], + "tcp": { + "parent_frames": dict(tcp_parent_frames), + "transform_direction": "parent_link_to_tcp", + "matrix_layout": "row_major_homogeneous_4x4", + "quaternion_order": "not_applicable", + "transform": [list(row) for row in self.tcp_transform], + }, + "grasp_model": self.grasp_model.as_mapping(), + } + + @staticmethod + def _validate_side(side: _Side) -> None: + if side not in {"left", "right"}: + raise ValueError("Gripper side must be 'left' or 'right'.") + + +_PGI_PROFILE = GripperProfile( + model=GripperModel.PGI, + asset_path="DH_PGI_140_80/DH_PGI_140_80.urdf", + assembly_name="dh_pgi_140_80", + tcp_transform=( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.121), + (0.0, 0.0, 0.0, 1.0), + ), + left_control_joints=("left_gripper_finger1_joint_1",), + right_control_joints=("right_gripper_finger1_joint_1",), + left_mimic_joints=("left_gripper_finger2_joint_1",), + right_mimic_joints=("right_gripper_finger2_joint_1",), + mimic_multipliers=(1.0,), + mimic_offsets=(0.0,), + simulated_joint_initial_positions=(0.0, 0.0), + open_positions=(0.0,), + close_positions=(0.04,), + control_limits=((0.0, 0.04),), + drive_stiffness=1.0e3, + drive_damping=1.0e2, + drive_max_effort=1.0e4, + grasp_model=GraspModelSpec( + model_id="dh_pgi_140_80", + min_opening_width=0.003, + max_opening_width=0.100, + finger_length=0.10, + finger_width=0.040, + finger_thickness=0.01, + palm_depth=0.096, + opening_margin=0.03, + ), +) + +_ROBOTIQ_MIMIC_MULTIPLIERS = (-1.0, 1.0, -1.0, -1.0, 1.0) +_ROBOTIQ_PROFILE = GripperProfile( + model=GripperModel.ROBOTIQ, + asset_path="Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf", + assembly_name="robotiq_arg2f_140", + tcp_transform=( + (0.0, -1.0, 0.0, 0.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.2), + (0.0, 0.0, 0.0, 1.0), + ), + left_control_joints=( + "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_control_joints=( + "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", + ), + left_mimic_joints=( + "left_inner_knuckle_joint", + "left_inner_finger_joint", + "left_right_outer_knuckle_joint", + "left_right_inner_knuckle_joint", + "left_right_inner_finger_joint", + ), + right_mimic_joints=( + "right_left_inner_knuckle_joint", + "right_left_inner_finger_joint", + "right_outer_knuckle_joint", + "right_inner_knuckle_joint", + "right_inner_finger_joint", + ), + mimic_multipliers=_ROBOTIQ_MIMIC_MULTIPLIERS, + mimic_offsets=(0.0, 0.0, 0.0, 0.0, 0.0), + simulated_joint_initial_positions=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + open_positions=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + close_positions=(0.7, -0.7, 0.7, -0.7, -0.7, 0.7), + control_limits=( + (0.0, 0.7), + (-0.8757, 0.8757), + (-0.8757, 0.8757), + (-0.725, 0.725), + (-0.8757, 0.8757), + (-0.8757, 0.8757), + ), + drive_stiffness=50.0, + drive_damping=5.0, + drive_max_effort=500.0, + grasp_model=GraspModelSpec( + model_id="robotiq_arg2f_140", + min_opening_width=0.01, + max_opening_width=0.15, + finger_length=0.13, + finger_width=0.03, + finger_thickness=0.01, + palm_depth=0.08, + opening_margin=0.01, + ), +) + +_GRIPPER_PROFILES = { + GripperModel.PGI: _PGI_PROFILE, + GripperModel.ROBOTIQ: _ROBOTIQ_PROFILE, +} + + +def get_gripper_profile(model: GripperModel | str) -> GripperProfile: + """Return one strictly selected GenSim gripper profile.""" + if isinstance(model, GripperModel): + selected = model + elif isinstance(model, str): + try: + selected = GripperModel(model) + except ValueError as exc: + raise ValueError( + f"Unsupported gripper model {model!r}; expected one of: pgi, robotiq." + ) from exc + else: + raise TypeError( + f"Gripper model must be a string; expected one of: pgi, robotiq, got " + f"{type(model).__name__}." + ) + return _GRIPPER_PROFILES[selected] diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index b0f88e1d8..19f9df38d 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -31,6 +31,7 @@ build_atomic_capability_registry, ) from embodichain.gen_sim.action_engine.config import default_runtime_policy +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.lab.sim.atomic_actions import ( ActionBinding, ActionInvocation, @@ -168,6 +169,9 @@ def __init__( self.env = env self.num_envs = int(env.num_envs) self.device = env.device + self.gripper_profile = get_gripper_profile( + getattr(env, "agent_gripper_model", "pgi") + ) if grasp_policy is None: profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) grasp_policy = default_runtime_policy(profile).grasp @@ -980,6 +984,7 @@ def _planner_trace( trace = { "action_class": grounded.action_class, "arm": grounded.arm, + "gripper_model": self.gripper_profile.model.value, "planner": str(self.planner_policy["backend"]), "primary_strategy": invocation.motion_policy.strategy, "dynamic_collision_mode": invocation.motion_policy.dynamic_collision_mode.value, @@ -1857,14 +1862,15 @@ def _grasp_pose_generators( if not isinstance(filter_ground_collision, bool): raise TypeError("filter_ground_collision must be a boolean.") options = self.grasp_policy + geometry = self.gripper_profile.grasp_model model = ParallelJawGripperModelCfg( - model_id="gen_sim_parallel_jaw", - min_opening_width=float(options["min_open_length"]), - max_opening_width=float(options["max_open_length"]), - finger_length=float(options["finger_length"]), - finger_width=0.03, - finger_thickness=0.01, - palm_depth=0.08, + model_id=geometry.model_id, + min_opening_width=geometry.min_opening_width, + max_opening_width=geometry.max_opening_width, + finger_length=geometry.finger_length, + finger_width=geometry.finger_width, + finger_thickness=geometry.finger_thickness, + palm_depth=geometry.palm_depth, ) algorithm = AntipodalGraspPoseGeneratorCfg( sample_count=int(options["antipodal_n_sample"]), @@ -1876,7 +1882,7 @@ def _grasp_pose_generators( collision = ParallelJawGraspCollisionCfg( point_sample_density=float(options["point_sample_dense"]), max_decomposition_hulls=int(options["max_decomposition_hulls"]), - opening_margin=0.01, + opening_margin=geometry.opening_margin, filter_ground_collision=filter_ground_collision, ) annotation = GraspAnnotationCfg( diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py index b1fd60550..71c7ff368 100644 --- a/embodichain/gen_sim/task_engine/cli.py +++ b/embodichain/gen_sim/task_engine/cli.py @@ -19,6 +19,7 @@ from __future__ import annotations import argparse +from dataclasses import replace import json from pathlib import Path import sys @@ -100,6 +101,12 @@ def _add_workflow_arguments(parser: argparse.ArgumentParser) -> None: choices=_ROBOT_PROFILES, default="franka", ) + parser.add_argument( + "--gripper-model", + choices=("pgi", "robotiq"), + default=None, + help="Override the Task Engine planning gripper profile.", + ) _add_failure_policy_argument(parser) @@ -152,6 +159,9 @@ def _run_workflow( instruction = _instruction(args) adapter = SceneAdapter(model=args.model, robot_profile=args.robot_profile) workflow = TaskEngineWorkflow(scene_adapter=adapter) + workflow_cfg, planning_cfg, execution_cfg = load_task_engine_config(args.config) + if args.gripper_model is not None: + planning_cfg = replace(planning_cfg, gripper_model=args.gripper_model) with reserve_run_directory(args.output_root) as allocation: if scene is not None: validate_scene_output_separation(scene, allocation.path) @@ -166,6 +176,9 @@ def _run_workflow( "output_dir": allocation.path.as_posix(), }, config_path=args.config, + workflow_cfg=workflow_cfg, + planning_cfg=planning_cfg, + execution_cfg=execution_cfg, model=args.model, vlm_model=args.vlm_model, base_seed=args.base_seed, diff --git a/embodichain/gen_sim/task_engine/config.py b/embodichain/gen_sim/task_engine/config.py index 1a6df29bd..b2caa8ca0 100644 --- a/embodichain/gen_sim/task_engine/config.py +++ b/embodichain/gen_sim/task_engine/config.py @@ -107,6 +107,7 @@ class TaskEnginePlanningCfg: candidate_count: int = 3 planning_mode: str = "offline" + gripper_model: str = "pgi" max_episodes: int = 1 max_episode_steps: int = 6000 @@ -121,6 +122,11 @@ def __post_init__(self) -> None: raise ValueError(f"{field_name} must be a positive integer.") if self.planning_mode not in {"offline", "ab"}: raise ValueError("planning_mode must be offline or ab.") + if self.gripper_model not in {"pgi", "robotiq"}: + raise ValueError( + f"Unsupported gripper model {self.gripper_model!r}; expected one " + "of: pgi, robotiq." + ) def load_task_engine_config( @@ -167,6 +173,7 @@ def load_task_engine_config( if set(planning) != { "candidate_count", "planning_mode", + "gripper_model", "max_episodes", "max_episode_steps", }: diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml index e2ef9e5ad..c7b02ea17 100644 --- a/embodichain/gen_sim/task_engine/defaults.yaml +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -24,6 +24,7 @@ workflow: planning: candidate_count: 3 planning_mode: offline + gripper_model: robotiq max_episodes: 1 max_episode_steps: 6000 diff --git a/embodichain/gen_sim/task_engine/orchestration/coordinator.py b/embodichain/gen_sim/task_engine/orchestration/coordinator.py index dab074430..6196c1d60 100644 --- a/embodichain/gen_sim/task_engine/orchestration/coordinator.py +++ b/embodichain/gen_sim/task_engine/orchestration/coordinator.py @@ -31,6 +31,7 @@ generate_action_engine_config, ) from embodichain.gen_sim.action_engine.generation.artifacts import artifact_paths +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.gen_sim.action_engine.agent import ActionAgent from embodichain.gen_sim.action_engine.unbound import ActionCapabilityError from embodichain.gen_sim.action_engine.domain.task_contracts import ( @@ -180,6 +181,7 @@ def prepare( candidate_count: int = 3, overwrite: bool = False, planning_mode: str = "offline", + gripper_model: str = "pgi", vlm_model: str | None = None, max_episodes: int | None = None, max_episode_steps: int | None = None, @@ -196,6 +198,7 @@ def prepare( They publish the complete audit hand-off but never publish a TaskSpec, SeedGraph, Gym configuration, or GroundedTaskPlan. """ + gripper_model = get_gripper_profile(gripper_model).model.value normalized_source = self._coerce_source(source) validate_scene_output_separation(normalized_source.path, output_dir) with ArtifactTransaction(output_dir, overwrite=overwrite) as transaction: @@ -354,6 +357,7 @@ def prepare( "task_name": grounded_plan["task_id"], "task_spec": grounded_plan["task_spec"], "robot_profile": robot_profile, + "gripper_model": gripper_model, "source_scene_z_rotation_degrees": ( adaptation.prepared_scene.z_rotation_degrees ), diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py index 88968d021..feb0af5e3 100644 --- a/embodichain/gen_sim/task_engine/workflow.py +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -651,6 +651,7 @@ def run( model=model, candidate_count=effective_candidate_count, planning_mode=planning_cfg.planning_mode, + gripper_model=planning_cfg.gripper_model, vlm_model=vlm_model, max_episodes=planning_cfg.max_episodes, max_episode_steps=planning_cfg.max_episode_steps, @@ -1035,6 +1036,7 @@ def _publish( "planning": { "candidate_count": planning_cfg.candidate_count, "planning_mode": planning_cfg.planning_mode, + "gripper_model": planning_cfg.gripper_model, "max_episodes": planning_cfg.max_episodes, "max_episode_steps": planning_cfg.max_episode_steps, }, diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py index 6581c2707..a9b33084d 100644 --- a/tests/gen_sim/action_engine/config/test_runtime_policy.py +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -83,6 +83,9 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: } assert runtime.grounding["arrangement"]["row_search_radius"] == 0.25 assert runtime.grasp["antipodal_n_sample"] == 10000 + assert "max_open_length" not in runtime.grasp + assert "min_open_length" not in runtime.grasp + assert "finger_length" not in runtime.grasp assert runtime.motion_modifiers["orientation"]["upright"]["MoveHeldObject"][ "surface_clearance" ] == pytest.approx(0.05) @@ -117,6 +120,7 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: 0.2617993877991494 ) assert generation["physics"]["rigid_object"]["mass"] == pytest.approx(0.1) + assert generation["task"]["default_gripper_model"] == "pgi" assert generation["environment"]["arm_aim_yaw_offset"] == { "left": pytest.approx(0.0), "right": pytest.approx(0.0), @@ -299,11 +303,42 @@ def test_v6_policy_snapshot_adds_axis_align_defaults_without_rewriting_e1() -> N } ) - assert resolved.schema_version == "action_engine_runtime_policy_v7" + assert resolved.schema_version == "action_engine_runtime_policy_v8" assert resolved.motion_defaults["AxisAlign"]["sample_interval"] == 180 assert resolved.motion_defaults["PickUp"]["lift_height"] == pytest.approx(0.11) +def test_v7_policy_snapshot_drops_legacy_gripper_geometry() -> None: + snapshot = default_runtime_policy("dual_ur10").as_mapping() + snapshot["schema_version"] = "action_engine_runtime_policy_v7" + snapshot["grasp"].update( + { + "min_open_length": 0.01, + "max_open_length": 0.15, + "finger_length": 0.13, + } + ) + payload = json.dumps( + snapshot, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode("utf-8") + + resolved = resolve_agent_runtime_policy( + { + "robot_profile": "dual_ur10", + "runtime_policy": snapshot, + "runtime_policy_hash": hashlib.sha256(payload).hexdigest(), + } + ) + + assert resolved.schema_version == "action_engine_runtime_policy_v8" + assert "min_open_length" not in resolved.grasp + assert "max_open_length" not in resolved.grasp + assert "finger_length" not in resolved.grasp + + def test_narrow_v1_policy_snapshot_is_migrated_to_complete_runtime_policy() -> None: snapshot = { "schema_version": "action_engine_runtime_policy_v1", @@ -354,7 +389,7 @@ def test_v3_policy_snapshot_is_migrated_with_default_planner_policy() -> None: } ) - assert resolved.schema_version == "action_engine_runtime_policy_v7" + assert resolved.schema_version == "action_engine_runtime_policy_v8" assert resolved.planner == expected.planner diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py index 4f5116272..fc474ff25 100644 --- a/tests/gen_sim/action_engine/generation/test_generation.py +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -48,6 +48,7 @@ build_agent_config, build_fast_gym_config, ) +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.gen_sim.action_engine.generation.generator import ( _add_ab_camera_requirements, _scene_requirements_from_bindings, @@ -390,7 +391,7 @@ def test_fast_gym_config_has_runnable_franka_contract(gym_export: Path) -> None: ) assert config["env"]["observations"]["norm_robot_eef_joint"]["params"][ "joint_ids" - ] == list(range(14, 26)) + ] == [14, 16] def test_offline_recording_can_be_disabled_but_ab_keeps_audience_recorder( @@ -669,12 +670,7 @@ def test_fast_gym_config_supports_all_robot_profiles( robot_uid: str, solver_type: str | None, ) -> None: - rotated_tcp = [ - [0.0, -1.0, 0.0, 0.0], - [1.0, 0.0, 0.0, 0.0], - [0.0, 0.0, 1.0, 0.2], - [0.0, 0.0, 0.0, 1.0], - ] + pgi = get_gripper_profile("pgi") identity_hand_mount = [ [1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], @@ -694,8 +690,11 @@ def test_fast_gym_config_supports_all_robot_profiles( assert config["robot"]["uid"] == robot_uid assert config["env"]["extensions"]["agent_robot_profile"] == profile + assert config["env"]["extensions"]["agent_gripper_model"] == "pgi" for arm in ("left_arm", "right_arm"): - assert config["robot"]["solver_cfg"][arm]["tcp"] == rotated_tcp + assert config["robot"]["solver_cfg"][arm]["tcp"] == [ + list(row) for row in pgi.tcp_transform + ] components = { component["component_type"]: component for component in config["robot"]["urdf_cfg"]["components"] @@ -706,6 +705,83 @@ def test_fast_gym_config_supports_all_robot_profiles( assert config["robot"]["solver_cfg"]["left_arm"]["ur_type"] == solver_type +@pytest.mark.parametrize("gripper_model", ["pgi", "robotiq"]) +def test_fast_gym_config_applies_one_complete_gripper_profile( + gym_export: Path, + gripper_model: str, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="gripper_profile_task", + task_description="Profile smoke test.", + robot_profile="ur10", + gripper_model=gripper_model, + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=20, + ) + profile = get_gripper_profile(gripper_model) + robot = config["robot"] + extensions = config["env"]["extensions"] + components = { + component["component_type"]: component + for component in robot["urdf_cfg"]["components"] + } + + assert extensions["agent_gripper_model"] == gripper_model + assert extensions["gripper_open_state"] == list(profile.open_positions) + assert extensions["gripper_close_state"] == list(profile.close_positions) + assert extensions["gripper_profile"]["model"] == gripper_model + assert robot["control_parts"]["left_eef"] == list( + profile.control_joint_names("left") + ) + assert robot["control_parts"]["right_eef"] == list( + profile.control_joint_names("right") + ) + assert components["left_hand"]["urdf_path"] == profile.asset_path + assert components["right_hand"]["urdf_path"] == profile.asset_path + assert robot["solver_cfg"]["left_arm"]["tcp"] == [ + list(row) for row in profile.tcp_transform + ] + assert robot["solver_cfg"]["right_arm"]["tcp"] == [ + list(row) for row in profile.tcp_transform + ] + assert robot["urdf_cfg"]["fname"].endswith(f"{profile.assembly_name}_basket") + + +def test_config_builders_reject_unknown_gripper_before_materialization( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + + with pytest.raises(ValueError, match="pgi.*robotiq"): + build_fast_gym_config( + scene, + task_name="invalid_gripper", + task_description="Invalid profile.", + robot_profile="ur10", + gripper_model="parallel_jaw", + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + +def test_agent_config_serializes_selected_gripper(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + config = build_agent_config( + task_name="gripper_profile_task", + robot_profile="ur10", + gripper_model="robotiq", + execution_program_hash="a" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + ) + + assert config["gripper_model"] == "robotiq" + + @pytest.mark.parametrize( ( "profile", @@ -998,7 +1074,7 @@ def capture_writer(*args, **kwargs): 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_v7" + "action_engine_runtime_policy_v8" ) assert agent_config["runtime_policy"]["planner"]["dynamic_collision"] is True assert agent_config["runtime_policy"]["planner"]["static_obstacle_uids"] == [ @@ -1415,7 +1491,7 @@ def test_agent_config_uses_relative_program_paths(gym_export: Path) -> None: assert config["runtime_policy"]["motion_defaults"]["PickUp"][ "lift_height" ] == pytest.approx(0.16) - assert config["runtime_policy"]["grasp"]["max_open_length"] == pytest.approx(0.15) + assert "max_open_length" not in config["runtime_policy"]["grasp"] assert len(config["runtime_policy_hash"]) == 64 diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 5a70dc228..902ae54f8 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -35,6 +35,7 @@ GroundedAction, ) from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.lab.sim.atomic_actions import ( Affordance, ActionBinding, @@ -419,6 +420,7 @@ def _planner_env( *, table: Any | None = None, rigid_objects: dict[str, Any] | None = None, + gripper_model: str = "pgi", ) -> SimpleNamespace: entities = dict(rigid_objects or {}) if table is not None: @@ -434,6 +436,7 @@ def _planner_env( right_eef_joints=[6, 7], open_state=torch.zeros(2), close_state=torch.ones(2), + agent_gripper_model=gripper_model, get_agent_arm_control_part=lambda is_left: ( "physical_left_arm" if is_left else "physical_right_arm" ), @@ -661,7 +664,7 @@ def test_coordinated_pickment_geometry_candidates_are_deterministic_for_tray() - def test_grasp_generators_follow_mainline_service_contract() -> None: - adapter = AtomicActionAdapter(_planner_env()) + adapter = AtomicActionAdapter(_planner_env(gripper_model="pgi")) generators = adapter._grasp_pose_generators() @@ -669,6 +672,10 @@ def test_grasp_generators_follow_mainline_service_contract() -> None: generator = generators["physical_left_eef"] assert generators["physical_right_eef"] is generator assert isinstance(generator, AntipodalGraspPoseGenerator) + assert generator.gripper_model.model_id == "dh_pgi_140_80" + assert generator.gripper_model.max_opening_width == pytest.approx(0.100) + assert generator.gripper_model.finger_length == pytest.approx(0.10) + assert generator.collision_cfg.opening_margin == pytest.approx(0.03) assert generator.algorithm_cfg.sample_count == 10000 assert generator.algorithm_cfg.approach_direction_samples == 4 assert generator.algorithm_cfg.max_candidates == 500 @@ -676,6 +683,73 @@ def test_grasp_generators_follow_mainline_service_contract() -> None: assert generator.collision_cfg.filter_ground_collision is True +def test_robotiq_grasp_generator_preserves_existing_geometry() -> None: + adapter = AtomicActionAdapter(_planner_env(gripper_model="robotiq")) + + generators = adapter._grasp_pose_generators() + generator = generators["physical_left_eef"] + + assert generators["physical_right_eef"] is generator + assert generator.gripper_model.model_id == "robotiq_arg2f_140" + assert generator.gripper_model.max_opening_width == pytest.approx(0.15) + assert generator.gripper_model.finger_length == pytest.approx(0.13) + assert generator.collision_cfg.opening_margin == pytest.approx(0.01) + + +@pytest.mark.parametrize("gripper_model", ["pgi", "robotiq"]) +def test_control_profiles_use_selected_gripper_joint_semantics( + gripper_model: str, +) -> None: + selected = get_gripper_profile(gripper_model) + hand_dof = len(selected.open_positions) + joint_ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": list(range(2, 2 + hand_dof)), + "physical_right_arm": [2 + hand_dof, 3 + hand_dof], + "physical_right_eef": list(range(4 + hand_dof, 4 + 2 * hand_dof)), + } + robot = SimpleNamespace( + uid="profile_robot", + dof=4 + 2 * hand_dof, + control_parts=joint_ids, + get_joint_ids=lambda *, name: list(joint_ids[name]), + ) + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + robot=robot, + sim=SimpleNamespace(get_rigid_object=lambda _uid: None), + agent_gripper_model=gripper_model, + open_state=torch.tensor(selected.open_positions), + close_state=torch.tensor(selected.close_positions), + get_agent_arm_control_part=lambda is_left: ( + "physical_left_arm" if is_left else "physical_right_arm" + ), + get_agent_eef_control_part=lambda is_left: ( + "physical_left_eef" if is_left else "physical_right_eef" + ), + ) + + profiles = AtomicActionAdapter(env)._control_profiles() + + assert set(profiles) == {"physical_left_eef", "physical_right_eef"} + for command_profile in profiles.values(): + open_qpos = command_profile.commands["open"].resolve( + num_envs=1, + control_dof=hand_dof, + device="cpu", + ) + grasp_qpos = command_profile.commands["grasp"].resolve( + num_envs=1, + control_dof=hand_dof, + device="cpu", + ) + torch.testing.assert_close(open_qpos[0], torch.tensor(selected.open_positions)) + torch.testing.assert_close( + grasp_qpos[0], torch.tensor(selected.close_positions) + ) + + def test_coordinated_grasp_generator_honors_ground_filter_policy() -> None: adapter = AtomicActionAdapter(_planner_env()) @@ -1173,6 +1247,7 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: ) assert outcome.planner_trace["primary_action_diagnostics"]["marker"] == 1.0 assert outcome.planner_trace["fallback_action_diagnostics"]["marker"] == 2.0 + assert outcome.planner_trace["gripper_model"] == "pgi" def test_collision_required_cleanup_does_not_use_unsafe_fallback( diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index bc48cb88f..c078a4fb5 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -682,7 +682,7 @@ def test_runtime_policy_v4_migrates_grasp_direction_count() -> None: } ) - assert policy.schema_version == "action_engine_runtime_policy_v7" + assert policy.schema_version == "action_engine_runtime_policy_v8" assert policy.grasp["n_deviated_approach_directions"] == 4 @@ -728,7 +728,7 @@ def test_runtime_policy_v5_migrates_support_geometry_thresholds() -> None: } ) - assert policy.schema_version == "action_engine_runtime_policy_v7" + assert policy.schema_version == "action_engine_runtime_policy_v8" assert policy.predicate_fallbacks["support_min_overlap_ratio"] == 0.25 assert policy.grounding["placement"]["clearance"] == 0.019 assert policy.grounding["placement"]["candidate_count"] == 5 diff --git a/tests/gen_sim/action_engine/test_gripper_profiles.py b/tests/gen_sim/action_engine/test_gripper_profiles.py new file mode 100644 index 000000000..672d65bde --- /dev/null +++ b/tests/gen_sim/action_engine/test_gripper_profiles.py @@ -0,0 +1,102 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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.gripper_profiles import ( + GripperModel, + get_gripper_profile, +) + + +def test_gripper_models_are_strictly_validated() -> None: + assert get_gripper_profile("pgi").model is GripperModel.PGI + assert get_gripper_profile("robotiq").model is GripperModel.ROBOTIQ + + for invalid in ("", "PGI", "robotiq ", "unknown", None): + with pytest.raises((TypeError, ValueError), match="pgi.*robotiq"): + get_gripper_profile(invalid) # type: ignore[arg-type] + + +def test_pgi_profile_owns_asset_control_mimic_tcp_and_grasp_geometry() -> None: + profile = get_gripper_profile("pgi") + + assert profile.asset_path == "DH_PGI_140_80/DH_PGI_140_80.urdf" + assert profile.control_joint_names("left") == ("left_gripper_finger1_joint_1",) + assert profile.mimic_joint_names("left") == ("left_gripper_finger2_joint_1",) + assert profile.simulated_joint_names("left") == ( + "left_gripper_finger1_joint_1", + "left_gripper_finger2_joint_1", + ) + assert profile.open_positions == (0.0,) + assert profile.close_positions == (0.04,) + assert profile.control_limits == ((0.0, 0.04),) + assert profile.tcp_transform == ( + (1.0, 0.0, 0.0, 0.0), + (0.0, 1.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.121), + (0.0, 0.0, 0.0, 1.0), + ) + assert profile.grasp_model.model_id == "dh_pgi_140_80" + assert profile.grasp_model.max_opening_width == pytest.approx(0.100) + assert profile.grasp_model.finger_length == pytest.approx(0.10) + assert profile.grasp_model.opening_margin == pytest.approx(0.03) + + +def test_robotiq_profile_preserves_existing_rotation_and_joint_semantics() -> None: + profile = get_gripper_profile("robotiq") + + assert profile.asset_path == ("Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf") + assert profile.control_joint_names("left") == ( + "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", + ) + assert profile.open_positions == (0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + assert profile.close_positions == (0.7, -0.7, 0.7, -0.7, -0.7, 0.7) + assert profile.tcp_transform == ( + (0.0, -1.0, 0.0, 0.0), + (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 1.0, 0.2), + (0.0, 0.0, 0.0, 1.0), + ) + assert profile.grasp_model.model_id == "robotiq_arg2f_140" + assert profile.grasp_model.max_opening_width == pytest.approx(0.15) + assert profile.grasp_model.finger_length == pytest.approx(0.13) + assert profile.grasp_model.opening_margin == pytest.approx(0.01) + + +def test_profile_manifest_records_tcp_frame_and_transform_conventions() -> None: + profile = get_gripper_profile("pgi") + + manifest = profile.runtime_manifest( + tcp_parent_frames={"left": "left_ee_link", "right": "right_ee_link"} + ) + + assert manifest["model"] == "pgi" + assert manifest["tcp"]["parent_frames"] == { + "left": "left_ee_link", + "right": "right_ee_link", + } + assert manifest["tcp"]["transform_direction"] == "parent_link_to_tcp" + assert manifest["tcp"]["matrix_layout"] == "row_major_homogeneous_4x4" + assert manifest["tcp"]["quaternion_order"] == "not_applicable" + assert manifest["grasp_model"]["model_id"] == "dh_pgi_140_80" diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py index 8096eab98..6c2b2bca5 100644 --- a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -600,6 +600,7 @@ def generator(_scene, output, **kwargs): assert result.bound assert generator_calls + assert generator_calls[0]["gripper_model"] == "pgi" assert not (result.output_dir / ".task_engine_input").exists() grounded = json.loads( (result.output_dir / "grounded_task_plan.json").read_text(encoding="utf-8") diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py index f135c5ce0..d7068c126 100644 --- a/tests/gen_sim/task_engine/test_parallel_workflow.py +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -404,11 +404,13 @@ def test_parallel_workflow_accepts_one_success_and_publishes_all_graphs( assert manifest["configuration"]["planning"] == { "candidate_count": 3, "planning_mode": "offline", + "gripper_model": "pgi", "max_episodes": 1, "max_episode_steps": 6000, } assert manifest["configuration"]["execution"]["dataset_saving"] is False assert coordinator.kwargs[0]["max_episode_steps"] == 6000 + assert coordinator.kwargs[0]["gripper_model"] == "pgi" assert coordinator.kwargs[0]["final_inspection"]["scene_revision_id"] == "0" * 64 assert ( coordinator.kwargs[0]["unbound_action_plan"]["candidate_id"] == "candidate_01" diff --git a/tests/gen_sim/task_engine/test_workflow.py b/tests/gen_sim/task_engine/test_workflow.py index e9e40fa35..2a9df8092 100644 --- a/tests/gen_sim/task_engine/test_workflow.py +++ b/tests/gen_sim/task_engine/test_workflow.py @@ -272,6 +272,7 @@ def test_packaged_workflow_configuration_uses_recovery_defaults() -> None: assert workflow.max_action_attempts == 3 assert planning.candidate_count == 3 assert planning.planning_mode == "offline" + assert planning.gripper_model == "pgi" assert planning.max_episodes == 1 assert planning.max_episode_steps == 6000 assert execution.num_envs == 1 @@ -290,6 +291,7 @@ def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: planning: candidate_count: 7 planning_mode: offline + gripper_model: robotiq max_episodes: 2 max_episode_steps: 5000 execution: @@ -306,6 +308,7 @@ def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: assert workflow.max_scene_attempts == 4 assert workflow.max_action_attempts == 5 assert planning.candidate_count == 7 + assert planning.gripper_model == "robotiq" assert planning.max_episodes == 2 assert planning.max_episode_steps == 5000 assert execution.num_envs == 6 @@ -335,3 +338,5 @@ def test_planning_configuration_rejects_invalid_values() -> None: TaskEnginePlanningCfg(candidate_count=0) with pytest.raises(ValueError, match="planning_mode"): TaskEnginePlanningCfg(planning_mode="unsupported") + with pytest.raises(ValueError, match="pgi.*robotiq"): + TaskEnginePlanningCfg(gripper_model="unsupported") From 62ab860c4233b96a1283889ae85bc60243a6201f Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:03:03 +0800 Subject: [PATCH 71/85] fix(gen-sim): preserve source videos and reduce recording frequency --- embodichain/gen_sim/action_engine/config/defaults.yaml | 2 +- embodichain/gen_sim/video_archive.py | 9 ++++++--- tests/gen_sim/test_video_archive.py | 10 +++++----- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index c5442dd04..4643b4ee3 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -30,7 +30,7 @@ generation: recording: enabled: true resolution: [640, 360] - interval_step: 1 + interval_step: 5 arm_aim_yaw_offset: left: 0.0 right: 0.0 diff --git a/embodichain/gen_sim/video_archive.py b/embodichain/gen_sim/video_archive.py index 9d49ee19a..706952d7b 100644 --- a/embodichain/gen_sim/video_archive.py +++ b/embodichain/gen_sim/video_archive.py @@ -14,12 +14,13 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Rename one completed GenSim recording to its task ID.""" +"""Copy one completed GenSim recording to its task ID.""" from __future__ import annotations import argparse from pathlib import Path +import shutil import sys from typing import Any, Sequence @@ -89,7 +90,7 @@ def _archive_task_video( source_stem: str, task_id: str, ) -> Path: - """Move a completed recording to ``.``. + """Copy a completed recording to ``.``. Args: video_directory: Directory containing the completed recording. @@ -132,7 +133,9 @@ def _archive_task_video( source = candidates[0] extension = source.name[len(source_stem) :] destination = directory / f"{task_id}{extension}" - source.replace(destination) + if destination.exists() or destination.is_symlink(): + destination.unlink() + shutil.copy2(source, destination) return destination diff --git a/tests/gen_sim/test_video_archive.py b/tests/gen_sim/test_video_archive.py index 658ea9120..feaa770d7 100644 --- a/tests/gen_sim/test_video_archive.py +++ b/tests/gen_sim/test_video_archive.py @@ -54,7 +54,7 @@ def _write_source(directory: Path, extension: str, content: bytes = b"video") -> return source -def test_archive_task_video_renames_source_and_preserves_extension( +def test_archive_task_video_copies_source_and_preserves_extension( tmp_path: Path, ) -> None: source = _write_source(tmp_path, ".webm") @@ -67,7 +67,7 @@ def test_archive_task_video_renames_source_and_preserves_extension( assert destination == tmp_path / "task2_1.webm" assert destination.read_bytes() == b"video" - assert not source.exists() + assert source.read_bytes() == b"video" @pytest.mark.parametrize("task_id", ["../task2_1", "task2/1", r"task2\1", ".."]) @@ -113,7 +113,7 @@ def test_archive_task_video_overwrites_existing_target( assert result == destination assert destination.read_bytes() == b"new" - assert not source.exists() + assert source.read_bytes() == b"new" def test_consecutive_tasks_keep_independent_videos(tmp_path: Path) -> None: @@ -127,7 +127,7 @@ def test_consecutive_tasks_keep_independent_videos(tmp_path: Path) -> None: assert (tmp_path / "task2_1.mp4").read_bytes() == b"first" assert (tmp_path / "task2_2.mp4").read_bytes() == b"second" - assert not (tmp_path / f"{SOURCE_STEM}.mp4").exists() + assert (tmp_path / f"{SOURCE_STEM}.mp4").read_bytes() == b"second" def test_task_recording_uses_runtime_recorder_path(tmp_path: Path) -> None: @@ -138,7 +138,7 @@ def test_task_recording_uses_runtime_recorder_path(tmp_path: Path) -> None: assert destination == tmp_path / "task2_1.mkv" assert destination.read_bytes() == b"video" - assert not source.exists() + assert source.read_bytes() == b"video" def test_task_recording_is_noop_when_recording_is_disabled() -> None: From f611c0fec295f1c007f27dfffa4ad3cdf7a48363 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Wed, 26 Aug 2026 19:55:43 +0800 Subject: [PATCH 72/85] fix(action-engine): align E2 upright orientation to the local z-axis --- embodichain/gen_sim/action_engine/runtime/actions.py | 2 +- embodichain/gen_sim/action_engine/runtime/grounding.py | 9 +++++++++ .../gen_sim/action_engine/tasks/interpretation.py | 2 +- embodichain/gen_sim/action_engine/tasks/recipes.py | 2 +- .../action_engine/runtime/test_runtime_contracts.py | 3 ++- tests/gen_sim/action_engine/tasks/test_factory.py | 2 ++ tests/gen_sim/action_engine/tasks/test_interpretation.py | 3 +++ 7 files changed, 19 insertions(+), 4 deletions(-) diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 19f9df38d..a27e6cd5f 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -1492,7 +1492,7 @@ def _build_single_arm_config( ) elif approach_mode is not None: raise ValueError(f"Unknown approach_direction_mode {approach_mode!r}.") - for name in ("approach_direction", "obj_upright_direction"): + for name in ("approach_direction", "obj_upright_direction", "target_axis"): if name in policy and not isinstance(policy[name], torch.Tensor): policy[name] = torch.as_tensor( policy[name], dtype=torch.float32, device=self.device diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index b1855c1f9..7c8892f29 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -76,6 +76,8 @@ __all__ = ["ActionGrounder", "LiveArrangementPlan", "LivePlacementPlan"] _E2_CLEARANCE_RETREAT_DISTANCE = 0.20 +DEFAULT_INTERNAL_AXIS = (0.0, 0.0, 1.0) +DEFAULT_TARGET_AXIS = (0.0, 0.0, 1.0) def _batched_pose(value: Any, env: Any) -> torch.Tensor: @@ -812,6 +814,7 @@ def ground( internal_axis=self._upright_local_direction(step), ), ) + policy.setdefault("target_axis", DEFAULT_TARGET_AXIS) policy.setdefault( "surface_clearance", float(self._policy_value(policy, "surface_clearance")), @@ -2542,6 +2545,12 @@ def _relative_object_spacing( def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: axis = self._upright_local_axis(step) + if axis == "z": + return torch.tensor( + DEFAULT_INTERNAL_AXIS, + dtype=torch.float32, + device=self.env.device, + ) entity = _object(self.env, step.object_uid) vertices = _local_vertices(entity, self.env, 0) if axis in {"long_axis", "short_axis"}: diff --git a/embodichain/gen_sim/action_engine/tasks/interpretation.py b/embodichain/gen_sim/action_engine/tasks/interpretation.py index f84f6d2bd..c4d002b34 100644 --- a/embodichain/gen_sim/action_engine/tasks/interpretation.py +++ b/embodichain/gen_sim/action_engine/tasks/interpretation.py @@ -262,7 +262,7 @@ def _emit_step( { "orientation_goal": "upright", "support_role": "table", - "upright_local_axis": "long_axis", + "upright_local_axis": "z", } ) elif task_type == "E3": diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index 8966eb82c..2ced064e7 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -406,7 +406,7 @@ def _recipe( "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")), + "upright_local_axis": str(params.get("upright_local_axis", "z")), **_orientation_extensions(params), } if terminal_behavior == "hold": diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index c078a4fb5..1cccaa34c 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -3816,8 +3816,9 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - ) assert torch.equal( orient_alignment.target.semantics.affordance.internal_axis, - torch.tensor([1.0, 0.0, 0.0]), + torch.tensor([0.0, 0.0, 1.0]), ) + assert orient_alignment.cfg["target_axis"] == (0.0, 0.0, 1.0) assert handover_pickup.target.grasp_xpos is None assert isinstance( handover_pickup.target.semantics.affordance, diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index 07c84fbf3..a3a483c14 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -238,6 +238,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "MoveEndEffector", "MoveJoints", ] + assert orient["goal"]["upright_local_axis"] == "z" + assert orient_nodes[0]["postcondition"]["local_axis"] == "z" assert orient_nodes[0]["motion_policy"] == {"modifiers": []} assert [node["role"] for node in orient_nodes] == [ "primary", diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index 9c7541ed9..263790883 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -310,6 +310,9 @@ def caller(**kwargs): "object_01": "purple_can", "object_02": "orange_can", } + assert ( + grounded.task_spec["task_instances"][0]["params"]["upright_local_axis"] == "z" + ) placement_actions = [ node["atomic_action"] for node in graph["nodes"] if node["task_type"] == "E1" ] From fa26c340aa02551495d2224c6bce287648864d96 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Thu, 27 Aug 2026 05:20:48 +0800 Subject: [PATCH 73/85] feat(gen-sim): add configurable planner modes and runtime diagnostics --- .../cli/generate_action_agent_config.py | 47 ++++++ .../config/planner_modes/curobo.yaml | 18 ++ .../config/planner_modes/ik_interp.yaml | 18 ++ .../config/planner_modes/toppra.yaml | 18 ++ .../action_engine/config/runtime_policy.py | 106 +++++++++++- .../generation/config_builder.py | 19 ++- .../action_engine/generation/generator.py | 2 + .../gen_sim/action_engine/runtime/actions.py | 74 ++++++++- .../gen_sim/action_engine/runtime/executor.py | 23 ++- embodichain/gen_sim/task_engine/cli.py | 19 +++ embodichain/gen_sim/task_engine/config.py | 18 +- embodichain/gen_sim/task_engine/defaults.yaml | 2 + .../task_engine/orchestration/coordinator.py | 2 + embodichain/gen_sim/task_engine/workflow.py | 2 + .../generation/test_generation.py | 154 ++++++++++++++++++ .../action_engine/runtime/test_actions.py | 111 ++++++++++++- .../orchestration/test_coordinator_cli.py | 53 ++++++ tests/gen_sim/task_engine/test_workflow.py | 11 ++ 18 files changed, 676 insertions(+), 21 deletions(-) create mode 100644 embodichain/gen_sim/action_engine/config/planner_modes/curobo.yaml create mode 100644 embodichain/gen_sim/action_engine/config/planner_modes/ik_interp.yaml create mode 100644 embodichain/gen_sim/action_engine/config/planner_modes/toppra.yaml 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 index 65b4665b8..2f46d4864 100644 --- a/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py +++ b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py @@ -19,9 +19,17 @@ from __future__ import annotations import argparse +from collections.abc import Mapping from pathlib import Path +import yaml + from embodichain.gen_sim.action_engine.config import generation_defaults +from embodichain.gen_sim.action_engine.config.runtime_policy import ( + _PLANNER_MODES, + _planner_policy_with_mode, + _resolve_planner_policy, +) from embodichain.gen_sim.action_engine.generation import ( generate_action_engine_config, ) @@ -123,6 +131,23 @@ def build_parser() -> argparse.ArgumentParser: default="offline", help="Generate one offline bundle or an offline/online A/B bundle.", ) + parser.add_argument( + "--planner-config", + default=None, + help=( + "Optional YAML file containing a planner mapping. The canonical " + "policy and runtime_policy_hash are regenerated into agent_config.json." + ), + ) + parser.add_argument( + "--planner-mode", + choices=_PLANNER_MODES, + default=None, + help=( + "Explicit planner mode override. When omitted, planner YAML/defaults " + "remain authoritative." + ), + ) parser.add_argument( "--source_scene_z_rotation_degrees", "--source-scene-z-rotation-degrees", @@ -183,6 +208,12 @@ def cli() -> None: """Generate and report the canonical Action Engine artifact bundle.""" args = build_parser().parse_args() task_description = _resolve_task_description(args) + planner_policy = _load_planner_config(args.planner_config) + if args.planner_mode is not None: + planner_policy = _planner_policy_with_mode( + planner_policy, + args.planner_mode, + ) paths = generate_action_engine_config( args.gym_project, args.output_dir, @@ -202,6 +233,7 @@ def cli() -> None: randomize_table_material=args.randomize_table_material, planning_mode=args.planning_mode, vlm_model=args.vlm_model, + planner_policy=planner_policy, ) print(f"Generated gym config: {paths.gym_config}") @@ -220,6 +252,21 @@ def cli() -> None: ) +def _load_planner_config(path: str | None) -> dict[str, object] | None: + if path is None: + return None + config_path = Path(path).expanduser().resolve() + raw = yaml.safe_load(config_path.read_text(encoding="utf-8")) + if not isinstance(raw, Mapping): + raise TypeError("Planner YAML must contain a mapping.") + planner = raw.get("planner") if set(raw) == {"planner"} else raw + if not isinstance(planner, Mapping): + raise TypeError("Planner YAML 'planner' must be a mapping.") + resolved = dict(planner) + _resolve_planner_policy(resolved) + return resolved + + def _resolve_task_description(args: argparse.Namespace) -> str: task_spec = getattr(args, "task_spec", None) if task_spec: diff --git a/embodichain/gen_sim/action_engine/config/planner_modes/curobo.yaml b/embodichain/gen_sim/action_engine/config/planner_modes/curobo.yaml new file mode 100644 index 000000000..b322cd707 --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/planner_modes/curobo.yaml @@ -0,0 +1,18 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +planner: + mode: curobo diff --git a/embodichain/gen_sim/action_engine/config/planner_modes/ik_interp.yaml b/embodichain/gen_sim/action_engine/config/planner_modes/ik_interp.yaml new file mode 100644 index 000000000..2acd3e2b3 --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/planner_modes/ik_interp.yaml @@ -0,0 +1,18 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +planner: + mode: ik_interp diff --git a/embodichain/gen_sim/action_engine/config/planner_modes/toppra.yaml b/embodichain/gen_sim/action_engine/config/planner_modes/toppra.yaml new file mode 100644 index 000000000..20d9b7433 --- /dev/null +++ b/embodichain/gen_sim/action_engine/config/planner_modes/toppra.yaml @@ -0,0 +1,18 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +planner: + mode: toppra diff --git a/embodichain/gen_sim/action_engine/config/runtime_policy.py b/embodichain/gen_sim/action_engine/config/runtime_policy.py index cfcba3c22..0e0775a9c 100644 --- a/embodichain/gen_sim/action_engine/config/runtime_policy.py +++ b/embodichain/gen_sim/action_engine/config/runtime_policy.py @@ -128,6 +128,32 @@ "max_attempts", "collision_activation_distance", } +_PLANNER_MODES = ("curobo", "toppra", "ik_interp") +_PLANNER_MODE_FIELDS = frozenset( + { + "backend", + "single_arm_strategy", + "coordinated_strategy", + "dynamic_collision", + } +) +_PLANNER_MODE_PATCHES: dict[str, dict[str, Any]] = { + "curobo": { + "backend": "curobo", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + }, + "toppra": { + "backend": "toppra", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + }, + "ik_interp": { + "backend": "toppra", + "single_arm_strategy": "ik_interp", + "coordinated_strategy": "ik_interp", + }, +} _MOTION_DEFAULT_ACTIONS = { "AxisAlign", "CoordinatedPickment", @@ -445,6 +471,80 @@ def generation_defaults() -> dict[str, Any]: return deepcopy(dict(value)) +def _resolve_planner_policy( + planner_policy: Mapping[str, Any] | None = None, + *, + robot_profile: str = "dual_ur10", +) -> dict[str, Any]: + """Merge and validate a partial generation-time planner policy. + + ``mode`` is a generation-only shorthand for the backend and strategy + fields. The returned mapping is the canonical runtime snapshot fragment. + Backend-specific overrides are rejected when that backend is not selected, + while package defaults remain present but dormant for stable hashes. + + Args: + planner_policy: Optional partial planner mapping loaded from YAML. + robot_profile: Runtime profile used to resolve package defaults. + + Returns: + A detached, fully materialized planner policy mapping. + """ + base = default_runtime_policy(robot_profile).planner + if planner_policy is None: + return deepcopy(base) + if not isinstance(planner_policy, Mapping): + raise TypeError("planner policy must be a mapping.") + override = deepcopy(dict(planner_policy)) + mode = override.pop("mode", None) + if mode is not None: + mode = _validate_planner_mode(mode) + conflicts = set(override).intersection(_PLANNER_MODE_FIELDS) + if conflicts: + raise ValueError( + "planner.mode cannot be combined with mode-owned fields: " + f"{sorted(conflicts)}." + ) + override = _deep_merge(override, _PLANNER_MODE_PATCHES[mode]) + unknown = set(override).difference(_PLANNER_KEYS) + if unknown: + raise ValueError(f"planner policy contains unknown fields: {sorted(unknown)}.") + resolved = _deep_merge(base, override) + backend = resolved.get("backend") + if backend != "curobo" and "curobo" in override: + raise ValueError("planner.curobo options require the cuRobo backend.") + _validate_finite_numbers(resolved, "planner") + _validate_planner(resolved) + return resolved + + +def _planner_policy_with_mode( + planner_policy: Mapping[str, Any] | None, + planner_mode: str, +) -> dict[str, Any]: + """Apply an explicit CLI mode after YAML planner configuration.""" + mode = _validate_planner_mode(planner_mode) + if planner_policy is not None and not isinstance(planner_policy, Mapping): + raise TypeError("planner policy must be a mapping.") + result = deepcopy(dict(planner_policy or {})) + result.pop("mode", None) + for field_name in _PLANNER_MODE_FIELDS: + result.pop(field_name, None) + if mode != "curobo": + result.pop("curobo", None) + result["mode"] = mode + _resolve_planner_policy(result) + return result + + +def _validate_planner_mode(value: Any) -> str: + if not isinstance(value, str) or value not in _PLANNER_MODES: + raise ValueError( + f"planner mode must be one of {list(_PLANNER_MODES)}, got {value!r}." + ) + return value + + def _load_defaults() -> dict[str, Any]: document = load_config(_DEFAULTS_PATH) if not isinstance(document, dict) or set(document) != { @@ -515,10 +615,8 @@ def _validate_planner(value: Mapping[str, Any]) -> None: raise ValueError(f"planner.{name} must be 'motion_gen' or 'ik_interp'.") if value.get("fallback_strategy") != "ik_interp": raise ValueError("planner.fallback_strategy must be 'ik_interp'.") - if value.get("coordinated_strategy") == "motion_gen" and backend == "curobo": - raise ValueError( - "planner.coordinated_strategy must be 'ik_interp' with cuRobo." - ) + if value.get("coordinated_strategy") != "ik_interp": + raise ValueError("planner.coordinated_strategy must be 'ik_interp'.") for name in ("allow_fallback", "dynamic_collision"): if not isinstance(value.get(name), bool): raise ValueError(f"planner.{name} must be a boolean.") diff --git a/embodichain/gen_sim/action_engine/generation/config_builder.py b/embodichain/gen_sim/action_engine/generation/config_builder.py index 5ca8501b9..965151fb2 100644 --- a/embodichain/gen_sim/action_engine/generation/config_builder.py +++ b/embodichain/gen_sim/action_engine/generation/config_builder.py @@ -33,6 +33,9 @@ generation_defaults, runtime_policy_hash, ) +from embodichain.gen_sim.action_engine.config.runtime_policy import ( + _resolve_planner_policy, +) from embodichain.gen_sim.action_engine.gripper_profiles import ( GripperProfile, get_gripper_profile, @@ -112,11 +115,22 @@ def build_agent_config( seed_task_graph_path: str | Path | None = EXECUTION_PROGRAM_FILENAME, vlm_model: str | None = None, vlm_camera_uids: Sequence[str] | None = None, + planner_policy: Mapping[str, Any] | None = None, ) -> dict[str, Any]: """Build the small manifest consumed by ``run_agent``.""" profile = canonical_robot_profile(robot_profile) selected_gripper = canonical_gripper_model(gripper_model) runtime_policy = default_runtime_policy(profile) + explicit_dynamic_collision = ( + planner_policy is not None and "dynamic_collision" in planner_policy + ) + if planner_policy is not None: + policy = runtime_policy.as_mapping() + policy["planner"] = _resolve_planner_policy( + planner_policy, + robot_profile=profile, + ) + runtime_policy = RuntimePolicyCfg.from_mapping(policy) if ( static_obstacle_uids is not None or dynamic_obstacle_uids is not None @@ -130,7 +144,10 @@ def build_agent_config( planner["dynamic_obstacle_uids"] = [ str(uid) for uid in dynamic_obstacle_uids ] - planner["dynamic_collision"] = bool(dynamic_obstacle_uids) + if not explicit_dynamic_collision: + planner["dynamic_collision"] = bool(dynamic_obstacle_uids) and ( + planner["backend"] == "curobo" + ) if table_top_z is not None: tabletop = float(table_top_z) if not math.isfinite(tabletop): diff --git a/embodichain/gen_sim/action_engine/generation/generator.py b/embodichain/gen_sim/action_engine/generation/generator.py index b9f07fda7..dd1d4f82c 100644 --- a/embodichain/gen_sim/action_engine/generation/generator.py +++ b/embodichain/gen_sim/action_engine/generation/generator.py @@ -76,6 +76,7 @@ def generate_action_engine_config( randomize_table_material: bool = False, planning_mode: str = "offline", vlm_model: str | None = None, + planner_policy: Mapping[str, Any] | None = None, ) -> GeneratedConfigPaths: """Generate the complete Action Engine input bundle. @@ -250,6 +251,7 @@ def generate_action_engine_config( seed_task_graph_path=graph_relative_path, vlm_model=vlm_model, vlm_camera_uids=vlm_camera_uids, + planner_policy=planner_policy, ) gym_config = build_fast_gym_config( scene, diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index a27e6cd5f..6504a5043 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -521,6 +521,7 @@ def plan( context=context, state=state, primary_success=primary_success, + primary_diagnostics=plan.diagnostics, fallback_allowed=fallback_allowed, fallback_strategy=( str(fallback_strategy) @@ -963,6 +964,7 @@ def _planner_trace( context: PlanningContext, state: ExecutionState, primary_success: torch.Tensor, + primary_diagnostics: Any, fallback_allowed: bool, fallback_strategy: str | None, fallback_attempted: torch.Tensor, @@ -981,25 +983,81 @@ def _planner_trace( dtype=torch.int64, device=self.device, ) + requested_backend = str(self.planner_policy["backend"]) + primary_strategy = invocation.motion_policy.strategy + diagnostic_backend = str(primary_diagnostics.backend) + if ( + primary_strategy == "motion_gen" + and diagnostic_backend in {"curobo", "toppra"} + and diagnostic_backend != requested_backend + ): + raise RuntimeError( + "Planner backend mismatch: runtime policy requested " + f"{requested_backend!r}, but the action plan reported " + f"{diagnostic_backend!r}." + ) + primary_effective_backend = ( + diagnostic_backend if primary_strategy == "motion_gen" else "ik_interp" + ) + if bool(fallback_used.all()) and bool(fallback_used.any()): + effective_backend = "ik_interp" + effective_strategy = "ik_interp" + elif bool(fallback_used.any()): + effective_backend = "mixed" + effective_strategy = "mixed" + else: + effective_backend = primary_effective_backend + effective_strategy = primary_strategy + planner_failure_reason = None + if not bool(primary_success.all()): + if primary_diagnostics.messages: + planner_failure_reason = "; ".join(primary_diagnostics.messages) + else: + metadata = primary_diagnostics.metadata + for key in ("failure_reason", "reason", "error"): + if metadata.get(key) is not None: + planner_failure_reason = str(metadata[key]) + break + if planner_failure_reason is None: + planner_failure_reason = "planner_reported_failure" + collision_planning_capable = ( + effective_backend == "curobo" and effective_strategy == "motion_gen" + ) + search_budget: dict[str, Any] = { + "requested_backend": requested_backend, + "fallback_enabled": bool(fallback_allowed), + } + if requested_backend == "curobo": + search_budget["primary_max_attempts"] = int( + self.planner_policy.get("curobo", {}).get("max_attempts", 1) + ) trace = { "action_class": grounded.action_class, "arm": grounded.arm, "gripper_model": self.gripper_profile.model.value, - "planner": str(self.planner_policy["backend"]), - "primary_strategy": invocation.motion_policy.strategy, + "planner": requested_backend, + "requested_backend": requested_backend, + "effective_backend": effective_backend, + "effective_strategy": effective_strategy, + "primary_effective_backend": primary_effective_backend, + "primary_strategy": primary_strategy, "dynamic_collision_mode": invocation.motion_policy.dynamic_collision_mode.value, + "collision_planning_capable": collision_planning_capable, + "collision_check_scope": ( + "static_and_dynamic" + if collision_planning_capable + and invocation.motion_policy.dynamic_collision_mode.value != "off" + else ("static" if collision_planning_capable else "not_supported") + ), "primary_success": primary_success.detach().clone(), + "planner_failure_reason": planner_failure_reason, "fallback_allowed": fallback_allowed, "fallback_strategy": fallback_strategy, "fallback_attempted": fallback_attempted.detach().clone(), "fallback_success": fallback_success.detach().clone(), "fallback_used": fallback_used.detach().clone(), - "search_budget": { - "primary_max_attempts": int( - self.planner_policy.get("curobo", {}).get("max_attempts", 1) - ), - "fallback_enabled": bool(fallback_allowed), - }, + "fallback_occurred": fallback_attempted.detach().clone(), + "search_budget": search_budget, "collision_world_revision": revisions, "collision_obstacle_positions": obstacle_positions, "collision_exclusions": { diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index 6c263045d..c39dae017 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -872,6 +872,10 @@ def _edge_exception_result( trace = { "action_class": str(action.get("atomic_action_class")), "arm": arm, + "requested_backend": self._planner_search_budget()["requested_backend"], + "effective_backend": "none", + "effective_strategy": "planner_exception", + "planner_failure_reason": f"{type(exc).__name__}: {exc}", "primary_strategy": "planner_exception", "primary_success": torch.zeros_like(inherited_failed), "fallback_attempted": torch.zeros_like(inherited_failed), @@ -2503,14 +2507,20 @@ def _candidate_exception_blockers( ] def _planner_search_budget(self) -> dict[str, Any]: - """Return the configured finite search budget used by motion planning.""" + """Return backend-neutral configured motion-planning search limits.""" runtime_policy = getattr(self, "runtime_policy", None) planner = getattr(runtime_policy, "planner", {}) - curobo = planner.get("curobo", {}) if isinstance(planner, Mapping) else {} - return { - "primary_max_attempts": int(curobo.get("max_attempts", 1)), + planner = planner if isinstance(planner, Mapping) else {} + backend = str(planner.get("backend", "unknown")) + budget = { + "requested_backend": backend, "fallback_enabled": bool(planner.get("allow_fallback", False)), } + if backend == "curobo": + curobo = planner.get("curobo", {}) + curobo = curobo if isinstance(curobo, Mapping) else {} + budget["primary_max_attempts"] = int(curobo.get("max_attempts", 1)) + return budget def _planner_failure_details( self, @@ -2528,6 +2538,9 @@ def _planner_failure_details( ) attempts = reachability.get("attempts", ()) evidence: dict[str, Any] = { + "requested_backend": str(trace.get("requested_backend", "unknown")), + "effective_backend": str(trace.get("effective_backend", "unknown")), + "effective_strategy": str(trace.get("effective_strategy", strategy)), "primary_success": bool( self._row_trace_value(trace.get("primary_success", False), env_id) ), @@ -2538,6 +2551,8 @@ def _planner_failure_details( self._row_trace_value(trace.get("fallback_success", False), env_id) ), } + if trace.get("planner_failure_reason") is not None: + evidence["planner_failure_reason"] = str(trace["planner_failure_reason"]) if trace.get("exception") is not None: evidence["exception"] = str(trace["exception"]) if isinstance(attempts, Sequence) and not isinstance( diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py index 71c7ff368..fd214d4e0 100644 --- a/embodichain/gen_sim/task_engine/cli.py +++ b/embodichain/gen_sim/task_engine/cli.py @@ -25,6 +25,11 @@ import sys from typing import Any, Final, Sequence +from embodichain.gen_sim.action_engine.config.runtime_policy import ( + _PLANNER_MODES, + _planner_policy_with_mode, +) + from .config import load_task_engine_config from .orchestration.scene_adapter import SceneAdapter from .run_directory import reserve_run_directory @@ -107,6 +112,12 @@ def _add_workflow_arguments(parser: argparse.ArgumentParser) -> None: default=None, help="Override the Task Engine planning gripper profile.", ) + parser.add_argument( + "--planner-mode", + choices=_PLANNER_MODES, + default=None, + help="Override planning.planner.mode from the Task Engine YAML.", + ) _add_failure_policy_argument(parser) @@ -162,6 +173,14 @@ def _run_workflow( workflow_cfg, planning_cfg, execution_cfg = load_task_engine_config(args.config) if args.gripper_model is not None: planning_cfg = replace(planning_cfg, gripper_model=args.gripper_model) + if args.planner_mode is not None: + planning_cfg = replace( + planning_cfg, + planner=_planner_policy_with_mode( + planning_cfg.planner, + args.planner_mode, + ), + ) with reserve_run_directory(args.output_root) as allocation: if scene is not None: validate_scene_output_separation(scene, allocation.path) diff --git a/embodichain/gen_sim/task_engine/config.py b/embodichain/gen_sim/task_engine/config.py index b2caa8ca0..2b5b94f70 100644 --- a/embodichain/gen_sim/task_engine/config.py +++ b/embodichain/gen_sim/task_engine/config.py @@ -19,6 +19,7 @@ from __future__ import annotations from collections.abc import Mapping +from copy import deepcopy from importlib.resources import files from pathlib import Path from typing import Any, Final @@ -110,6 +111,7 @@ class TaskEnginePlanningCfg: gripper_model: str = "pgi" max_episodes: int = 1 max_episode_steps: int = 6000 + planner: dict[str, Any] = {} def __post_init__(self) -> None: for field_name in ( @@ -127,6 +129,14 @@ def __post_init__(self) -> None: f"Unsupported gripper model {self.gripper_model!r}; expected one " "of: pgi, robotiq." ) + if not isinstance(self.planner, Mapping): + raise TypeError("planner must be a mapping.") + from embodichain.gen_sim.action_engine.config.runtime_policy import ( + _resolve_planner_policy, + ) + + _resolve_planner_policy(self.planner) + self.planner = deepcopy(dict(self.planner)) def load_task_engine_config( @@ -170,14 +180,20 @@ def load_task_engine_config( "max_action_attempts", }: raise ValueError("Task Engine workflow configuration fields are invalid.") - if set(planning) != { + required_planning = { "candidate_count", "planning_mode", "gripper_model", "max_episodes", "max_episode_steps", + } + if set(planning) not in { + frozenset(required_planning), + frozenset((*required_planning, "planner")), }: raise ValueError("Task Engine planning configuration fields are invalid.") + if "planner" in planning: + planning["planner"] = _mapping(planning["planner"], "planning.planner") if set(execution) != { "num_envs", "success_policy", diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml index c7b02ea17..6e03b85e5 100644 --- a/embodichain/gen_sim/task_engine/defaults.yaml +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -27,6 +27,8 @@ planning: gripper_model: robotiq max_episodes: 1 max_episode_steps: 6000 + planner: + mode: ik_interp execution: num_envs: 1 diff --git a/embodichain/gen_sim/task_engine/orchestration/coordinator.py b/embodichain/gen_sim/task_engine/orchestration/coordinator.py index 6196c1d60..a7ef2de55 100644 --- a/embodichain/gen_sim/task_engine/orchestration/coordinator.py +++ b/embodichain/gen_sim/task_engine/orchestration/coordinator.py @@ -185,6 +185,7 @@ def prepare( vlm_model: str | None = None, max_episodes: int | None = None, max_episode_steps: int | None = None, + planner_policy: Mapping[str, Any] | None = None, randomize_scene: bool = False, randomize_table_material: bool = False, candidate_set: TaskCandidateSet | Mapping[str, Any] | None = None, @@ -371,6 +372,7 @@ def prepare( "randomize_table_material": randomize_table_material, "planning_mode": planning_mode, "vlm_model": vlm_model, + "planner_policy": planner_policy, } if max_episodes is not None: generator_kwargs["max_episodes"] = max_episodes diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py index feb0af5e3..b3a4bca09 100644 --- a/embodichain/gen_sim/task_engine/workflow.py +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -655,6 +655,7 @@ def run( vlm_model=vlm_model, max_episodes=planning_cfg.max_episodes, max_episode_steps=planning_cfg.max_episode_steps, + planner_policy=planning_cfg.planner, candidate_set=candidate_set, force_most_likely=True, final_inspection=final_inspection, @@ -1039,6 +1040,7 @@ def _publish( "gripper_model": planning_cfg.gripper_model, "max_episodes": planning_cfg.max_episodes, "max_episode_steps": planning_cfg.max_episode_steps, + "planner": deepcopy(planning_cfg.planner), }, "execution": { "num_envs": execution_cfg.num_envs, diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py index fc474ff25..56fc67564 100644 --- a/tests/gen_sim/action_engine/generation/test_generation.py +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -35,6 +35,7 @@ generate_action_agent_config as cli_module, ) from embodichain.gen_sim.action_engine.cli.generate_action_agent_config import ( + _load_planner_config, build_parser, ) from embodichain.gen_sim.action_engine.generation.artifacts import ( @@ -492,6 +493,104 @@ def test_fast_gym_config_preserves_unicode_instruction_and_uses_task_name_label( assert params["extra"]["task_description"] == task_name +@pytest.mark.parametrize( + ("name", "planner_policy", "expected"), + [ + ( + "curobo", + {"mode": "curobo"}, + ("curobo", "motion_gen", "ik_interp", True), + ), + ( + "toppra", + {"mode": "toppra"}, + ("toppra", "motion_gen", "ik_interp", False), + ), + ( + "ik_interp", + {"mode": "ik_interp"}, + ("toppra", "ik_interp", "ik_interp", False), + ), + ], +) +def test_agent_config_materializes_yaml_planner_policy_and_hash( + tmp_path: Path, + name: str, + planner_policy: dict[str, object], + expected: tuple[str, str, str, bool], +) -> None: + config = build_agent_config( + task_name=name, + robot_profile="ur10", + execution_program_hash="d" * 64, + source_config_path=tmp_path / "gym_config.json", + uid_map={"table": "table", "cube": "cube"}, + static_obstacle_uids=["table"], + dynamic_obstacle_uids=["cube"], + planner_policy=planner_policy, + ) + + planner = config["runtime_policy"]["planner"] + assert ( + planner["backend"], + planner["single_arm_strategy"], + planner["coordinated_strategy"], + planner["dynamic_collision"], + ) == expected + assert planner["static_obstacle_uids"] == ["table"] + assert planner["dynamic_obstacle_uids"] == ["cube"] + assert len(config["runtime_policy_hash"]) == 64 + + from embodichain.gen_sim.action_engine.config import resolve_agent_runtime_policy + + resolved = resolve_agent_runtime_policy(config) + assert resolved.planner == planner + + +def test_agent_config_rejects_toppra_dynamic_collision_before_writing( + tmp_path: Path, +) -> None: + with pytest.raises(ValueError, match="dynamic_collision.*cuRobo"): + build_agent_config( + task_name="invalid_toppra", + robot_profile="ur10", + execution_program_hash="d" * 64, + source_config_path=tmp_path / "gym_config.json", + uid_map={"cube": "cube"}, + dynamic_obstacle_uids=["cube"], + planner_policy={"backend": "toppra", "dynamic_collision": True}, + ) + + +def test_planner_yaml_loader_accepts_wrapped_policy_and_rejects_backend_leaks( + tmp_path: Path, +) -> None: + config_path = tmp_path / "toppra.yaml" + config_path.write_text( + """\ +planner: + mode: toppra +""", + encoding="utf-8", + ) + + assert _load_planner_config(str(config_path)) == { + "mode": "toppra", + } + + config_path.write_text( + """\ +planner: + mode: toppra + curobo: + max_attempts: 2 +""", + encoding="utf-8", + ) + with pytest.raises(ValueError, match="curobo.*cuRobo backend"): + _load_planner_config(str(config_path)) + + def test_ab_config_uses_offline_branch_and_four_vlm_cameras( gym_export: Path, tmp_path: Path, @@ -1552,6 +1651,7 @@ def test_generation_cli_defaults_to_mature_robot_without_scene_randomization() - assert args.robot_profile == "ur10" assert args.randomize_scene is False assert args.planning_mode == "offline" + assert args.planner_mode is None assert not hasattr(args, "instruction_parser") assert not hasattr(args, "task_agent") @@ -1634,6 +1734,60 @@ def test_generation_cli_reports_seed_png_path( ) +def test_generation_cli_explicit_mode_overrides_planner_yaml( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + paths = artifact_paths(tmp_path / "output") + planner_path = tmp_path / "planner.yaml" + planner_path.write_text( + """\ +planner: + backend: curobo + single_arm_strategy: motion_gen + coordinated_strategy: ik_interp + dynamic_collision: true + allow_fallback: false + curobo: + max_attempts: 2 +""", + encoding="utf-8", + ) + captured: dict[str, object] = {} + + def generate(*_args, **kwargs): + captured.update(kwargs) + return paths + + monkeypatch.setattr(cli_module, "generate_action_engine_config", generate) + monkeypatch.setattr( + sys, + "argv", + [ + "generate_action_agent_config", + "--gym-project", + "gym_export", + "--output-dir", + str(tmp_path / "output"), + "--task-name", + "task", + "--task-description", + "test-instruction", + "--planner-config", + str(planner_path), + "--planner-mode", + "ik_interp", + ], + ) + + cli_module.cli() + + assert captured["planner_policy"] == { + "allow_fallback": False, + "mode": "ik_interp", + } + + @pytest.mark.parametrize( "removed_args", [ diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 902ae54f8..7a49d7dc4 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -60,7 +60,7 @@ TimedTrajectory, TrackingPolicy, ) -from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.lab.sim.planners import CuroboPlannerCfg, ToppraPlannerCfg from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator @@ -539,6 +539,96 @@ def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None assert hand.motion_policy.strategy == "ik_interp" +@pytest.mark.parametrize( + ("planner_policy", "single_strategy"), + [ + ( + { + "backend": "toppra", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "dynamic_collision": False, + }, + "motion_gen", + ), + ( + { + "backend": "toppra", + "single_arm_strategy": "ik_interp", + "coordinated_strategy": "ik_interp", + "dynamic_collision": False, + }, + "ik_interp", + ), + ], +) +def test_toppra_and_ik_interp_policies_reach_runtime_factory_and_strategy( + planner_policy: dict[str, object], + single_strategy: str, +) -> None: + adapter = AtomicActionAdapter(_planner_env(), planner_policy=planner_policy) + adapter._atomic_engine = _FakeEngine() + goal = JointPositionGoal(target=torch.zeros(2, 2)) + + invocation = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "arm", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + + assert isinstance(adapter._motion_generator_cfg().planner_cfg, ToppraPlannerCfg) + assert invocation.motion_policy.strategy == single_strategy + assert invocation.motion_policy.dynamic_collision_mode.value == "off" + + +def test_runtime_rejects_requested_and_effective_backend_mismatch( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter( + env, + planner_policy={ + "backend": "toppra", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "dynamic_collision": False, + }, + ) + trajectory = TimedTrajectory.from_uniform_step( + torch.zeros(2, 2, 8), + env_ids=torch.arange(2), + step_dt=0.01, + ) + wrong_backend_plan = ActionPlan( + skill_id="move_joints", + plan_success=torch.ones(2, dtype=torch.bool), + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="curobo"), + expected_effects=StateDelta(), + ) + monkeypatch.setattr( + adapter, + "_engine", + lambda: _FakeEngine(lambda *_args: wrong_backend_plan), + ) + + with pytest.raises(RuntimeError, match="requested 'toppra'.*reported 'curobo'"): + adapter.plan( + GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ), + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + def test_coordinated_pickment_uses_engine_scoped_grasp_generator() -> None: adapter = AtomicActionAdapter(_planner_env()) adapter._atomic_engine = _FakeEngine() @@ -1173,6 +1263,8 @@ def action_plan( success: torch.Tensor, terminal: float, held: HeldObjectState, + *, + messages: tuple[str, ...] = (), ) -> ActionPlan: positions = torch.full((2, 2, 8), terminal) trajectory = TimedTrajectory.from_uniform_step( @@ -1190,7 +1282,8 @@ def action_plan( planned_scene_version=0, planned_collision_world_revision=(0, 0), diagnostics=PlannerDiagnostics( - backend="fake", + backend="curobo", + messages=messages, metadata={"marker": terminal}, ), expected_effects=StateDelta( @@ -1200,7 +1293,12 @@ def action_plan( plans = iter( ( - action_plan(torch.tensor([True, False]), 1.0, held_at(1.0)), + action_plan( + torch.tensor([True, False]), + 1.0, + held_at(1.0), + messages=("IK unreachable",), + ), action_plan(torch.tensor([True, True]), 2.0, held_at(2.0)), ) ) @@ -1248,6 +1346,13 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: assert outcome.planner_trace["primary_action_diagnostics"]["marker"] == 1.0 assert outcome.planner_trace["fallback_action_diagnostics"]["marker"] == 2.0 assert outcome.planner_trace["gripper_model"] == "pgi" + assert outcome.planner_trace["requested_backend"] == "curobo" + assert outcome.planner_trace["effective_backend"] == "mixed" + assert outcome.planner_trace["effective_strategy"] == "mixed" + assert outcome.planner_trace["primary_effective_backend"] == "curobo" + assert bool(outcome.planner_trace["fallback_occurred"].any()) + assert outcome.planner_trace["planner_failure_reason"] == "IK unreachable" + assert outcome.planner_trace["collision_planning_capable"] is False def test_collision_required_cleanup_does_not_use_unsafe_fallback( diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py index 6c2b2bca5..e4ca7d5d5 100644 --- a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -843,6 +843,58 @@ def run(self, request, **kwargs): assert Path(payload["output_dir"]).parent == tmp_path / "history" +def test_unified_cli_explicit_planner_mode_overrides_packaged_yaml( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + captured: dict[str, object] = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **kwargs): + captured.update(kwargs) + output = Path(request["output_dir"]) + return SimpleNamespace( + status="prepared", + succeeded=False, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + + assert ( + cli.main( + [ + "prepare", + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + str(tmp_path / "input.png"), + "--output-root", + str(tmp_path / "history"), + "--planner-mode", + "ik_interp", + ] + ) + == 0 + ) + + planning_cfg = captured["planning_cfg"] + assert planning_cfg.planner == {"mode": "ik_interp"} + assert json.loads(capsys.readouterr().out)["status"] == "prepared" + + def test_unified_cli_reuses_history_root_without_modifying_prior_scene( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -978,6 +1030,7 @@ def test_public_cli_exposes_prepare_run_and_run_all_modes() -> None: assert arguments.command == "prepare" assert arguments.dataset_saving is True assert arguments.failure_policy == "stop" + assert arguments.planner_mode is None def test_prepare_cli_stops_before_simulator_execution( diff --git a/tests/gen_sim/task_engine/test_workflow.py b/tests/gen_sim/task_engine/test_workflow.py index 2a9df8092..0ea7ba9db 100644 --- a/tests/gen_sim/task_engine/test_workflow.py +++ b/tests/gen_sim/task_engine/test_workflow.py @@ -279,6 +279,12 @@ def test_packaged_workflow_configuration_uses_recovery_defaults() -> None: assert execution.required_successes == 1 +def test_packaged_planner_mode_defaults_to_curobo() -> None: + _, planning, _ = load_task_engine_config() + + assert planning.planner == {"mode": "curobo"} + + def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: config = tmp_path / "task_engine.yaml" config.write_text( @@ -294,6 +300,8 @@ def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: gripper_model: robotiq max_episodes: 2 max_episode_steps: 5000 + planner: + mode: toppra execution: num_envs: 6 success_policy: at_least @@ -311,6 +319,7 @@ def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: assert planning.gripper_model == "robotiq" assert planning.max_episodes == 2 assert planning.max_episode_steps == 5000 + assert planning.planner == {"mode": "toppra"} assert execution.num_envs == 6 assert execution.required_successes == 2 @@ -340,3 +349,5 @@ def test_planning_configuration_rejects_invalid_values() -> None: TaskEnginePlanningCfg(planning_mode="unsupported") with pytest.raises(ValueError, match="pgi.*robotiq"): TaskEnginePlanningCfg(gripper_model="unsupported") + with pytest.raises(ValueError, match="mode cannot be combined"): + TaskEnginePlanningCfg(planner={"mode": "toppra", "dynamic_collision": True}) From 8caf7f745b87959bdb6d98e7621e3458ecba5cda Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Thu, 27 Aug 2026 11:22:30 +0800 Subject: [PATCH 74/85] feat(action-engine): extend E4 handover with hold and placement outcomes --- .../action_engine/capabilities/atomic.py | 4 + .../action_engine/config/defaults.yaml | 2 +- .../action_engine/domain/task_contracts.py | 18 +- .../gen_sim/action_engine/domain/v2.py | 72 +++--- .../gen_sim/action_engine/runtime/actions.py | 2 +- .../gen_sim/action_engine/runtime/executor.py | 156 +++++-------- .../action_engine/runtime/grounding.py | 151 +++++-------- .../action_engine/runtime/predicates.py | 2 +- .../gen_sim/action_engine/runtime/recovery.py | 50 ++--- .../gen_sim/action_engine/tasks/assembly.py | 8 +- .../action_engine/tasks/interpretation.py | 6 + .../gen_sim/action_engine/tasks/recipes.py | 212 ++++++++++++++++-- embodichain/gen_sim/task_engine/agent.py | 6 +- .../gen_sim/task_engine/interpretation.py | 149 ++++-------- embodichain/gen_sim/task_engine/ontology.py | 18 +- .../domain/test_task_contracts.py | 12 + .../action_engine/planning/test_linker.py | 4 +- .../action_engine/planning/test_online_v2.py | 2 +- .../action_engine/runtime/test_recovery_v2.py | 64 +++--- .../runtime/test_runtime_contracts.py | 25 +-- tests/gen_sim/action_engine/task_fixtures.py | 1 + .../action_engine/tasks/test_e3_pour.py | 2 +- .../action_engine/tasks/test_factory.py | 163 +++++++++++++- .../tasks/test_interpretation.py | 78 +++++-- .../tasks/test_payload_contracts.py | 45 ++++ .../action_engine/test_architecture.py | 42 ++++ 26 files changed, 805 insertions(+), 489 deletions(-) diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 86404ea9e..92cbf6987 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -628,6 +628,9 @@ def _resolve_default_contract( ) if capability.name == "MoveHeldObject": required_arm = _required_arm(arm, capability.name) + terminal_hold = binding.get("terminal_hold", False) + if not isinstance(terminal_hold, bool): + raise TypeError("MoveHeldObject terminal_hold must be a boolean.") return ResolvedActionContract( requires=( StateAtom("object_held", object_uid=object_uid, arm=required_arm), @@ -637,6 +640,7 @@ def _resolve_default_contract( ResourceClaim(f"object:{object_uid}", lifetime="until_release"), ) + payload_claims, + completion="terminal_barrier" if terminal_hold else "ordinary", ) if capability.name == "Place": required_arm = _required_arm(arm, capability.name) diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index 4643b4ee3..41d1eb21e 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -171,7 +171,7 @@ runtime: sample_interval: 180 pre_grasp_distance: 0.15 lift_height: 0.16 - lower_distance: 0.03 + lower_distance: 0.16 hand_interp_steps: 12 PickUp: pre_grasp_distance: 0.15 diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py index 5ee1228d3..3d1b268aa 100644 --- a/embodichain/gen_sim/action_engine/domain/task_contracts.py +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -72,7 +72,7 @@ def normalize_placement_relation(value: Any) -> str: "E1": ("PickUp", "MoveHeldObject", "Place"), "E2": ("AxisAlign",), "E3": ("PickUp", "MoveHeldObject", "Pour", "Place"), - "E4": ("PickUp", "MoveHeldObject", "HandOver"), + "E4": ("PickUp", "MoveHeldObject", "HandOver", "Place"), "E5": ("CoordinatedPickment",), "E6": ("PullArticulatedPart",), "E7": ("PushArticulatedPart",), @@ -81,6 +81,20 @@ def normalize_placement_relation(value: Any) -> str: } ) +_SIGNATURE_ACTIONS: Mapping[str, frozenset[str]] = MappingProxyType( + { + "E1": frozenset(), + "E2": frozenset(), + "E3": frozenset({"Pour"}), + "E4": frozenset({"HandOver"}), + "E5": frozenset(), + "E6": frozenset({"PullArticulatedPart"}), + "E7": frozenset({"PushArticulatedPart"}), + "E8": frozenset({"TurnKnob"}), + "E9": frozenset({"Press"}), + } +) + @dataclass(frozen=True, slots=True) class TaskContract: @@ -101,6 +115,7 @@ class TaskContract: direct_payload_relations: frozenset[str] accepts_incoming_hold: bool terminal_success_types: tuple[tuple[str, str], ...] + signature_actions: frozenset[str] = frozenset() def _action_contract(value: SemanticTaskContract) -> TaskContract: @@ -108,6 +123,7 @@ def _action_contract(value: SemanticTaskContract) -> TaskContract: task_type=value.task_type, semantics=value.semantics, core_actions=_CORE_ACTIONS[value.task_type], + signature_actions=_SIGNATURE_ACTIONS[value.task_type], applicable_intent_fields=value.applicable_intent_fields, source_structure=value.source_structure, required_affordances=value.required_affordances, diff --git a/embodichain/gen_sim/action_engine/domain/v2.py b/embodichain/gen_sim/action_engine/domain/v2.py index d8f581a59..f14f278a2 100644 --- a/embodichain/gen_sim/action_engine/domain/v2.py +++ b/embodichain/gen_sim/action_engine/domain/v2.py @@ -623,21 +623,32 @@ def _validate_task_group_semantics( nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]], ) -> None: + from .task_contracts import task_contract + node_by_id = {str(node["id"]): node for node in nodes} - required_actions = { - "E1": set(), - "E2": set(), - "E3": {"Pour"}, - "E4": {"HandOver"}, - "E5": set(), - "E6": {"PullArticulatedPart"}, - "E7": {"PushArticulatedPart"}, - "E8": {"TurnKnob"}, - "E9": {"Press"}, - } for group in groups: task_type = str(group["task_type"]) - if task_type == "E3": + contract = task_contract(task_type) + group_nodes = [node_by_id[node_id] for node_id in group["node_ids"]] + actions = {str(node["atomic_action"]) for node in group_nodes} + if group.get("role") == "recovery": + goal = group.get("goal", {}) + if goal.get("recovery_capability") != "object_upright": + raise ValueError( + f"Recovery TaskGroup {group['id']!r} requires a registered " + "recovery_capability." + ) + required = {"PickUp", "MoveHeldObject"} + if goal.get("terminal_behavior") == "place": + required.add("Place") + missing = required - actions + if missing: + raise ValueError( + f"Recovery TaskGroup {group['id']!r} is missing capability " + f"actions: {sorted(missing)}." + ) + continue + if contract.success_type == "poured": goal = group.get("goal", {}) unsupported = sorted( {"pour_mode", "pouring_arm", "holding_arm"} & set(goal) @@ -645,13 +656,15 @@ def _validate_task_group_semantics( if unsupported: raise ValueError( f"SeedGraph TaskGroup {group['id']!r} uses unsupported " - f"dual-arm E3 fields {unsupported}; regenerate it as a " + f"dual-arm pour fields {unsupported}; regenerate it as a " "single-arm pour over a fixed target container." ) - group_nodes = [node_by_id[node_id] for node_id in group["node_ids"]] - actions = {str(node["atomic_action"]) for node in group_nodes} - missing = required_actions[task_type] - actions - if task_type == "E1": + missing = set(contract.signature_actions) - actions + if ( + contract.resource_mode == "single_arm" + and contract.success_type == "semantic_goal" + and not contract.terminal_success_types + ): if not actions.intersection({"MoveHeldObject", "Place"}): missing = {"MoveHeldObject|Place"} elif "PickUp" not in actions: @@ -659,24 +672,25 @@ def _validate_task_group_semantics( precondition = first.get("precondition", {}) if precondition.get("type") != "object_held": missing = {"PickUp|object_held precondition"} - if task_type == "E2" and "AxisAlign" not in actions: + if contract.success_type == "object_upright" and "AxisAlign" not in actions: missing = {"MoveHeldObject", "Place"} - actions if "PickUp" not in actions: first = group_nodes[0] precondition = first.get("precondition", {}) if precondition.get("type") != "object_held": missing.add("PickUp|object_held precondition") - # Recovery may explicitly preserve a verified downstream hold. Ordinary - # E2 groups always complete their supported world state with Place. - if ( - task_type == "E2" - and "Place" not in actions - and group.get("goal", {}).get("terminal_behavior") == "hold" - and "MoveHeldObject" in actions - and group.get("role") == "recovery" - ): - missing.discard("Place") - if task_type == "E5" and not actions.intersection( + if contract.resource_mode == "handover": + terminal_behavior = str( + group.get("goal", {}).get("terminal_behavior", "hold") + ) + if terminal_behavior == "place" and "Place" not in actions: + missing.add("Place") + if terminal_behavior == "hold" and "Place" in actions: + raise ValueError( + f"SeedGraph TaskGroup {group['id']!r} cannot release during " + "terminal_behavior='hold'." + ) + if contract.resource_mode == "coordinated" and not actions.intersection( {"CoordinatedPickment", "CoordinatedPlacement"} ): missing = {"CoordinatedPickment|CoordinatedPlacement"} diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 6504a5043..48cfd30ba 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -626,7 +626,7 @@ def _adapt_coordinated_pickment_grasps( grounded: GroundedAction, capability: AtomicCapability, ) -> tuple[GroundedAction, ...]: - """Build deterministic geometry-ranked E5 partition candidates. + """Build deterministic geometry-ranked coordinated-grasp candidates. ``left_to_right_arm_direction`` remains the live base-to-base direction: it labels the two participant regions and is not the transport direction. diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index c39dae017..600f813c7 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -259,9 +259,7 @@ def __init__( for step_id in group.get("semantic_step_ids", ()) } arrangement_steps = [ - step - for step in program.semantic_steps - if step.operator in {"arrange_line", "place_in_line"} + step for step in program.semantic_steps if step.goal.get("layout") == "line" ] arrangement_groups: dict[str, list[SemanticStep]] = {} for step in arrangement_steps: @@ -286,8 +284,8 @@ def __init__( placement_groups: dict[str, list[SemanticStep]] = {} for step in program.semantic_steps: if ( - step.operator == "place_relative" - and step.goal.get("relation") == "inside" + step.goal.get("relation") == "inside" + and step.postcondition.get("type") == "semantic_goal" and isinstance(step.goal.get("reference_object"), str) ): placement_groups.setdefault( @@ -563,7 +561,7 @@ def run( if ( self.placement_recovery_attempts and bool(postcondition_failed.any()) - and step.operator == "place_relative" + and step.postcondition.get("type") == "semantic_goal" and normalize_placement_relation( step.goal.get("relation", "on") ) @@ -1023,7 +1021,7 @@ def _recover_object_fallen( fallen_transition: torch.Tensor, recorder: RuntimeRecorder, ) -> _EdgeResult: - """Run the bounded E2 repair and replay only the failed vector rows.""" + """Run bounded upright recovery and replay only the failed vector rows.""" if self.runtime_graph is None or len(edge.actions) != 1: return result node_id = edge.actions[0].get("seed_node_id") @@ -1770,7 +1768,15 @@ def _parallel_pickup_candidate(self, edge: ExecutionEdge) -> bool: capability.state_effect == "hold" and capability.resource_mode == "single_arm_object" and step.actor.get("mode") in {"auto", "required"} - and step.operator != "orient_object" + and not self._is_in_place_orientation(step) + ) + + @staticmethod + def _is_in_place_orientation(step: SemanticStep) -> bool: + return ( + step.goal.get("position_anchor") in {"initial_xy", "live_xy"} + and isinstance(step.goal.get("support_object"), str) + and "upright_local_axis" in step.goal ) def _preferred_in_place_arm( @@ -1779,7 +1785,7 @@ def _preferred_in_place_arm( env_id: int, ) -> str | None: """Map a clearly sided in-place object to the robot-view arm slot.""" - if step.operator != "orient_object": + if not self._is_in_place_orientation(step): return None initial = getattr(self.env, "agent_initial_object_poses", {}).get( step.object_uid @@ -1878,12 +1884,20 @@ def _ensure_assignment( first_capability = self.adapter.capabilities.get( str(first_action.get("atomic_action_class")) ) - if step.operator == "handover" and first_capability.state_effect != "hold": + has_transfer = any( + self.adapter.capabilities.get( + str(action.get("atomic_action_class")) + ).state_effect + == "transfer_hold" + for edge_id in step.edge_ids + for action in self.edges[edge_id].actions + ) + if has_transfer and first_capability.state_effect != "hold": # A coordinated handover has an internal, multi-arm planner. # Do not let a speculative single-arm suffix plan veto the # real execution (or create a misleading downstream pickup # error) once a predecessor already established the transfer - # hold. A standalone E4 starts with PickUp and still needs its + # hold. A standalone transfer starts with PickUp and still needs its # cached candidate plan for that first action. source_state = self._state_for(step, arm) has_hold = ( @@ -2112,11 +2126,8 @@ def _candidate( capability = self.adapter.capabilities.get( str(action.get("atomic_action_class")) ) - if ( - step.operator == "handover" - and capability.resource_mode == "coordinated_object" - ): - # A standalone E4 needs a speculative PickUp/staging + if capability.state_effect == "transfer_hold": + # A standalone transfer needs a speculative PickUp/staging # prefix to choose and cache its transfer arm. The # actual HandOver is coordinated, however, and must # only be planned from the live post-staging state. @@ -2613,15 +2624,7 @@ def _with_downstream_targets( state: ExecutionState, grounded: GroundedAction, ) -> GroundedAction: - """Screen grasp poses against every later held-object target. - - A handover is split across semantic steps: its staging ``MoveHeldObject`` - edge is not part of the pickup step's local edge suffix. Include that - first exchange pose here so ``PickUp`` can reject a grasp whose - ``object_to_eef`` transform makes the later transfer arm unreachable. - This keeps the screening speculative and bounded; no simulator steps - are sent while a candidate is being built. - """ + """Screen grasp poses against later held-object targets in this task.""" targets: list[torch.Tensor] = [] start = step.edge_ids.index(pickup_edge_id) + 1 for edge_id in step.edge_ids[start:]: @@ -2629,6 +2632,13 @@ def _with_downstream_targets( if len(edge.actions) != 1: continue action = edge.actions[0] + action_actor = action.get("actor", {}) + if ( + isinstance(action_actor, Mapping) + and action_actor.get("mode") == "required" + and action_actor.get("arm") != arm + ): + continue if ( self.adapter.capabilities.get( str(action.get("atomic_action_class")) @@ -2645,7 +2655,6 @@ def _with_downstream_targets( ) if future.target_object_pose is not None: targets.append(future.target_object_pose) - targets.extend(self._handover_successor_targets(step, arm, state)) if not targets: return grounded existing = tuple(grounded.cfg.get("downstream_object_target_poses", ())) @@ -2657,78 +2666,6 @@ def _with_downstream_targets( }, ) - def _handover_successor_targets( - self, - step: SemanticStep, - arm: str, - state: ExecutionState, - ) -> list[torch.Tensor]: - """Return staging poses for handovers downstream of a pickup. - - ``SemanticStep.depends_on`` contains semantic IDs rather than edge IDs, - so walk the small dependency graph instead of assuming the handover is - an immediate child. Only a handover that transfers this object from - the selected pickup arm is relevant to the grasp screen. - """ - reachable = {step.id} - changed = True - while changed: - changed = False - for candidate in self.steps.values(): - if candidate.id in reachable: - continue - if any(dependency in reachable for dependency in candidate.depends_on): - reachable.add(candidate.id) - changed = True - - targets: list[torch.Tensor] = [] - for successor in self.steps.values(): - if ( - successor.id not in reachable - or successor.id == step.id - or successor.operator != "handover" - or successor.object_uid != step.object_uid - ): - continue - for edge_id in successor.edge_ids: - edge = self.edges[edge_id] - if len(edge.actions) != 1: - continue - action = edge.actions[0] - binding = action.get("target_binding", {}) - if not isinstance(binding, Mapping): - continue - if binding.get("kind") != "handover_staging": - continue - transfer_arm = str( - binding.get( - "transfer_arm", - successor.goal.get("transfer_arm", ""), - ) - ) - if transfer_arm != arm: - break - try: - grounded = self.grounder.ground( - action, - successor, - arm=arm, - state=state, - orientation_reference_pose=self._orientation_references.get( - successor.id, - self._orientation_references.get(step.id), - ), - ) - except (AttributeError, KeyError, ValueError): - # A malformed/incomplete successor must not make an - # otherwise valid pickup candidate disappear. The normal - # successor execution will report that grounding error. - break - if grounded.target_object_pose is not None: - targets.append(grounded.target_object_pose) - break - return targets - def _eef_target(self, outcome: ActionOutcome) -> torch.Tensor | None: state = outcome.next_state held_object = state.get_held_object( @@ -2888,6 +2825,17 @@ def _execute_edge( f"Edge {edge.id!r} must contain one action or an explicit dual pair." ) assignments = self._assignments[step.id] + action_actor = edge.actions[0].get("actor", {}) + if ( + isinstance(action_actor, Mapping) + and action_actor.get("mode") == "required" + and action_actor.get("arm") in {"left_arm", "right_arm"} + ): + required_arm = str(action_actor["arm"]) + assignments = [ + required_arm if assignment is not None else None + for assignment in assignments + ] outcomes: dict[str, ActionOutcome | None] = { "left_arm": None, "right_arm": None, @@ -4124,13 +4072,11 @@ def _verify_step( success = torch.zeros_like(failed) log_info(f"Skipped verification for {step.id}: no active environments.") return failed, success, observed - relation = ( - normalize_placement_relation(step.goal.get("relation", "on")) - if step.operator == "place_relative" - else str(step.goal.get("relation", "")) - ) - reference = self._support_reference_uid(step) postcondition_type = step.postcondition.get("type") + relation = str(step.goal.get("relation", "")) + if postcondition_type == "semantic_goal" and relation not in {"", "none"}: + relation = normalize_placement_relation(relation) + reference = self._support_reference_uid(step) if postcondition_type in {"object_held", "handover_complete"}: # A planned hover target is not evidence that the object remains # grasped. Verify live TCP/object geometry and gripper closure. @@ -4189,7 +4135,7 @@ def _verify_step( reference, active, ) - elif step.operator == "orient_object": + elif self._is_in_place_orientation(step): position_anchor = str(step.goal.get("position_anchor", "initial_xy")) anchor_pose = None if position_anchor == "initial_xy": diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 7c8892f29..37c98bc2d 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -36,7 +36,6 @@ from embodichain.gen_sim.action_engine.orientation import ( AlignAxisConstraint, MatchRotationConstraint, - OrientationConstraint, compile_orientation_constraint, ) from embodichain.lab.sim.atomic_actions import ( @@ -108,6 +107,19 @@ def _live_pose(env: Any, uid: str) -> torch.Tensor: return _batched_pose(entity.get_local_pose(to_matrix=True), env) +def _placement_relation(step: SemanticStep) -> str: + """Normalize a placement relation from the step's own goal contract.""" + relation = str(step.goal.get("relation", "none")) + if relation in {"none", "handover", "held_above_initial"}: + return relation + if ( + step.postcondition.get("type") == "semantic_goal" + or step.goal.get("terminal_behavior") == "place" + ): + return normalize_placement_relation(relation) + return relation + + def _local_vertices(entity: Any, env: Any, env_id: int = 0) -> torch.Tensor: value = entity.get_vertices(env_ids=[env_id], scale=True) if isinstance(value, (list, tuple)): @@ -695,30 +707,7 @@ def ground( raise ValueError("target_binding must be a mapping.") kind = str(binding.get("kind", "")) orientation = compile_orientation_constraint(step.goal) - is_handover_continuation = self._is_handover_continuation(step) - uses_handover_staging = ( - kind == "handover_staging" - and capability.target_materializer == "semantic_held_object" - ) - use_upright_transport_policy = ( - is_handover_continuation or uses_handover_staging - ) and self._uses_upright_transport_policy( - step, - orientation, - ) - extra_modifiers: tuple[tuple[str, str], ...] = () - if ( - is_handover_continuation - and use_upright_transport_policy - and capability.target_materializer - in { - "semantic_held_object", - "current_held_pose", - "eef_pose", - } - ): - extra_modifiers = (("orientation", "upright"),) - policy = self.policy(action, extra_modifiers=extra_modifiers) + policy = self.policy(action) if kind == "joint_state": joint_defaults = self.runtime_policy.grounding["joint_state"] source = binding.get("source") @@ -751,7 +740,7 @@ def ground( # replace it with collision-unaware joint interpolation. policy["collision_safety"] = "required" object_pose = _live_pose(self.env, step.object_uid) - if step.operator == "orient_object": + if "upright_local_axis" in step.goal: policy["upright_local_axis"] = self._upright_local_axis(step) if capability.target_materializer == "object_grasp": policy["obj_upright_direction"] = self._upright_local_direction(step) @@ -1107,11 +1096,7 @@ def ground_candidates( ), ) placement_support_uid = self._placement_support_uid(step) - placement_relation = ( - normalize_placement_relation(step.goal.get("relation", "on")) - if step.operator == "place_relative" - else str(step.goal.get("relation", "none")) - ) + placement_relation = _placement_relation(step) is_on_placement = ( binding.get("kind") == "semantic_goal" and binding.get("phase", "final") != "staging" @@ -1286,19 +1271,6 @@ def _placement_support_uid(step: SemanticStep) -> str | None: return "table" return None - def _is_handover_continuation(self, step: SemanticStep) -> bool: - if step.operator != "place_relative": - return False - predecessors = { - candidate.id: candidate for candidate in self.program.semantic_steps - } - return any( - (predecessor := predecessors.get(dependency)) is not None - and predecessor.operator == "handover" - and predecessor.object_uid == step.object_uid - for dependency in step.depends_on - ) - def _visual_target( self, binding: Mapping[str, Any], @@ -1718,6 +1690,17 @@ def _semantic_target( phase: str, orientation_reference_pose: torch.Tensor | None = None, ) -> torch.Tensor: + if phase == "handover_exit": + receive_arm = str(step.goal.get("receive_arm", "")) + if receive_arm not in {"left_arm", "right_arm"}: + raise ValueError("handover_exit requires a concrete receive_arm.") + target = self._handover_receiver_exit(object_pose, receive_arm, policy) + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + return target if step.operator in {"arrange_line", "place_in_line"}: arrangement = self.arrangements.get(step.id) if arrangement is None: @@ -1763,7 +1746,9 @@ def _semantic_target( if phase == "staging": target[:, 2, 3] += float(policy["transport_clearance"]) return target - if step.operator == "orient_object": + if step.goal.get("position_anchor") in {"initial_xy", "live_xy"} and isinstance( + step.goal.get("support_object"), str + ): initial = None if step.goal.get("position_anchor", "initial_xy") == "initial_xy": initial = getattr(self.env, "agent_initial_object_poses", {}).get( @@ -1830,11 +1815,7 @@ def _semantic_target( # Operators without a relational goal (for example press or a # direction-only coordinated transport) must preserve the live origin # instead of being silently projected onto a synthetic table support. - relation = ( - normalize_placement_relation(step.goal.get("relation", "on")) - if step.operator == "place_relative" - else str(step.goal.get("relation", "none")) - ) + relation = _placement_relation(step) distance = float(self._policy_value(policy, "relation_distance")) relation_frame = str(step.goal.get("relation_frame", "world")) forward_distance = distance @@ -1945,7 +1926,8 @@ def _semantic_target( orientation_reference_pose=orientation_reference_pose, ) if ( - step.operator == "coordinated_transport" + step.actor.get("mode") == "coordinated" + and "terminal_behavior" in step.goal and relation not in {"on", "on_top", "on_top_of", "inside"} and direction not in {"up", "down"} ): @@ -1994,29 +1976,30 @@ def _semantic_target( supported_pose = object_pose supported_pose = _batched_pose(supported_pose, self.env) target[:, 2, 3] = supported_pose[:, 2, 3] + elif isinstance(step.goal.get("support_object"), str) and relation not in { + "none", + "handover", + "above", + "held_above_initial", + }: + support = _object(self.env, str(step.goal["support_object"])) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + + float(self._policy_value(policy, "surface_clearance")) + - bottom + ) if phase == "staging": # Staging is a runtime waypoint, not a persisted coordinate. This # keeps in-place orientation robust to the object's live height. target[:, 2, 3] += float(self._policy_value(policy, "transport_clearance")) - elif self._is_handover_continuation(step) and relation not in { - "on", - "on_top", - "on_top_of", - "inside", - }: - # A handover can leave the live rigid-body center a few centimetres - # below the original table-supported height. Reusing that drifted - # height for the lateral placement target makes the can intersect - # the table during release and it may tip or slide. Preserve the - # predecessor's supported height for the final held-object pose. - supported_pose = orientation_reference_pose - if supported_pose is None: - supported_pose = object_pose - supported_pose = _batched_pose(supported_pose, self.env) - target[:, 2, 3] = torch.maximum( - target[:, 2, 3], - supported_pose[:, 2, 3], - ) return target def _pour_source_semantics( @@ -2564,34 +2547,6 @@ def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: direction[axis_index] = 1.0 return direction - def _uses_upright_transport_policy( - self, - step: SemanticStep, - constraint: OrientationConstraint, - ) -> bool: - """Return whether transport should retain upright-specific tuning. - - A preceding upright operation may use higher-clearance motion settings - without turning that live posture into a hard terminal constraint. - """ - if constraint.requires_upright_axis_alignment: - return True - if ( - constraint.terms - or constraint.planning_preference != "minimize_rotation_from_current" - ): - return False - entity = _object(self.env, step.object_uid) - vertices = _local_vertices(entity, self.env, 0) - try: - axis_index = analyze_local_geometry_axes(vertices).long_axis_index - except ValueError: - return False - pose = _live_pose(self.env, step.object_uid) - cosine = pose[:, 2, axis_index].abs().clamp(0.0, 1.0) - tolerance = float(self.runtime_policy.predicate_fallbacks["upright_max_tilt"]) - return bool(torch.all(torch.arccos(cosine) <= tolerance).item()) - @staticmethod def _upright_local_axis(step: SemanticStep) -> str: align_terms = tuple( diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index c7ef5c0e3..f50e131c6 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -742,7 +742,7 @@ def evaluate_predicate( return _constant(env, False) if kind == "poured": if spec.get("verification") == "action_completion": - # Reaching semantic-step verification means every required E3 edge + # Reaching semantic-step verification means every required pour edge # already completed without a fatal planning or execution failure. return _constant(env, True) diff --git a/embodichain/gen_sim/action_engine/runtime/recovery.py b/embodichain/gen_sim/action_engine/runtime/recovery.py index cbe3caa7a..1855f5cfa 100644 --- a/embodichain/gen_sim/action_engine/runtime/recovery.py +++ b/embodichain/gen_sim/action_engine/runtime/recovery.py @@ -221,32 +221,16 @@ def insert_recovery_subgraph( for node_id in descendants if str(node_by_id[node_id]["task_instance_id"]) == failed_group_id } - cleanup_suffix_ids: set[str] = set() + abandoned_suffix_ids: set[str] = set() + failed_capability = self.registry.get(str(failed_node["atomic_action"])) if ( - str(failed_node["atomic_action"]) == "HandOver" + failed_capability.state_effect == "transfer_hold" and not preserve_failed_group_suffix ): - # A failed handover leaves ownership indeterminate. Its - # transfer-arm retreat/home tail must not execute from a stale - # handover pose; recovery owns the cleanup before replanning. - cleanup_suffix_ids = { - node_id - for node_id in same_group_descendants - if node_by_id[node_id]["role"] == "cleanup" - } - non_cleanup_dependents = [ - node_id - for node_id in same_group_descendants - cleanup_suffix_ids - if any( - dependency in cleanup_suffix_ids - for dependency in node_by_id[node_id]["depends_on"] - ) - ] - if non_cleanup_dependents: - raise ValueError( - "Cannot remove the HandOver cleanup suffix because it feeds " - f"same-group non-cleanup nodes: {sorted(non_cleanup_dependents)}." - ) + # A failed ownership transfer leaves every unfinished action in the + # same task invalid, including receiver-side continuation. Recovery + # replaces that suffix before downstream tasks are replanned. + abandoned_suffix_ids = set(same_group_descendants) first = [ node for node in nodes @@ -277,16 +261,16 @@ def insert_recovery_subgraph( dependency for dependency in node["depends_on"] if dependency != failed_node_id - and dependency not in cleanup_suffix_ids + and dependency not in abandoned_suffix_ids ] + terminal_ids ) ) - if cleanup_suffix_ids: + if abandoned_suffix_ids: patched["nodes"] = [ node for node in patched["nodes"] - if str(node["id"]) not in cleanup_suffix_ids + if str(node["id"]) not in abandoned_suffix_ids ] failed_group = next( item @@ -296,7 +280,7 @@ def insert_recovery_subgraph( failed_group["node_ids"] = [ node_id for node_id in failed_group["node_ids"] - if node_id not in cleanup_suffix_ids + if node_id not in abandoned_suffix_ids ] group["depends_on"] = list( dict.fromkeys([failed_group_id, *group.get("depends_on", [])]) @@ -475,10 +459,11 @@ def build_upright_recovery( revision: int, resume_failed_group: bool = False, ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Build a coordinate-free E2 recovery group for a fallen rigid object.""" + """Build a coordinate-free capability recovery for a fallen rigid object.""" failed = _node(graph, failed_node_id) object_uid = str(failed["object_uid"]) - group_id = f"recovery_e2_{int(revision):02d}_{failed_node_id}" + task_type = str(failed["task_type"]) + group_id = f"recovery_upright_{int(revision):02d}_{failed_node_id}" actor = _recovery_actor(graph, failed) held_consumer_arm = None if not resume_failed_group: @@ -531,7 +516,7 @@ def build_upright_recovery( "target_binding": binding, "depends_on": dependencies, "task_instance_id": group_id, - "task_type": "E2", + "task_type": task_type, "role": "recovery" if index <= 3 else "cleanup", "precondition": {}, "postcondition": {}, @@ -547,9 +532,9 @@ def build_upright_recovery( dependencies = [node_id] group = { "id": group_id, - "task_type": "E2", + "task_type": task_type, "role": "recovery", - "operator": "orient_object", + "operator": "recover_object_upright", "object_uid": object_uid, "actor": actor, "goal": { @@ -561,6 +546,7 @@ def build_upright_recovery( "support_object": "table", "upright_local_axis": "long_axis", "terminal_behavior": "hold" if hold_for_downstream else "place", + "recovery_capability": "object_upright", }, "depends_on": [], "parent_task_instance_id": str(failed["task_instance_id"]), diff --git a/embodichain/gen_sim/action_engine/tasks/assembly.py b/embodichain/gen_sim/action_engine/tasks/assembly.py index c4afa5f65..f4f954488 100644 --- a/embodichain/gen_sim/action_engine/tasks/assembly.py +++ b/embodichain/gen_sim/action_engine/tasks/assembly.py @@ -337,14 +337,12 @@ def validate_target_compatibility( relation: str, ) -> None: """Reject only structural or explicitly declared target contradictions.""" - if task_type == "E1" and relation == "on" and target is not None: + if relation == "on" and target is not None: # Support is a relation between two concrete bodies at a candidate # pose. A positive affordance list is not a closed-world inventory, so # omission of ``support_surface`` cannot prove incompatibility here. return - requires_container = task_type == "E3" or ( - task_type == "E1" and relation == "inside" - ) + requires_container = task_type == "E3" or relation == "inside" if requires_container and target is None: raise ValueError( f"{task_type} {relation} relation requires a target container." @@ -366,7 +364,7 @@ def validate_target_compatibility( def _target_affordances(task_type: str, relation: str) -> tuple[str, ...]: - if task_type == "E3" or (task_type == "E1" and relation == "inside"): + if task_type == "E3" or relation == "inside": return ("container",) return () diff --git a/embodichain/gen_sim/action_engine/tasks/interpretation.py b/embodichain/gen_sim/action_engine/tasks/interpretation.py index c4d002b34..be136183b 100644 --- a/embodichain/gen_sim/action_engine/tasks/interpretation.py +++ b/embodichain/gen_sim/action_engine/tasks/interpretation.py @@ -268,11 +268,17 @@ def _emit_step( elif task_type == "E3": params.update({"relation": "above", "relation_frame": "robot"}) elif task_type == "E4": + terminal_behavior = str(step["terminal_behavior"]) + if terminal_behavior == "none": + terminal_behavior = "place" if target is not None else "hold" params.update( { "transfer_arm": step["transfer_arm"], "receive_arm": step["receive_arm"], "orientation_goal": step["orientation_goal"], + "terminal_behavior": terminal_behavior, + "relation": step["relation"], + "relation_frame": "robot", } ) elif task_type == "E5": diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index 2ced064e7..c31930c18 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -352,6 +352,7 @@ def _recipe( role=role, already_held=incoming_held_arm is not None, payloads=payloads, + orientation_goal=str(params.get("orientation_goal", "none")), ), "arrange_line", goal, @@ -359,6 +360,11 @@ def _recipe( ) goal = { "reference_object": target, + "support_object": ( + target + if relation in {"on", "inside"} + else str(params.get("support_role", "table")) + ), "reference_state": "live", "relation": relation, "relation_frame": str(params.get("relation_frame", "world")), @@ -387,6 +393,7 @@ def _recipe( role=role, already_held=incoming_held_arm is not None, payloads=payloads, + orientation_goal=str(params.get("orientation_goal", "none")), ), "place_relative", goal, @@ -502,6 +509,7 @@ def _recipe( role=role, already_held=incoming_held_arm is not None, leave_held=terminal_behavior == "hold", + orientation_goal=str(params.get("orientation_goal", "upright")), ), "orient_object", goal, @@ -618,6 +626,20 @@ def _recipe( if task_type == "E4": transfer = str(params.get("transfer_arm", "left_arm")) receive = str(params.get("receive_arm", "right_arm")) + terminal_behavior = str(params.get("terminal_behavior", "hold")) + if terminal_behavior not in TERMINAL_BEHAVIORS - {"none"}: + raise ValueError("E4 terminal_behavior must be 'hold' or 'place'.") + target = params.get("target_role") + relation = str(params.get("relation", "none")) + if terminal_behavior == "place": + if not isinstance(target, str) or not target or relation == "none": + raise ValueError( + "E4 terminal_behavior='place' requires target_role and relation." + ) + elif target is not None or relation != "none": + raise ValueError( + "E4 terminal_behavior='hold' cannot carry target_role or relation." + ) if incoming_held_arm == "coordinated": raise ValueError( "E4 cannot consume a coordinated hold; an explicit single-arm " @@ -628,30 +650,36 @@ def _recipe( f"E4 transfer_arm {transfer!r} conflicts with the predecessor " f"holder {incoming_held_arm!r}." ) - pickup_actor = {"mode": "required", "arm": transfer} - pickup = None + transfer_actor = {"mode": "required", "arm": transfer} + receive_actor = {"mode": "required", "arm": receive} + nodes: list[dict[str, Any]] = [] + previous = list(dependencies) + next_index = 1 if incoming_held_arm is None: pickup = _node( group_id, - 1, + next_index, "PickUp", task_type, object_uid, - pickup_actor, + transfer_actor, "arm", {"kind": "object", "object": object_uid}, - dependencies, + previous, role, {"type": "object_held", "object": object_uid, "arm": transfer}, motion_policy(("handover_role", "transfer")), ) + nodes.append(pickup) + previous = [pickup["id"]] + next_index += 1 staging = _node( group_id, - 1 if pickup is None else 2, + next_index, "MoveHeldObject", task_type, object_uid, - pickup_actor, + transfer_actor, "arm", { "kind": "handover_staging", @@ -659,14 +687,17 @@ def _recipe( "transfer_arm": transfer, "receive_arm": receive, }, - dependencies if pickup is None else [pickup["id"]], + previous, role, {"type": "object_held", "object": object_uid, "arm": transfer}, motion_policy(), ) + nodes.append(staging) + previous = [staging["id"]] + next_index += 1 handover = _node( group_id, - 2 if pickup is None else 3, + next_index, "HandOver", task_type, object_uid, @@ -678,67 +709,199 @@ def _recipe( "transfer_arm": transfer, "receive_arm": receive, }, - [staging["id"]], + previous, role, {"type": "handover_complete", "object": object_uid, "arm": receive}, motion_policy(), ) + nodes.append(handover) + previous = [handover["id"]] + next_index += 1 # 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, + next_index, "MoveEndEffector", task_type, object_uid, - pickup_actor, + transfer_actor, "arm", { "kind": "policy_pose", "source": "handover", "operation": "retreat", }, - [handover["id"]], + previous, "cleanup", {}, motion_policy(), ) + nodes.append(retreat) + previous = [retreat["id"]] + next_index += 1 home = _node( group_id, - 4 if pickup is None else 5, + next_index, "MoveJoints", task_type, object_uid, - pickup_actor, + transfer_actor, "arm", { "kind": "joint_state", "source": "initial", "operation": "handover_home", }, - [retreat["id"]], + previous, + "cleanup", + {}, + motion_policy(), + ) + nodes.append(home) + previous = [home["id"]] + next_index += 1 + + orientation_modifiers: tuple[tuple[str, str], ...] = ( + (("orientation", "upright"),) + if params.get("orientation_goal") == "upright" + else () + ) + receiver_policy = motion_policy(*orientation_modifiers) + if terminal_behavior == "hold": + receiver_exit = _node( + group_id, + next_index, + "MoveHeldObject", + task_type, + object_uid, + receive_actor, + "arm", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "handover_exit", + "receive_arm": receive, + "terminal_hold": True, + }, + previous, + role, + {"type": "handover_complete", "object": object_uid, "arm": receive}, + receiver_policy, + ) + nodes.append(receiver_exit) + goal = { + "relation": "handover", + "orientation_goal": str(params.get("orientation_goal", "none")), + "orientation_axis": "none", + **_orientation_extensions(params), + "transfer_arm": transfer, + "receive_arm": receive, + "terminal_behavior": "hold", + } + return ( + nodes, + "handover", + goal, + {"type": "handover_complete", "object": object_uid, "arm": receive}, + ) + + assert isinstance(target, str) + for phase in ("staging", "final"): + receiver_move = _node( + group_id, + next_index, + "MoveHeldObject", + task_type, + object_uid, + receive_actor, + "arm", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": phase, + }, + previous, + role, + {}, + receiver_policy, + ) + nodes.append(receiver_move) + previous = [receiver_move["id"]] + next_index += 1 + release = _node( + group_id, + next_index, + "Place", + task_type, + object_uid, + receive_actor, + "arm", + {"kind": "current_held_pose"}, + previous, + role, + {}, + receiver_policy, + ) + nodes.append(release) + previous = [release["id"]] + next_index += 1 + receiver_retreat = _node( + group_id, + next_index, + "MoveEndEffector", + task_type, + object_uid, + receive_actor, + "arm", + {"kind": "policy_pose", "source": "release", "operation": "retreat"}, + previous, + "cleanup", + {}, + receiver_policy, + ) + nodes.append(receiver_retreat) + previous = [receiver_retreat["id"]] + next_index += 1 + receiver_home = _node( + group_id, + next_index, + "MoveJoints", + task_type, + object_uid, + receive_actor, + "arm", + {"kind": "joint_state", "source": "initial"}, + previous, "cleanup", {}, motion_policy(), ) + nodes.append(receiver_home) + success = { + "type": task_success_type(task_type, params), + "relation": relation, + "reference_object": target, + } return ( - [ - item - for item in (pickup, staging, handover, retreat, home) - if item is not None - ], + nodes, "handover", { - "relation": "handover", + "reference_object": target, + "support_object": (target if relation in {"on", "inside"} else "table"), + "reference_state": "live", + "relation": relation, + "relation_frame": str(params.get("relation_frame", "robot")), "orientation_goal": str(params.get("orientation_goal", "none")), "orientation_axis": "none", **_orientation_extensions(params), "transfer_arm": transfer, "receive_arm": receive, + "terminal_behavior": "place", }, - {"type": "handover_complete", "object": object_uid, "arm": receive}, + success, ) if task_type == "E5": terminal_behavior = str(params.get("terminal_behavior", "hold")) @@ -916,9 +1079,10 @@ def _single_arm_manipulation( already_held: bool = False, leave_held: bool = False, payloads: Sequence[Mapping[str, Any]] = (), + orientation_goal: str = "none", ) -> list[dict[str, Any]]: orientation_modifiers: tuple[tuple[str, str], ...] = ( - (("orientation", "upright"),) if task_type == "E2" else () + (("orientation", "upright"),) if orientation_goal == "upright" else () ) payload_binding = deepcopy(list(payloads)) specs = ( diff --git a/embodichain/gen_sim/task_engine/agent.py b/embodichain/gen_sim/task_engine/agent.py index f5ab773f0..1d183e9d2 100644 --- a/embodichain/gen_sim/task_engine/agent.py +++ b/embodichain/gen_sim/task_engine/agent.py @@ -276,14 +276,14 @@ def _canonicalize_intent(intent: Mapping[str, Any]) -> dict[str, Any]: def _target_affordances(task_type: str, relation: str) -> list[str]: - if task_type == "E3" or (task_type == "E1" and relation == "inside"): + if task_type == "E3" or relation == "inside": return ["container"] return [] def _target_structure(task_type: str, relation: str) -> str: - if task_type == "E1" and relation == "on": + if relation == "on": return "physical_entity" - if task_type == "E3" or (task_type == "E1" and relation == "inside"): + if task_type == "E3" or relation == "inside": return "rigid_object" return "spatial_reference" diff --git a/embodichain/gen_sim/task_engine/interpretation.py b/embodichain/gen_sim/task_engine/interpretation.py index 2b971344c..f1d5484ab 100644 --- a/embodichain/gen_sim/task_engine/interpretation.py +++ b/embodichain/gen_sim/task_engine/interpretation.py @@ -365,111 +365,22 @@ def _normalize_instruction_intent_fields( "reason": "e5_hold_defaults_to_lift", } ) - _normalize_handover_arm_continuity(raw_steps, changes) - return result, changes - - -def _normalize_handover_arm_continuity( - steps: Sequence[Any], - changes: list[dict[str, Any]], -) -> None: - """Repair a same-arm E4 only when adjacent ownership fixes both roles.""" - by_id: dict[str, Mapping[str, Any]] = {} - for step in steps: - if not isinstance(step, dict) or set(step) != _STEP_KEYS: - return - step_id = step.get("id") - if not isinstance(step_id, str) or not step_id or step_id in by_id: - return - by_id[step_id] = step - - explicit_arms = {"left_arm", "right_arm"} - for index, step in enumerate(steps): - assert isinstance(step, dict) - transfer = step.get("transfer_arm") - receive = step.get("receive_arm") - if ( - step.get("task_type") != "E4" - or transfer not in explicit_arms - or transfer != receive - ): - continue - - object_key = _object_lineage_key(step, by_id) - upstream_arm: str | None = None - for producer in reversed(steps[:index]): - assert isinstance(producer, Mapping) - if object_key is None or _object_lineage_key(producer, by_id) != object_key: - continue - candidate = ( - producer.get("receive_arm") - if producer.get("task_type") == "E4" - else producer.get("required_arm") + if task_type == "E4" and raw_step.get("terminal_behavior") == "none": + terminal = ( + "place" + if isinstance(target, Mapping) and target.get("kind") != "none" + else "hold" ) - if candidate in explicit_arms: - upstream_arm = str(candidate) - break - - downstream_arm: str | None = None - for consumer in steps[index + 1 :]: - assert isinstance(consumer, Mapping) - if object_key is None or _object_lineage_key(consumer, by_id) != object_key: - continue - candidate = ( - consumer.get("transfer_arm") - if consumer.get("task_type") == "E4" - else consumer.get("required_arm") - ) - if candidate in explicit_arms: - downstream_arm = str(candidate) - break - - desired_transfer = upstream_arm or str(transfer) - desired_receive = downstream_arm or str(receive) - if desired_transfer == desired_receive: - continue - for field, desired in ( - ("transfer_arm", desired_transfer), - ("receive_arm", desired_receive), - ): - if step[field] == desired: - continue - previous = step[field] - step[field] = desired + raw_step["terminal_behavior"] = terminal changes.append( { - "path": f"steps[{index}].{field}", - "from": previous, - "to": desired, - "reason": "handover_arm_continuity", + "path": f"steps[{index}].terminal_behavior", + "from": "none", + "to": terminal, + "reason": "e4_terminal_inferred_from_own_target", } ) - - -def _object_lineage_key( - step: Mapping[str, Any], - by_id: Mapping[str, Mapping[str, Any]], - seen: frozenset[str] = frozenset(), -) -> tuple[str, str] | None: - """Resolve object identity only through explicit step-result lineage.""" - selector = step.get("object") - if not isinstance(selector, Mapping): - return None - kind = selector.get("kind") - if kind == "scene_ref": - step_id = step.get("id") - if not isinstance(step_id, str) or not step_id: - return None - return ("step_result", step_id) - if kind != "step_result": - return None - producer_id = selector.get("step_id") - if not isinstance(producer_id, str) or producer_id in seen: - return None - producer = by_id.get(producer_id) - if producer is None: - return None - return _object_lineage_key(producer, by_id, seen | {producer_id}) + return result, changes def _empty_selector() -> dict[str, Any]: @@ -639,7 +550,7 @@ def _validate_task_fields(step: Mapping[str, Any], context: str) -> None: f"{context} {task_type} requires an object selector." ) target_kind = str(step["target"]["kind"]) - if task_type not in {"E1", "E3", "E5"} and step["relation"] != "none": + if task_type not in {"E1", "E3", "E4", "E5"} and step["relation"] != "none": raise ValueError(f"{context} {task_type} does not accept relation.") if task_type == "E3" and step["relation"] != "above": raise ValueError(f"{context} E3 relation must be above.") @@ -675,6 +586,24 @@ def _validate_task_fields(step: Mapping[str, Any], context: str) -> None: ) if step["relation"] == "none" and task_type == "E3": raise ValueError(f"{context} {task_type} requires a symbolic relation.") + elif task_type == "E4": + terminal = str(step["terminal_behavior"]) + effective_terminal = ( + "place" if terminal == "none" and target_kind != "none" else terminal + ) + if effective_terminal == "none": + effective_terminal = "hold" + if effective_terminal not in _TERMINAL_BEHAVIORS - {"none"}: + raise ValueError(f"{context} E4 requires terminal_behavior hold/place.") + if effective_terminal == "place": + if target_kind == "none" or step["relation"] == "none": + raise ValueError( + f"{context} E4 terminal_behavior=place requires target and relation." + ) + elif target_kind != "none" or step["relation"] != "none": + raise ValueError( + f"{context} E4 terminal_behavior=hold cannot carry target or relation." + ) elif task_type == "E5": direction = str(step["direction"]) terminal = str(step["terminal_behavior"]) @@ -696,7 +625,7 @@ def _validate_task_fields(step: Mapping[str, Any], context: str) -> None: ) elif target_kind != "none": raise ValueError(f"{context} {task_type} does not accept a target selector.") - if task_type != "E5": + if task_type not in {"E4", "E5"}: if step["direction"] != "none": raise ValueError(f"{context} direction is only valid for E5.") if step["terminal_behavior"] != "none": @@ -742,9 +671,11 @@ def _instruction_prompt(instruction: str) -> str: "paths, or reasoning. Encode explicit ordering with depends_on; same-action set " "members may remain independent. Use empty strings and 'none' for " "inapplicable required fields. A request to retract the transfer arm " - "immediately after an E4 handover is a mandatory runtime retreat/home " - "barrier for that E4; do " - "not emit a separate task step for it. The exact output keys are steps -> id, " + "E4 owns the complete transfer. For a handover followed by placement in " + "the same user intent, emit one E4 with target, relation, and " + "terminal_behavior=place; do not emit a trailing E1. Use " + "terminal_behavior=hold only when the receiver should keep holding the " + "object. The exact output keys are steps -> id, " "task_type, object, target, relation, required_arm, transfer_arm, " "receive_arm, orientation_goal, target_state, target_setting, layout, " "axis, direction, terminal_behavior, depends_on; each selector has kind, " @@ -755,7 +686,7 @@ def _instruction_prompt(instruction: str) -> str: "Emptying, dumping, or pouring contents from one container into another " "is exactly one E3 step: object selects the source container, target " "selects the receiving container, and relation=above. Pickup and staging " - "are internal to E3; do not emit a separate E1 for an explicit grab. " + "are internal to that E3 step. " "Opening or pulling out a drawer is E6 with object selecting that drawer " "and target_state=open. Closing or pushing in a drawer is E7 with object " "selecting that drawer and target_state=closed. " @@ -772,11 +703,11 @@ def _instruction_prompt(instruction: str) -> str: "and terminal_behavior=hold. Use hold unless the instruction explicitly " "says to put/release the object. For pick " "and release at the original location, use direction=none and place. A dual-arm " - "pick/move/transport request is E5, not E1. Final checklist: every step " + "pick/move/transport request uses E5. Final checklist: every step " "has all 16 step keys; every object and target " "has all 5 selector keys. For an inapplicable field use the canonical " "default shown in the example, never omit the field. E4 must explicitly " - "state transfer_arm and receive_arm. E1/E3 must explicitly state target " + "state transfer_arm, receive_arm, and terminal_behavior. E1/E3 must explicitly state target " "and relation (except E1 layout=line)." ) @@ -872,7 +803,7 @@ def _instruction_repair_guidance(error: Exception) -> str: "\nMissing-target repair rule for E3: keep task_type=E3. object is " "the source container whose contents are poured, target is the " "receiving container, and relation must be above. An explicit grab " - "is part of E3 and must not be reclassified as E1.\n" + "is part of the same E3 task.\n" ) return ( "\nMissing-target repair rule: for a non-line E1 placement, object is " diff --git a/embodichain/gen_sim/task_engine/ontology.py b/embodichain/gen_sim/task_engine/ontology.py index 4ae707ca7..1f497b645 100644 --- a/embodichain/gen_sim/task_engine/ontology.py +++ b/embodichain/gen_sim/task_engine/ontology.py @@ -188,14 +188,28 @@ def _contract( ), "E4": _contract( "E4", - "Transfer one held object from one arm to the other.", - frozenset({"transfer_arm", "receive_arm", "orientation_goal"}), + "Transfer one object between arms, then either leave the receiver " + "holding it safely or place it at a symbolic relation.", + frozenset( + { + "target", + "relation", + "transfer_arm", + "receive_arm", + "orientation_goal", + "terminal_behavior", + } + ), "rigid_object", frozenset({"graspable", "handover"}), "handover_complete", resource_mode="handover", moves_primary_object=True, accepts_incoming_hold=True, + terminal_success_types=( + ("hold", "handover_complete"), + ("place", "semantic_goal"), + ), ), "E5": _contract( "E5", diff --git a/tests/gen_sim/action_engine/domain/test_task_contracts.py b/tests/gen_sim/action_engine/domain/test_task_contracts.py index a1b76fb88..daa54a53c 100644 --- a/tests/gen_sim/action_engine/domain/test_task_contracts.py +++ b/tests/gen_sim/action_engine/domain/test_task_contracts.py @@ -67,6 +67,18 @@ def test_e5_success_depends_only_on_terminal_behavior() -> None: task_success_type("E5", {"terminal_behavior": "none"}) +def test_e4_success_depends_on_its_own_terminal_behavior() -> None: + contract = task_contract("E4") + + assert {"target", "relation", "terminal_behavior"} <= set( + contract.applicable_intent_fields + ) + assert task_success_type("E4", {"terminal_behavior": "hold"}) == ( + "handover_complete" + ) + assert task_success_type("E4", {"terminal_behavior": "place"}) == "semantic_goal" + + def test_symbolic_transport_values_are_language_neutral_protocol_enums() -> None: assert {"on", "inside", "behind", "left_of"} <= RELATIONS assert {"none", "up", "left", "world_y"} <= TRANSPORT_DIRECTIONS diff --git a/tests/gen_sim/action_engine/planning/test_linker.py b/tests/gen_sim/action_engine/planning/test_linker.py index 2548a43a3..dbc711602 100644 --- a/tests/gen_sim/action_engine/planning/test_linker.py +++ b/tests/gen_sim/action_engine/planning/test_linker.py @@ -172,9 +172,9 @@ def test_handover_ownership_flows_through_home_terminal_barrier() -> None: if node["task_instance_id"] == "task_03" and node["atomic_action"] == "HandOver" ) - assert terminal["atomic_action"] == "MoveJoints" + assert terminal["atomic_action"] == "MoveHeldObject" assert terminal["contract"]["completion"] == "terminal_barrier" - assert terminal["contract"]["failure_policy"] == "best_effort" + assert terminal["contract"]["failure_policy"] == "task_required" retreat = next( node for node in graph["nodes"] diff --git a/tests/gen_sim/action_engine/planning/test_online_v2.py b/tests/gen_sim/action_engine/planning/test_online_v2.py index 290f667bd..bfa7b8938 100644 --- a/tests/gen_sim/action_engine/planning/test_online_v2.py +++ b/tests/gen_sim/action_engine/planning/test_online_v2.py @@ -109,7 +109,7 @@ def caller(**kwargs): 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 "Transfer one object between arms" in prompts[0] assert graph["metadata"]["oracle_exposed"] is False assert any( node["target_binding"]["kind"] == "visual_constraint" for node in graph["nodes"] diff --git a/tests/gen_sim/action_engine/runtime/test_recovery_v2.py b/tests/gen_sim/action_engine/runtime/test_recovery_v2.py index 5544aea45..a112c7c24 100644 --- a/tests/gen_sim/action_engine/runtime/test_recovery_v2.py +++ b/tests/gen_sim/action_engine/runtime/test_recovery_v2.py @@ -221,38 +221,13 @@ def test_recovery_insertion_revises_runtime_graph_not_seed_graph() -> None: failed_node = next( node for node in graph["nodes"] if node["atomic_action"] == "HandOver" ) - recovery_source = _graph("E2") - source_group = recovery_source["task_groups"][0] - recovery_group_id = "recovery_upright_01" - recovery_nodes = [] - id_map = { - node["id"]: f"recovery_{index:02d}" - for index, node in enumerate(recovery_source["nodes"], start=1) - } - for node in recovery_source["nodes"]: - item = deepcopy(node) - item["id"] = id_map[node["id"]] - item["object_uid"] = failed_node["object_uid"] - item["target_binding"] = deepcopy(item["target_binding"]) - if item["target_binding"].get("kind") == "object": - item["target_binding"]["object"] = failed_node["object_uid"] - item["depends_on"] = [id_map.get(dep, dep) for dep in node["depends_on"]] - recovery_nodes.append(item) - recovery_group = deepcopy(source_group) - recovery_group.update( - { - "id": recovery_group_id, - "role": "recovery", - "object_uid": failed_node["object_uid"], - "node_ids": [node["id"] for node in recovery_nodes], - "depends_on": [], - "parent_task_instance_id": failed_node["task_instance_id"], - } - ) - recovery_group["success"] = { - "type": "object_upright", - "object": failed_node["object_uid"], - } + recovery_nodes, recovery_group = build_upright_recovery( + graph, + failed_node_id=failed_node["id"], + revision=1, + resume_failed_group=True, + ) + recovery_group_id = recovery_group["id"] patched = runtime.insert_recovery_subgraph( failed_node_id=failed_node["id"], @@ -322,6 +297,9 @@ def test_failed_group_resume_recovery_places_before_prefix_replay() -> None: and node["role"] == "cleanup" } assert group["goal"]["terminal_behavior"] == "place" + assert group["task_type"] == "E4" + assert group["operator"] == "recover_object_upright" + assert group["goal"]["recovery_capability"] == "object_upright" assert [node["atomic_action"] for node in nodes] == [ "PickUp", "MoveHeldObject", @@ -332,6 +310,26 @@ def test_failed_group_resume_recovery_places_before_prefix_replay() -> None: assert original_cleanup <= {node["id"] for node in patched["nodes"]} +@pytest.mark.parametrize("task_type", ("E1", "E3", "E4")) +def test_upright_recovery_inherits_failed_task_identity_without_e2_dispatch( + task_type: str, +) -> None: + graph = _graph(task_type) + failed = next(node for node in graph["nodes"] if node["role"] != "cleanup") + + nodes, recovery_group = build_upright_recovery( + graph, + failed_node_id=failed["id"], + revision=1, + resume_failed_group=True, + ) + + assert recovery_group["task_type"] == task_type + assert recovery_group["operator"] == "recover_object_upright" + assert recovery_group["goal"]["recovery_capability"] == "object_upright" + assert {node["task_type"] for node in nodes} == {task_type} + + @pytest.mark.parametrize( "actor", ( @@ -399,7 +397,7 @@ def test_local_recovery_replays_failed_group_prefix_and_preserves_seed_graph( recovery_edges = [ event for event in recorder.edge_events - if event["step_id"].startswith("recovery_e2_") + if event["step_id"].startswith("recovery_upright_") ] replay_edges = [ event for event in recorder.edge_events if event["step_id"] == step.id diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index 1cccaa34c..e7c1eacd6 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -1893,8 +1893,9 @@ def test_handover_continuation_uses_stable_upright_policies() -> None: assert grounded_staging.allow_yaw_search assert grounded_final.allow_yaw_search assert grounded_final_with_reference.target_object_pose is not None + assert grounded_final.target_object_pose is not None assert grounded_final_with_reference.target_object_pose[0, 2, 3] == pytest.approx( - 0.90 + float(grounded_final.target_object_pose[0, 2, 3]) ) assert ( grounded_release.cfg["sample_interval"] == release_defaults["sample_interval"] @@ -4357,10 +4358,10 @@ def test_live_pickup_planning_exception_is_a_retryable_edge_failure( assert result.planner_traces[0]["exception"] == "RuntimeError: no IK" -def test_pickup_candidate_screens_handover_successor_target( +def test_pickup_candidate_does_not_scan_a_successor_task_operator( monkeypatch: pytest.MonkeyPatch, ) -> None: - """A later handover staging pose participates in pickup grasp screening.""" + """A task screens only its own target bindings, never a successor recipe.""" entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)) env = _FakeEnv({"can": entity}) executor = ProgramExecutor( @@ -4404,7 +4405,6 @@ def test_pickup_candidate_screens_handover_successor_target( executor.edges[handover_edge.id] = handover_edge existing_target = _pose(0.0, 0.2, 0.85) - handover_target = _pose(0.0, 0.0, 1.15) grounded = GroundedAction( action_class="PickUp", arm="left_arm", @@ -4413,27 +4413,21 @@ def test_pickup_candidate_screens_handover_successor_target( cfg={"downstream_object_target_poses": (existing_target,)}, ) - def ground( + def own_task_ground( _action: Any, candidate: SemanticStep, - *, - arm: str, - state: ExecutionState, - reference_eef_pose: torch.Tensor | None = None, - orientation_reference_pose: torch.Tensor | None = None, + **_kwargs: Any, ) -> GroundedAction: - del arm, state, reference_eef_pose, orientation_reference_pose - target = handover_target if candidate.id == handover_step.id else None + assert candidate.id == pickup_step.id return GroundedAction( action_class="MoveHeldObject", arm="left_arm", control="arm", target=SimpleNamespace(), cfg={}, - target_object_pose=target, ) - monkeypatch.setattr(executor.grounder, "ground", ground) + monkeypatch.setattr(executor.grounder, "ground", own_task_ground) result = executor._with_downstream_targets( pickup_step, pickup_step.edge_ids[0], @@ -4443,9 +4437,8 @@ def ground( ) targets = result.cfg["downstream_object_target_poses"] - assert len(targets) == 2 + assert len(targets) == 1 assert torch.equal(targets[0], existing_target) - assert torch.equal(targets[1], handover_target) def test_object_held_predicate_checks_live_gripper_and_tcp_geometry() -> None: diff --git a/tests/gen_sim/action_engine/task_fixtures.py b/tests/gen_sim/action_engine/task_fixtures.py index f7308dc1a..001b4a91f 100644 --- a/tests/gen_sim/action_engine/task_fixtures.py +++ b/tests/gen_sim/action_engine/task_fixtures.py @@ -313,6 +313,7 @@ def make_task_spec( "transfer_arm": "left_arm", "receive_arm": "right_arm", "orientation_goal": "none", + "terminal_behavior": "hold", } ) elif task_type == "E5": diff --git a/tests/gen_sim/action_engine/tasks/test_e3_pour.py b/tests/gen_sim/action_engine/tasks/test_e3_pour.py index 54c61d8b9..19c9e0f7e 100644 --- a/tests/gen_sim/action_engine/tasks/test_e3_pour.py +++ b/tests/gen_sim/action_engine/tasks/test_e3_pour.py @@ -111,7 +111,7 @@ def test_seed_graph_rejects_legacy_dual_arm_e3_goal() -> None: graph = _graph() graph["task_groups"][0]["goal"]["pour_mode"] = "dual_arm" - with pytest.raises(ValueError, match="unsupported dual-arm E3 fields"): + with pytest.raises(ValueError, match="unsupported dual-arm pour fields"): validate_seed_graph(graph) diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index a3a483c14..e7a4f3cf6 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -19,7 +19,11 @@ from __future__ import annotations from copy import deepcopy +from dataclasses import replace +import embodichain.gen_sim.action_engine.domain.task_contracts as task_contracts_module +import embodichain.gen_sim.action_engine.domain.v2 as domain_v2_module +from embodichain.gen_sim.action_engine.runtime import load_execution_program from embodichain.gen_sim.action_engine.tasks import ( ground_instruction_draft, instantiate_seed_graph, @@ -68,7 +72,7 @@ def _intent_step( "layout": "none", "axis": "none", "direction": "none", - "terminal_behavior": "none", + "terminal_behavior": "hold" if task_type == "E4" else "none", "depends_on": [], } step.update(updates) @@ -273,6 +277,7 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "HandOver", "MoveEndEffector", "MoveJoints", + "MoveHeldObject", ] assert handover_nodes[0]["motion_policy"] == { "modifiers": [{"type": "handover_role", "mode": "transfer"}] @@ -428,6 +433,7 @@ def test_handover_to_place_uses_receiver_hold_without_repickup() -> None: "HandOver", "MoveEndEffector", "MoveJoints", + "MoveHeldObject", ] assert handover["actor"] == {"mode": "required", "arm": "left_arm"} assert graph["nodes"][0]["motion_policy"] == { @@ -469,6 +475,151 @@ def test_handover_to_place_uses_receiver_hold_without_repickup() -> None: assert placement_nodes[0]["depends_on"] == [handover["node_ids"][-1]] +def test_e4_hold_owns_receiver_safe_exit() -> None: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "handover_hold", + "level": "L1", + "instruction": "Hand the can to the right arm and keep holding it.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "handover", + "task_type": "E4", + "params": { + "object_role": "can", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "terminal_behavior": "hold", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {}, + "oracle": {}, + "metadata": {}, + } + + graph = instantiate_seed_graph(task, {"can": "can_uid"}) + group = graph["task_groups"][0] + receiver_exit = graph["nodes"][-1] + + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + "MoveHeldObject", + ] + assert receiver_exit["actor"] == {"mode": "required", "arm": "right_arm"} + assert receiver_exit["target_binding"]["phase"] == "handover_exit" + assert receiver_exit["target_binding"]["terminal_hold"] is True + assert group["success"]["type"] == "handover_complete" + assert group["contract"]["completion"] == "terminal_barrier" + + +def test_e4_place_owns_receiver_placement_without_e1(monkeypatch) -> None: + task = { + "schema_version": "action_engine_task_spec_v2", + "task_id": "handover_place", + "level": "L1", + "instruction": "Hand the can to the right arm and place it on the notebook.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "handover_place", + "task_type": "E4", + "params": { + "object_role": "can", + "target_role": "notebook", + "relation": "on", + "transfer_arm": "left_arm", + "receive_arm": "right_arm", + "terminal_behavior": "place", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {}, + "oracle": {}, + "metadata": {}, + } + + graph = instantiate_seed_graph( + task, + {"can": "can_uid", "notebook": "notebook_uid"}, + ) + group = graph["task_groups"][0] + + assert {node["task_type"] for node in graph["nodes"]} == {"E4"} + assert len(graph["task_groups"]) == 1 + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + "MoveHeldObject", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + assert all( + node["actor"] == {"mode": "required", "arm": "right_arm"} + for node in graph["nodes"][5:] + ) + assert group["goal"]["terminal_behavior"] == "place" + assert group["goal"]["reference_object"] == "notebook_uid" + assert group["success"] == { + "type": "semantic_goal", + "relation": "on", + "reference_object": "notebook_uid", + } + + renamed = deepcopy(graph) + renamed["task_groups"][0]["operator"] = "equivalent_transfer_and_place" + program = load_execution_program(renamed) + assert program.semantic_steps[0].operator == "equivalent_transfer_and_place" + assert [ + action["atomic_action_class"] + for edge in program.edges + for action in edge.actions + ] == [node["atomic_action"] for node in graph["nodes"]] + + alias_contract = replace( + task_contracts_module.task_contract("E4"), + task_type="transfer_and_place", + ) + monkeypatch.setattr( + domain_v2_module, + "TASK_TYPES", + frozenset({*domain_v2_module.TASK_TYPES, "transfer_and_place"}), + ) + monkeypatch.setattr( + task_contracts_module, + "TASK_CONTRACTS", + { + **dict(task_contracts_module.TASK_CONTRACTS), + "transfer_and_place": alias_contract, + }, + ) + renamed_type = deepcopy(graph) + for node in renamed_type["nodes"]: + node["task_type"] = "transfer_and_place" + renamed_type["task_groups"][0]["task_type"] = "transfer_and_place" + + alias_program = load_execution_program(renamed_type) + assert [ + action["atomic_action_class"] + for edge in alias_program.edges + for action in edge.actions + ] == [node["atomic_action"] for node in graph["nodes"]] + + def test_structured_draft_grounds_handover_then_receiver_placement() -> None: scene = [ { @@ -536,6 +687,7 @@ def test_structured_draft_grounds_handover_then_receiver_placement() -> None: "MoveJoints", "MoveHeldObject", "MoveHeldObject", + "MoveHeldObject", "Place", "MoveEndEffector", "MoveJoints", @@ -617,7 +769,14 @@ def test_seed_graph_adds_missing_same_object_e2_handover_dependency() -> None: node["atomic_action"] for node in graph["nodes"] if node["task_instance_id"] == "task_03" - ] == ["PickUp", "MoveHeldObject", "HandOver", "MoveEndEffector", "MoveJoints"] + ] == [ + "PickUp", + "MoveHeldObject", + "HandOver", + "MoveEndEffector", + "MoveJoints", + "MoveHeldObject", + ] pickup = next( node for node in graph["nodes"] diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index 263790883..928c339d5 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -93,7 +93,7 @@ def _step(step_id: str, task_type: str, object_selector: dict, **values): "layout": "none", "axis": "none", "direction": "none", - "terminal_behavior": "none", + "terminal_behavior": "hold" if task_type == "E4" else "none", "depends_on": [], } result.update(values) @@ -332,6 +332,7 @@ def caller(**kwargs): "HandOver", "MoveEndEffector", "MoveJoints", + "MoveHeldObject", ] assert handover_nodes[0]["motion_policy"] == { "modifiers": [{"type": "handover_role", "mode": "transfer"}] @@ -362,6 +363,50 @@ def caller(**kwargs): assert calls[0]["model"] == "test-model" +def test_complete_handover_place_intent_emits_one_e4_task() -> None: + intent = { + "steps": [ + _step( + "handover_place", + "E4", + _selector("scene_ref", reference="object-alpha"), + target=_selector("scene_ref", reference="object-beta"), + relation="on", + transfer_arm="left_arm", + receive_arm="right_arm", + terminal_behavior="place", + ) + ] + } + prompts: list[str] = [] + + def caller(**kwargs): + prompts.append(kwargs["prompt"]) + return deepcopy(intent) + + grounded = interpret_and_ground_task_spec( + "complete_handover_place", + "Use the left arm to hand the can to the right arm, then place it on the notebook.", + _scene(), + robot_profile="ur10", + caller=caller, + grounding_caller=_grounding_caller( + **{ + "handover_place.object": "purple_can", + "handover_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"]] == [ + "E4" + ] + assert {node["task_type"] for node in graph["nodes"]} == {"E4"} + assert graph["task_groups"][0]["goal"]["terminal_behavior"] == "place" + assert "do not emit a trailing E1" in prompts[0] + + def test_six_step_repeated_objects_preserve_two_handover_continuations() -> None: intent = { "steps": [ @@ -1044,7 +1089,7 @@ def test_interpreter_normalizes_registry_inapplicable_e4_required_arm() -> None: ] -def test_interpreter_resolves_same_arm_handover_from_step_result_ownership() -> None: +def test_interpreter_does_not_infer_handover_arms_from_adjacent_tasks() -> None: intent = { "steps": [ _step( @@ -1083,26 +1128,12 @@ def test_interpreter_resolves_same_arm_handover_from_step_result_ownership() -> 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", - }, - ) + with pytest.raises(ValueError, match="failed validation after one repair"): + task_interpretation_module.interpret_instruction_draft( + "test-instruction-invalid-handover", + model="test-model", + caller=lambda **_kwargs: deepcopy(intent), + ) def test_interpreter_repairs_direct_reference_handover_from_later_arm_semantics() -> ( @@ -1315,7 +1346,7 @@ def caller(**kwargs): assert "exactly one E3 step" in prompts[0] assert "Missing-target repair rule for E3" in prompts[1] - assert "must not be reclassified as E1" in prompts[1] + assert "part of the same E3 task" in prompts[1] def test_missing_e6_object_reaches_targeted_repair() -> None: @@ -1838,6 +1869,7 @@ def test_multi_object_handover_keeps_both_order_and_holder_dependencies() -> Non "HandOver", "MoveEndEffector", "MoveJoints", + "MoveHeldObject", ] 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") diff --git a/tests/gen_sim/action_engine/tasks/test_payload_contracts.py b/tests/gen_sim/action_engine/tasks/test_payload_contracts.py index 11e537bd3..c921bb46c 100644 --- a/tests/gen_sim/action_engine/tasks/test_payload_contracts.py +++ b/tests/gen_sim/action_engine/tasks/test_payload_contracts.py @@ -31,6 +31,7 @@ import embodichain.gen_sim.action_engine.runtime.executor as executor_module from embodichain.gen_sim.action_engine.domain import task_contract from embodichain.gen_sim.action_engine.protocol import TASK_SPEC_SCHEMA +from embodichain.gen_sim.action_engine.runtime import load_execution_program from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph from ..task_fixtures import make_task_spec @@ -157,6 +158,28 @@ def test_loaded_e5_propagates_payloads_through_goal_binding_and_claims() -> None ] +def test_standalone_e5_binds_task_owned_payload_without_a_producer() -> None: + task, _requirements = make_task_spec("E5") + task["task_instances"][0]["params"]["payload_roles"] = ["payload"] + + graph = instantiate_seed_graph( + task, + {"object_01": "tray_uid", "payload": "can_uid"}, + ) + node = graph["nodes"][0] + + assert graph["metadata"]["direct_payload_links"] == [] + assert graph["task_groups"][0]["goal"]["payloads"] == [ + {"object": "can_uid", "slot": "center"} + ] + assert node["target_binding"]["payloads"] == [ + {"object": "can_uid", "slot": "center"} + ] + assert "object:can_uid" in { + claim["resource"] for claim in node["contract"]["claims"] + } + + def test_loaded_carrier_graph_and_payload_order_are_deterministic() -> None: task, bindings = _loaded_carrier_task() @@ -234,6 +257,25 @@ def test_payload_infrastructure_does_not_branch_on_task_numbers() -> None: assert literals.isdisjoint(task_numbers), function.__qualname__ +def test_e2_incoming_hold_reads_only_generic_ownership_state() -> None: + holder = recipes_module._incoming_held_arm( + "E2", + "object_uid", + ["anonymous_producer"], + {"anonymous_producer": ("object_uid", "right_arm")}, + ) + tree = ast.parse(dedent(inspect.getsource(recipes_module._incoming_held_arm))) + task_numbers = {f"E{index}" for index in range(1, 10)} + literals = { + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + + assert holder == "right_arm" + assert literals.isdisjoint(task_numbers) + + @pytest.mark.parametrize("task_type", [f"E{index}" for index in range(1, 10)]) def test_standalone_tasks_gain_no_payload_dependency(task_type: str) -> None: task, _requirements = make_task_spec(task_type) @@ -248,6 +290,9 @@ def test_standalone_tasks_gain_no_payload_dependency(task_type: str) -> None: bindings = {role: f"runtime_{role}" for role in role_names} graph = instantiate_seed_graph(task, bindings) + program = load_execution_program(graph) assert graph["task_groups"][0]["depends_on"] == [] assert graph["metadata"]["direct_payload_links"] == [] + assert len(program.semantic_steps) == 1 + assert program.semantic_steps[0].id == "task_01" diff --git a/tests/gen_sim/action_engine/test_architecture.py b/tests/gen_sim/action_engine/test_architecture.py index 689743318..5ecc21393 100644 --- a/tests/gen_sim/action_engine/test_architecture.py +++ b/tests/gen_sim/action_engine/test_architecture.py @@ -153,3 +153,45 @@ def test_runtime_core_has_no_action_name_dispatch_branches() -> None: if duplicated: offenders[relative] = duplicated assert offenders == {} + + +def test_runtime_never_dispatches_on_task_operator_names() -> None: + forbidden = {"handover", "orient_object", "place_relative"} + offenders: list[str] = [] + for relative in ( + "runtime/executor.py", + "runtime/grounding.py", + "runtime/recovery.py", + ): + path = _PACKAGE_ROOT / relative + source = path.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(path)) + for node in ast.walk(tree): + if not isinstance(node, ast.Compare): + continue + segment = ast.get_source_segment(source, node) or "" + literals = { + item.value + for item in ast.walk(node) + if isinstance(item, ast.Constant) and isinstance(item.value, str) + } + if ".operator" in segment and literals & forbidden: + offenders.append(f"{relative}:{node.lineno}") + assert offenders == [] + + +def test_runtime_recovery_contains_no_e_task_identity() -> None: + runtime_root = _PACKAGE_ROOT / "runtime" + offenders = {} + task_types = {f"E{index}" for index in range(1, 10)} + for path in sorted(runtime_root.rglob("*.py")): + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + literals = { + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Constant) and isinstance(node.value, str) + } + duplicated = sorted(literals & task_types) + if duplicated: + offenders[path.relative_to(_PACKAGE_ROOT).as_posix()] = duplicated + assert offenders == {} From ad283d5436e75d9ec0e70a739dcac2bd8a56ca72 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:46:45 +0800 Subject: [PATCH 75/85] fix(action-engine): harden E2 and E5 post-action cleanup safety --- .../action_engine/capabilities/atomic.py | 90 +++++- .../action_engine/capabilities/builtins.py | 43 ++- .../gen_sim/action_engine/runtime/actions.py | 276 ++++++++++++++++- .../gen_sim/action_engine/runtime/executor.py | 27 +- .../action_engine/runtime/grounding.py | 45 ++- .../gen_sim/action_engine/tasks/recipes.py | 101 ++++++- .../action_engine/runtime/test_actions.py | 230 ++++++++++++++ .../runtime/test_runtime_contracts.py | 284 +++++++++++++++++- .../action_engine/tasks/test_factory.py | 29 +- .../tasks/test_interpretation.py | 91 +++++- .../tasks/test_language_decoupling.py | 12 +- .../tasks/test_payload_contracts.py | 7 +- 12 files changed, 1188 insertions(+), 47 deletions(-) diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 92cbf6987..0f483e7b7 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -435,6 +435,7 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: "control_part", "preserve", "joint_state", + verifier_hook=_verify_required_home, contract_resolver_hook=_resolve_joints_contract, ), AtomicCapability( @@ -795,7 +796,9 @@ def _resolve_end_effector_contract( StateAtom( ( "arm_clear" - if binding.get("operation") == "retreat_after_lift" + if binding.get("requires_arm_clear", False) + or binding.get("operation") + in {"reorient_tool_down", "retreat_after_lift"} else "arm_free" ), arm=arm, @@ -1012,20 +1015,83 @@ def _verify_arm_clearance( ) clear = distance >= float(minimum_clearance) role_axis = policy.get("transfer_role_axis") - if role_axis is None: - return attempted & clear - role_axis = torch.as_tensor( - role_axis, - dtype=offset.dtype, - device=offset.device, - ) - if role_axis.ndim == 1: - role_axis = role_axis.unsqueeze(0).repeat(int(executor.env.num_envs), 1) - lateral = torch.sum(offset * role_axis, dim=1) - clear &= lateral >= float(policy.get("minimum_transfer_lateral_clearance", 0.06)) + if role_axis is not None: + role_axis = torch.as_tensor( + role_axis, + dtype=offset.dtype, + device=offset.device, + ) + if role_axis.ndim == 1: + role_axis = role_axis.unsqueeze(0).repeat(int(executor.env.num_envs), 1) + lateral = torch.sum(offset * role_axis, dim=1) + clear &= lateral >= float( + policy.get("minimum_transfer_lateral_clearance", 0.06) + ) + if bool(policy.get("verify_lift_clear", False)): + target = getattr(outcome.grounded.target, "xpos", None) + if not isinstance(target, torch.Tensor): + return torch.zeros_like(attempted) + if target.ndim == 4: + target = target[:, -1] + if target.shape != eef.shape: + return torch.zeros_like(attempted) + target = target.to(dtype=eef.dtype, device=eef.device) + tolerance = float( + policy.get( + "postcondition_tolerance", + executor.runtime_policy.predicate_fallbacks["position_tolerance"], + ) + ) + clear &= ( + torch.linalg.vector_norm( + eef[:, :3, 3] - target[:, :3, 3], + dim=1, + ) + <= tolerance + ) return attempted & clear +def _verify_required_home( + *, + executor: Any, + step: Any, + arm: str, + outcome: Any, + attempted: torch.Tensor, +) -> torch.Tensor: + """Verify required cleanup against the live arm joint state.""" + del step + policy = outcome.grounded.motion_policy + if not bool(policy.get("verify_required_home", False)): + return attempted + env = executor.env + get_part = getattr(env, "get_agent_arm_control_part", None) + if not callable(get_part): + return torch.zeros_like(attempted) + control_part = get_part(arm == "left_arm") + if not isinstance(control_part, str) or not control_part: + return torch.zeros_like(attempted) + target = getattr(outcome.grounded.target, "target", None) + if not isinstance(target, torch.Tensor): + return torch.zeros_like(attempted) + joint_ids = env.robot.get_joint_ids(name=control_part) + current = env.robot.get_qpos()[:, joint_ids] + target = target.to(dtype=current.dtype, device=current.device) + if target.ndim == 1: + target = target.unsqueeze(0).repeat(int(env.num_envs), 1) + if target.shape != current.shape: + return torch.zeros_like(attempted) + tolerance = float( + policy.get( + "postcondition_tolerance", + executor.runtime_policy.predicate_fallbacks["arm_initial_qpos_tolerance"], + ) + ) + reached = torch.all(torch.abs(current - target) <= tolerance, dim=1) + return attempted & reached + + def _actor_arms(actor: Mapping[str, Any]) -> tuple[str, ...]: mode = str(actor.get("mode", "auto")) if mode == "coordinated": diff --git a/embodichain/gen_sim/action_engine/capabilities/builtins.py b/embodichain/gen_sim/action_engine/capabilities/builtins.py index 1af8b9950..4011083c8 100644 --- a/embodichain/gen_sim/action_engine/capabilities/builtins.py +++ b/embodichain/gen_sim/action_engine/capabilities/builtins.py @@ -649,7 +649,7 @@ def _build_coordinated_transport_phases( ) if step["goal"]["terminal_behavior"] != "place": return phases - return phases + ( + release = ( PhaseTemplate( name="dual_release", state_semantic="Both grippers release the transported object", @@ -672,6 +672,47 @@ def _build_coordinated_transport_phases( ), ), ) + lifts = tuple( + PhaseTemplate( + name=f"{side}_lift_clear", + state_semantic=f"The {side} end effector lifts clear of the object", + actions=( + ActionTemplate( + "MoveEndEffector", + { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + "verify_lift_clear": True, + }, + build_motion_policy(), + actor={"mode": "required", "arm": f"{side}_arm"}, + ), + ), + ) + for side in ("left", "right") + ) + homes = tuple( + PhaseTemplate( + name=f"{side}_home", + state_semantic=f"The {side} arm returns to its initial state", + actions=( + ActionTemplate( + "MoveJoints", + { + "kind": "joint_state", + "source": "initial", + "operation": "e5_home", + "required_home": True, + }, + build_motion_policy(), + actor={"mode": "required", "arm": f"{side}_arm"}, + ), + ), + ) + for side in ("left", "right") + ) + return phases + release + lifts + homes def _build_press_phases(step: Mapping[str, Any]) -> tuple[PhaseTemplate, ...]: diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 48cfd30ba..c3c0373dd 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -18,10 +18,13 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager from copy import deepcopy from dataclasses import replace +import logging import math +from threading import RLock from typing import Any import torch @@ -69,7 +72,9 @@ GraspAnnotationCfg, ParallelJawGraspCollisionCfg, ) -from embodichain.utils.logger import log_info +from embodichain.utils import logger as project_logger +from embodichain.utils.logger import log_info, log_warning +from embodichain.utils.math import matrix_from_quat, quat_from_matrix, quat_slerp from .body_grasp import AxisAlignBodyGraspAdapter from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache @@ -104,6 +109,29 @@ _BODY_GRASP_CANDIDATE_LIMIT = 500 _BODY_GRASP_SEED = 17_392 _FREE_YAW_SAMPLE_COUNT = 8 +_RETREAT_LOG_LOCK = RLock() + + +@contextmanager +def _capture_retreat_warnings(enabled: bool) -> Iterator[list[str]]: + """Capture candidate-level planner warnings during bounded retreat search.""" + messages: list[str] = [] + if not enabled: + yield messages + return + collector = logging.Handler(level=logging.WARNING) + collector.emit = lambda record: messages.append(record.getMessage()) + logger = project_logger.logger + with _RETREAT_LOG_LOCK: + handlers = list(logger.handlers) + propagate = logger.propagate + try: + logger.handlers[:] = [collector] + logger.propagate = False + yield messages + finally: + logger.handlers[:] = handlers + logger.propagate = propagate def _collision_cache_for_world( @@ -358,13 +386,14 @@ def plan( capability, ) grounded_candidates = tuple( - candidate + reorientation for coordinated in coordinated_candidates for candidate in self._adapt_axis_align_body_grasps( coordinated, context, capability, ) + for reorientation in self._adapt_tool_down_candidates(candidate) ) selected: ( tuple[ @@ -375,6 +404,9 @@ def plan( ] | None ) = None + selected_warnings: tuple[str, ...] = () + candidate_search_warnings: list[str] = [] + candidate_search_attempts = 0 best_failure_count = self.num_envs + 1 for candidate in grounded_candidates: candidate_engine = self._engine_for(candidate, capability) @@ -383,7 +415,15 @@ def plan( capability, engine=candidate_engine, ) - candidate_plan = candidate_engine.plan(candidate_invocation, context) + capture_warnings = bool( + candidate.motion_policy.get("retreat_reachability_search", False) + or candidate.motion_policy.get("reorient_tool_down", False) + ) + with _capture_retreat_warnings(capture_warnings) as warnings: + candidate_plan = candidate_engine.plan(candidate_invocation, context) + if capture_warnings: + candidate_search_warnings.extend(warnings) + candidate_search_attempts += 1 failure_count = int((~candidate_plan.plan_success).sum().item()) if selected is None or failure_count < best_failure_count: selected = ( @@ -392,12 +432,26 @@ def plan( candidate_plan, candidate_engine, ) + selected_warnings = tuple(warnings) best_failure_count = failure_count if failure_count == 0: break if selected is None: raise RuntimeError("Atomic action adaptation produced no plan candidate.") grounded, invocation, plan, selected_engine = selected + if bool(grounded.motion_policy.get("reorient_tool_down", False)): + summary = ( + "Tool-down reorientation search: " + f"resolved={int(plan.plan_success.sum())}/{plan.plan_success.numel()}, " + f"attempts={candidate_search_attempts}, " + f"selected_yaw_degrees=" + f"{grounded.motion_policy.get('reorient_selected_yaw_degrees')}, " + f"suppressed_warnings={len(candidate_search_warnings)}." + ) + if bool(plan.plan_success.all()): + log_info(summary) + else: + log_warning(summary) selected_positions = self._positions_with_agent_holds( plan, grounded, @@ -419,6 +473,7 @@ def plan( invocation=invocation, initial_positions=selected_positions, initial_success=primary_success, + initial_warnings=selected_warnings, ) invocation = replace(invocation, goal=grounded.target) combined_success = primary_success.clone() @@ -621,6 +676,117 @@ def _adapt_axis_align_body_grasps( ) return tuple(candidates) + def _adapt_tool_down_candidates( + self, + grounded: GroundedAction, + ) -> tuple[GroundedAction, ...]: + """Build fixed-position, downward-TCP yaw candidates after E2 lift-clear.""" + if not bool(grounded.motion_policy.get("reorient_tool_down", False)): + return (grounded,) + reference = grounded.motion_policy.get("reorient_reference_pose") + if not isinstance(reference, torch.Tensor): + raise ValueError("Tool-down reorientation requires a reference TCP pose.") + reference = reference.to(device=self.device, dtype=torch.float32) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + if reference.shape != (self.num_envs, 4, 4): + raise ValueError( + "Tool-down reorientation reference must have shape " + f"({self.num_envs}, 4, 4)." + ) + waypoint_count = int(grounded.cfg.get("reorient_waypoint_count", 5)) + if not 2 <= waypoint_count <= 16: + raise ValueError("reorient_waypoint_count must be in [2, 16].") + raw_yaws = grounded.cfg.get( + "reorient_yaw_degrees", (0.0, 45.0, -45.0, 90.0, -90.0, 180.0) + ) + if not isinstance(raw_yaws, Sequence) or isinstance(raw_yaws, (str, bytes)): + raise TypeError("reorient_yaw_degrees must be a sequence.") + yaw_degrees = [float(value) for value in raw_yaws] + if not yaw_degrees or any(not math.isfinite(value) for value in yaw_degrees): + raise ValueError("reorient_yaw_degrees must contain finite values.") + + base_rotation = self._downward_tcp_rotation(reference[:, :3, :3]) + candidates = [] + for yaw_degrees_value in yaw_degrees: + yaw = math.radians(yaw_degrees_value) + yaw_rotation = reference.new_tensor( + [ + [math.cos(yaw), -math.sin(yaw), 0.0], + [math.sin(yaw), math.cos(yaw), 0.0], + [0.0, 0.0, 1.0], + ] + ) + target_rotation = torch.matmul(yaw_rotation, base_rotation) + waypoints = self._rotation_waypoints( + reference, + target_rotation, + waypoint_count=waypoint_count, + ) + policy = { + **grounded.motion_policy, + "reorient_selected_yaw_degrees": yaw_degrees_value, + } + candidates.append( + replace( + grounded, + target=EndEffectorPoseGoal(xpos=waypoints), + cfg={**grounded.cfg, **policy}, + motion_policy=policy, + ) + ) + return tuple(candidates) + + def _downward_tcp_rotation(self, rotation: torch.Tensor) -> torch.Tensor: + """Project the current TCP heading while aligning local +Z with world -Z.""" + x_axis = rotation[:, :3, 0].clone() + x_axis[:, 2] = 0.0 + norm = torch.linalg.vector_norm(x_axis, dim=1, keepdim=True) + fallback = rotation[:, :3, 1].clone() + fallback[:, 2] = 0.0 + fallback_norm = torch.linalg.vector_norm(fallback, dim=1, keepdim=True) + world_x = rotation.new_tensor([1.0, 0.0, 0.0]).expand_as(x_axis) + fallback = torch.where( + fallback_norm > 1.0e-6, + fallback / fallback_norm.clamp_min(1.0e-6), + world_x, + ) + x_axis = torch.where( + norm > 1.0e-6, + x_axis / norm.clamp_min(1.0e-6), + fallback, + ) + z_axis = rotation.new_tensor([0.0, 0.0, -1.0]).expand_as(x_axis) + y_axis = torch.linalg.cross(z_axis, x_axis, dim=1) + return torch.stack((x_axis, y_axis, z_axis), dim=2) + + def _rotation_waypoints( + self, + reference: torch.Tensor, + target_rotation: torch.Tensor, + *, + waypoint_count: int, + ) -> torch.Tensor: + """Interpolate fixed-position TCP rotations, excluding the observed start.""" + start_quat = quat_from_matrix(reference[:, :3, :3]) + end_quat = quat_from_matrix(target_rotation) + end_quat = torch.where( + torch.sum(start_quat * end_quat, dim=1, keepdim=True) < 0.0, + -end_quat, + end_quat, + ) + poses = reference[:, None].repeat(1, waypoint_count, 1, 1) + for index in range(waypoint_count): + fraction = float(index + 1) / float(waypoint_count) + interpolated = torch.stack( + [ + quat_slerp(start_quat[env_id], end_quat[env_id], tau=fraction) + for env_id in range(self.num_envs) + ] + ) + poses[:, index, :3, :3] = matrix_from_quat(interpolated) + return poses + def _adapt_coordinated_pickment_grasps( self, grounded: GroundedAction, @@ -788,8 +954,9 @@ def _search_reachable_retreat( invocation: ActionInvocation, initial_positions: torch.Tensor, initial_success: torch.Tensor, + initial_warnings: Sequence[str] = (), ) -> tuple[GroundedAction, torch.Tensor, torch.Tensor, dict[str, Any]]: - """Select the highest row-local retreat accepted by the live planner.""" + """Select a row-local retreat candidate accepted by the live planner.""" candidates = self._retreat_search_targets(grounded) target = getattr(grounded.target, "xpos", None) if not isinstance(target, torch.Tensor) or len(candidates) <= 1: @@ -811,10 +978,27 @@ def _search_reachable_retreat( selected_target = candidates[0][1].clone() selected_positions = initial_positions success = initial_success.clone() + suppressed_warnings = list(initial_warnings) + reference = grounded.motion_policy["retreat_reference_pose"].to( + device=selected_target.device, + dtype=selected_target.dtype, + ) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + selected_candidates = [ + candidates[0][0] if bool(success[env_id]) else "unresolved" + for env_id in range(self.num_envs) + ] attempts: list[dict[str, Any]] = [ { "candidate": candidates[0][0], "target_z": candidates[0][1][:, 2, 3].detach().clone(), + "target_distance": torch.linalg.vector_norm( + candidates[0][1][:, :2, 3] - reference[:, :2, 3], + dim=1, + ) + .detach() + .clone(), "success": initial_success.detach().clone(), } ] @@ -835,7 +1019,9 @@ def _search_reachable_retreat( invocation, goal=candidate_grounded.target, ) - candidate_plan = self._engine().plan(candidate_invocation, context) + with _capture_retreat_warnings(True) as warnings: + candidate_plan = self._engine().plan(candidate_invocation, context) + suppressed_warnings.extend(warnings) candidate_positions = self._positions_with_agent_holds( candidate_plan, candidate_grounded, @@ -854,19 +1040,45 @@ def _search_reachable_retreat( candidate_target, selected_target, ) + for env_id in ( + torch.nonzero(selected_rows, as_tuple=False).flatten().tolist() + ): + selected_candidates[env_id] = label success |= candidate_success attempts.append( { "candidate": label, "target_z": candidate_target[:, 2, 3].detach().clone(), + "target_distance": torch.linalg.vector_norm( + candidate_target[:, :2, 3] - reference[:, :2, 3], + dim=1, + ) + .detach() + .clone(), "success": candidate_success.detach().clone(), } ) + selected_distance = torch.linalg.vector_norm( + selected_target[:, :2, 3] - reference[:, :2, 3], dim=1 + ) metadata = { "retreat_selected_target_z": selected_target[:, 2, 3].detach().clone(), + "retreat_selected_target_distance": selected_distance.detach().clone(), "retreat_reachability_found": success.detach().clone(), } + summary = ( + "Retreat reachability search: " + f"mode={grounded.motion_policy.get('retreat_search_mode', 'vertical_then_baseward')}, " + f"resolved={int(success.sum())}/{success.numel()}, " + f"attempts={len(attempts)}, " + f"selected={selected_candidates}, " + f"suppressed_warnings={len(suppressed_warnings)}." + ) + if bool(success.all()): + log_info(summary) + else: + log_warning(summary) selected_grounded = replace( grounded, target=EndEffectorPoseGoal(xpos=selected_target), @@ -881,6 +1093,9 @@ def _search_reachable_retreat( "strategy": "bounded_motion_planner", "attempts": attempts, "selected_target_z": selected_target[:, 2, 3].detach().clone(), + "selected_target_distance": selected_distance.detach().clone(), + "selected_candidates": selected_candidates, + "suppressed_warnings": len(suppressed_warnings), }, ) @@ -908,6 +1123,52 @@ def _retreat_search_targets( sample_count = int(grounded.cfg.get("retreat_search_samples", 6)) if not 2 <= sample_count <= 16: raise ValueError("retreat_search_samples must be in [2, 16].") + search_mode = str( + grounded.motion_policy.get("retreat_search_mode", "vertical_then_baseward") + ) + allowed_modes = {"horizontal_only", "vertical_only", "vertical_then_baseward"} + if search_mode not in allowed_modes: + raise ValueError( + "retreat_search_mode must be 'horizontal_only', 'vertical_only', " + "or 'vertical_then_baseward'." + ) + if search_mode == "horizontal_only": + direction = target[:, :2, 3] - reference[:, :2, 3] + requested_distance = torch.linalg.vector_norm( + direction, dim=1, keepdim=True + ) + if bool((requested_distance <= 1.0e-6).any()): + return [("requested", target.clone())] + direction = direction / requested_distance + minimum_distance = float(grounded.cfg.get("minimum_retreat_distance", 0.05)) + if not math.isfinite(minimum_distance) or minimum_distance < 0.0: + raise ValueError( + "minimum_retreat_distance must be finite and non-negative." + ) + minimum = torch.minimum( + requested_distance[:, 0], + torch.full_like(requested_distance[:, 0], minimum_distance), + ) + fractions = torch.linspace( + 1.0, + 0.0, + sample_count, + dtype=target.dtype, + device=target.device, + ) + distances = ( + minimum[:, None] + + (requested_distance[:, 0] - minimum)[:, None] * fractions[None] + ) + candidates = [("requested", target.clone())] + for index in range(1, sample_count): + candidate = target.clone() + candidate[:, :2, 3] = ( + reference[:, :2, 3] + direction * distances[:, index, None] + ) + candidates.append((f"distance_{index}", candidate)) + return candidates + minimum_height = float(grounded.cfg.get("minimum_retreat_height", 0.05)) if not math.isfinite(minimum_height) or minimum_height < 0.0: raise ValueError("minimum_retreat_height must be finite and non-negative.") @@ -935,6 +1196,9 @@ def _retreat_search_targets( candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] candidates.append((f"height_{index}", candidate)) + if search_mode == "vertical_only": + return candidates + from .frames import arm_base_poses left_base, right_base = arm_base_poses(self.env) diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index 600f813c7..678952852 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -393,11 +393,15 @@ def run( edge.id: self._dependency_failures(edge, edge_failures) for edge in ready } + safety_dependency_failed = { + edge.id: self._safety_dependency_failures(edge, edge_failures) + for edge in ready + } scheduling_blocked = { edge_id: ( failed if self.failure_policy == "stop" - else torch.zeros_like(failed) + else safety_dependency_failed[edge_id] ) for edge_id, failed in dependency_failed.items() } @@ -1541,6 +1545,22 @@ def _dependency_failures( result |= failures[dependency] return result + def _safety_dependency_failures( + self, + edge: ExecutionEdge, + failures: Mapping[str, torch.Tensor], + ) -> torch.Tensor: + """Keep failed safety cleanup prerequisites blocked in continue mode.""" + result = torch.zeros( + int(self.env.num_envs), dtype=torch.bool, device=self.env.device + ) + for dependency in edge.depends_on: + if (dependency, edge.id) in self._completion_only_dependencies: + continue + if self._edge_failure_policy(self.edges[dependency]) == "safety_required": + result |= failures[dependency] + return result + def _completion_only_dependency_edges(self) -> frozenset[tuple[str, str]]: """Resolve linker-added resource ordering to executable edge pairs.""" graph = self.program.seed_graph @@ -2681,7 +2701,10 @@ def _eef_target(self, outcome: ActionOutcome) -> torch.Tensor | None: if held_object is not None: return held_object.grasp_xpos target = outcome.grounded.target - return getattr(target, "xpos", None) + target_pose = getattr(target, "xpos", None) + if isinstance(target_pose, torch.Tensor) and target_pose.dim() == 4: + return target_pose[:, -1] + return target_pose def _state_for(self, step: SemanticStep, arm: str) -> ExecutionState: """Refresh qpos while retaining holds across TaskGroup boundaries.""" diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 37c98bc2d..b9d22ee07 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -75,6 +75,7 @@ __all__ = ["ActionGrounder", "LiveArrangementPlan", "LivePlacementPlan"] _E2_CLEARANCE_RETREAT_DISTANCE = 0.20 +_E2_REORIENT_MINIMUM_CLEARANCE = 0.15 DEFAULT_INTERNAL_AXIS = (0.0, 0.0, 1.0) DEFAULT_TARGET_AXIS = (0.0, 0.0, 1.0) @@ -949,11 +950,35 @@ def ground( arm, reference_eef_pose, ) - if operation == "retreat": + if operation in {"retreat", "lift_clear", "retreat_after_lift"}: policy["retreat_reachability_search"] = True policy["retreat_reference_pose"] = retreat_reference.clone() + if operation == "lift_clear": + policy["retreat_search_mode"] = "vertical_only" + if bool(binding.get("verify_lift_clear", False)): + policy["verify_lift_clear"] = True + policy["minimum_clearance"] = max( + float(policy.get("minimum_clearance", 0.0)), + _E2_REORIENT_MINIMUM_CLEARANCE, + ) + if operation == "reorient_tool_down": + policy["reorient_tool_down"] = True + policy["reorient_reference_pose"] = retreat_reference.clone() + policy["reorient_waypoint_count"] = 5 + policy["reorient_yaw_degrees"] = [ + 0.0, + 45.0, + -45.0, + 90.0, + -90.0, + 180.0, + ] + policy["minimum_clearance"] = _E2_REORIENT_MINIMUM_CLEARANCE if operation == "retreat_after_lift": policy["retreat_distance"] = _E2_CLEARANCE_RETREAT_DISTANCE + policy["minimum_retreat_distance"] = 0.05 + policy["retreat_search_mode"] = "horizontal_only" + policy["retreat_search_samples"] = 4 if source in {"release", "handover"}: policy["clearance_object_uid"] = step.object_uid policy["collision_safety"] = "required" @@ -971,12 +996,16 @@ def ground( device=object_pose.device, ) target = EndEffectorPoseGoal( - xpos=self._retreat_pose( - arm, - policy, - retreat_reference, - clear_exchange=source == "handover", - retreat_after_lift=operation == "retreat_after_lift", + xpos=( + retreat_reference.clone() + if operation == "reorient_tool_down" + else self._retreat_pose( + arm, + policy, + retreat_reference, + clear_exchange=source == "handover", + retreat_after_lift=operation == "retreat_after_lift", + ) ) ) elif kind == "visual_constraint": @@ -993,6 +1022,8 @@ def ground( "cannot resolve a visual_constraint." ) elif kind == "joint_state": + if bool(binding.get("required_home", False)): + policy["verify_required_home"] = True target = JointPositionGoal( target=self._joint_target( arm, diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index c31930c18..0b22dd10e 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -456,7 +456,7 @@ def _recipe( {}, motion_policy(), ) - retreat = _node( + reorient = _node( group_id, 3, "MoveEndEffector", @@ -467,16 +467,53 @@ def _recipe( { "kind": "policy_pose", "source": "release", - "operation": "retreat_after_lift", + "operation": "reorient_tool_down", }, [lift_clear["id"]], "cleanup", {}, motion_policy(("orientation", "upright")), ) - home = _node( + post_reorient_lift = _node( group_id, 4, + "MoveEndEffector", + task_type, + object_uid, + actor, + "arm", + { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + "requires_arm_clear": True, + }, + [reorient["id"]], + "cleanup", + {}, + motion_policy(), + ) + retreat = _node( + group_id, + 5, + "MoveEndEffector", + task_type, + object_uid, + actor, + "arm", + { + "kind": "policy_pose", + "source": "release", + "operation": "retreat_after_lift", + }, + [post_reorient_lift["id"]], + "cleanup", + {}, + motion_policy(("orientation", "upright")), + ) + home = _node( + group_id, + 6, "MoveJoints", task_type, object_uid, @@ -494,7 +531,14 @@ def _recipe( motion_policy(), ) return ( - [alignment, lift_clear, retreat, home], + [ + alignment, + lift_clear, + reorient, + post_reorient_lift, + retreat, + home, + ], "orient_object", goal, success, @@ -957,6 +1001,7 @@ def _recipe( nodes = [pick] if terminal_behavior == "place": release_sync_group = f"{group_id}__dual_release" + releases = [] for index, arm, release_role in ( (2, "left_arm", "participant"), (3, "right_arm", "commit"), @@ -981,6 +1026,54 @@ def _recipe( ) release["sync_group"] = release_sync_group nodes.append(release) + releases.append(release) + release_ids = [release["id"] for release in releases] + lifts = [] + for index, arm in ((4, "left_arm"), (5, "right_arm")): + lift = _node( + group_id, + index, + "MoveEndEffector", + task_type, + object_uid, + {"mode": "required", "arm": arm}, + "arm", + { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + "verify_lift_clear": True, + }, + release_ids, + "cleanup", + {}, + motion_policy(), + ) + nodes.append(lift) + lifts.append(lift) + lift_ids = [lift["id"] for lift in lifts] + for index, arm in ((6, "left_arm"), (7, "right_arm")): + nodes.append( + _node( + group_id, + index, + "MoveJoints", + task_type, + object_uid, + {"mode": "required", "arm": arm}, + "arm", + { + "kind": "joint_state", + "source": "initial", + "operation": "e5_home", + "required_home": True, + }, + lift_ids, + "cleanup", + {}, + motion_policy(), + ) + ) success_type = task_success_type(task_type, params) success = ( {"type": success_type, "object": object_uid} diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 7a49d7dc4..2d03d3e25 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -62,6 +62,7 @@ ) from embodichain.lab.sim.planners import CuroboPlannerCfg, ToppraPlannerCfg from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator +from embodichain.utils.logger import log_warning class _MeshEntity: @@ -926,6 +927,235 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: assert len(search["attempts"]) == len(attempted_targets) +def test_lift_clear_reachability_search_uses_only_vertical_candidates( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.05 + requested = reference.clone() + requested[:, 2, 3] = 1.35 + height_thresholds = torch.tensor([1.24, 1.00]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + height_reachable = target[:, 2, 3] <= height_thresholds + baseward_reachable = target[:, 1, 3] < -0.05 + success = height_reachable | baseward_reachable + terminal = target[:, 2, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "retreat_height": 0.30, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_search_mode": "vertical_only", + "retreat_reference_pose": reference, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert all( + torch.equal(target[:, 1, 3], reference[:, 1, 3]) for target in attempted_targets + ) + assert outcome.success.tolist() == [True, False] + assert outcome.grounded.target.xpos[:, 1, 3].tolist() == pytest.approx([0.0, 0.0]) + + +def test_retreat_after_lift_search_reduces_only_horizontal_distance( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.35 + requested = reference.clone() + requested[:, 1, 3] = -0.20 + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + success = target[:, 1, 3].abs() <= 0.10 + 1.0e-6 + if not bool(success.all()): + log_warning("Synthetic unreachable horizontal retreat candidate.") + terminal = target[:, 1, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "minimum_retreat_distance": 0.05, + "retreat_search_samples": 4, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_search_mode": "horizontal_only", + "retreat_reference_pose": reference, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + attempted_distances = torch.stack( + [ + torch.linalg.vector_norm(target[:, :2, 3] - reference[:, :2, 3], dim=1) + for target in attempted_targets + ] + ) + torch.testing.assert_close( + attempted_distances, + torch.tensor([[0.20, 0.20], [0.15, 0.15], [0.10, 0.10]]), + ) + assert all( + torch.equal(target[:, 2, 3], reference[:, 2, 3]) for target in attempted_targets + ) + assert bool(outcome.success.all()) + search = outcome.planner_trace["reachability_search"] + assert search["selected_candidates"] == ["distance_2", "distance_2"] + assert search["selected_target_distance"].tolist() == pytest.approx([0.10, 0.10]) + assert search["suppressed_warnings"] == 2 + + +def test_tool_down_reorientation_uses_waypoints_and_yaw_candidates( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, :3, :3] = _rotation_x(55.0) + reference[:, :3, 3] = torch.tensor([0.1, -0.2, 1.35]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + success = torch.full((2,), len(attempted_targets) >= 2, dtype=torch.bool) + if not bool(success.all()): + log_warning("Synthetic unreachable tool-down yaw candidate.") + positions = torch.zeros(2, 30, 8) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=reference), + { + "sample_interval": 30, + "reorient_tool_down": True, + "reorient_reference_pose": reference, + "reorient_waypoint_count": 5, + "reorient_yaw_degrees": [0.0, 45.0, -45.0], + }, + motion_policy={ + "reorient_tool_down": True, + "reorient_reference_pose": reference, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert len(attempted_targets) == 2 + for target in attempted_targets: + assert target.shape == (2, 5, 4, 4) + torch.testing.assert_close( + target[:, :, :3, 3], + reference[:, None, :3, 3].expand(-1, 5, -1), + ) + torch.testing.assert_close( + target[:, -1, :3, 2], + torch.tensor([0.0, 0.0, -1.0]).repeat(2, 1), + atol=1.0e-6, + rtol=1.0e-6, + ) + assert bool(outcome.success.all()) + assert outcome.grounded.motion_policy["reorient_selected_yaw_degrees"] == 45.0 + + def test_curobo_generator_receives_generated_static_obstacles( monkeypatch: Any, ) -> None: diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index e7c1eacd6..dc321c53e 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -101,6 +101,7 @@ EndEffectorPoseGoal, HeldObjectPoseGoal, HeldObjectState, + JointPositionGoal, ObjectSemantics, PickUpOptions, PourGoal, @@ -634,6 +635,20 @@ def test_e5_synchronized_release_uses_physics_verified_opening_time() -> None: if action.get("target_binding", {}).get("coordinated_release_role") == "participant" ) + left_lift = next( + action + for edge in program.edges + for action in edge.actions + if action.get("actor", {}).get("arm") == "left_arm" + and action.get("target_binding", {}).get("operation") == "lift_clear" + ) + left_home = next( + action + for edge in program.edges + for action in edge.actions + if action.get("actor", {}).get("arm") == "left_arm" + and action.get("target_binding", {}).get("operation") == "e5_home" + ) grounder = ActionGrounder(program, env, lambda _uid: None) grounded = grounder.ground( @@ -642,8 +657,121 @@ def test_e5_synchronized_release_uses_physics_verified_opening_time() -> None: arm="left_arm", state=_coordinated_held_state(env, entity), ) + lift_grounded = grounder.ground( + left_lift, + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + home_grounded = grounder.ground( + left_home, + step, + arm="left_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) assert grounded.cfg["sample_interval"] == 60 + assert lift_grounded.motion_policy["retreat_search_mode"] == "vertical_only" + assert lift_grounded.motion_policy["verify_lift_clear"] is True + assert home_grounded.motion_policy["verify_required_home"] is True + + +def test_e5_lift_clear_verifier_requires_the_live_tcp_to_reach_its_target() -> None: + entity = _FakeEntity("tray", _pose(0.0, 0.0, 0.75), _box_vertices(0.20)) + env = _FakeEnv({"tray": entity}) + target_pose = _pose(0.0, -0.20, 1.05) + live_pose = target_pose.clone() + env.get_current_xpos_agent = lambda: (live_pose, _pose(0.0, 0.20, 1.05)) + grounded = GroundedAction( + action_class="MoveEndEffector", + arm="left_arm", + control="arm", + target=EndEffectorPoseGoal(xpos=target_pose), + cfg={}, + motion_policy={ + "clearance_object_uid": "tray", + "minimum_clearance": 0.15, + "postcondition_tolerance": 0.05, + "verify_lift_clear": True, + }, + object_uid="tray", + ) + outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.ones(1, dtype=torch.bool), + next_state=ExecutionState(last_qpos=env.robot.get_qpos()), + grounded=grounded, + ) + verifier = build_atomic_capability_registry().get("MoveEndEffector").verifier_hook + assert verifier is not None + executor = SimpleNamespace( + env=env, + runtime_policy=default_runtime_policy("dual_ur10"), + ) + + reached = verifier( + executor=executor, + step=SimpleNamespace(object_uid="tray"), + arm="left_arm", + outcome=outcome, + attempted=torch.ones(1, dtype=torch.bool), + ) + live_pose[:, 2, 3] -= 0.10 + missed = verifier( + executor=executor, + step=SimpleNamespace(object_uid="tray"), + arm="left_arm", + outcome=outcome, + attempted=torch.ones(1, dtype=torch.bool), + ) + + assert reached.tolist() == [True] + assert missed.tolist() == [False] + + +def test_required_home_verifier_checks_live_arm_qpos() -> None: + env = _FakeEnv() + target = torch.tensor([[0.20, -0.20]]) + grounded = GroundedAction( + action_class="MoveJoints", + arm="left_arm", + control="arm", + target=JointPositionGoal(target=target), + cfg={}, + motion_policy={"verify_required_home": True}, + ) + outcome = ActionOutcome( + trajectory=torch.zeros(1, 1, env.robot.dof), + success=torch.ones(1, dtype=torch.bool), + next_state=ExecutionState(last_qpos=env.robot.get_qpos()), + grounded=grounded, + ) + verifier = build_atomic_capability_registry().get("MoveJoints").verifier_hook + assert verifier is not None + executor = SimpleNamespace( + env=env, + runtime_policy=default_runtime_policy("dual_ur10"), + ) + env.robot._qpos[:, env.left_arm_joints] = target + + reached = verifier( + executor=executor, + step=SimpleNamespace(object_uid="tray"), + arm="left_arm", + outcome=outcome, + attempted=torch.ones(1, dtype=torch.bool), + ) + env.robot._qpos[:, env.left_arm_joints] += 0.10 + missed = verifier( + executor=executor, + step=SimpleNamespace(object_uid="tray"), + arm="left_arm", + outcome=outcome, + attempted=torch.ones(1, dtype=torch.bool), + ) + + assert reached.tolist() == [True] + assert missed.tolist() == [False] def test_runtime_policy_discards_legacy_support_z_fallbacks() -> None: @@ -3715,12 +3843,14 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - if candidate.id in orient_step.edge_ids and candidate.actions[0]["atomic_action_class"] == "AxisAlign" ) - orient_lift_edge = next( + orient_lift_edges = [ candidate for candidate in program.edges if candidate.id in orient_step.edge_ids and candidate.actions[0]["target_binding"].get("operation") == "lift_clear" - ) + ] + assert len(orient_lift_edges) == 2 + orient_lift_edge, orient_post_reorient_lift_edge = orient_lift_edges orient_retreat_edge = next( candidate for candidate in program.edges @@ -3728,6 +3858,13 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - and candidate.actions[0]["target_binding"].get("operation") == "retreat_after_lift" ) + orient_reorient_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["target_binding"].get("operation") + == "reorient_tool_down" + ) handover_edge = next( candidate for candidate in program.edges @@ -3760,12 +3897,26 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - state=ExecutionState(last_qpos=env.robot.get_qpos()), reference_eef_pose=release_pose, ) + orient_reorient = grounder.ground( + orient_reorient_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + reference_eef_pose=orient_lift.target.xpos, + ) + orient_post_reorient_lift = grounder.ground( + orient_post_reorient_lift_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + reference_eef_pose=orient_reorient.target.xpos, + ) orient_retreat = grounder.ground( orient_retreat_edge.actions[0], orient_step, arm="right_arm", state=ExecutionState(last_qpos=env.robot.get_qpos()), - reference_eef_pose=orient_lift.target.xpos, + reference_eef_pose=orient_post_reorient_lift.target.xpos, ) handover_pickup = grounder.ground( handover_edge.actions[0], @@ -3784,25 +3935,49 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - release_pose[:, :2, 3], ) assert bool((orient_lift.target.xpos[:, 2, 3] > release_pose[:, 2, 3]).all()) - assert "retreat_reachability_search" not in orient_lift.motion_policy + assert orient_lift.motion_policy["retreat_reachability_search"] is True + assert orient_lift.motion_policy["retreat_search_mode"] == "vertical_only" + assert isinstance(orient_reorient.target, EndEffectorPoseGoal) + assert orient_reorient.motion_policy["reorient_tool_down"] is True + assert orient_reorient.motion_policy["reorient_waypoint_count"] == 5 + assert orient_reorient.motion_policy["minimum_clearance"] == pytest.approx(0.15) + assert isinstance(orient_post_reorient_lift.target, EndEffectorPoseGoal) + assert ( + orient_post_reorient_lift.motion_policy["retreat_search_mode"] + == "vertical_only" + ) + assert bool( + ( + orient_post_reorient_lift.target.xpos[:, 2, 3] + > orient_reorient.target.xpos[:, 2, 3] + ).all() + ) assert isinstance(orient_retreat.target, EndEffectorPoseGoal) + assert orient_retreat.motion_policy["retreat_reachability_search"] is True + assert orient_retreat.motion_policy["retreat_search_mode"] == "horizontal_only" + assert orient_retreat.motion_policy["retreat_search_samples"] == 4 retreat_distance = torch.linalg.vector_norm( - orient_retreat.target.xpos[:, :2, 3] - orient_lift.target.xpos[:, :2, 3], + orient_retreat.target.xpos[:, :2, 3] + - orient_post_reorient_lift.target.xpos[:, :2, 3], dim=1, ) torch.testing.assert_close( retreat_distance, torch.full_like(retreat_distance, 0.20) ) assert bool( - (orient_retreat.target.xpos[:, 0, 3] > orient_lift.target.xpos[:, 0, 3]).all() + ( + orient_retreat.target.xpos[:, 0, 3] + > orient_post_reorient_lift.target.xpos[:, 0, 3] + ).all() ) assert torch.equal( orient_retreat.target.xpos[:, 2, 3], - orient_lift.target.xpos[:, 2, 3], + orient_post_reorient_lift.target.xpos[:, 2, 3], ) assert bool( ( - orient_retreat.target.xpos[:, :2, 3] != orient_lift.target.xpos[:, :2, 3] + orient_retreat.target.xpos[:, :2, 3] + != orient_post_reorient_lift.target.xpos[:, :2, 3] ).any() ) assert orient_alignment.target.grasp_xpos is None @@ -5219,6 +5394,99 @@ def test_continue_failure_policy_executes_downstream_without_clearing_history( assert result.success.tolist() == [True, False] +def test_continue_failure_policy_blocks_failed_safety_cleanup_chain( + monkeypatch: pytest.MonkeyPatch, +) -> None: + task = { + "schema_version": TASK_SPEC_SCHEMA, + "task_id": "safe_orient_cleanup", + "level": "L1", + "instruction": "Stand the can upright.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "orient", + "task_type": "E2", + "params": { + "object_role": "can", + "required_arm": "right_arm", + "orientation_goal": "upright", + "support_role": "table", + }, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "object_upright", "task_instance_id": "orient"}, + "oracle": {}, + "metadata": {}, + } + entities = { + "can": _FakeEntity("can", _pose(0.0, 0.2, 0.75), _box_vertices(0.03)), + "table": _FakeEntity("table", _pose(0.0, 0.0, 0.5), _box_vertices(0.5)), + } + env = _FakeEnv(entities) + program = load_execution_program(instantiate_seed_graph(task, {"can": "can"})) + executor = ProgramExecutor( + program, + env, + settle_steps=0, + record_runtime=False, + failure_policy="continue", + ) + executor.runtime_graph = None + monkeypatch.setattr( + executor, + "_ensure_assignment", + lambda step, failed: executor._assignments.setdefault( + step.id, + [None if bool(failed[0]) else str(step.actor["arm"])], + ), + ) + lift_edge = next( + edge + for edge in program.edges + if edge.actions[0]["target_binding"].get("operation") == "lift_clear" + ) + active_by_edge: dict[str, list[bool]] = {} + + def execute( + edge: ExecutionEdge, + _step: SemanticStep, + *, + failed: torch.Tensor, + ) -> _EdgeResult: + active = ~failed + active_by_edge[edge.id] = active.tolist() + result_failed = failed.clone() + if edge.id == lift_edge.id: + result_failed |= active + return _EdgeResult( + actions=[], + failed=result_failed, + grounded=[], + executed=active, + ) + + monkeypatch.setattr(executor, "_execute_edge", execute) + monkeypatch.setattr( + executor, + "_verify_step", + lambda _step, failed: ( + failed, + ~failed, + torch.zeros(env.num_envs, 3), + ), + ) + + result = executor.run() + + assert active_by_edge[lift_edge.id] == [True] + for edge_id in program.semantic_steps[0].edge_ids[2:]: + assert active_by_edge[edge_id] == [False] + assert not bool(result.success[0]) + + def test_continue_failure_policy_records_failed_then_executed_checkpoints( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index e7a4f3cf6..c40b894a9 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -143,6 +143,8 @@ def test_historical_task2_1_uses_axis_align_and_explicit_handover_arms() -> None "AxisAlign", "MoveEndEffector", "MoveEndEffector", + "MoveEndEffector", + "MoveEndEffector", "MoveJoints", ] assert actions["task_01"] == expected_orient_actions @@ -240,6 +242,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "AxisAlign", "MoveEndEffector", "MoveEndEffector", + "MoveEndEffector", + "MoveEndEffector", "MoveJoints", ] assert orient["goal"]["upright_local_axis"] == "z" @@ -250,6 +254,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "cleanup", "cleanup", "cleanup", + "cleanup", + "cleanup", ] assert orient_nodes[1]["target_binding"] == { "kind": "policy_pose", @@ -260,9 +266,20 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[2]["target_binding"] == { "kind": "policy_pose", "source": "release", - "operation": "retreat_after_lift", + "operation": "reorient_tool_down", } assert orient_nodes[3]["target_binding"] == { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + "requires_arm_clear": True, + } + assert orient_nodes[4]["target_binding"] == { + "kind": "policy_pose", + "source": "release", + "operation": "retreat_after_lift", + } + assert orient_nodes[5]["target_binding"] == { "kind": "joint_state", "source": "initial", "operation": "e2_home", @@ -271,6 +288,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[1]["depends_on"] == [orient_nodes[0]["id"]] assert orient_nodes[2]["depends_on"] == [orient_nodes[1]["id"]] assert orient_nodes[3]["depends_on"] == [orient_nodes[2]["id"]] + assert orient_nodes[4]["depends_on"] == [orient_nodes[3]["id"]] + assert orient_nodes[5]["depends_on"] == [orient_nodes[4]["id"]] assert [node["atomic_action"] for node in handover_nodes] == [ "PickUp", "MoveHeldObject", @@ -289,6 +308,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[0]["contract"]["failure_policy"] == "task_required" assert orient_nodes[1]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[2]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[3]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[4]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[-1]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[1]["contract"]["requires"] == [ {"predicate": "arm_free", "arm": "left_arm"} @@ -296,6 +317,12 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[2]["contract"]["requires"] == [ {"predicate": "arm_clear", "arm": "left_arm"} ] + assert orient_nodes[3]["contract"]["requires"] == [ + {"predicate": "arm_clear", "arm": "left_arm"} + ] + assert orient_nodes[4]["contract"]["requires"] == [ + {"predicate": "arm_clear", "arm": "left_arm"} + ] assert any( effect["atom"]["predicate"] == "arm_home" for effect in orient["contract"]["exit_effects"] diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index 928c339d5..71d4656b8 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -599,8 +599,12 @@ def test_e5_relative_transport_emits_pickment_and_optional_release() -> None: "CoordinatedPickment", "MoveJoints", "MoveJoints", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + "MoveJoints", ] - release_nodes = released["nodes"][1:] + release_nodes = released["nodes"][1:3] assert all( node["depends_on"] == [released["nodes"][0]["id"]] for node in release_nodes ) @@ -638,6 +642,70 @@ def test_e5_relative_transport_emits_pickment_and_optional_release() -> None: ("add", "arm_free", "left_arm"), ("add", "arm_free", "right_arm"), } + lift_nodes = released["nodes"][3:5] + release_ids = {node["id"] for node in release_nodes} + assert {node["actor"]["arm"] for node in lift_nodes} == { + "left_arm", + "right_arm", + } + assert all(node["role"] == "cleanup" for node in lift_nodes) + assert all(node["control"] == "arm" for node in lift_nodes) + assert all(set(node["depends_on"]) == release_ids for node in lift_nodes) + assert all( + node["target_binding"] + == { + "kind": "policy_pose", + "source": "release", + "operation": "lift_clear", + "verify_lift_clear": True, + } + for node in lift_nodes + ) + assert all( + node["contract"]["requires"] + == [{"predicate": "arm_free", "arm": node["actor"]["arm"]}] + for node in lift_nodes + ) + assert all( + node["contract"]["effects"] + == [ + { + "op": "add", + "atom": {"predicate": "arm_clear", "arm": node["actor"]["arm"]}, + } + ] + for node in lift_nodes + ) + assert all( + node["contract"]["failure_policy"] == "safety_required" for node in lift_nodes + ) + home_nodes = released["nodes"][5:] + lift_ids = {node["id"] for node in lift_nodes} + assert {node["actor"]["arm"] for node in home_nodes} == { + "left_arm", + "right_arm", + } + assert all(node["role"] == "cleanup" for node in home_nodes) + assert all(node["control"] == "arm" for node in home_nodes) + assert all(set(node["depends_on"]) == lift_ids for node in home_nodes) + assert all( + node["target_binding"] + == { + "kind": "joint_state", + "source": "initial", + "operation": "e5_home", + "required_home": True, + } + for node in home_nodes + ) + assert all( + node["contract"]["requires"] + == [{"predicate": "arm_clear", "arm": node["actor"]["arm"]}] + for node in home_nodes + ) + assert all( + node["contract"]["failure_policy"] == "safety_required" for node in home_nodes + ) from embodichain.gen_sim.action_engine.runtime import load_execution_program program = load_execution_program(released) @@ -645,8 +713,17 @@ def test_e5_relative_transport_emits_pickment_and_optional_release() -> None: action["atomic_action_class"] for edge in program.edges for action in edge.actions - ] == ["CoordinatedPickment", "MoveJoints", "MoveJoints"] - assert len(program.edges[-1].actions) == 2 + ] == [ + "CoordinatedPickment", + "MoveJoints", + "MoveJoints", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + "MoveJoints", + ] + assert len(program.edges[1].actions) == 2 + assert all(len(edge.actions) == 1 for edge in program.edges[2:]) in_place_spec = deepcopy(released_spec) in_place_params = in_place_spec["task_instances"][0]["params"] @@ -657,6 +734,10 @@ def test_e5_relative_transport_emits_pickment_and_optional_release() -> None: "CoordinatedPickment", "MoveJoints", "MoveJoints", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + "MoveJoints", ] assert "reference_object" not in in_place["task_groups"][0]["goal"] @@ -769,6 +850,10 @@ def test_task1_2_open_reference_generates_coordinated_pick_move_and_release() -> "CoordinatedPickment", "MoveJoints", "MoveJoints", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + "MoveJoints", ] assert grounded.scene_requirements["objects"][0]["category"] == "rigid_object" assert grounded.task_spec["metadata"]["instruction_call_count"] == 1 diff --git a/tests/gen_sim/action_engine/tasks/test_language_decoupling.py b/tests/gen_sim/action_engine/tasks/test_language_decoupling.py index b9f0768d3..519a5e526 100644 --- a/tests/gen_sim/action_engine/tasks/test_language_decoupling.py +++ b/tests/gen_sim/action_engine/tasks/test_language_decoupling.py @@ -356,7 +356,15 @@ def test_unfamiliar_wording_and_categories_flow_through_injected_llm_stages() -> terminal_behavior="place", ), [_binding("move_fixture.object", "aerogel_fixture_7")], - ["CoordinatedPickment", "MoveJoints", "MoveJoints"], + [ + "CoordinatedPickment", + "MoveJoints", + "MoveJoints", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + "MoveJoints", + ], "semantic_goal", ), ( @@ -412,7 +420,7 @@ def test_e5_symbolic_intent_reaches_the_seed_graph( ) assert graph["task_groups"][0]["goal"]["relation"] == "behind" if name == "dual_move_place": - release_nodes = graph["nodes"][1:] + release_nodes = graph["nodes"][1:3] assert {node["actor"]["arm"] for node in release_nodes} == { "left_arm", "right_arm", diff --git a/tests/gen_sim/action_engine/tasks/test_payload_contracts.py b/tests/gen_sim/action_engine/tasks/test_payload_contracts.py index c921bb46c..e1c81fa17 100644 --- a/tests/gen_sim/action_engine/tasks/test_payload_contracts.py +++ b/tests/gen_sim/action_engine/tasks/test_payload_contracts.py @@ -201,8 +201,13 @@ def test_loaded_e5_place_keeps_payload_contract_through_synchronized_release() - "CoordinatedPickment", "MoveJoints", "MoveJoints", + "MoveEndEffector", + "MoveEndEffector", + "MoveJoints", + "MoveJoints", ] - assert len({node.get("sync_group") for node in nodes[1:]}) == 1 + assert len({node.get("sync_group") for node in nodes[1:3]}) == 1 + assert all(node.get("sync_group") is None for node in nodes[3:]) def test_consumer_without_direct_payload_capability_is_rejected() -> None: From d8c5d2c73f52e6f9a39cb335ce9e444f3d4373ae Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Thu, 27 Aug 2026 13:13:32 +0800 Subject: [PATCH 76/85] refactor(action-engine): use stable entity IDs for object semantics --- .../gen_sim/action_engine/runtime/actions.py | 11 ++--- .../action_engine/runtime/grounding.py | 45 ++----------------- .../action_engine/runtime/predicates.py | 12 ++--- .../action_engine/runtime/test_actions.py | 43 +++++++++++++----- .../action_engine/runtime/test_body_grasp.py | 1 + 5 files changed, 43 insertions(+), 69 deletions(-) diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index c3c0373dd..4b96364ad 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -301,7 +301,7 @@ def _build_scene_provider(self) -> SceneProvider | None: ) def semantics(self, uid: str) -> ObjectSemantics: - """Build object semantics once while retaining the live entity handle.""" + """Build object semantics once for the stable scene entity ID.""" cached = self._semantics.get(uid) if cached is not None: return cached @@ -360,7 +360,7 @@ def semantics(self, uid: str) -> ObjectSemantics: semantics = ObjectSemantics( label=uid, - entity=entity, + entity_id=uid, geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, affordance=AntipodalAffordance( object_label=uid, @@ -629,11 +629,6 @@ def _adapt_axis_align_body_grasps( "AxisAlign body grasp requires a finite grounded live object pose " f"with shape ({self.num_envs}, 4, 4)." ) - if goal.semantics.entity_id is not None: - goal = replace( - goal, - semantics=replace(goal.semantics, entity_id=None), - ) _, hand_part, _ = self._parts(grounded.arm) if hand_part is None: raise ValueError("AxisAlign body grasp requires a configured hand part.") @@ -1590,7 +1585,7 @@ def include(uid: str | None, env_mask: torch.Tensor | None = None) -> None: include(target_uid) for held in state.held_objects.values(): - include(held.semantics.label, held.env_mask) + include(held.semantics.entity_id, held.env_mask) collision_exclusion_uids = grounded.motion_policy.get( "collision_exclusion_uids", () ) diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index b9d22ee07..03ecf1f5b 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -58,15 +58,10 @@ TwistAffordance, TwistGoal, ) -from embodichain.toolkits.graspkit.pg_grasp import ( - AntipodalSamplerCfg, - GraspGeneratorCfg, - GripperCollisionCfg, -) from embodichain.utils.logger import log_info + from .frames import arm_base_poses, relation_offset, robot_frame_axes from .geometry_axes import analyze_local_geometry_axes -from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache from .models import ExecutionProgram, GroundedAction, SemanticStep from .motion_policy import resolve_motion_policy, with_motion_modifiers from .robot_parts import arm_control_part @@ -798,9 +793,6 @@ def ground( custom_config=dict(affordance.custom_config), mesh_vertices=affordance.mesh_vertices, mesh_triangles=affordance.mesh_triangles, - generator_cfg=affordance.generator_cfg, - gripper_collision_cfg=affordance.gripper_collision_cfg, - force_reannotate=affordance.force_reannotate, internal_axis=self._upright_local_direction(step), ), ) @@ -2078,9 +2070,6 @@ def _pour_source_semantics( custom_config=dict(affordance.custom_config), mesh_vertices=affordance.mesh_vertices, mesh_triangles=affordance.mesh_triangles, - generator_cfg=affordance.generator_cfg, - gripper_collision_cfg=affordance.gripper_collision_cfg, - force_reannotate=affordance.force_reannotate, internal_axis=local_axes[0], ), ) @@ -2274,40 +2263,10 @@ def _slide_target( f"eef_position={current_eef[0, :3, 3].detach().cpu().tolist()}, " f"arm_base_position={arm_base[0, :3, 3].detach().cpu().tolist()}." ) - grasp_options = self.runtime_policy.grasp - sampler = AntipodalSamplerCfg( - n_sample=int(grasp_options["antipodal_n_sample"]), - max_angle=float(grasp_options["antipodal_max_angle"]), - max_length=float(grasp_options["max_open_length"]), - min_length=float(grasp_options["min_open_length"]), - ) - generator = GraspGeneratorCfg( - viser_port=int(grasp_options["viser_port"]), - antipodal_sampler_cfg=sampler, - max_deviation_angle=float(grasp_options["max_deviation_angle"]), - n_deviated_approach_directions=int( - grasp_options["n_deviated_approach_directions"] - ), - ) - max_hulls = int(grasp_options["max_decomposition_hulls"]) - collision = GripperCollisionCfg( - max_open_length=float(grasp_options["max_open_length"]), - finger_length=float(grasp_options["finger_length"]), - point_sample_dense=float(grasp_options["point_sample_dense"]), - max_decomposition_hulls=max_hulls, - ) - ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=max_hulls, - ) affordance = SlideAffordance( object_label=f"{step.object_uid}:{child_link}", mesh_vertices=vertices, mesh_triangles=triangles, - generator_cfg=generator, - gripper_collision_cfg=collision, - force_reannotate=bool(grasp_options["force_grasp_reannotate"]), translation_axis=push_local[0], joint_name=joint_name, joint_limits=(float(limits[0, 0]), float(limits[0, 1])), @@ -2315,6 +2274,7 @@ def _slide_target( semantics = ObjectSemantics( affordance=affordance, geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + entity_id=f"{step.object_uid}:{child_link}", label=f"{step.object_uid}:{child_link}", ) return ( @@ -2505,6 +2465,7 @@ def _twist_target( semantics = ObjectSemantics( affordance=affordance, geometry={"mesh_vertices": vertices}, + entity_id=f"{step.object_uid}:{child_link}", label=f"{step.object_uid}:{child_link}", ) scoped_policy = dict(policy) diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index f50e131c6..45feef127 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -346,10 +346,8 @@ def _object_held( gripper = gripper_values[arm_index] if held is None or actual_eef is None or gripper is None: continue - label = getattr(held.semantics, "label", None) - if not label and held.semantics.entity is not None: - label = getattr(held.semantics.entity, "uid", None) - if label != uid: + entity_id = getattr(held.semantics, "entity_id", None) + if entity_id != uid: continue actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) expected_eef = torch.bmm( @@ -395,10 +393,8 @@ def _coordinated_held( return result for held in held_relations: assert held is not None - label = getattr(held.semantics, "label", None) - if not label and getattr(held.semantics, "entity", None) is not None: - label = getattr(held.semantics.entity, "uid", None) - if label != uid: + entity_id = getattr(held.semantics, "entity_id", None) + if entity_id != uid: return result eef_values = _arm_values(env, "xpos") gripper_values = _arm_values(env, "gripper_state") diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 2d03d3e25..289960fe7 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -52,6 +52,7 @@ JointPositionGoal, ObjectSemantics, PlannerDiagnostics, + PlanningFailure, RecoveryPolicy, RuntimeCommandFrame, SceneSnapshot, @@ -175,6 +176,11 @@ def _commands_for(trajectory: TimedTrajectory) -> TimedCommandSequence: return TimedCommandSequence(frames=frames, env_ids=trajectory.env_ids) +def _planner_diagnostics(success: torch.Tensor) -> PlannerDiagnostics: + failure = None if bool(success.all()) else PlanningFailure("planning_failed") + return PlannerDiagnostics(backend="fake", failure=failure) + + class _FakeEngine: """Minimal endpoint-binding and planning surface for adapter unit tests.""" @@ -226,7 +232,12 @@ def test_free_yaw_search_uses_an_internal_reachability_sample_set( adapter = AtomicActionAdapter(env) target = torch.eye(4).repeat(2, 1, 1) target[:, :3, 3] = torch.tensor([0.1, -0.2, 0.9]) - semantics = ObjectSemantics(label="can", geometry={}, affordance=Affordance()) + semantics = ObjectSemantics( + label="can", + entity_id="can", + geometry={}, + affordance=Affordance(), + ) held = HeldObjectState( semantics=semantics, object_to_eef=torch.eye(4).repeat(2, 1, 1), @@ -361,7 +372,7 @@ def get_valid_grasp_poses(self, **kwargs: Any): assert len(adapted_items) == 1 assert isinstance(adapted.target, AxisAlignGoal) - assert adapted.target.semantics.entity_id is None + assert adapted.target.semantics.entity_id == "can" assert adapted.target.grasp_xpos is not None assert adapted.motion_policy["body_grasp"]["long_axis_index"] == 1 assert adapted.motion_policy["body_grasp"]["candidate_counts"] == [1, 1] @@ -387,6 +398,7 @@ def test_axis_align_body_grasp_does_not_fall_back_to_scene_snapshot_pose() -> No AxisAlignGoal( semantics=ObjectSemantics( label="can", + entity_id="can", geometry={}, affordance=AxisAlignAffordance( mesh_vertices=vertices, @@ -485,7 +497,7 @@ def fake_affordance(**kwargs: Any) -> Affordance: second = adapter.semantics("cube") assert first is second - assert first.entity_id is None + assert first.entity_id == "cube" assert events == ["cache", "affordance"] assert observed["max_decomposition_hulls"] == 8 assert observed["mesh_vertices"].dtype == torch.float32 @@ -509,6 +521,7 @@ def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None coordinated_goal = CoordinatedPickGoal( semantics=ObjectSemantics( label="tray", + entity_id="tray", geometry={}, affordance=AntipodalAffordance(), ), @@ -637,6 +650,7 @@ def test_coordinated_pickment_uses_engine_scoped_grasp_generator() -> None: goal = CoordinatedPickGoal( semantics=ObjectSemantics( label="tray", + entity_id="tray", geometry={}, affordance=affordance, ), @@ -677,6 +691,7 @@ def _coordinated_grounded( goal = CoordinatedPickGoal( semantics=ObjectSemantics( label="test_object", + entity_id="test_object", geometry={}, affordance=AntipodalAffordance( object_label="test_object", @@ -886,7 +901,7 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: tracking_policy=TrackingPolicy.timed(), planned_scene_version=0, planned_collision_world_revision=(0, 0), - diagnostics=PlannerDiagnostics(backend="fake"), + diagnostics=_planner_diagnostics(success), expected_effects=StateDelta(), ) @@ -961,7 +976,7 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: tracking_policy=TrackingPolicy.timed(), planned_scene_version=0, planned_collision_world_revision=(0, 0), - diagnostics=PlannerDiagnostics(backend="fake"), + diagnostics=_planner_diagnostics(success), expected_effects=StateDelta(), ) @@ -1032,7 +1047,7 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: tracking_policy=TrackingPolicy.timed(), planned_scene_version=0, planned_collision_world_revision=(0, 0), - diagnostics=PlannerDiagnostics(backend="fake"), + diagnostics=_planner_diagnostics(success), expected_effects=StateDelta(), ) @@ -1111,7 +1126,7 @@ def plan(invocation: Any, _context: Any) -> ActionPlan: tracking_policy=TrackingPolicy.timed(), planned_scene_version=0, planned_collision_world_revision=(0, 0), - diagnostics=PlannerDiagnostics(backend="fake"), + diagnostics=_planner_diagnostics(success), expected_effects=StateDelta(), ) @@ -1225,7 +1240,7 @@ def test_dynamic_scene_parks_contact_target_and_held_rows() -> None: ) held_semantics = ObjectSemantics( label="held", - entity=entities["held"], + entity_id="held", geometry={}, affordance=Affordance(), ) @@ -1426,7 +1441,7 @@ def test_retreat_parks_intentional_contact_objects() -> None: def test_action_outcome_commits_state_delta_only_for_verified_rows() -> None: semantics = ObjectSemantics( label="cube", - entity=object(), + entity_id="cube", geometry={}, affordance=Affordance(), ) @@ -1475,7 +1490,7 @@ def test_fallback_rows_keep_the_fallback_plan_effects(monkeypatch: Any) -> None: adapter = AtomicActionAdapter(env) semantics = ObjectSemantics( label="cube", - entity=object(), + entity_id="cube", geometry={}, affordance=Affordance(), ) @@ -1515,6 +1530,9 @@ def action_plan( backend="curobo", messages=messages, metadata={"marker": terminal}, + failure=( + None if bool(success.all()) else PlanningFailure("planning_failed") + ), ), expected_effects=StateDelta( held_object_updates={"physical_left_arm": held} @@ -1610,7 +1628,10 @@ def test_collision_required_cleanup_does_not_use_unsafe_fallback( tracking_policy=TrackingPolicy.timed(), planned_scene_version=1, planned_collision_world_revision=(1, 1), - diagnostics=PlannerDiagnostics(backend="fake"), + diagnostics=PlannerDiagnostics( + backend="fake", + failure=PlanningFailure("planning_failed"), + ), expected_effects=StateDelta(), ) strategies: list[str] = [] diff --git a/tests/gen_sim/action_engine/runtime/test_body_grasp.py b/tests/gen_sim/action_engine/runtime/test_body_grasp.py index 3eb9a8cb0..7c3ec8040 100644 --- a/tests/gen_sim/action_engine/runtime/test_body_grasp.py +++ b/tests/gen_sim/action_engine/runtime/test_body_grasp.py @@ -123,6 +123,7 @@ def test_axis_align_adapter_injects_the_selected_body_grasp_unchanged() -> None: goal = AxisAlignGoal( semantics=ObjectSemantics( label="can", + entity_id="can", geometry={}, affordance=AxisAlignAffordance( mesh_vertices=vertices, From 7def96b71de0b5dcf46a5b5007d33912b0d5fb9e Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Thu, 27 Aug 2026 14:31:13 +0800 Subject: [PATCH 77/85] refactor(action-engine): use stable entity IDs for object semantics --- .../gen_sim/action_engine/runtime/actions.py | 11 - .../runtime/grasp_collision_cache.py | 330 ---------------- .../action_engine/runtime/grounding.py | 45 +-- .../capabilities/test_atomic_v2.py | 2 + .../config/test_runtime_policy.py | 6 +- .../generation/test_generation.py | 4 +- .../action_engine/runtime/test_actions.py | 39 +- .../runtime/test_grasp_collision_cache.py | 354 ------------------ .../runtime/test_runtime_contracts.py | 68 +++- .../tasks/test_interpretation.py | 2 + 10 files changed, 81 insertions(+), 780 deletions(-) delete mode 100644 embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py delete mode 100644 tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 4b96364ad..604dc3245 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -77,7 +77,6 @@ from embodichain.utils.math import matrix_from_quat, quat_from_matrix, quat_slerp from .body_grasp import AxisAlignBodyGraspAdapter -from .grasp_collision_cache import ensure_vhacd_grasp_collision_cache from .models import ActionOutcome, GroundedAction from .state import ExecutionState @@ -348,16 +347,6 @@ def semantics(self, uid: str) -> ObjectSemantics: if triangles.ndim != 2 or triangles.shape[-1] != 3 or triangles.numel() == 0: raise ValueError(f"Object {uid!r} has invalid mesh triangles.") - grasp_options = self.grasp_policy - max_hulls = int(grasp_options["max_decomposition_hulls"]) - cache_result = ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=max_hulls, - ) - if cache_result.status != "hit": - log_info(f"Prepared V-HACD grasp cache for {uid!r}: {cache_result.status}.") - semantics = ObjectSemantics( label=uid, entity_id=uid, diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py b/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py deleted file mode 100644 index 252c388d6..000000000 --- a/embodichain/gen_sim/action_engine/runtime/grasp_collision_cache.py +++ /dev/null @@ -1,330 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Prepare checksummed V-HACD caches for the shared grasp collision checker. - -The sidecar identifies the backend without changing Main's cache key or pickle -payload, so an unlabelled CoACD cache is never silently reused as V-HACD. -""" - -from __future__ import annotations - -from dataclasses import dataclass -import hashlib -import io -import json -import operator -import os -from pathlib import Path -import pickle -import stat -import tempfile -from typing import Literal - -import numpy as np -import torch - -__all__ = [ - "GraspCollisionCacheError", - "GraspCollisionCacheResult", - "ensure_vhacd_grasp_collision_cache", - "grasp_collision_cache_path", -] - -_CACHE_SCHEMA_VERSION = 1 -_METADATA_SUFFIX = ".action_engine.json" -_DEFAULT_CACHE_DIR = ( - Path.home() / ".cache" / "embodichain_cache" / "convex_decomposition" -) - -CacheStatus = Literal["hit", "generated", "replaced"] - - -class GraspCollisionCacheError(RuntimeError): - """Raised when a safe, Main-compatible V-HACD cache cannot be prepared.""" - - -@dataclass(frozen=True) -class GraspCollisionCacheResult: - """Describe the prepared cache files and whether decomposition ran.""" - - status: CacheStatus - cache_path: Path - metadata_path: Path - - -def grasp_collision_cache_path( - mesh_vertices: torch.Tensor | np.ndarray, - mesh_triangles: torch.Tensor | np.ndarray, - max_decomposition_hulls: int, - *, - cache_dir: str | Path | None = None, -) -> Path: - """Return Main's exact ``_.pkl`` cache path.""" - vertices, triangles = _validate_mesh(mesh_vertices, mesh_triangles) - hull_limit = _validate_hull_limit(max_decomposition_hulls) - mesh_hash = hashlib.md5(vertices.tobytes() + triangles.tobytes()).hexdigest() - return _resolve_cache_dir(cache_dir) / f"{mesh_hash}_{hull_limit}.pkl" - - -def ensure_vhacd_grasp_collision_cache( - *, - mesh_vertices: torch.Tensor | np.ndarray, - mesh_triangles: torch.Tensor | np.ndarray, - max_decomposition_hulls: int, - cache_dir: str | Path | None = None, -) -> GraspCollisionCacheResult: - """Create or validate a V-HACD cache and its checksummed backend sidecar.""" - vertices, triangles = _validate_mesh(mesh_vertices, mesh_triangles) - hull_limit = _validate_hull_limit(max_decomposition_hulls) - mesh_hash = hashlib.md5(vertices.tobytes() + triangles.tobytes()).hexdigest() - cache_path = _resolve_cache_dir(cache_dir) / f"{mesh_hash}_{hull_limit}.pkl" - metadata_path = cache_path.with_name(f"{cache_path.name}{_METADATA_SUFFIX}") - expected_metadata: dict[str, object] = { - "schema_version": _CACHE_SCHEMA_VERSION, - "backend": "vhacd", - "mesh_hash": mesh_hash, - "max_decomposition_hulls": hull_limit, - } - - _prepare_private_directory(cache_path.parent) - _refuse_symlink(cache_path) - _refuse_symlink(metadata_path) - if _cache_matches_metadata(cache_path, metadata_path, expected_metadata): - return GraspCollisionCacheResult("hit", cache_path, metadata_path) - - exists = cache_path.exists() or metadata_path.exists() - status: CacheStatus = "replaced" if exists else "generated" - try: - plane_equations = _compute_vhacd_plane_equations( - vertices, - triangles, - hull_limit, - ) - cache_bytes = _serialize_checker_payload(plane_equations) - metadata = { - **expected_metadata, - "cache_sha256": hashlib.sha256(cache_bytes).hexdigest(), - } - - # Publish the complete pickle before its sidecar. A crash between the - # two replaces leaves a cache miss on retry, never a partial pickle. - _write_bytes_atomic(cache_path, cache_bytes) - metadata_bytes = (json.dumps(metadata, sort_keys=True) + "\n").encode() - _write_bytes_atomic(metadata_path, metadata_bytes) - except GraspCollisionCacheError: - raise - except Exception as exc: - raise GraspCollisionCacheError( - f"Failed to prepare V-HACD grasp collision cache {cache_path}: {exc}" - ) from exc - - return GraspCollisionCacheResult(status, cache_path, metadata_path) - - -def _compute_vhacd_plane_equations( - vertices: np.ndarray, - triangles: np.ndarray, - max_decomposition_hulls: int, -) -> list[tuple[np.ndarray, np.ndarray]]: - """Run DexSim V-HACD and convert its hulls to checker plane equations.""" - import open3d as o3d - from dexsim.kit.meshproc import convex_decomposition_vhacd - - from embodichain.toolkits.graspkit.pg_grasp.collision_checker import ( - extract_plane_equations, - ) - - mesh = o3d.t.geometry.TriangleMesh() - mesh.vertex.positions = o3d.core.Tensor(vertices.astype(np.float32, copy=False)) - mesh.triangle.indices = o3d.core.Tensor(triangles.astype(np.int32, copy=False)) - is_success, hull_meshes = convex_decomposition_vhacd( - mesh, - max_convex_hull_num=max_decomposition_hulls, - ) - if not is_success or not hull_meshes: - raise GraspCollisionCacheError( - "V-HACD returned no convex hulls for the grasp collision mesh." - ) - - convex_parts = [ - ( - np.asarray(hull.vertex.positions.numpy()), - np.asarray(hull.triangle.indices.numpy()), - ) - for hull in hull_meshes - ] - plane_equations = extract_plane_equations(convex_parts) - if not plane_equations: - raise GraspCollisionCacheError( - "V-HACD hulls produced no grasp collision plane equations." - ) - return plane_equations - - -def _serialize_checker_payload( - plane_equations: list[tuple[np.ndarray, np.ndarray]], -) -> bytes: - """Pack plane equations in the exact tensor dictionary Main unpickles.""" - if not plane_equations: - raise ValueError("V-HACD must produce at least one convex hull.") - - normalized: list[tuple[np.ndarray, np.ndarray]] = [] - for normals_value, offsets_value in plane_equations: - normals = np.asarray(normals_value, dtype=np.float32) - offsets = np.asarray(offsets_value, dtype=np.float32) - if normals.ndim != 2 or normals.shape[1:] != (3,) or not len(normals): - raise ValueError("Each V-HACD hull must have normals shaped [K, 3].") - if offsets.shape != (len(normals),): - raise ValueError("Each hull needs one offset per plane normal.") - if not np.isfinite(normals).all() or not np.isfinite(offsets).all(): - raise ValueError("V-HACD plane equations must contain finite values.") - normalized.append((normals, offsets)) - - max_plane_count = max(normals.shape[0] for normals, _ in normalized) - equations = torch.zeros((len(normalized), max_plane_count, 4)) - counts = torch.zeros(len(normalized), dtype=torch.int32) - for index, (normals, offsets) in enumerate(normalized): - plane_count = normals.shape[0] - equations[index, :plane_count, :3] = torch.from_numpy(normals) - equations[index, :plane_count, 3] = torch.from_numpy(offsets) - counts[index] = plane_count - - stream = io.BytesIO() - payload = {"plane_equations": equations, "plane_equation_counts": counts} - pickle.dump(payload, stream, protocol=pickle.HIGHEST_PROTOCOL) - return stream.getvalue() - - -def _cache_matches_metadata( - cache_path: Path, - metadata_path: Path, - expected_metadata: dict[str, object], -) -> bool: - if not cache_path.is_file() or not metadata_path.is_file(): - return False - try: - metadata = json.loads(metadata_path.read_text(encoding="utf-8")) - checksum = metadata.get("cache_sha256") - expected_checksum = hashlib.sha256(cache_path.read_bytes()).hexdigest() - return ( - isinstance(metadata, dict) - and all( - metadata.get(key) == value for key, value in expected_metadata.items() - ) - and isinstance(checksum, str) - and checksum == expected_checksum - ) - except (AttributeError, OSError, UnicodeDecodeError, json.JSONDecodeError): - return False - - -def _validate_mesh( - mesh_vertices: torch.Tensor | np.ndarray, - mesh_triangles: torch.Tensor | np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - if isinstance(mesh_vertices, torch.Tensor): - mesh_vertices = mesh_vertices.detach().cpu().numpy() - if isinstance(mesh_triangles, torch.Tensor): - mesh_triangles = mesh_triangles.detach().cpu().numpy() - if not isinstance(mesh_vertices, np.ndarray): - raise TypeError("mesh_vertices must be a torch.Tensor or numpy.ndarray.") - if not isinstance(mesh_triangles, np.ndarray): - raise TypeError("mesh_triangles must be a torch.Tensor or numpy.ndarray.") - vertices = np.ascontiguousarray(mesh_vertices) - triangles = np.ascontiguousarray(mesh_triangles) - if vertices.ndim != 2 or vertices.shape[1:] != (3,) or len(vertices) == 0: - raise ValueError("mesh_vertices must have non-empty shape [N, 3].") - if triangles.ndim != 2 or triangles.shape[1:] != (3,) or len(triangles) == 0: - raise ValueError("mesh_triangles must have non-empty shape [M, 3].") - if not np.issubdtype(vertices.dtype, np.number): - raise TypeError("mesh_vertices must contain numeric values.") - if not np.isfinite(vertices).all(): - raise ValueError("mesh_vertices must contain only finite values.") - if not np.issubdtype(triangles.dtype, np.integer): - raise TypeError("mesh_triangles must contain integer indices.") - if triangles.min() < 0 or triangles.max() >= len(vertices): - raise ValueError("mesh_triangles contains out-of-range vertex indices.") - return vertices, triangles - - -def _validate_hull_limit(value: int) -> int: - if isinstance(value, (bool, np.bool_)): - raise TypeError("max_decomposition_hulls must be an integer.") - try: - hull_limit = operator.index(value) - except TypeError as exc: - raise TypeError("max_decomposition_hulls must be an integer.") from exc - if hull_limit <= 0: - raise ValueError("max_decomposition_hulls must be positive.") - return hull_limit - - -def _resolve_cache_dir(cache_dir: str | Path | None) -> Path: - if cache_dir is not None: - return Path(cache_dir).expanduser().resolve() - try: - from embodichain.lab.sim import CONVEX_DECOMP_DIR - except Exception: - return _DEFAULT_CACHE_DIR - return Path(CONVEX_DECOMP_DIR).expanduser().resolve() - - -def _prepare_private_directory(path: Path) -> None: - try: - path.mkdir(parents=True, exist_ok=True, mode=0o700) - path.chmod(0o700) - except OSError as exc: - raise GraspCollisionCacheError( - f"Cannot secure grasp collision cache directory: {path}" - ) from exc - if path.stat().st_mode & (stat.S_IWGRP | stat.S_IWOTH): - raise GraspCollisionCacheError(f"Refusing writable cache directory: {path}") - - -def _refuse_symlink(path: Path) -> None: - if path.is_symlink(): - raise GraspCollisionCacheError( - f"Refusing symlinked grasp collision cache path: {path}" - ) - - -def _write_bytes_atomic(path: Path, payload: bytes) -> None: - """Publish one complete file with a same-directory atomic replacement.""" - _refuse_symlink(path) - file_descriptor, temporary_name = tempfile.mkstemp( - prefix=f".{path.name}.", - suffix=".tmp", - dir=path.parent, - ) - temporary_path = Path(temporary_name) - try: - os.fchmod(file_descriptor, 0o600) - with os.fdopen(file_descriptor, "wb") as output: - file_descriptor = -1 - output.write(payload) - output.flush() - os.fsync(output.fileno()) - _refuse_symlink(path) - os.replace(temporary_path, path) - path.chmod(0o600) - finally: - if file_descriptor >= 0: - os.close(file_descriptor) - try: - temporary_path.unlink(missing_ok=True) - except OSError: - pass diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 03ecf1f5b..b0f8e8c24 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -2152,24 +2152,8 @@ def _slide_target( "Slide grounding requires exactly one fixed contact endpoint " f"on prismatic child link {child_link!r}; found {contact_links}." ) - grasp_pose = _batched_pose( - articulation.get_link_pose(contact_links[0], to_matrix=True), self.env - ) - # The bundled drawer contact frame owns TCP +Z as its approach axis. - # A local quarter turn aligns the parallel-gripper closing direction - # with the narrow handle, matching scripts/tutorials/sim/open_drawer.py. - grasp_roll = torch.eye( - 4, - dtype=grasp_pose.dtype, - device=grasp_pose.device, - ) - grasp_roll[0, 0] = 0.0 - grasp_roll[0, 1] = -1.0 - grasp_roll[1, 0] = 1.0 - grasp_roll[1, 1] = 0.0 - grasp_pose = torch.matmul(grasp_pose, grasp_roll) - - vertices, triangles = articulation.get_link_vert_face(child_link) + contact_link = contact_links[0] + vertices, triangles = articulation.get_link_vert_face(contact_link) vertices = torch.as_tensor( vertices, dtype=torch.float32, device=self.env.device ) @@ -2177,12 +2161,12 @@ def _slide_target( triangles, dtype=torch.int64, device=self.env.device ) if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): - raise ValueError("Slide child link has no valid grasp geometry.") + raise ValueError("Slide contact link has no valid grasp geometry.") if triangles.ndim != 2 or triangles.shape[-1] != 3 or not triangles.numel(): - raise ValueError("Slide child link has no valid triangle geometry.") + raise ValueError("Slide contact link has no valid triangle geometry.") - child_pose = _batched_pose( - articulation.get_link_pose(child_link, to_matrix=True), self.env + contact_pose = _batched_pose( + articulation.get_link_pose(contact_link, to_matrix=True), self.env ) parent_pose = _batched_pose( articulation.get_link_pose(parent_link, to_matrix=True), self.env @@ -2204,7 +2188,7 @@ def _slide_target( ) push_world = -torch.nn.functional.normalize(opening_world, dim=1) push_local = torch.bmm( - child_pose[:, :3, :3].transpose(1, 2), + contact_pose[:, :3, :3].transpose(1, 2), push_world.unsqueeze(2), ).squeeze(2) push_local = torch.nn.functional.normalize(push_local, dim=1) @@ -2249,7 +2233,6 @@ def _slide_target( "articulation_joint_id": joint_id, "articulation_initial_qpos": qpos, "articulation_target_qpos": target_qpos, - "articulation_grasp_position": grasp_pose[:, :3, 3], "articulation_push_axis_world": push_world, } ) @@ -2257,14 +2240,15 @@ def _slide_target( arm_base = left_base if arm == "left_arm" else right_base current_eef = self._current_eef_pose(arm) log_info( - f"Slide grounding {step.id}/{arm}: grasp_position=" - f"{grasp_pose[0, :3, 3].detach().cpu().tolist()}, " + f"Slide grounding {step.id}/{arm}: contact_link={contact_link!r}, " + f"contact_position={contact_pose[0, :3, 3].detach().cpu().tolist()}, " f"push_axis={push_world[0].detach().cpu().tolist()}, " f"eef_position={current_eef[0, :3, 3].detach().cpu().tolist()}, " f"arm_base_position={arm_base[0, :3, 3].detach().cpu().tolist()}." ) + contact_entity_id = f"{step.object_uid}:{contact_link}" affordance = SlideAffordance( - object_label=f"{step.object_uid}:{child_link}", + object_label=contact_entity_id, mesh_vertices=vertices, mesh_triangles=triangles, translation_axis=push_local[0], @@ -2274,14 +2258,13 @@ def _slide_target( semantics = ObjectSemantics( affordance=affordance, geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, - entity_id=f"{step.object_uid}:{child_link}", - label=f"{step.object_uid}:{child_link}", + entity_id=contact_entity_id, + label=contact_entity_id, ) return ( SlideGoal( semantics=semantics, - target_pose=child_pose, - grasp_xpos=grasp_pose, + target_pose=contact_pose, ), scoped_policy, ) diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py index 09e00b6d4..276c815c3 100644 --- a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -45,6 +45,7 @@ RuntimeCommandFrame, TimedTrajectory, TimedCommandSequence, + TrackingPolicy, ) @@ -91,6 +92,7 @@ def plan(self, invocation, context): ), joint_trajectory=trajectory, recovery_policy=invocation.recovery_policy, + tracking_policy=TrackingPolicy.timed(), planned_scene_version=context.scene.version, planned_collision_world_revision=(0,) * context.batch_size, diagnostics=PlannerDiagnostics(backend="test"), diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py index a9b33084d..bbdfc7e0d 100644 --- a/tests/gen_sim/action_engine/config/test_runtime_policy.py +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -128,7 +128,7 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: assert generation["environment"]["recording"] == { "enabled": True, "resolution": [640, 360], - "interval_step": 1, + "interval_step": 5, } assert generation["scene"]["object_length_sample_points"] == 5000 assert generation["dataset"]["control_frequency"] == 25 @@ -163,14 +163,14 @@ def test_e1_motion_defaults_match_atomic_action_tutorial_cadence() -> None: } -def test_axis_align_defaults_match_atomic_action_tutorial_cadence() -> None: +def test_axis_align_defaults_preserve_action_engine_clearance_policy() -> None: axis_align = default_runtime_policy("dual_franka").motion_defaults["AxisAlign"] assert axis_align == { "sample_interval": 180, "pre_grasp_distance": pytest.approx(0.15), "lift_height": pytest.approx(0.16), - "lower_distance": pytest.approx(0.03), + "lower_distance": pytest.approx(0.16), "hand_interp_steps": 12, } diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py index 56fc67564..b425bea15 100644 --- a/tests/gen_sim/action_engine/generation/test_generation.py +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -378,7 +378,7 @@ def test_fast_gym_config_has_runnable_franka_contract(gym_export: Path) -> None: 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["interval_step"] == 5 assert recorder["params"]["resolution"] == [640, 360] assert recorder["params"]["intrinsics"] == pytest.approx( [280.0, 280.0, 320.0, 180.0] @@ -433,7 +433,7 @@ def test_offline_recording_can_be_disabled_but_ab_keeps_audience_recorder( assert ab["env"]["events"]["record_camera"]["params"]["name"] == ( "record_cam_audience_view" ) - assert ab["env"]["events"]["record_camera"]["interval_step"] == 1 + assert ab["env"]["events"]["record_camera"]["interval_step"] == 5 @pytest.mark.parametrize( diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 289960fe7..c918cfe9d 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -459,12 +459,7 @@ def _planner_env( ) -def test_semantics_prewarms_vhacd_cache_before_affordance( - monkeypatch: Any, -) -> None: - """The lazy shared checker must see V-HACD's pickle, never create CoACD.""" - events: list[str] = [] - observed: dict[str, Any] = {} +def test_semantics_builds_one_mesh_only_affordance() -> None: entity = _MeshEntity() env = SimpleNamespace( num_envs=1, @@ -472,25 +467,7 @@ def test_semantics_prewarms_vhacd_cache_before_affordance( sim=SimpleNamespace( get_rigid_object=lambda uid: entity if uid == "cube" else None ), - agent_grasp_runtime_defaults={"max_decomposition_hulls": 8}, - ) - - def fake_prepare(**kwargs: Any) -> SimpleNamespace: - events.append("cache") - observed.update(kwargs) - return SimpleNamespace(status="hit") - - def fake_affordance(**kwargs: Any) -> Affordance: - events.append("affordance") - observed["affordance_kwargs"] = kwargs - return Affordance() - - monkeypatch.setattr( - actions, - "ensure_vhacd_grasp_collision_cache", - fake_prepare, ) - monkeypatch.setattr(actions, "AntipodalAffordance", fake_affordance) adapter = AtomicActionAdapter(env) first = adapter.semantics("cube") @@ -498,15 +475,11 @@ def fake_affordance(**kwargs: Any) -> Affordance: assert first is second assert first.entity_id == "cube" - assert events == ["cache", "affordance"] - assert observed["max_decomposition_hulls"] == 8 - assert observed["mesh_vertices"].dtype == torch.float32 - assert observed["mesh_triangles"].dtype == torch.int64 - assert set(observed["affordance_kwargs"]) == { - "object_label", - "mesh_vertices", - "mesh_triangles", - } + assert isinstance(first.affordance, AntipodalAffordance) + assert first.affordance.mesh_vertices is not None + assert first.affordance.mesh_triangles is not None + assert first.affordance.mesh_vertices.dtype == torch.float32 + assert first.affordance.mesh_triangles.dtype == torch.int64 def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None: diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py b/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py deleted file mode 100644 index a08771054..000000000 --- a/tests/gen_sim/action_engine/runtime/test_grasp_collision_cache.py +++ /dev/null @@ -1,354 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import hashlib -import json -import os -from pathlib import Path -import pickle -from typing import Callable - -import numpy as np -import pytest -import torch - -from embodichain.gen_sim.action_engine.runtime import grasp_collision_cache -from embodichain.gen_sim.action_engine.runtime.grasp_collision_cache import ( - GraspCollisionCacheError, - ensure_vhacd_grasp_collision_cache, - grasp_collision_cache_path, -) - - -def _tetrahedron() -> tuple[torch.Tensor, torch.Tensor]: - vertices = torch.tensor( - [ - [0.0, 0.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 1.0, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=torch.float32, - ) - triangles = torch.tensor( - [ - [0, 2, 1], - [0, 1, 3], - [0, 3, 2], - [1, 2, 3], - ], - dtype=torch.int64, - ) - return vertices, triangles - - -def _plane_equations() -> list[tuple[np.ndarray, np.ndarray]]: - return [ - ( - np.asarray( - [ - [1.0, 0.0, 0.0], - [0.0, 1.0, 0.0], - [0.0, 0.0, 1.0], - ], - dtype=np.float32, - ), - np.asarray([-1.0, -1.0, -1.0], dtype=np.float32), - ), - ( - np.asarray([[1.0, 1.0, 1.0]], dtype=np.float32), - np.asarray([-1.0], dtype=np.float32), - ), - ] - - -def _install_fake_decomposer( - monkeypatch: pytest.MonkeyPatch, -) -> list[tuple[tuple[int, ...], tuple[int, ...], int]]: - calls: list[tuple[tuple[int, ...], tuple[int, ...], int]] = [] - - def fake_decompose( - vertices: np.ndarray, - triangles: np.ndarray, - max_decomposition_hulls: int, - ) -> list[tuple[np.ndarray, np.ndarray]]: - calls.append( - ( - tuple(vertices.shape), - tuple(triangles.shape), - max_decomposition_hulls, - ) - ) - return _plane_equations() - - monkeypatch.setattr( - grasp_collision_cache, - "_compute_vhacd_plane_equations", - fake_decompose, - ) - return calls - - -def test_cache_key_and_payload_match_main_checker_contract( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - vertices, triangles = _tetrahedron() - _install_fake_decomposer(monkeypatch) - - result = ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=16, - cache_dir=tmp_path, - ) - - expected_hash = hashlib.md5( - vertices.numpy().tobytes() + triangles.numpy().tobytes() - ).hexdigest() - assert result.cache_path == tmp_path / f"{expected_hash}_16.pkl" - with result.cache_path.open("rb") as cache_file: - payload = pickle.load(cache_file) - assert set(payload) == {"plane_equations", "plane_equation_counts"} - assert payload["plane_equations"].shape == (2, 3, 4) - assert payload["plane_equations"].dtype == torch.float32 - assert payload["plane_equation_counts"].tolist() == [3, 1] - assert payload["plane_equation_counts"].dtype == torch.int32 - - -def test_main_checker_loads_prepared_cache_without_running_coacd( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - import embodichain.lab.sim - from embodichain.toolkits.graspkit.pg_grasp import collision_checker - - vertices, triangles = _tetrahedron() - _install_fake_decomposer(monkeypatch) - result = ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=16, - cache_dir=tmp_path, - ) - - def fail_coacd(*args: object, **kwargs: object) -> None: - raise AssertionError("The prepared V-HACD cache must bypass CoACD.") - - monkeypatch.setattr(embodichain.lab.sim, "CONVEX_DECOMP_DIR", tmp_path) - monkeypatch.setattr(collision_checker, "convex_decomposition_coacd", fail_coacd) - checker = collision_checker.ConvexCollisionChecker( - vertices, - triangles, - max_decomposition_hulls=16, - ) - - assert checker.cache_path == result.cache_path.as_posix() - assert checker.plane_equations["plane_equation_counts"].tolist() == [3, 1] - - -def test_matching_vhacd_metadata_returns_cache_hit( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - vertices, triangles = _tetrahedron() - calls = _install_fake_decomposer(monkeypatch) - kwargs = { - "mesh_vertices": vertices, - "mesh_triangles": triangles, - "max_decomposition_hulls": 16, - "cache_dir": tmp_path, - } - - first = ensure_vhacd_grasp_collision_cache(**kwargs) - second = ensure_vhacd_grasp_collision_cache(**kwargs) - - assert first.status == "generated" - assert second.status == "hit" - assert len(calls) == 1 - - -def test_non_vhacd_metadata_forces_cache_replacement( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - vertices, triangles = _tetrahedron() - calls = _install_fake_decomposer(monkeypatch) - kwargs = { - "mesh_vertices": vertices, - "mesh_triangles": triangles, - "max_decomposition_hulls": 16, - "cache_dir": tmp_path, - } - first = ensure_vhacd_grasp_collision_cache(**kwargs) - metadata = json.loads(first.metadata_path.read_text(encoding="utf-8")) - metadata["backend"] = "coacd" - first.metadata_path.write_text(json.dumps(metadata), encoding="utf-8") - - replaced = ensure_vhacd_grasp_collision_cache(**kwargs) - - assert replaced.status == "replaced" - assert len(calls) == 2 - repaired = json.loads(replaced.metadata_path.read_text(encoding="utf-8")) - assert repaired["backend"] == "vhacd" - - -def test_modified_cache_fails_checksum_and_is_rebuilt( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - vertices, triangles = _tetrahedron() - calls = _install_fake_decomposer(monkeypatch) - kwargs = { - "mesh_vertices": vertices, - "mesh_triangles": triangles, - "max_decomposition_hulls": 16, - "cache_dir": tmp_path, - } - first = ensure_vhacd_grasp_collision_cache(**kwargs) - first.cache_path.write_bytes(b"not a valid collision cache") - - replaced = ensure_vhacd_grasp_collision_cache(**kwargs) - - assert replaced.status == "replaced" - assert len(calls) == 2 - with replaced.cache_path.open("rb") as cache_file: - assert "plane_equations" in pickle.load(cache_file) - - -def test_cache_and_metadata_are_published_by_atomic_replace( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - vertices, triangles = _tetrahedron() - _install_fake_decomposer(monkeypatch) - replacements: list[tuple[Path, Path]] = [] - real_replace: Callable[[os.PathLike[str], os.PathLike[str]], None] = os.replace - - def recording_replace( - source: os.PathLike[str], - destination: os.PathLike[str], - ) -> None: - replacements.append((Path(source), Path(destination))) - real_replace(source, destination) - - monkeypatch.setattr(grasp_collision_cache.os, "replace", recording_replace) - - result = ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=16, - cache_dir=tmp_path, - ) - - assert [destination for _, destination in replacements] == [ - result.cache_path, - result.metadata_path, - ] - assert all( - source.parent == destination.parent for source, destination in replacements - ) - assert all(not source.exists() for source, _ in replacements) - - -@pytest.mark.parametrize( - ("vertices", "triangles", "message"), - [ - ( - torch.empty((0, 3), dtype=torch.float32), - torch.tensor([[0, 1, 2]], dtype=torch.int64), - "mesh_vertices", - ), - ( - torch.zeros((3, 3), dtype=torch.float32), - torch.tensor([[0, 1]], dtype=torch.int64), - "mesh_triangles", - ), - ( - torch.tensor([[0.0, 0.0, 0.0], [1.0, float("nan"), 0.0], [0.0, 1.0, 0.0]]), - torch.tensor([[0, 1, 2]], dtype=torch.int64), - "finite", - ), - ( - torch.zeros((3, 3), dtype=torch.float32), - torch.tensor([[0, 1, 3]], dtype=torch.int64), - "indices", - ), - ], -) -def test_invalid_mesh_is_rejected_before_decomposition( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, - vertices: torch.Tensor, - triangles: torch.Tensor, - message: str, -) -> None: - calls = _install_fake_decomposer(monkeypatch) - - with pytest.raises(ValueError, match=message): - ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=16, - cache_dir=tmp_path, - ) - - assert calls == [] - - -@pytest.mark.parametrize("max_decomposition_hulls", [True, 0, -1, 1.5]) -def test_invalid_hull_limit_is_rejected( - tmp_path: Path, - max_decomposition_hulls: object, -) -> None: - vertices, triangles = _tetrahedron() - - with pytest.raises((TypeError, ValueError), match="max_decomposition_hulls"): - ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=max_decomposition_hulls, # type: ignore[arg-type] - cache_dir=tmp_path, - ) - - -def test_symlinked_cache_path_is_refused( - monkeypatch: pytest.MonkeyPatch, - tmp_path: Path, -) -> None: - vertices, triangles = _tetrahedron() - _install_fake_decomposer(monkeypatch) - cache_path = grasp_collision_cache_path( - vertices, - triangles, - 16, - cache_dir=tmp_path, - ) - victim = tmp_path / "victim.pkl" - victim.write_bytes(b"do not overwrite") - cache_path.symlink_to(victim) - - with pytest.raises(GraspCollisionCacheError, match="symlink"): - ensure_vhacd_grasp_collision_cache( - mesh_vertices=vertices, - mesh_triangles=triangles, - max_decomposition_hulls=16, - cache_dir=tmp_path, - ) - - assert victim.read_bytes() == b"do not overwrite" diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index dc321c53e..55af776db 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -191,7 +191,12 @@ def __init__(self, uid: str, qpos: float) -> None: self._qpos = torch.tensor([[qpos]], dtype=torch.float32) self._limits = torch.tensor([[[0.0, 0.2]]], dtype=torch.float32) self._pose = _pose(0.0, 0.0, 0.7) + self._handle_pose = _pose(0.05, 0.0, 0.72) + self._handle_pose[:, :3, :3] = torch.tensor( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] + ) self._vertices = _box_vertices(0.05) + self._handle_vertices = _box_vertices(0.02) self._triangles = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int64) self._joint_info = SimpleNamespace( joint_type=SimpleNamespace(name="PRISMATIC"), @@ -222,9 +227,13 @@ def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: def get_link_pose(self, link_name: str, *, to_matrix: bool) -> torch.Tensor: assert link_name in self.link_names assert to_matrix + if link_name == "handle": + return self._handle_pose.clone() return self._pose.clone() def get_link_vert_face(self, link_name: str) -> tuple[torch.Tensor, torch.Tensor]: + if link_name == "handle": + return self._handle_vertices.clone(), self._triangles.clone() assert link_name == "drawer_link" return self._vertices.clone(), self._triangles.clone() @@ -399,8 +408,8 @@ def test_press_grounding_adapts_top_surface_and_depth_to_mainline_contract() -> semantics = ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id="button", label="button", - entity=entity, ) step = program.semantic_steps[0] action = program.edges[0].actions[0] @@ -437,7 +446,12 @@ def test_press_grounding_uses_calibrated_prismatic_button_state() -> None: env.agent_config = { "articulation_settings": {"button": {"slide_joint": [0.0, 0.02]}} } - semantics = ObjectSemantics(affordance=Affordance(), geometry={}) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + entity_id="button", + label="button", + ) step = program.semantic_steps[0] grounded = ActionGrounder(program, env, lambda _uid: semantics).ground( @@ -1526,8 +1540,8 @@ def _held_state( semantics = ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id=entity.uid, label=entity.uid, - entity=entity, ) left_eef, right_eef = env.get_current_xpos_agent() eef = left_eef if arm == "left_arm" else right_eef @@ -1551,8 +1565,8 @@ def _coordinated_held_state( semantics = ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id=entity.uid, label=entity.uid, - entity=entity, ) left_eef, right_eef = env.get_current_xpos_agent() object_pose = entity.get_local_pose(to_matrix=True) @@ -3876,10 +3890,11 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - affordance=AntipodalAffordance( object_label="can", mesh_vertices=vertices, + mesh_triangles=entity.get_triangles(env_ids=[0]), ), geometry={"mesh_vertices": vertices}, + entity_id="can", label="can", - entity=entity, ) grounder = ActionGrounder(program, env, lambda _uid: semantics) @@ -4048,10 +4063,11 @@ def test_pour_grounding_targets_receiver_without_physical_contents() -> None: affordance=AntipodalAffordance( object_label="source", mesh_vertices=_box_vertices(0.05), + mesh_triangles=source.get_triangles(env_ids=[0]), ), geometry={}, + entity_id="source", label="source", - entity=source, ) grounder = ActionGrounder(program, env, lambda _uid: semantics) state = ExecutionState(last_qpos=env.robot.get_qpos()) @@ -4101,7 +4117,12 @@ def test_articulation_grounding_reuses_slide_and_observes_joint_state( grounder = ActionGrounder( program, env, - lambda _uid: ObjectSemantics(affordance=Affordance(), geometry={}), + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + entity_id=uid, + label=uid, + ), ) step = program.semantic_steps[0] edge = program.edges[0] @@ -4114,8 +4135,13 @@ def test_articulation_grounding_reuses_slide_and_observes_joint_state( ) assert isinstance(grounded.target, SlideGoal) - assert grounded.target.grasp_xpos is not None + assert grounded.target.semantics.entity_id == "drawer:handle" + assert torch.equal(grounded.target.target_pose, articulation._handle_pose) assert isinstance(grounded.target.semantics.affordance, SlideAffordance) + assert torch.equal( + grounded.target.semantics.affordance.mesh_vertices, + articulation._handle_vertices, + ) assert grounded.cfg["direction"] == direction assert grounded.cfg["translation_distance"] == pytest.approx(0.2) assert grounded.cfg["articulation_joint_name"] == "slide_joint" @@ -4124,7 +4150,7 @@ def test_articulation_grounding_reuses_slide_and_observes_joint_state( ) assert torch.allclose( grounded.target.semantics.affordance.translation_axis, - torch.tensor([-1.0, 0.0, 0.0]), + torch.tensor([0.0, 1.0, 0.0]), ) assert edge.actions[0]["failure_policy"] == "task_required" @@ -4173,7 +4199,12 @@ def test_turn_knob_requires_setting_map_and_reuses_twist() -> None: grounder = ActionGrounder( program, env, - lambda _uid: ObjectSemantics(affordance=Affordance(), geometry={}), + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + entity_id=uid, + label=uid, + ), ) step = program.semantic_steps[0] @@ -4206,7 +4237,12 @@ def test_turn_knob_requires_setting_map_and_reuses_twist() -> None: ActionGrounder( program, env, - lambda _uid: ObjectSemantics(affordance=Affordance(), geometry={}), + lambda uid: ObjectSemantics( + affordance=Affordance(), + geometry={}, + entity_id=uid, + label=uid, + ), ).ground( program.edges[0].actions[0], program.semantic_steps[0], @@ -6131,8 +6167,8 @@ def semantics(uid: str) -> ObjectSemantics: "mesh_vertices": entity.get_vertices(env_ids=[0], scale=True), "mesh_triangles": entity.get_triangles(env_ids=[0]), }, + entity_id=uid, label=uid, - entity=entity, ) step = program.semantic_steps[0] @@ -6273,8 +6309,8 @@ def test_lay_flat_surface_height_uses_rotated_live_mesh() -> None: lambda uid: ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id=uid, label=uid, - entity=entities[uid], ), ) grounded = grounder.ground( @@ -6339,8 +6375,8 @@ def test_orient_object_anchors_final_pose_to_support_not_live_lift_height() -> N lambda uid: ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id=uid, label=uid, - entity=entities[uid], ), ) @@ -6414,8 +6450,8 @@ def test_orient_grounding_uses_mature_robot_profile_policy() -> None: lambda uid: ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id=uid, label=uid, - entity=entities[uid], ), ) @@ -6785,8 +6821,8 @@ def semantics(uid: str) -> ObjectSemantics: return ObjectSemantics( affordance=Affordance(), geometry={}, + entity_id=uid, label=uid, - entity=entities[uid], ) step = program.semantic_steps[0] diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index 71d4656b8..b9051a2c0 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -324,6 +324,8 @@ def caller(**kwargs): "AxisAlign", "MoveEndEffector", "MoveEndEffector", + "MoveEndEffector", + "MoveEndEffector", "MoveJoints", ] assert [node["atomic_action"] for node in handover_nodes] == [ From 7304a92c7070a2996e11852afb9c0d2e35c067d5 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Fri, 28 Aug 2026 00:41:05 +0800 Subject: [PATCH 78/85] feat(gen-sim): add IK solver profiles and E5 grasp diagnostics --- .../cli/generate_action_agent_config.py | 7 + .../gen_sim/action_engine/cli/run_agent.py | 23 ++ .../action_engine/config/defaults.yaml | 3 + .../action_engine/config/runtime_policy.py | 5 + .../action_engine/environment/agent_env.py | 8 + .../action_engine/generation/__init__.py | 2 + .../generation/config_builder.py | 48 +++ .../action_engine/generation/generator.py | 17 + .../gen_sim/action_engine/gripper_profiles.py | 35 ++ .../gen_sim/action_engine/runtime/actions.py | 323 ++++++++++++--- .../gen_sim/action_engine/runtime/executor.py | 380 ++++++++++++++++++ .../action_engine/runtime/grasp_debug.py | 231 +++++++++++ .../runtime/grasp_diagnostics.py | 203 ++++++++++ .../action_engine/runtime/predicates.py | 32 +- .../gen_sim/action_engine/solver_profiles.py | 84 ++++ embodichain/gen_sim/task_engine/cli.py | 16 + embodichain/gen_sim/task_engine/config.py | 8 + embodichain/gen_sim/task_engine/defaults.yaml | 1 + .../task_engine/orchestration/coordinator.py | 2 + embodichain/gen_sim/task_engine/workflow.py | 25 +- .../action_engine/cli/test_run_agent.py | 53 +++ .../config/test_runtime_policy.py | 1 + .../generation/test_generation.py | 80 ++++ .../action_engine/runtime/test_actions.py | 86 +++- .../action_engine/runtime/test_grasp_debug.py | 133 ++++++ .../runtime/test_grasp_diagnostics.py | 124 ++++++ .../runtime/test_runtime_contracts.py | 71 +++- .../action_engine/test_gripper_profiles.py | 27 +- .../action_engine/test_solver_profiles.py | 42 ++ .../orchestration/test_coordinator_cli.py | 27 ++ .../task_engine/test_parallel_workflow.py | 40 ++ tests/gen_sim/task_engine/test_workflow.py | 5 + 32 files changed, 2079 insertions(+), 63 deletions(-) create mode 100644 embodichain/gen_sim/action_engine/runtime/grasp_debug.py create mode 100644 embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py create mode 100644 embodichain/gen_sim/action_engine/solver_profiles.py create mode 100644 tests/gen_sim/action_engine/runtime/test_grasp_debug.py create mode 100644 tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py create mode 100644 tests/gen_sim/action_engine/test_solver_profiles.py 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 index 2f46d4864..0660b7f2f 100644 --- a/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py +++ b/embodichain/gen_sim/action_engine/cli/generate_action_agent_config.py @@ -112,6 +112,12 @@ def build_parser() -> argparse.ArgumentParser: default=str(_TASK_DEFAULTS["default_gripper_model"]), help="Gripper asset, control, TCP, and grasp profile used by both arms.", ) + parser.add_argument( + "--ik-solver", + choices=("auto", "ur", "pytorch"), + default=str(_TASK_DEFAULTS["default_ik_solver"]), + help="Generation-time IK solver used by both arms.", + ) parser.add_argument( "--llm_model", "--llm-model", @@ -222,6 +228,7 @@ def cli() -> None: task_spec=args.task_spec, robot_profile=args.robot_profile, gripper_model=args.gripper_model, + ik_solver=args.ik_solver, llm_model=args.llm_model, source_scene_z_rotation_degrees=args.source_scene_z_rotation_degrees, body_scale_policy=args.body_scale_policy, diff --git a/embodichain/gen_sim/action_engine/cli/run_agent.py b/embodichain/gen_sim/action_engine/cli/run_agent.py index 06bd72e8d..9d6edcab4 100644 --- a/embodichain/gen_sim/action_engine/cli/run_agent.py +++ b/embodichain/gen_sim/action_engine/cli/run_agent.py @@ -78,6 +78,11 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Show physical collision geometry after every reset.", ) + parser.add_argument( + "--show-grasp-poses", + action="store_true", + help="Write one static PNG for the valid grasp pair selected by E5.", + ) parser.add_argument( "--seed", type=int, @@ -127,6 +132,9 @@ def _validate_run_contract( ) -> None: """Validate the small cross-artifact contract before simulator startup.""" from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + from embodichain.gen_sim.action_engine.solver_profiles import ( + validate_robot_ik_solver_contract, + ) configured_task = agent_config.get("task_name") if configured_task != task_name: @@ -157,6 +165,18 @@ def _validate_run_contract( "Gym and agent configs have different gripper models: " f"gym={gym_gripper!r}, agent={agent_gripper!r}." ) + extensions = gym_config.get("env", {}).get("extensions", {}) + gym_solver = extension.get("ik_solver") + env_solver = extensions.get("agent_ik_solver") + agent_solver = agent_config.get("ik_solver") + if gym_solver is None and env_solver is None and agent_solver is None: + return + if gym_solver != agent_solver or env_solver != agent_solver: + raise ValueError( + "Gym and agent configs have different IK solvers: " + f"gym={gym_solver!r}, env={env_solver!r}, agent={agent_solver!r}." + ) + validate_robot_ik_solver_contract(gym_config.get("robot", {}), str(agent_solver)) def cli() -> int | None: @@ -210,6 +230,8 @@ def cli() -> int | None: "runtime_backend": str(args.runtime_backend), "task_name": str(args.task_name), } + if args.show_grasp_poses: + runtime_arguments["show_grasp_poses"] = True any_failed = False task_engine_reports: list[ExecutionReport] = [] episode_index = 0 @@ -234,6 +256,7 @@ def cli() -> int | None: result = execute( regenerate=bool(args.regenerate), failure_policy=str(args.failure_policy), + show_grasp_poses=bool(args.show_grasp_poses), runtime_run_id=run_id, episode_index=episode_index, ) diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index 41d1eb21e..93ed15678 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -22,6 +22,7 @@ generation: task: default_robot_profile: ur10 default_gripper_model: pgi + default_ik_solver: auto max_episodes: 1 max_episode_steps: 2000 environment: @@ -227,6 +228,8 @@ runtime: pre_grasp_distance: 0.10 lift_height: 0.08 middle_empty_ratio: 0.4 + grasp_opening_margin: 0.02 + grasp_seed: 17393 is_filter_ground_collision: false release_sample_interval: 60 release_gripper_tolerance: 0.08 diff --git a/embodichain/gen_sim/action_engine/config/runtime_policy.py b/embodichain/gen_sim/action_engine/config/runtime_policy.py index 0e0775a9c..ee725f08b 100644 --- a/embodichain/gen_sim/action_engine/config/runtime_policy.py +++ b/embodichain/gen_sim/action_engine/config/runtime_policy.py @@ -28,6 +28,7 @@ from embodichain.gen_sim.action_engine.domain.motion import MOTION_MODIFIER_MODES from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile +from embodichain.gen_sim.action_engine.solver_profiles import resolve_ik_solver_mode from embodichain.utils import configclass from embodichain.utils.utility import load_config @@ -468,6 +469,10 @@ def generation_defaults() -> dict[str, Any]: if not isinstance(task, Mapping): raise ValueError("Generation task defaults must be a mapping.") get_gripper_profile(task.get("default_gripper_model")) + resolve_ik_solver_mode( + task.get("default_ik_solver"), + "dual_ur10", + ) return deepcopy(dict(value)) diff --git a/embodichain/gen_sim/action_engine/environment/agent_env.py b/embodichain/gen_sim/action_engine/environment/agent_env.py index d970b38e0..bbd4125f1 100644 --- a/embodichain/gen_sim/action_engine/environment/agent_env.py +++ b/embodichain/gen_sim/action_engine/environment/agent_env.py @@ -96,6 +96,10 @@ def __init__(self, cfg: EmbodiedEnvCfg | None = None, **kwargs: Any) -> None: self.agent_config.get("gripper_model", "pgi"), ) ) + self.agent_gripper_state_joint_indices = { + side: selected_gripper.state_joint_indices(side) + for side in ("left", "right") + } agent_gripper = self.agent_config.get( "gripper_model", selected_gripper.model.value ) @@ -402,9 +406,12 @@ def create_demo_action_list( self, regenerate: bool = False, failure_policy: str = "stop", + show_grasp_poses: bool = False, **kwargs: Any, ) -> Any: """Compile in memory when requested, then execute the program online.""" + if not isinstance(show_grasp_poses, bool): + raise TypeError("show_grasp_poses must be a boolean.") program = load_agent_execution_program( self.agent_config, agent_config_path=self.agent_config_path, @@ -427,6 +434,7 @@ def create_demo_action_list( record_root=getattr(self, "action_engine_record_root", None), runtime_policy=self.runtime_policy, failure_policy=failure_policy, + show_grasp_poses=show_grasp_poses, ) self.last_execution = executor.run( run_id=kwargs.get("runtime_run_id"), diff --git a/embodichain/gen_sim/action_engine/generation/__init__.py b/embodichain/gen_sim/action_engine/generation/__init__.py index 733b84406..086d479f7 100644 --- a/embodichain/gen_sim/action_engine/generation/__init__.py +++ b/embodichain/gen_sim/action_engine/generation/__init__.py @@ -21,6 +21,7 @@ from .config_builder import ( VLM_CAMERA_UIDS, canonical_gripper_model, + canonical_ik_solver, canonical_robot_profile, ) from .assets import normalize_scene_assets @@ -32,6 +33,7 @@ "PreparedScene", "VLM_CAMERA_UIDS", "canonical_gripper_model", + "canonical_ik_solver", "canonical_robot_profile", "generate_action_engine_config", "normalize_scene_assets", diff --git a/embodichain/gen_sim/action_engine/generation/config_builder.py b/embodichain/gen_sim/action_engine/generation/config_builder.py index 965151fb2..0b2054eeb 100644 --- a/embodichain/gen_sim/action_engine/generation/config_builder.py +++ b/embodichain/gen_sim/action_engine/generation/config_builder.py @@ -40,6 +40,10 @@ GripperProfile, get_gripper_profile, ) +from embodichain.gen_sim.action_engine.solver_profiles import ( + resolve_ik_solver_mode, + validate_robot_ik_solver_contract, +) from embodichain.gen_sim.action_engine.protocol import ( ACTION_ENGINE_CONFIG_SCHEMA, ACTION_ENGINE_ENV_ID, @@ -54,6 +58,7 @@ "build_agent_config", "build_fast_gym_config", "canonical_gripper_model", + "canonical_ik_solver", "canonical_robot_profile", "VLM_CAMERA_UIDS", "validate_fast_gym_config", @@ -63,6 +68,7 @@ _GENERATION_DEFAULTS = generation_defaults() _DEFAULT_TABLETOP_Z = float(_GENERATION_DEFAULTS["scene"]["default_tabletop_z"]) _DEFAULT_GRIPPER_MODEL = str(_GENERATION_DEFAULTS["task"]["default_gripper_model"]) +_DEFAULT_IK_SOLVER = str(_GENERATION_DEFAULTS["task"]["default_ik_solver"]) _ARM_SLOTS = { "left": {"arm": "left_arm", "eef": "left_eef"}, @@ -99,6 +105,11 @@ def canonical_gripper_model(model: str) -> str: return get_gripper_profile(model).model.value +def canonical_ik_solver(mode: str, robot_profile: str) -> str: + """Resolve one generation-time IK solver mode for a robot profile.""" + return resolve_ik_solver_mode(mode, canonical_robot_profile(robot_profile)) + + def build_agent_config( *, task_name: str, @@ -107,6 +118,7 @@ def build_agent_config( source_config_path: Path, uid_map: dict[str, str], gripper_model: str = _DEFAULT_GRIPPER_MODEL, + ik_solver: str = _DEFAULT_IK_SOLVER, static_obstacle_uids: Sequence[str] | None = None, dynamic_obstacle_uids: Sequence[str] | None = None, table_top_z: float | None = None, @@ -120,6 +132,7 @@ def build_agent_config( """Build the small manifest consumed by ``run_agent``.""" profile = canonical_robot_profile(robot_profile) selected_gripper = canonical_gripper_model(gripper_model) + selected_ik_solver = resolve_ik_solver_mode(ik_solver, profile) runtime_policy = default_runtime_policy(profile) explicit_dynamic_collision = ( planner_policy is not None and "dynamic_collision" in planner_policy @@ -174,6 +187,7 @@ def build_agent_config( "task_name": task_name, "robot_profile": profile, "gripper_model": selected_gripper, + "ik_solver": selected_ik_solver, "planning_mode": planning_mode, "task_spec": TASK_SPEC_FILENAME, "scene_requirements": SCENE_REQUIREMENTS_FILENAME, @@ -245,6 +259,7 @@ def build_fast_gym_config( max_episodes: int, max_episode_steps: int, gripper_model: str = _DEFAULT_GRIPPER_MODEL, + ik_solver: str = _DEFAULT_IK_SOLVER, randomize_scene: bool = False, randomize_table_material: bool = False, planning_mode: str = "offline", @@ -261,6 +276,7 @@ def build_fast_gym_config( graph_path = f"offline/{EXECUTION_PROGRAM_FILENAME}" profile = canonical_robot_profile(robot_profile) gripper_profile = get_gripper_profile(gripper_model) + selected_ik_solver = resolve_ik_solver_mode(ik_solver, profile) profile_config = _profile(profile) robot = _make_robot( @@ -268,6 +284,7 @@ def build_fast_gym_config( profile_config, scene.table_top_z, gripper_profile=gripper_profile, + ik_solver=selected_ik_solver, ) observations = _make_observations(robot, gripper_profile) # These two template fields describe serialization order to generation, not @@ -299,6 +316,7 @@ def build_fast_gym_config( "task_name": task_name, "robot_profile": profile, "gripper_model": gripper_profile.model.value, + "ik_solver": selected_ik_solver, "planning_mode": planning_mode, "task_spec": TASK_SPEC_FILENAME, "scene_requirements": SCENE_REQUIREMENTS_FILENAME, @@ -317,6 +335,7 @@ def build_fast_gym_config( "action_engine": engine_extension, "agent_robot_profile": profile, "agent_gripper_model": gripper_profile.model.value, + "agent_ik_solver": selected_ik_solver, "agent_arm_slots": deepcopy(_ARM_SLOTS), "agent_static_obstacle_uids": background_uids, "agent_dynamic_obstacle_uids": rigid_uids, @@ -426,6 +445,10 @@ def validate_fast_gym_config(config: dict[str, Any]) -> None: if extensions.get("agent_gripper_model") != gripper_profile.model.value: raise ValueError("Gym config gripper model fields do not match.") _validate_robot_gripper_contract(config["robot"], gripper_profile) + ik_solver = action_engine.get("ik_solver") + if extensions.get("agent_ik_solver") != ik_solver: + raise ValueError("Gym config IK solver fields do not match.") + validate_robot_ik_solver_contract(config["robot"], str(ik_solver)) graph_path = action_engine.get("seed_task_graph") if ( not isinstance(graph_path, str) @@ -466,6 +489,7 @@ def _make_robot( table_top_z: float | None, *, gripper_profile: GripperProfile, + ik_solver: str, ) -> dict[str, Any]: robot = _load_template(str(profile["template"])) tabletop_z = _DEFAULT_TABLETOP_Z if table_top_z is None else float(table_top_z) @@ -503,11 +527,35 @@ def _make_robot( "fname" ] = f"dual_{family}_{gripper_profile.assembly_name}_basket" _apply_gripper_profile(robot, gripper_profile) + _apply_ik_solver(robot, ik_solver) if profile_id != canonical_robot_profile(profile_id): raise ValueError(f"Invalid canonical robot profile {profile_id!r}.") return robot +def _apply_ik_solver(robot: dict[str, Any], mode: str) -> None: + """Materialize one concrete solver mode while preserving frames and TCPs.""" + solvers = robot.get("solver_cfg") + if not isinstance(solvers, dict): + raise ValueError("Robot template requires solver_cfg.") + if mode == "pytorch": + for arm in ("left_arm", "right_arm"): + current = solvers.get(arm) + if not isinstance(current, dict): + raise ValueError(f"Robot template requires solver_cfg.{arm}.") + if current.get("class_type") == "PytorchSolver": + continue + solvers[arm] = { + "class_type": "PytorchSolver", + "urdf_path": current.get("urdf_path"), + "end_link_name": current["end_link_name"], + "root_link_name": current["root_link_name"], + "tcp": deepcopy(current["tcp"]), + "num_samples": 30, + } + validate_robot_ik_solver_contract(robot, mode) + + def _apply_gripper_profile( robot: dict[str, Any], profile: GripperProfile, diff --git a/embodichain/gen_sim/action_engine/generation/generator.py b/embodichain/gen_sim/action_engine/generation/generator.py index dd1d4f82c..bd85978a7 100644 --- a/embodichain/gen_sim/action_engine/generation/generator.py +++ b/embodichain/gen_sim/action_engine/generation/generator.py @@ -43,6 +43,7 @@ VLM_CAMERA_UIDS, build_agent_config, build_fast_gym_config, + canonical_robot_profile, ) from .models import GeneratedConfigPaths from .source_scene import prepare_scene @@ -64,6 +65,7 @@ def generate_action_engine_config( task_spec: Mapping[str, Any] | str | Path | None = None, robot_profile: str = str(_TASK_DEFAULTS["default_robot_profile"]), gripper_model: str = str(_TASK_DEFAULTS["default_gripper_model"]), + ik_solver: str = str(_TASK_DEFAULTS["default_ik_solver"]), llm_model: str | None = None, source_scene_z_rotation_degrees: float | None = None, source_scene_xy_translation: Sequence[float] | None = None, @@ -95,8 +97,15 @@ def generate_action_engine_config( if planning_mode not in {"offline", "ab"}: raise ValueError("planning_mode must be 'offline' or 'ab'.") from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + from embodichain.gen_sim.action_engine.solver_profiles import ( + resolve_ik_solver_mode, + ) gripper_model = get_gripper_profile(gripper_model).model.value + ik_solver = resolve_ik_solver_mode( + ik_solver, + canonical_robot_profile(robot_profile), + ) _raise_if_outputs_exist( output_dir, overwrite=overwrite, @@ -232,6 +241,7 @@ def generate_action_engine_config( task_name=task_name, robot_profile=robot_profile, gripper_model=gripper_model, + ik_solver=ik_solver, execution_program_hash=program_hash, source_config_path=scene.source_config_path, uid_map=scene.uid_map, @@ -259,6 +269,7 @@ def generate_action_engine_config( task_description=task_description, robot_profile=robot_profile, gripper_model=gripper_model, + ik_solver=ik_solver, execution_program_hash=program_hash, max_episodes=max_episodes, max_episode_steps=max_episode_steps, @@ -744,8 +755,14 @@ def _validate_agent_config(config: Mapping[str, Any]) -> None: if config.get("scene_requirements") != SCENE_REQUIREMENTS_FILENAME: raise ValueError("Agent config must point to canonical SceneRequirements.") from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + from embodichain.gen_sim.action_engine.solver_profiles import ( + resolve_ik_solver_mode, + ) get_gripper_profile(config.get("gripper_model")) + solver = config.get("ik_solver") + if resolve_ik_solver_mode(solver, str(config.get("robot_profile"))) != solver: + raise ValueError("Agent config must store a concrete IK solver mode.") graph_path = config.get("seed_task_graph") if ( not isinstance(graph_path, str) diff --git a/embodichain/gen_sim/action_engine/gripper_profiles.py b/embodichain/gen_sim/action_engine/gripper_profiles.py index 69e921f69..b9b3e7bc2 100644 --- a/embodichain/gen_sim/action_engine/gripper_profiles.py +++ b/embodichain/gen_sim/action_engine/gripper_profiles.py @@ -87,6 +87,8 @@ class GripperProfile: tcp_transform: _Transform left_control_joints: tuple[str, ...] right_control_joints: tuple[str, ...] + left_state_joints: tuple[str, ...] + right_state_joints: tuple[str, ...] left_mimic_joints: tuple[str, ...] right_mimic_joints: tuple[str, ...] mimic_multipliers: tuple[float, ...] @@ -123,12 +125,28 @@ def __post_init__(self) -> None: == mimic_count ): raise ValueError("Gripper mimic metadata must have matching lengths.") + state_count = len(self.left_state_joints) + if not state_count or len(self.right_state_joints) != state_count: + raise ValueError( + "Gripper profiles require matching non-empty state joints." + ) if len(self.simulated_joint_initial_positions) != len( self.simulated_joint_names("left") ): raise ValueError( "Gripper simulated initial positions must match physical joints." ) + for side in ("left", "right"): + controls = set(self.control_joint_names(side)) + states = set(self.state_joint_names(side)) + if not states <= controls: + raise ValueError(f"Gripper {side} state joints must be control joints.") + overlap = states & set(self.mimic_joint_names(side)) + if overlap: + raise ValueError( + "Gripper state and mimic joints must be disjoint; " + f"{side} overlaps: {sorted(overlap)}." + ) def control_joint_names(self, side: _Side) -> tuple[str, ...]: """Return the exact assembled control-joint names for one hand.""" @@ -140,6 +158,16 @@ def mimic_joint_names(self, side: _Side) -> tuple[str, ...]: self._validate_side(side) return self.left_mimic_joints if side == "left" else self.right_mimic_joints + def state_joint_names(self, side: _Side) -> tuple[str, ...]: + """Return joints that define semantic open/closed state for one hand.""" + self._validate_side(side) + return self.left_state_joints if side == "left" else self.right_state_joints + + def state_joint_indices(self, side: _Side) -> tuple[int, ...]: + """Return semantic-state indices within one hand's command vector.""" + controls = self.control_joint_names(side) + return tuple(controls.index(name) for name in self.state_joint_names(side)) + def simulated_joint_names(self, side: _Side) -> tuple[str, ...]: """Return physical movable joints in assembled qpos order for one hand.""" return tuple( @@ -166,6 +194,9 @@ def runtime_manifest( "control_joints": { side: list(self.control_joint_names(side)) for side in ("left", "right") }, + "state_joints": { + side: list(self.state_joint_names(side)) for side in ("left", "right") + }, "mimic_joints": { side: [ { @@ -209,6 +240,8 @@ def _validate_side(side: _Side) -> None: ), left_control_joints=("left_gripper_finger1_joint_1",), right_control_joints=("right_gripper_finger1_joint_1",), + left_state_joints=("left_gripper_finger1_joint_1",), + right_state_joints=("right_gripper_finger1_joint_1",), left_mimic_joints=("left_gripper_finger2_joint_1",), right_mimic_joints=("right_gripper_finger2_joint_1",), mimic_multipliers=(1.0,), @@ -259,6 +292,8 @@ def _validate_side(side: _Side) -> None: "right_inner_knuckle_joint", "right_inner_finger_joint", ), + left_state_joints=("left_finger_joint",), + right_state_joints=("right_finger_joint",), left_mimic_joints=( "left_inner_knuckle_joint", "left_inner_finger_joint", diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 604dc3245..ac035b64f 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -19,7 +19,7 @@ from __future__ import annotations from collections.abc import Callable, Iterator, Mapping, Sequence -from contextlib import contextmanager +from contextlib import contextmanager, nullcontext from copy import deepcopy from dataclasses import replace import logging @@ -35,6 +35,9 @@ ) from embodichain.gen_sim.action_engine.config import default_runtime_policy from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile +from embodichain.gen_sim.action_engine.solver_profiles import ( + expected_ik_solver_class, +) from embodichain.lab.sim.atomic_actions import ( ActionBinding, ActionInvocation, @@ -77,6 +80,7 @@ from embodichain.utils.math import matrix_from_quat, quat_from_matrix, quat_slerp from .body_grasp import AxisAlignBodyGraspAdapter +from .grasp_diagnostics import _TracingAntipodalGraspPoseGenerator from .models import ActionOutcome, GroundedAction from .state import ExecutionState @@ -107,6 +111,7 @@ _COLLISION_PARKING_Z_OFFSET = -100.0 _BODY_GRASP_CANDIDATE_LIMIT = 500 _BODY_GRASP_SEED = 17_392 +_COORDINATED_GRASP_SEED = 17_393 _FREE_YAW_SAMPLE_COUNT = 8 _RETREAT_LOG_LOCK = RLock() @@ -199,6 +204,7 @@ def __init__( self.gripper_profile = get_gripper_profile( getattr(env, "agent_gripper_model", "pgi") ) + self.ik_solver, self.ik_solver_classes = self._resolve_runtime_ik_solver() if grasp_policy is None: profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) grasp_policy = default_runtime_policy(profile).grasp @@ -223,13 +229,40 @@ def __init__( self.capabilities = capability_registry or build_atomic_capability_registry() self._motion_generator: MotionGenerator | None = None self._atomic_engine: AtomicActionEngine | None = None - self._coordinated_engines: dict[bool, AtomicActionEngine] = {} + self._coordinated_engines: dict[tuple[bool, float], AtomicActionEngine] = {} self._semantics: dict[str, ObjectSemantics] = {} self._scene_time = 0.0 if scene_provider is not None and not isinstance(scene_provider, SceneProvider): raise TypeError("scene_provider must implement SceneProvider.") self.scene_provider = scene_provider or self._build_scene_provider() + def _resolve_runtime_ik_solver(self) -> tuple[str, dict[str, str]]: + """Validate a declared bundle solver against the initialized robot.""" + declared = getattr(self.env, "agent_ik_solver", None) + robot = getattr(self.env, "robot", None) + get_solver = getattr(robot, "get_solver", None) + classes: dict[str, str] = {} + if callable(get_solver): + for arm in ("left_arm", "right_arm"): + part = self.env.get_agent_arm_control_part(arm == "left_arm") + classes[arm] = type(get_solver(name=part)).__name__ + if declared is None: + inferred = { + "URSolver": "ur", + "PytorchSolver": "pytorch", + } + modes = {inferred[name] for name in classes.values() if name in inferred} + return (modes.pop() if len(modes) == 1 else "unknown"), classes + declared = str(declared) + expected = expected_ik_solver_class(declared) + for arm, actual in classes.items(): + if actual != expected: + raise ValueError( + f"Runtime {arm} must use {expected} for " + f"agent_ik_solver={declared!r}, got {actual!r}." + ) + return declared, classes + @staticmethod def _merge_planner_policy( target: dict[str, Any], @@ -396,6 +429,7 @@ def plan( selected_warnings: tuple[str, ...] = () candidate_search_warnings: list[str] = [] candidate_search_attempts = 0 + coordinated_search_attempts: list[dict[str, Any]] = [] best_failure_count = self.num_envs + 1 for candidate in grounded_candidates: candidate_engine = self._engine_for(candidate, capability) @@ -408,8 +442,40 @@ def plan( candidate.motion_policy.get("retreat_reachability_search", False) or candidate.motion_policy.get("reorient_tool_down", False) ) - with _capture_retreat_warnings(capture_warnings) as warnings: + grasp_seed = candidate.motion_policy.get("grasp_seed") + seed_context = ( + nullcontext() + if grasp_seed is None + else self._isolated_random_seed(int(grasp_seed)) + ) + with seed_context, _capture_retreat_warnings(capture_warnings) as warnings: candidate_plan = candidate_engine.plan(candidate_invocation, context) + coordinated_trace = candidate.motion_policy.get("coordinated_grasp") + if isinstance(coordinated_trace, dict): + grasp_stages = self._latest_coordinated_grasp_trace(candidate_engine) + if grasp_stages is not None: + coordinated_trace["stages"] = grasp_stages + coordinated_search_attempts.append( + { + "candidate_index": coordinated_trace.get("candidate_index"), + "approach_candidate_label": coordinated_trace.get( + "approach_candidate_label" + ), + "approach_direction": coordinated_trace.get( + "approach_direction" + ), + "middle_empty_ratio": coordinated_trace.get( + "selected_middle_empty_ratio" + ), + "plan_success": candidate_plan.plan_success.detach() + .cpu() + .tolist(), + "planner_messages": list( + candidate_plan.diagnostics.messages + ), + "stages": deepcopy(grasp_stages), + } + ) if capture_warnings: candidate_search_warnings.extend(warnings) candidate_search_attempts += 1 @@ -428,6 +494,11 @@ def plan( if selected is None: raise RuntimeError("Atomic action adaptation produced no plan candidate.") grounded, invocation, plan, selected_engine = selected + selected_coordinated_trace = grounded.motion_policy.get("coordinated_grasp") + if isinstance(selected_coordinated_trace, dict): + selected_coordinated_trace["search_attempts"] = deepcopy( + coordinated_search_attempts + ) if bool(grounded.motion_policy.get("reorient_tool_down", False)): summary = ( "Tool-down reorientation search: " @@ -776,12 +847,7 @@ def _adapt_coordinated_pickment_grasps( grounded: GroundedAction, capability: AtomicCapability, ) -> tuple[GroundedAction, ...]: - """Build deterministic geometry-ranked coordinated-grasp candidates. - - ``left_to_right_arm_direction`` remains the live base-to-base direction: - it labels the two participant regions and is not the transport direction. - Object geometry only adjusts how much of the projected middle is excluded. - """ + """Build deterministic live-geometry approach and partition candidates.""" if capability.config_materializer != "coordinated_pickment": return (grounded,) target = self._validate_coordinated_pickment_goal(grounded) @@ -837,10 +903,8 @@ def _adapt_coordinated_pickment_grasps( covariance = centered.transpose(0, 1) @ centered / float(vertices.shape[0]) eigenvalues, eigenvectors = torch.linalg.eigh(covariance) principal_local = eigenvectors[:, -1] - principal_world = torch.matmul( - live_pose[:, :3, :3], - principal_local, - ) + world_axes = torch.matmul(live_pose[:, :3, :3], eigenvectors) + principal_world = world_axes[:, :, -1] principal_world = principal_world / torch.linalg.vector_norm( principal_world, dim=1, @@ -873,61 +937,179 @@ def _adapt_coordinated_pickment_grasps( if not any(abs(ratio - existing) <= 1.0e-6 for existing in ratios): ratios.append(ratio) - approach = grounded.cfg.get("approach_direction", (0.0, 0.0, -1.0)) - approach = torch.as_tensor( - approach, + requested_approach = grounded.cfg.get("approach_direction", (0.0, 0.0, -1.0)) + requested_approach = torch.as_tensor( + requested_approach, dtype=torch.float32, device=self.device, ) - if approach.shape != (3,) or not bool(torch.isfinite(approach).all()): + if requested_approach.shape != (3,) or not bool( + torch.isfinite(requested_approach).all() + ): raise ValueError("approach_direction must be a finite vector shaped (3,).") - approach_norm = torch.linalg.vector_norm(approach) + approach_norm = torch.linalg.vector_norm(requested_approach) if float(approach_norm) <= 1.0e-6: raise ValueError("approach_direction must be non-zero.") - approach = approach / approach_norm + requested_approach = requested_approach / approach_norm + + horizontal_arm = shared_direction.clone() + horizontal_arm[2] = 0.0 + horizontal_arm_norm = torch.linalg.vector_norm(horizontal_arm) + if float(horizontal_arm_norm) <= 1.0e-6: + raise ValueError( + "CoordinatedPickment arm bases must be separated in the world XY plane." + ) + horizontal_arm = horizontal_arm / horizontal_arm_norm + robot_forward = torch.stack( + (-horizontal_arm[1], horizontal_arm[0], horizontal_arm.new_tensor(0.0)) + ) + base_midpoint = 0.5 * (left_base[:, :3, 3] + right_base[:, :3, 3]) + object_from_bases = live_pose[:, :3, 3] - base_midpoint + reach_alignment = torch.sum(object_from_bases[:, :2] * robot_forward[:2], dim=1) + if float(reach_alignment.mean()) < 0.0: + robot_forward = -robot_forward + reach_alignment = -reach_alignment + + down = requested_approach.new_tensor([0.0, 0.0, -1.0]) + approach_candidates: list[tuple[str, torch.Tensor]] = [] + + def add_approach(label: str, direction: torch.Tensor) -> None: + direction = direction.to(device=self.device, dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if not bool(torch.isfinite(direction).all()) or float(norm) <= 1.0e-6: + return + direction = direction / norm + if any( + float(torch.dot(direction, existing)) >= 1.0 - 1.0e-5 + for _, existing in approach_candidates + ): + return + approach_candidates.append((label, direction)) + + add_approach("robot_forward_down", robot_forward + down) + add_approach("current", requested_approach) + add_approach("world_down", down) + add_approach("robot_forward", robot_forward) + + horizontal_axis_index: int | None = None + for axis_index in torch.argsort(eigenvalues).tolist(): + axes = world_axes[:, :, axis_index] + if float(torch.mean(torch.abs(axes[:, 2]))) > 0.75: + continue + reference = axes[0] + consistency = torch.abs(torch.matmul(axes, reference)) + if bool((consistency < 0.90).any()): + continue + horizontal_axis_index = int(axis_index) + horizontal_axis = reference.clone() + if float(torch.dot(horizontal_axis[:2], robot_forward[:2])) < 0.0: + horizontal_axis = -horizontal_axis + add_approach("short_axis_forward_down", horizontal_axis + down) + add_approach("short_axis_forward", horizontal_axis) + add_approach("short_axis_reverse_down", -horizontal_axis + down) + add_approach("short_axis_reverse", -horizontal_axis) + break + + grasp_seed = grounded.cfg.get("grasp_seed", _COORDINATED_GRASP_SEED) + if type(grasp_seed) is not int or grasp_seed < 0: + raise ValueError("grasp_seed must be a non-negative integer.") trace = { - "strategy": "live_geometry_partition_search", + "strategy": "live_geometry_approach_search", + "local_pca_axes": eigenvectors.detach().cpu().tolist(), + "world_pca_axes": world_axes.detach().cpu().tolist(), + "pca_eigenvalues": eigenvalues.detach().cpu().tolist(), "local_principal_axis": principal_local.detach().cpu().tolist(), "world_principal_axes": principal_world.detach().cpu().tolist(), "elongation_ratio": float(elongation_ratio), "elongation_confidence": confidence, "arm_axis_alignment": arm_alignment.detach().cpu().tolist(), "left_to_right_arm_direction": shared_direction.detach().cpu().tolist(), - "approach_direction": approach.detach().cpu().tolist(), + "robot_forward_direction": robot_forward.detach().cpu().tolist(), + "shared_reach_alignment": reach_alignment.detach().cpu().tolist(), + "requested_approach_direction": requested_approach.detach().cpu().tolist(), + "selected_horizontal_axis_index": horizontal_axis_index, + "approach_candidates": [ + { + "label": label, + "direction": direction.detach().cpu().tolist(), + } + for label, direction in approach_candidates + ], "candidate_middle_empty_ratios": list(ratios), + "grasp_seed": grasp_seed, } candidates: list[GroundedAction] = [] - for candidate_index, ratio in enumerate(ratios): - cfg = { - **grounded.cfg, - "left_to_right_arm_direction": shared_direction.clone(), - "approach_direction": approach.clone(), - "middle_empty_ratio": ratio, - } - motion_policy = { - **grounded.motion_policy, - "coordinated_grasp": { - **trace, - "candidate_index": candidate_index, - "selected_middle_empty_ratio": ratio, - }, - } - candidates.append( - replace( - grounded, - target=replace(target, object_initial_pose=live_pose.clone()), - cfg=cfg, - object_pose=live_pose.clone(), - motion_policy=motion_policy, + candidate_index = 0 + for approach_index, (approach_label, approach) in enumerate( + approach_candidates + ): + for ratio in ratios: + cfg = { + **grounded.cfg, + "left_to_right_arm_direction": shared_direction.clone(), + "approach_direction": approach.clone(), + "middle_empty_ratio": ratio, + "grasp_seed": grasp_seed, + } + motion_policy = { + **grounded.motion_policy, + "grasp_seed": grasp_seed, + "coordinated_grasp": { + **trace, + "candidate_index": candidate_index, + "approach_candidate_index": approach_index, + "approach_candidate_label": approach_label, + "approach_direction": approach.detach().cpu().tolist(), + "selected_middle_empty_ratio": ratio, + }, + } + candidates.append( + replace( + grounded, + target=replace(target, object_initial_pose=live_pose.clone()), + cfg=cfg, + object_pose=live_pose.clone(), + motion_policy=motion_policy, + ) ) - ) + candidate_index += 1 return tuple(candidates) + @contextmanager + def _isolated_random_seed(self, seed: int) -> Iterator[None]: + """Run one stochastic grasp attempt without perturbing global RNG state.""" + if type(seed) is not int or seed < 0: + raise ValueError("seed must be a non-negative integer.") + device = torch.device(self.device) + cuda_devices: list[int] = [] + if device.type == "cuda": + cuda_devices.append( + torch.cuda.current_device() if device.index is None else device.index + ) + with torch.random.fork_rng(devices=cuda_devices): + torch.manual_seed(seed) + if cuda_devices: + torch.cuda.manual_seed_all(seed) + yield + def _coordinated_arm_bases(self) -> tuple[torch.Tensor, torch.Tensor]: from .frames import arm_base_poses return arm_base_poses(self.env) + def _latest_coordinated_grasp_trace( + self, + engine: AtomicActionEngine, + ) -> dict[str, Any] | None: + """Return the S1-S5 trace emitted by the shared E5 grasp service.""" + _, hand_part, _ = self._parts("left_arm") + if hand_part is None: + return None + generator = engine.grasp_pose_generators.get(hand_part) + if not isinstance(generator, _TracingAntipodalGraspPoseGenerator): + return None + return generator.last_dual_trace + def _search_reachable_retreat( self, *, @@ -1283,6 +1465,9 @@ def _planner_trace( "action_class": grounded.action_class, "arm": grounded.arm, "gripper_model": self.gripper_profile.model.value, + "ik_solver": self.ik_solver, + "left_solver_class": self.ik_solver_classes.get("left_arm"), + "right_solver_class": self.ik_solver_classes.get("right_arm"), "planner": requested_backend, "requested_backend": requested_backend, "effective_backend": effective_backend, @@ -2055,15 +2240,36 @@ def _engine_for( ) if not isinstance(filter_ground_collision, bool): raise TypeError("is_filter_ground_collision must be a boolean.") - if filter_ground_collision: + opening_margin = grounded.cfg.get( + "grasp_opening_margin", + self.gripper_profile.grasp_model.opening_margin, + ) + if isinstance(opening_margin, bool) or not isinstance( + opening_margin, (int, float) + ): + raise TypeError("grasp_opening_margin must be a real number.") + opening_margin = float(opening_margin) + if ( + not math.isfinite(opening_margin) + or opening_margin < 0.0 + or opening_margin >= self.gripper_profile.grasp_model.max_opening_width + ): + raise ValueError( + "grasp_opening_margin must be finite, non-negative, and smaller " + "than the selected gripper's maximum opening width." + ) + profile_margin = self.gripper_profile.grasp_model.opening_margin + if filter_ground_collision and opening_margin == profile_margin: return self._engine() - cached = self._coordinated_engines.get(filter_ground_collision) + cache_key = (filter_ground_collision, opening_margin) + cached = self._coordinated_engines.get(cache_key) if cached is None: cached = self._new_engine( MotionGenerator(cfg=self._motion_generator_cfg()), filter_ground_collision=filter_ground_collision, + opening_margin=opening_margin, ) - self._coordinated_engines[filter_ground_collision] = cached + self._coordinated_engines[cache_key] = cached return cached def _new_engine( @@ -2071,6 +2277,7 @@ def _new_engine( motion_generator: MotionGenerator, *, filter_ground_collision: bool, + opening_margin: float | None = None, ) -> AtomicActionEngine: from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver @@ -2081,6 +2288,7 @@ def _new_engine( control_profiles=self._control_profiles(), grasp_pose_generators=self._grasp_pose_generators( filter_ground_collision=filter_ground_collision, + opening_margin=opening_margin, ), ) engine.register(ExactTargetMoveHeldObject(), replace=True) @@ -2163,12 +2371,29 @@ def _grasp_pose_generators( self, *, filter_ground_collision: bool = True, + opening_margin: float | None = None, ) -> dict[str, AntipodalGraspPoseGenerator]: """Build one mainline grasp service for each runtime hand endpoint.""" if not isinstance(filter_ground_collision, bool): raise TypeError("filter_ground_collision must be a boolean.") options = self.grasp_policy geometry = self.gripper_profile.grasp_model + if opening_margin is None: + opening_margin = geometry.opening_margin + if isinstance(opening_margin, bool) or not isinstance( + opening_margin, (int, float) + ): + raise TypeError("opening_margin must be a real number or None.") + opening_margin = float(opening_margin) + if ( + not math.isfinite(opening_margin) + or opening_margin < 0.0 + or opening_margin >= geometry.max_opening_width + ): + raise ValueError( + "opening_margin must be finite, non-negative, and smaller than " + "the selected gripper's maximum opening width." + ) model = ParallelJawGripperModelCfg( model_id=geometry.model_id, min_opening_width=geometry.min_opening_width, @@ -2188,7 +2413,7 @@ def _grasp_pose_generators( collision = ParallelJawGraspCollisionCfg( point_sample_density=float(options["point_sample_dense"]), max_decomposition_hulls=int(options["max_decomposition_hulls"]), - opening_margin=geometry.opening_margin, + opening_margin=opening_margin, filter_ground_collision=filter_ground_collision, ) annotation = GraspAnnotationCfg( @@ -2196,7 +2421,7 @@ def _grasp_pose_generators( viser_port=int(options["viser_port"]), force_refresh=bool(options["force_grasp_reannotate"]), ) - shared_generator = AntipodalGraspPoseGenerator( + shared_generator = _TracingAntipodalGraspPoseGenerator( model, algorithm_cfg=algorithm, collision_cfg=collision, diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index 678952852..d07ab4aa0 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -23,6 +23,7 @@ from copy import deepcopy from dataclasses import dataclass, field, replace import logging +from pathlib import Path from threading import RLock from typing import Any @@ -185,6 +186,7 @@ def __init__( capability_registry: Any | None = None, scene_provider: SceneProvider | None = None, failure_policy: str = "stop", + show_grasp_poses: bool = False, ) -> None: self.program = program self.env = env @@ -195,6 +197,10 @@ def __init__( "ProgramExecutor failure_policy must be 'stop' or 'continue'." ) self.failure_policy = str(failure_policy) + if not isinstance(show_grasp_poses, bool): + raise TypeError("ProgramExecutor show_grasp_poses must be a boolean.") + self.show_grasp_poses = show_grasp_poses + self._runtime_output_dir: Path | None = None if runtime_policy is None: profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) runtime_policy = default_runtime_policy(profile) @@ -365,6 +371,7 @@ def run( runtime_policy_hash=runtime_policy_hash(self.runtime_policy), failure_policy=self.failure_policy, ) + self._runtime_output_dir = recorder.output_dir aggregate_failed = torch.zeros( int(self.env.num_envs), dtype=torch.bool, @@ -3302,6 +3309,234 @@ def _physical_coordinated_hold( ) return attempted & held & reached + @staticmethod + def _rotation_error_radians( + actual: torch.Tensor, + expected: torch.Tensor, + ) -> torch.Tensor: + """Return the geodesic angle between batched rotation matrices.""" + relative = torch.matmul(actual.transpose(1, 2), expected) + cosine = (torch.diagonal(relative, dim1=1, dim2=2).sum(dim=1) - 1.0) * 0.5 + return torch.acos(torch.clamp(cosine, -1.0, 1.0)) + + def _coordinated_physical_trace( + self, + uid: str, + state: ExecutionState, + grounded: GroundedAction, + outcome: ActionOutcome, + attempted: torch.Tensor, + *, + object_pose_before: torch.Tensor, + execution_observations: Mapping[str, Any], + dual_hold: torch.Tensor, + accepted: torch.Tensor, + ) -> dict[str, Any]: + """Describe the measured E5 tracking, closure, lift, and acceptance state.""" + device = self.env.device + actual_object = self._entity_pose(uid).to(device=device, dtype=torch.float32) + before = object_pose_before.to(device=device, dtype=torch.float32) + target = grounded.target_object_pose + if not isinstance(target, torch.Tensor): + target_pose = torch.full_like(actual_object, torch.nan) + reached = torch.zeros_like(attempted) + else: + target_pose = torch.as_tensor( + target, + dtype=torch.float32, + device=device, + ) + if target_pose.shape == (4, 4): + target_pose = target_pose.unsqueeze(0).repeat( + int(self.env.num_envs), 1, 1 + ) + tolerance = float( + grounded.cfg.get( + "postcondition_tolerance", + self.runtime_policy.predicate_fallbacks["position_tolerance"], + ) + ) + reached = ( + torch.linalg.vector_norm( + actual_object[:, :3, 3] - target_pose[:, :3, 3], dim=1 + ) + <= tolerance + ) + + displacement = actual_object[:, :3, 3] - before[:, :3, 3] + intended = target_pose[:, :3, 3] - before[:, :3, 3] + intended_distance = torch.linalg.vector_norm(intended, dim=1) + intended_direction = intended / intended_distance[:, None].clamp_min(1.0e-8) + displacement_distance = torch.linalg.vector_norm(displacement, dim=1) + displacement_projection = torch.sum(displacement * intended_direction, dim=1) + direction_cosine = displacement_projection / displacement_distance.clamp_min( + 1.0e-8 + ) + direction_cosine = torch.where( + (intended_distance > 1.0e-8) & (displacement_distance > 1.0e-8), + direction_cosine, + torch.zeros_like(direction_cosine), + ) + + maximum_object_z = torch.as_tensor( + execution_observations.get("maximum_object_z", before[:, 2, 3]), + dtype=torch.float32, + device=device, + ) + eef_values = self.env.get_current_xpos_agent() + gripper_values = self.env.get_current_gripper_state_agent() + actual_qpos = self.env.robot.get_qpos().to(device=device, dtype=torch.float32) + terminal_qpos = outcome.trajectory[:, -1].to(device=device, dtype=torch.float32) + gripper_tolerance = float( + self.runtime_policy.predicate_fallbacks["held_gripper_tolerance"] + ) + arms: dict[str, Any] = {} + for arm_index, arm in enumerate(("left_arm", "right_arm")): + control_part = arm_control_part(self.env, arm) + held = state.get_held_object(control_part) + actual_eef = torch.as_tensor( + eef_values[arm_index], + dtype=torch.float32, + device=device, + ) + gripper = torch.as_tensor( + gripper_values[arm_index], + dtype=torch.float32, + device=device, + ) + open_state = torch.as_tensor( + self.env.open_state, + dtype=torch.float32, + device=device, + ).flatten() + close_state = torch.as_tensor( + self.env.close_state, + dtype=torch.float32, + device=device, + ).flatten() + repeats = (gripper.shape[-1] + open_state.numel() - 1) // open_state.numel() + expected_open = open_state.repeat(repeats)[: gripper.shape[-1]] + expected_close = close_state.repeat(repeats)[: gripper.shape[-1]] + open_distance = torch.linalg.vector_norm( + gripper - expected_open.unsqueeze(0), dim=1 + ) + close_distance = torch.linalg.vector_norm( + gripper - expected_close.unsqueeze(0), dim=1 + ) + arm_joint_ids = list(self.env.robot.get_joint_ids(name=control_part)) + hand_part = self.env.get_agent_eef_control_part(arm_index == 0) + hand_joint_ids = ( + [] + if hand_part is None + else list(self.env.robot.get_joint_ids(name=hand_part)) + ) + arm_trace: dict[str, Any] = { + "control_part": control_part, + "actual_eef_pose": actual_eef.detach().cpu().tolist(), + "gripper_qpos": gripper.detach().cpu().tolist(), + "gripper_open_distance": open_distance.detach().cpu().tolist(), + "gripper_close_distance": close_distance.detach().cpu().tolist(), + "gripper_closed": (open_distance > gripper_tolerance) + .detach() + .cpu() + .tolist(), + "arm_joint_tracking_error": torch.linalg.vector_norm( + actual_qpos[:, arm_joint_ids] - terminal_qpos[:, arm_joint_ids], + dim=1, + ) + .detach() + .cpu() + .tolist(), + "gripper_joint_tracking_error": ( + torch.linalg.vector_norm( + actual_qpos[:, hand_joint_ids] + - terminal_qpos[:, hand_joint_ids], + dim=1, + ) + if hand_joint_ids + else torch.zeros(int(self.env.num_envs), device=device) + ) + .detach() + .cpu() + .tolist(), + } + if held is not None: + planned_relation = held.object_to_eef.to( + device=device, dtype=torch.float32 + ) + expected_eef = torch.bmm(actual_object, planned_relation) + actual_relation = torch.bmm(torch.linalg.inv(actual_object), actual_eef) + planned_target_eef = held.grasp_xpos.to( + device=device, dtype=torch.float32 + ) + arm_trace.update( + { + "planned_grasp_pose": planned_target_eef.detach() + .cpu() + .tolist(), + "planned_object_to_eef": planned_relation.detach() + .cpu() + .tolist(), + "actual_object_to_eef": actual_relation.detach().cpu().tolist(), + "eef_target_position_error": torch.linalg.vector_norm( + actual_eef[:, :3, 3] - planned_target_eef[:, :3, 3], + dim=1, + ) + .detach() + .cpu() + .tolist(), + "eef_target_orientation_error": self._rotation_error_radians( + actual_eef[:, :3, :3], planned_target_eef[:, :3, :3] + ) + .detach() + .cpu() + .tolist(), + "eef_relation_position_error": torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], dim=1 + ) + .detach() + .cpu() + .tolist(), + "eef_relation_orientation_error": self._rotation_error_radians( + actual_eef[:, :3, :3], expected_eef[:, :3, :3] + ) + .detach() + .cpu() + .tolist(), + } + ) + arms[arm] = arm_trace + + return { + "attempted": attempted.detach().cpu().tolist(), + "dual_hold_predicate": dual_hold.detach().cpu().tolist(), + "semantic_target_reached": reached.detach().cpu().tolist(), + "accepted": accepted.detach().cpu().tolist(), + "object_pose_before": before.detach().cpu().tolist(), + "object_pose_after": actual_object.detach().cpu().tolist(), + "object_target_pose": target_pose.detach().cpu().tolist(), + "object_displacement": displacement.detach().cpu().tolist(), + "object_displacement_distance": displacement_distance.detach() + .cpu() + .tolist(), + "object_intended_distance": intended_distance.detach().cpu().tolist(), + "object_displacement_projection": displacement_projection.detach() + .cpu() + .tolist(), + "object_displacement_direction_cosine": direction_cosine.detach() + .cpu() + .tolist(), + "maximum_object_lift": (maximum_object_z - before[:, 2, 3]) + .detach() + .cpu() + .tolist(), + "maximum_object_z_waypoint": execution_observations.get( + "maximum_object_z_waypoint" + ), + "segment_observations": execution_observations.get("segments", {}), + "arms": arms, + } + def _commit_coordinated_ownership( self, uid: str, @@ -3445,6 +3680,11 @@ def _execute_coordinated( for message in dict.fromkeys(selected_warnings): log_warning(message) grounded, outcome = selected + if ( + self.show_grasp_poses + and capability.config_materializer == "coordinated_pickment" + ): + self._write_selected_coordinated_grasp_pose(step, outcome) if capability.state_effect == "transfer_hold": outcome = replace( outcome, @@ -3458,19 +3698,81 @@ def _execute_coordinated( ) self._remember_target(step, grounded) successful = active & outcome.success + object_pose_before = self._entity_pose(step.object_uid).detach().clone() + execution_observations: dict[str, Any] = { + "maximum_object_z": object_pose_before[:, 2, 3].detach().clone(), + "maximum_object_z_waypoint": None, + "segments": {}, + } + segment_stops = { + int(segment["stop"]) - 1: name + for name, segment in outcome.planner_trace.get( + "action_segments", {} + ).items() + if int(segment.get("stop", 0)) > 0 + } + + def observe_waypoint(waypoint_index: int) -> None: + object_pose = self._entity_pose(step.object_uid).detach().clone() + higher = object_pose[:, 2, 3] > execution_observations["maximum_object_z"] + execution_observations["maximum_object_z"] = torch.where( + higher, + object_pose[:, 2, 3], + execution_observations["maximum_object_z"], + ) + if bool(higher.any()): + execution_observations["maximum_object_z_waypoint"] = waypoint_index + segment_name = segment_stops.get(waypoint_index) + if segment_name is not None: + eef = self.env.get_current_xpos_agent() + gripper = self.env.get_current_gripper_state_agent() + execution_observations["segments"][segment_name] = { + "waypoint": waypoint_index, + "object_pose": object_pose.detach().cpu().tolist(), + "left_eef_pose": torch.as_tensor(eef[0]).detach().cpu().tolist(), + "right_eef_pose": torch.as_tensor(eef[1]).detach().cpu().tolist(), + "left_gripper_qpos": torch.as_tensor(gripper[0]) + .detach() + .cpu() + .tolist(), + "right_gripper_qpos": torch.as_tensor(gripper[1]) + .detach() + .cpu() + .tolist(), + } + actions = self.adapter.execute_trajectory( outcome.trajectory, active=successful, + waypoint_observer=observe_waypoint, ) physical_failed = torch.zeros_like(failed) committed_state = outcome.state_after(successful) if capability.state_effect == "coordinated_hold": + dual_hold = evaluate_predicate( + self.env, + {"type": "held_by_both_grippers", "object": step.object_uid}, + coordinated_state=committed_state, + ) physical = self._physical_coordinated_hold( step.object_uid, committed_state, grounded, successful, ) + outcome.planner_trace["physical_execution"] = ( + self._coordinated_physical_trace( + step.object_uid, + committed_state, + grounded, + outcome, + successful, + object_pose_before=object_pose_before, + execution_observations=execution_observations, + dual_hold=dual_hold, + accepted=physical, + ) + ) physical_failed |= successful & ~physical successful = physical committed_state = outcome.state_after(successful) @@ -3573,6 +3875,37 @@ def _execute_coordinated( successful, ) + def _write_selected_coordinated_grasp_pose( + self, + step: SemanticStep, + outcome: ActionOutcome, + ) -> None: + """Write the valid env-zero grasp pair selected by one E5 plan.""" + from .grasp_debug import ( + grasp_pose_image_path, + render_coordinated_grasp_pose_png, + selected_coordinated_grasp_scene, + ) + + if self._runtime_output_dir is None: + return + try: + scene = selected_coordinated_grasp_scene( + outcome, + left_control_part=arm_control_part(self.env, "left_arm"), + right_control_part=arm_control_part(self.env, "right_arm"), + ) + if scene is None: + return + path = grasp_pose_image_path(self._runtime_output_dir, step.id) + render_coordinated_grasp_pose_png(scene, path) + outcome.planner_trace["grasp_pose_visualization"] = path.as_posix() + log_info(f"E5 grasp pose visualization: {path}") + except Exception as exc: + message = f"{type(exc).__name__}: {exc}" + outcome.planner_trace["grasp_pose_visualization_error"] = message + log_warning(f"Unable to render E5 grasp poses: {message}") + def _rebase_held_state( self, uid: str, @@ -3719,6 +4052,53 @@ def _execute_explicit_dual( ) released = active & opened physical_failed = active & ~opened + gripper_values = self.env.get_current_gripper_state_agent() + open_state = torch.as_tensor( + self.env.open_state, + dtype=torch.float32, + device=self.env.device, + ).flatten() + gripper_trace: dict[str, Any] = {} + for arm_index, arm in enumerate(("left", "right")): + gripper = torch.as_tensor( + gripper_values[arm_index], + dtype=torch.float32, + device=self.env.device, + ) + repeats = ( + gripper.shape[-1] + open_state.numel() - 1 + ) // open_state.numel() + expected_open = open_state.repeat(repeats)[: gripper.shape[-1]] + gripper_trace[f"{arm}_gripper_qpos"] = gripper.detach().cpu().tolist() + gripper_trace[f"{arm}_gripper_open_error"] = ( + torch.linalg.vector_norm( + gripper - expected_open.unsqueeze(0), dim=1 + ) + .detach() + .cpu() + .tolist() + ) + release_trace = { + "active": active.detach().cpu().tolist(), + "both_grippers_open": opened.detach().cpu().tolist(), + "released": released.detach().cpu().tolist(), + "release_gripper_tolerance": float( + coordinated_policy.get( + "release_gripper_tolerance", + self.runtime_policy.predicate_fallbacks[ + "gripper_state_tolerance" + ], + ) + ), + "object_pose_after_release": self._entity_pose(step.object_uid) + .detach() + .cpu() + .tolist(), + **gripper_trace, + } + for outcome in outcomes.values(): + if outcome is not None: + outcome.planner_trace["physical_release"] = deepcopy(release_trace) control_parts = ( arm_control_part(self.env, "left_arm"), arm_control_part(self.env, "right_arm"), diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_debug.py b/embodichain/gen_sim/action_engine/runtime/grasp_debug.py new file mode 100644 index 000000000..2b5c3d704 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grasp_debug.py @@ -0,0 +1,231 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 static visualization of the grasp poses selected by E5.""" + +from __future__ import annotations + +from dataclasses import dataclass +import os +from pathlib import Path +import re +import tempfile + +import numpy as np +import torch + +from embodichain.lab.sim.atomic_actions import AntipodalAffordance + +from .models import ActionOutcome + +__all__ = [ + "CoordinatedGraspPoseScene", + "grasp_pose_image_path", + "render_coordinated_grasp_pose_png", + "selected_coordinated_grasp_scene", +] + + +_SAFE_NAME = re.compile(r"[^0-9A-Za-z._-]+") +_IMAGE_WIDTH = 960 +_IMAGE_HEIGHT = 720 + + +@dataclass(frozen=True, slots=True) +class CoordinatedGraspPoseScene: + """Detached env-zero inputs for one selected coordinated grasp image.""" + + object_label: str + mesh_vertices: torch.Tensor + mesh_triangles: torch.Tensor + object_pose: torch.Tensor + left_grasp_pose: torch.Tensor + right_grasp_pose: torch.Tensor + + +def _pose_row(value: torch.Tensor, env_id: int, *, name: str) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32) + if pose.shape == (4, 4): + if env_id != 0: + raise ValueError(f"{name} has no row for environment {env_id}.") + result = pose + elif pose.ndim == 3 and pose.shape[1:] == (4, 4) and pose.shape[0] > env_id: + result = pose[env_id] + else: + raise ValueError(f"{name} must contain a (4, 4) pose for environment {env_id}.") + if not bool(torch.isfinite(result).all().item()): + raise ValueError(f"{name} must contain only finite values.") + return result.detach().cpu().clone() + + +def selected_coordinated_grasp_scene( + outcome: ActionOutcome, + *, + left_control_part: str, + right_control_part: str, + env_id: int = 0, +) -> CoordinatedGraspPoseScene | None: + """Extract the exact valid grasp pair selected by E5 for one environment.""" + if not isinstance(outcome, ActionOutcome): + raise TypeError("outcome must be an ActionOutcome.") + if type(env_id) is not int or env_id < 0: + raise ValueError("env_id must be a non-negative integer.") + success = torch.as_tensor(outcome.success, dtype=torch.bool).reshape(-1) + if success.numel() <= env_id or not bool(success[env_id].item()): + return None + object_pose_value = outcome.grounded.object_pose + if object_pose_value is None: + raise ValueError("Selected E5 outcome does not retain its live object pose.") + object_pose = _pose_row(object_pose_value, env_id, name="live object pose") + left_held = outcome.next_state.get_held_object(left_control_part) + right_held = outcome.next_state.get_held_object(right_control_part) + if left_held is None or right_held is None: + raise ValueError("Selected E5 outcome does not retain both held-object states.") + if left_held.semantics.entity_id != right_held.semantics.entity_id: + raise ValueError("Selected E5 grasps must refer to the same object.") + affordance = left_held.semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise TypeError("Selected E5 object must retain an AntipodalAffordance.") + left_relation = _pose_row( + left_held.object_to_eef, + env_id, + name="left object-to-EEF pose", + ) + right_relation = _pose_row( + right_held.object_to_eef, + env_id, + name="right object-to-EEF pose", + ) + return CoordinatedGraspPoseScene( + object_label=left_held.semantics.label, + mesh_vertices=torch.as_tensor( + affordance.mesh_vertices, + dtype=torch.float32, + ) + .detach() + .cpu() + .clone(), + mesh_triangles=torch.as_tensor( + affordance.mesh_triangles, + dtype=torch.int64, + ) + .detach() + .cpu() + .clone(), + object_pose=object_pose, + left_grasp_pose=object_pose @ left_relation, + right_grasp_pose=object_pose @ right_relation, + ) + + +def grasp_pose_image_path(output_root: str | Path, step_id: str) -> Path: + """Return the stable env-zero image path for one E5 semantic step.""" + safe_step = _SAFE_NAME.sub("_", str(step_id)).strip("._") + if not safe_step: + raise ValueError("step_id must contain a usable path component.") + return Path(output_root) / "env_0000" / "grasp_poses" / f"{safe_step}.png" + + +def _material(rendering: object, color: tuple[float, float, float, float]) -> object: + material = rendering.MaterialRecord() + material.shader = "defaultUnlit" + material.base_color = color + return material + + +def render_coordinated_grasp_pose_png( + scene: CoordinatedGraspPoseScene, + output_path: str | Path, +) -> Path: + """Render one object mesh and its selected left/right grasp frames to PNG.""" + if not isinstance(scene, CoordinatedGraspPoseScene): + raise TypeError("scene must be a CoordinatedGraspPoseScene.") + path = Path(output_path).expanduser().resolve() + if path.suffix.lower() != ".png": + raise ValueError("Grasp pose visualization output must use a .png suffix.") + + import open3d as o3d + + renderer = o3d.visualization.rendering.OffscreenRenderer( + _IMAGE_WIDTH, + _IMAGE_HEIGHT, + ) + renderer.scene.set_background(np.array([1.0, 1.0, 1.0, 1.0])) + + mesh = o3d.geometry.TriangleMesh( + vertices=o3d.utility.Vector3dVector(scene.mesh_vertices.numpy()), + triangles=o3d.utility.Vector3iVector(scene.mesh_triangles.numpy()), + ) + mesh.compute_vertex_normals() + mesh.transform(scene.object_pose.numpy()) + renderer.scene.add_geometry( + "object", + mesh, + _material(o3d.visualization.rendering, (0.25, 0.62, 0.34, 1.0)), + ) + + world_vertices = np.asarray(mesh.vertices) + mesh_extent = float(np.ptp(world_vertices, axis=0).max()) + frame_size = max(0.035, min(0.12, 0.35 * mesh_extent)) + center_radius = max(0.006, 0.075 * frame_size) + for name, pose, color in ( + ("left", scene.left_grasp_pose, (0.82, 0.16, 0.58, 1.0)), + ("right", scene.right_grasp_pose, (0.05, 0.55, 0.80, 1.0)), + ): + pose_np = pose.numpy() + frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=frame_size) + frame.transform(pose_np) + renderer.scene.add_geometry( + f"{name}_frame", + frame, + _material(o3d.visualization.rendering, (1.0, 1.0, 1.0, 1.0)), + ) + center = o3d.geometry.TriangleMesh.create_sphere(radius=center_radius) + center.compute_vertex_normals() + center.translate(pose_np[:3, 3]) + renderer.scene.add_geometry( + f"{name}_center", + center, + _material(o3d.visualization.rendering, color), + ) + + bounds = renderer.scene.bounding_box + center = np.asarray(bounds.get_center(), dtype=np.float64) + extent = max(float(np.asarray(bounds.get_extent()).max()), 0.10) + view = np.array([1.4, -1.8, 1.2], dtype=np.float64) + view /= np.linalg.norm(view) + renderer.setup_camera( + 50.0, + center, + center + view * (2.2 * extent), + np.array([0.0, 0.0, 1.0]), + ) + image = renderer.render_to_image() + path.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".png", + ) + os.close(descriptor) + temporary_path = Path(temporary_name) + try: + if not o3d.io.write_image(temporary_path.as_posix(), image): + raise OSError(f"Open3D could not write {temporary_path}.") + os.replace(temporary_path, path) + finally: + temporary_path.unlink(missing_ok=True) + return path diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py new file mode 100644 index 000000000..ea273a8f9 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py @@ -0,0 +1,203 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Stage-separated diagnostics for GenSim antipodal grasp generation.""" + +from __future__ import annotations + +from copy import deepcopy +from typing import Any + +import torch +import torch.nn.functional as F + +from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator + +__all__: list[str] = [] + + +class _TracingAntipodalGraspPoseGenerator(AntipodalGraspPoseGenerator): + """Retain compact S1-S5 evidence from the concrete GenSim grasp backend.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._last_dual_trace: dict[str, Any] | None = None + + @property + def last_dual_trace(self) -> dict[str, Any] | None: + """Return an owned snapshot of the most recent dual-grasp trace.""" + return deepcopy(self._last_dual_trace) + + @staticmethod + def _transform_points(points: torch.Tensor, pose: torch.Tensor) -> torch.Tensor: + return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + @classmethod + def _filter_counts( + cls, + backend: Any, + *, + mesh_vertices: torch.Tensor, + object_pose: torch.Tensor, + arm_direction: torch.Tensor, + approach_direction: torch.Tensor, + middle_empty_ratio: float, + ) -> dict[str, int]: + pairs = backend.antipodal_pairs.to(dtype=torch.float32) + origin = cls._transform_points(pairs[:, 0], object_pose) + hit = cls._transform_points(pairs[:, 1], object_pose) + world_vertices = cls._transform_points(mesh_vertices, object_pose) + projection = torch.matmul(world_vertices, arm_direction) + extent = projection.max() - projection.min() + left_threshold = projection.min() + extent * (0.5 - middle_empty_ratio * 0.5) + right_threshold = projection.max() - extent * (0.5 - middle_empty_ratio * 0.5) + origin_projection = torch.matmul(origin, arm_direction) + hit_projection = torch.matmul(hit, arm_direction) + masks = { + "left": (origin_projection < left_threshold) + | (hit_projection < left_threshold), + "right": (origin_projection > right_threshold) + | (hit_projection > right_threshold), + } + counts: dict[str, int] = {} + for side, mask in masks.items(): + grasp_x = F.normalize(hit[mask] - origin[mask], dim=1) + cosine = torch.clamp( + torch.sum(grasp_x * approach_direction, dim=1), -1.0, 1.0 + ) + angle = torch.abs(torch.acos(cosine)) + angle_valid = torch.abs(angle - torch.pi * 0.5) <= float( + backend._max_deviation_angle + ) + counts[f"{side}_partition_pair_count"] = int(mask.sum().item()) + counts[f"{side}_angle_valid_pair_count"] = int(angle_valid.sum().item()) + return counts + + @staticmethod + def _candidate_count(result: dict[str, Any]) -> int: + poses = result.get("grasp_poses") + if not result.get("is_success", False) or not isinstance(poses, torch.Tensor): + return 0 + return int(poses.shape[0]) if poses.ndim == 3 else 0 + + def get_dual_arm_valid_grasp_poses(self, **kwargs: Any) -> list[dict | None]: + """Run the standard generator while observing its NMS/collision boundary.""" + vertices = kwargs["mesh_vertices"] + triangles = kwargs["mesh_triangles"] + backend = self._backend(vertices, triangles) + poses = self._object_poses(kwargs["obj_poses"], device=backend.device) + directions = self._approach_directions( + kwargs["approach_direction"], + batch_size=poses.shape[0], + device=backend.device, + ) + arm_direction = self._approach_directions( + kwargs["left_to_right_arm_direction"], + batch_size=1, + device=backend.device, + )[0] + ratio = float(kwargs.get("middle_empty_ratio", 0.4)) + collision_records: list[dict[str, Any]] = [] + checker = backend._collision_checker + original_query = checker.query + + def traced_query(*args: Any, **query_kwargs: Any): + colliding, distance = original_query(*args, **query_kwargs) + distance = torch.as_tensor(distance, dtype=torch.float32) + collision_records.append( + { + "nms_candidate_count": int(colliding.numel()), + "noncolliding_candidate_count": int((~colliding).sum().item()), + "minimum_signed_distance": float(distance.min().item()), + "maximum_signed_distance": float(distance.max().item()), + } + ) + return colliding, distance + + checker.query = traced_query + try: + results = super().get_dual_arm_valid_grasp_poses(**kwargs) + finally: + checker.query = original_query + + row_traces: list[dict[str, Any]] = [] + for row_index, (object_pose, approach, result) in enumerate( + zip(poses, directions, results, strict=True) + ): + filter_counts = self._filter_counts( + backend, + mesh_vertices=vertices.to(device=backend.device, dtype=torch.float32), + object_pose=object_pose, + arm_direction=arm_direction, + approach_direction=approach, + middle_empty_ratio=ratio, + ) + record_offset = 2 * row_index + records = collision_records[record_offset : record_offset + 2] + left_record = records[0] if len(records) >= 1 else {} + right_record = records[1] if len(records) >= 2 else {} + left = {} if result is None else result["left"] + right = {} if result is None else result["right"] + left_final = self._candidate_count(left) + right_final = self._candidate_count(right) + row_traces.append( + { + "environment_index": row_index, + "approach_direction": approach.detach().cpu().tolist(), + "middle_empty_ratio": ratio, + "S1_grasp_pair_generation": { + "antipodal_pair_count": int(backend.antipodal_pairs.shape[0]), + }, + "S2_approach_angle_filtering": filter_counts, + "S3_nms": { + "left_candidate_count": int( + left_record.get("nms_candidate_count", 0) + ), + "right_candidate_count": int( + right_record.get("nms_candidate_count", 0) + ), + }, + "S4_collision_filtering": { + "left_candidate_count": int( + left_record.get("noncolliding_candidate_count", 0) + ), + "right_candidate_count": int( + right_record.get("noncolliding_candidate_count", 0) + ), + "left_minimum_signed_distance": left_record.get( + "minimum_signed_distance" + ), + "left_maximum_signed_distance": left_record.get( + "maximum_signed_distance" + ), + "right_minimum_signed_distance": right_record.get( + "minimum_signed_distance" + ), + "right_maximum_signed_distance": right_record.get( + "maximum_signed_distance" + ), + }, + "S5_left_right_pairing": { + "left_final_count": left_final, + "right_final_count": right_final, + "paired": left_final > 0 and right_final > 0, + }, + } + ) + self._last_dual_trace = ( + row_traces[0] if len(row_traces) == 1 else {"environment_rows": row_traces} + ) + return results diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index 45feef127..afc86ef71 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -266,7 +266,7 @@ def _arm_values( if callable(getter): left, right = getter() values = [] - for value in (left, right): + for side, value in zip(("left", "right"), (left, right)): if value is None: values.append(None) continue @@ -275,6 +275,17 @@ def _arm_values( item = item.unsqueeze(0).repeat(int(env.num_envs), 1, 1) elif kind == "gripper_state" and item.ndim == 1: item = item.unsqueeze(0) + if kind == "gripper_state": + configured = getattr( + env, + "agent_gripper_state_joint_indices", + {}, + ) + indices = ( + configured.get(side) if isinstance(configured, Mapping) else None + ) + if indices is not None: + item = item[:, list(indices)] values.append(item) return values[0], values[1] if kind != "gripper_state": @@ -285,7 +296,12 @@ def _arm_values( ids = list(getattr(env, f"{side}_eef_joints", ())) if not ids: return None - values.append(qpos[:, ids]) + item = qpos[:, ids] + configured = getattr(env, "agent_gripper_state_joint_indices", {}) + indices = configured.get(side) if isinstance(configured, Mapping) else None + if indices is not None: + item = item[:, list(indices)] + values.append(item) return values[0], values[1] @@ -681,11 +697,13 @@ def evaluate_predicate( dim=-1, ) if kind in {"both_grippers_open", "grippers_open"}: - if not hasattr(env, "get_current_gripper_state_agent"): + gripper_values = _arm_values(env, "gripper_state") + if gripper_values is None: return _constant(env, False) - left, right = env.get_current_gripper_state_agent() results = [] - for side, value in zip(("left", "right"), (left, right)): + for side, value in zip(("left", "right"), gripper_values): + if value is None: + return _constant(env, False) value = torch.as_tensor(value, dtype=torch.float32, device=env.device) if value.ndim == 1: value = value.unsqueeze(0).repeat(int(env.num_envs), 1) @@ -697,6 +715,10 @@ def evaluate_predicate( ) if expected.ndim == 1: expected = expected.unsqueeze(0).repeat(int(env.num_envs), 1) + configured = getattr(env, "agent_gripper_state_joint_indices", {}) + indices = configured.get(side) if isinstance(configured, Mapping) else None + if indices is not None: + expected = expected[:, list(indices)] results.append( torch.linalg.vector_norm(value - expected, dim=-1) <= float(spec.get("tolerance", defaults["gripper_state_tolerance"])) diff --git a/embodichain/gen_sim/action_engine/solver_profiles.py b/embodichain/gen_sim/action_engine/solver_profiles.py new file mode 100644 index 000000000..120fbea77 --- /dev/null +++ b/embodichain/gen_sim/action_engine/solver_profiles.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. +# ---------------------------------------------------------------------------- + +"""Generation-time IK solver selection for GenSim robot bundles.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any, Final + +__all__ = [ + "IK_SOLVER_MODES", + "expected_ik_solver_class", + "resolve_ik_solver_mode", + "validate_robot_ik_solver_contract", +] + +IK_SOLVER_MODES: Final = ("auto", "ur", "pytorch") +_UR_PROFILES = frozenset({"dual_ur3", "dual_ur5", "dual_ur10"}) +_SUPPORTED_PROFILES = _UR_PROFILES | {"dual_franka"} +_CLASS_BY_MODE = {"ur": "URSolver", "pytorch": "PytorchSolver"} + + +def resolve_ik_solver_mode(mode: str, robot_profile: str) -> str: + """Resolve one requested mode to a concrete solver for a robot profile.""" + if not isinstance(mode, str): + raise TypeError( + "IK solver mode must be a string; expected one of: auto, ur, pytorch." + ) + if mode not in IK_SOLVER_MODES: + raise ValueError( + f"Unsupported IK solver mode {mode!r}; expected one of: " + "auto, ur, pytorch." + ) + profile = str(robot_profile) + if profile not in _SUPPORTED_PROFILES: + raise ValueError(f"Unsupported IK solver robot profile {profile!r}.") + resolved = "pytorch" if mode == "auto" and profile == "dual_franka" else mode + if resolved == "auto": + resolved = "ur" + if resolved == "ur" and profile == "dual_franka": + raise ValueError("Franka does not support the analytical URSolver.") + return resolved + + +def expected_ik_solver_class(mode: str) -> str: + """Return the serialized/runtime class name for one concrete mode.""" + if mode not in _CLASS_BY_MODE: + raise ValueError("Concrete IK solver mode must be 'ur' or 'pytorch'.") + return _CLASS_BY_MODE[mode] + + +def validate_robot_ik_solver_contract( + robot: Mapping[str, Any], + mode: str, +) -> None: + """Validate that both generated arms use the declared concrete solver.""" + expected = expected_ik_solver_class(mode) + solvers = robot.get("solver_cfg") + if not isinstance(solvers, Mapping): + raise ValueError("Generated robot requires a solver_cfg mapping.") + for arm in ("left_arm", "right_arm"): + solver = solvers.get(arm) + if not isinstance(solver, Mapping): + raise ValueError(f"Generated robot requires solver_cfg.{arm}.") + actual = solver.get("class_type") + if actual != expected: + raise ValueError( + f"Generated robot {arm} must use {expected} for ik_solver={mode!r}, " + f"got {actual!r}." + ) diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py index fd214d4e0..3b7ee42c3 100644 --- a/embodichain/gen_sim/task_engine/cli.py +++ b/embodichain/gen_sim/task_engine/cli.py @@ -78,6 +78,7 @@ def build_parser() -> argparse.ArgumentParser: run_parser.add_argument("--seed", type=int, default=0) run_parser.add_argument("--num-envs", type=int, default=None) run_parser.add_argument("--dataset-saving", action="store_true") + run_parser.add_argument("--show-grasp-poses", action="store_true") _add_failure_policy_argument(run_parser) return parser @@ -112,6 +113,17 @@ def _add_workflow_arguments(parser: argparse.ArgumentParser) -> None: default=None, help="Override the Task Engine planning gripper profile.", ) + parser.add_argument( + "--ik-solver", + choices=("auto", "ur", "pytorch"), + default=None, + help="Override the generation-time IK solver for both arms.", + ) + parser.add_argument( + "--show-grasp-poses", + action="store_true", + help="Write one static PNG for the valid grasp pair selected by E5.", + ) parser.add_argument( "--planner-mode", choices=_PLANNER_MODES, @@ -173,6 +185,8 @@ def _run_workflow( workflow_cfg, planning_cfg, execution_cfg = load_task_engine_config(args.config) if args.gripper_model is not None: planning_cfg = replace(planning_cfg, gripper_model=args.gripper_model) + if args.ik_solver is not None: + planning_cfg = replace(planning_cfg, ik_solver=args.ik_solver) if args.planner_mode is not None: planning_cfg = replace( planning_cfg, @@ -203,6 +217,7 @@ def _run_workflow( base_seed=args.base_seed, dataset_saving=args.dataset_saving, failure_policy=args.failure_policy, + show_grasp_poses=args.show_grasp_poses, run_id=allocation.run_id, created_at=allocation.created_at, execute=execute, @@ -236,6 +251,7 @@ def _run_prepared_bundle(args: argparse.Namespace) -> int: num_envs=num_envs, dataset_saving=bool(args.dataset_saving), failure_policy=args.failure_policy, + show_grasp_poses=bool(args.show_grasp_poses), ) environments = report.get("environments", ()) successes = [ diff --git a/embodichain/gen_sim/task_engine/config.py b/embodichain/gen_sim/task_engine/config.py index 2b5b94f70..0653bd322 100644 --- a/embodichain/gen_sim/task_engine/config.py +++ b/embodichain/gen_sim/task_engine/config.py @@ -37,6 +37,7 @@ ] TASK_ENGINE_DEFAULTS_SCHEMA: Final = "embodichain.task-engine-defaults/v1" +_IK_SOLVER_MODES: Final = ("auto", "ur", "pytorch") @configclass @@ -109,6 +110,7 @@ class TaskEnginePlanningCfg: candidate_count: int = 3 planning_mode: str = "offline" gripper_model: str = "pgi" + ik_solver: str = "auto" max_episodes: int = 1 max_episode_steps: int = 6000 planner: dict[str, Any] = {} @@ -129,6 +131,11 @@ def __post_init__(self) -> None: f"Unsupported gripper model {self.gripper_model!r}; expected one " "of: pgi, robotiq." ) + if self.ik_solver not in _IK_SOLVER_MODES: + raise ValueError( + f"Unsupported IK solver mode {self.ik_solver!r}; expected one of: " + "auto, ur, pytorch." + ) if not isinstance(self.planner, Mapping): raise TypeError("planner must be a mapping.") from embodichain.gen_sim.action_engine.config.runtime_policy import ( @@ -184,6 +191,7 @@ def load_task_engine_config( "candidate_count", "planning_mode", "gripper_model", + "ik_solver", "max_episodes", "max_episode_steps", } diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml index 6e03b85e5..4c58490d7 100644 --- a/embodichain/gen_sim/task_engine/defaults.yaml +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -25,6 +25,7 @@ planning: candidate_count: 3 planning_mode: offline gripper_model: robotiq + ik_solver: pytorch max_episodes: 1 max_episode_steps: 6000 planner: diff --git a/embodichain/gen_sim/task_engine/orchestration/coordinator.py b/embodichain/gen_sim/task_engine/orchestration/coordinator.py index a7ef2de55..cc1aaf0b7 100644 --- a/embodichain/gen_sim/task_engine/orchestration/coordinator.py +++ b/embodichain/gen_sim/task_engine/orchestration/coordinator.py @@ -182,6 +182,7 @@ def prepare( overwrite: bool = False, planning_mode: str = "offline", gripper_model: str = "pgi", + ik_solver: str = "auto", vlm_model: str | None = None, max_episodes: int | None = None, max_episode_steps: int | None = None, @@ -359,6 +360,7 @@ def prepare( "task_spec": grounded_plan["task_spec"], "robot_profile": robot_profile, "gripper_model": gripper_model, + "ik_solver": ik_solver, "source_scene_z_rotation_degrees": ( adaptation.prepared_scene.z_rotation_degrees ), diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py index b3a4bca09..b89ddfcd1 100644 --- a/embodichain/gen_sim/task_engine/workflow.py +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -107,6 +107,7 @@ def __call__( num_envs: int, dataset_saving: bool = False, failure_policy: str = "stop", + show_grasp_poses: bool = False, ) -> Mapping[str, Any]: """Run one simulator attempt and preserve its report and trajectory. @@ -118,6 +119,7 @@ def __call__( dataset_saving: Whether to enable the Gym project's dataset recorder. failure_policy: Whether failed dependencies stop or permit downstream diagnostic execution. + show_grasp_poses: Whether to write the valid E5 grasp pair as a PNG. Returns: Validated Action Engine execution report. @@ -126,6 +128,8 @@ def __call__( attempt_root = Path(output_root).expanduser().resolve() if failure_policy not in {"stop", "continue"}: raise ValueError("failure_policy must be 'stop' or 'continue'.") + if not isinstance(show_grasp_poses, bool): + raise TypeError("show_grasp_poses must be a boolean.") attempt_root.mkdir(parents=True, exist_ok=False) command = [ sys.executable, @@ -142,6 +146,8 @@ def __call__( ] if not dataset_saving: command.append("--filter_dataset_saving") + if show_grasp_poses: + command.append("--show-grasp-poses") command.extend(["--failure-policy", failure_policy]) log_path = attempt_root / "action.log" print( @@ -296,6 +302,7 @@ def run( base_seed: int = 0, dataset_saving: bool = False, failure_policy: str = "stop", + show_grasp_poses: bool = False, run_id: str | None = None, created_at: datetime | None = None, overwrite: bool = False, @@ -315,6 +322,7 @@ def run( dataset_saving: Whether Action attempts may initialize dataset recording. failure_policy: Whether failed dependencies stop or permit downstream diagnostic execution. + show_grasp_poses: Whether to write the valid E5 grasp pair as a PNG. run_id: Optional externally allocated run identifier. created_at: Optional timezone-aware run creation timestamp. overwrite: Whether to atomically replace an existing run directory. @@ -326,6 +334,8 @@ def run( normalized = validate_task_run_request(request) if not isinstance(dataset_saving, bool): raise TypeError("dataset_saving must be a boolean.") + if not isinstance(show_grasp_poses, bool): + raise TypeError("show_grasp_poses must be a boolean.") if failure_policy not in {"stop", "continue"}: raise ValueError("failure_policy must be 'stop' or 'continue'.") if workflow_cfg is None or planning_cfg is None or execution_cfg is None: @@ -652,6 +662,7 @@ def run( candidate_count=effective_candidate_count, planning_mode=planning_cfg.planning_mode, gripper_model=planning_cfg.gripper_model, + ik_solver=planning_cfg.ik_solver, vlm_model=vlm_model, max_episodes=planning_cfg.max_episodes, max_episode_steps=planning_cfg.max_episode_steps, @@ -860,13 +871,18 @@ def run( "error": None, } try: + execution_options = { + "seed": action_seed, + "num_envs": execution_cfg.num_envs, + "dataset_saving": bool(dataset_saving), + "failure_policy": failure_policy, + } + if show_grasp_poses: + execution_options["show_grasp_poses"] = True report = self.action_executor( preparation.output_dir, action_root, - seed=action_seed, - num_envs=execution_cfg.num_envs, - dataset_saving=bool(dataset_saving), - failure_policy=failure_policy, + **execution_options, ) successes = _environment_successes( report, @@ -1038,6 +1054,7 @@ def _publish( "candidate_count": planning_cfg.candidate_count, "planning_mode": planning_cfg.planning_mode, "gripper_model": planning_cfg.gripper_model, + "ik_solver": planning_cfg.ik_solver, "max_episodes": planning_cfg.max_episodes, "max_episode_steps": planning_cfg.max_episode_steps, "planner": deepcopy(planning_cfg.planner), diff --git a/tests/gen_sim/action_engine/cli/test_run_agent.py b/tests/gen_sim/action_engine/cli/test_run_agent.py index 3a83339db..40ffeaf46 100644 --- a/tests/gen_sim/action_engine/cli/test_run_agent.py +++ b/tests/gen_sim/action_engine/cli/test_run_agent.py @@ -31,6 +31,7 @@ _prepare_ab_branches, _publish_task_engine_report, _task_engine_exit_code, + _validate_run_contract, ) from embodichain.gen_sim.action_engine.runtime import ( ExecutionReport, @@ -65,6 +66,58 @@ def __init__(self, recorder=None) -> None: ) +def _solver_run_contract(ik_solver: str) -> tuple[dict, dict]: + class_type = "URSolver" if ik_solver == "ur" else "PytorchSolver" + gym_config = { + "env": { + "extensions": { + "action_engine": { + "task_name": "solver_task", + "seed_task_graph_hash": "a" * 64, + "planning_mode": "offline", + "gripper_model": "pgi", + "ik_solver": ik_solver, + }, + "agent_ik_solver": ik_solver, + } + }, + "robot": { + "solver_cfg": { + "left_arm": {"class_type": class_type}, + "right_arm": {"class_type": class_type}, + } + }, + } + agent_config = { + "task_name": "solver_task", + "seed_task_graph_hash": "a" * 64, + "planning_mode": "offline", + "gripper_model": "pgi", + "ik_solver": ik_solver, + } + return gym_config, agent_config + + +@pytest.mark.parametrize("ik_solver", ["ur", "pytorch"]) +def test_run_contract_accepts_matching_concrete_ik_solver(ik_solver: str) -> None: + gym_config, agent_config = _solver_run_contract(ik_solver) + + _validate_run_contract(gym_config, agent_config, "solver_task") + + +def test_run_contract_rejects_ik_solver_artifact_drift() -> None: + gym_config, agent_config = _solver_run_contract("pytorch") + agent_config["ik_solver"] = "ur" + + with pytest.raises(ValueError, match="different IK solvers"): + _validate_run_contract(gym_config, agent_config, "solver_task") + + agent_config["ik_solver"] = "pytorch" + gym_config["robot"]["solver_cfg"]["right_arm"]["class_type"] = "URSolver" + with pytest.raises(ValueError, match="right_arm.*PytorchSolver"): + _validate_run_contract(gym_config, agent_config, "solver_task") + + def test_capture_ab_initial_frame_invokes_only_audience_recorder() -> None: recorder = record_camera_data() env = _FakeEnv(recorder) diff --git a/tests/gen_sim/action_engine/config/test_runtime_policy.py b/tests/gen_sim/action_engine/config/test_runtime_policy.py index bbdfc7e0d..b2ffe9a7b 100644 --- a/tests/gen_sim/action_engine/config/test_runtime_policy.py +++ b/tests/gen_sim/action_engine/config/test_runtime_policy.py @@ -121,6 +121,7 @@ def test_defaults_cover_current_execution_and_generation_policy() -> None: ) assert generation["physics"]["rigid_object"]["mass"] == pytest.approx(0.1) assert generation["task"]["default_gripper_model"] == "pgi" + assert generation["task"]["default_ik_solver"] == "auto" assert generation["environment"]["arm_aim_yaw_offset"] == { "left": pytest.approx(0.0), "right": pytest.approx(0.0), diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py index b425bea15..d30f2e42b 100644 --- a/tests/gen_sim/action_engine/generation/test_generation.py +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -867,6 +867,85 @@ def test_config_builders_reject_unknown_gripper_before_materialization( ) +@pytest.mark.parametrize( + ("ik_solver", "class_type"), + (("ur", "URSolver"), ("pytorch", "PytorchSolver")), +) +def test_fast_gym_config_materializes_selected_ur10_ik_solver( + gym_export: Path, + ik_solver: str, + class_type: str, +) -> None: + scene = prepare_scene(gym_export) + config = build_fast_gym_config( + scene, + task_name="solver_profile_task", + task_description="Exercise one generated IK solver.", + robot_profile="ur10", + gripper_model="robotiq", + ik_solver=ik_solver, + execution_program_hash="d" * 64, + max_episodes=1, + max_episode_steps=20, + ) + robot = config["robot"] + extension = config["env"]["extensions"] + tcp = [list(row) for row in get_gripper_profile("robotiq").tcp_transform] + + assert extension["action_engine"]["ik_solver"] == ik_solver + assert extension["agent_ik_solver"] == ik_solver + for arm in ("left_arm", "right_arm"): + solver = robot["solver_cfg"][arm] + assert solver["class_type"] == class_type + assert solver["tcp"] == tcp + assert solver["end_link_name"] == f"{arm.split('_')[0]}_ee_link" + assert solver["root_link_name"] == f"{arm.split('_')[0]}_base_link" + if ik_solver == "pytorch": + assert solver["num_samples"] == 30 + + +def test_fast_gym_config_rejects_analytic_ur_solver_for_franka( + gym_export: Path, +) -> None: + scene = prepare_scene(gym_export) + + with pytest.raises(ValueError, match="Franka.*URSolver"): + build_fast_gym_config( + scene, + task_name="invalid_solver", + task_description="Reject incompatible IK.", + robot_profile="franka", + ik_solver="ur", + execution_program_hash="e" * 64, + max_episodes=1, + max_episode_steps=20, + ) + + +def test_agent_config_serializes_concrete_ik_solver(gym_export: Path) -> None: + scene = prepare_scene(gym_export) + + ur = build_agent_config( + task_name="ur_solver_task", + robot_profile="ur10", + ik_solver="auto", + execution_program_hash="f" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + ) + pytorch = build_agent_config( + task_name="pytorch_solver_task", + robot_profile="ur10", + ik_solver="pytorch", + execution_program_hash="1" * 64, + source_config_path=scene.source_config_path, + uid_map=scene.uid_map, + ) + + assert ur["ik_solver"] == "ur" + assert pytorch["ik_solver"] == "pytorch" + + def test_agent_config_serializes_selected_gripper(gym_export: Path) -> None: scene = prepare_scene(gym_export) config = build_agent_config( @@ -1652,6 +1731,7 @@ def test_generation_cli_defaults_to_mature_robot_without_scene_randomization() - assert args.randomize_scene is False assert args.planning_mode == "offline" assert args.planner_mode is None + assert args.ik_solver == "auto" assert not hasattr(args, "instruction_parser") assert not hasattr(args, "task_agent") diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index c918cfe9d..33300de93 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -705,11 +705,22 @@ def test_coordinated_pickment_geometry_candidates_are_live_and_continuous() -> N ) preferred = [ - candidates[0].cfg["middle_empty_ratio"] + next( + candidate.cfg["middle_empty_ratio"] + for candidate in candidates + if candidate.motion_policy["coordinated_grasp"]["approach_candidate_label"] + == "robot_forward_down" + ) for candidates in (vertical, tilted, horizontal) ] assert preferred[0] < preferred[1] < preferred[2] - assert yawed[0].cfg["middle_empty_ratio"] == pytest.approx(preferred[0]) + yawed_forward_down = next( + candidate + for candidate in yawed + if candidate.motion_policy["coordinated_grasp"]["approach_candidate_label"] + == "robot_forward_down" + ) + assert yawed_forward_down.cfg["middle_empty_ratio"] == pytest.approx(preferred[0]) for candidates in (vertical, tilted, horizontal, yawed): assert candidates assert torch.allclose( @@ -718,10 +729,47 @@ def test_coordinated_pickment_geometry_candidates_are_live_and_continuous() -> N ) assert torch.allclose( candidates[0].cfg["approach_direction"], - torch.tensor([0.0, 0.0, -1.0]), + torch.tensor([2**-0.5, 0.0, -(2**-0.5)]), + ) + assert any( + torch.allclose( + candidate.cfg["approach_direction"], + torch.tensor([0.0, 0.0, -1.0]), + ) + for candidate in candidates ) +def test_coordinated_pickment_approach_family_follows_live_shared_reach() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + grounded = _coordinated_grounded( + _rotation_z(90.0), + vertices=_cuboid_vertices(0.20, 0.08, 0.02), + ) + left = torch.eye(4).repeat(2, 1, 1) + right = torch.eye(4).repeat(2, 1, 1) + left[:, 0, 3] = -0.3 + right[:, 0, 3] = 0.3 + left[:, 1, 3] = right[:, 1, 3] = -0.5 + adapter._coordinated_arm_bases = lambda: (left, right) + + candidates = adapter._adapt_coordinated_pickment_grasps(grounded, capability) + + trace = candidates[0].motion_policy["coordinated_grasp"] + assert trace["approach_candidate_label"] == "robot_forward_down" + torch.testing.assert_close( + candidates[0].cfg["approach_direction"], + torch.tensor([0.0, 2**-0.5, -(2**-0.5)]), + ) + assert trace["grasp_seed"] == 17_393 + assert [item["label"] for item in trace["approach_candidates"][:3]] == [ + "robot_forward_down", + "current", + "robot_forward", + ] + + def test_coordinated_pickment_geometry_candidates_are_deterministic_for_tray() -> None: adapter = AtomicActionAdapter(_planner_env()) capability = adapter.capabilities.get("CoordinatedPickment") @@ -775,6 +823,38 @@ def test_robotiq_grasp_generator_preserves_existing_geometry() -> None: assert generator.collision_cfg.opening_margin == pytest.approx(0.01) +def test_coordinated_robotiq_generator_uses_e5_opening_margin_only() -> None: + adapter = AtomicActionAdapter(_planner_env(gripper_model="robotiq")) + + coordinated = adapter._grasp_pose_generators( + filter_ground_collision=False, + opening_margin=0.02, + )["physical_left_eef"] + ordinary = adapter._grasp_pose_generators()["physical_left_eef"] + + assert coordinated.collision_cfg.opening_margin == pytest.approx(0.02) + assert ordinary.collision_cfg.opening_margin == pytest.approx(0.01) + + +def test_adapter_validates_declared_runtime_ik_solver_classes() -> None: + pytorch_solver_type = type("PytorchSolver", (), {}) + env = _planner_env() + env.agent_ik_solver = "pytorch" + env.robot.get_solver = lambda **_kwargs: pytorch_solver_type() + + adapter = AtomicActionAdapter(env) + + assert adapter.ik_solver == "pytorch" + assert adapter.ik_solver_classes == { + "left_arm": "PytorchSolver", + "right_arm": "PytorchSolver", + } + + env.agent_ik_solver = "ur" + with pytest.raises(ValueError, match="left_arm.*URSolver.*PytorchSolver"): + AtomicActionAdapter(env) + + @pytest.mark.parametrize("gripper_model", ["pgi", "robotiq"]) def test_control_profiles_use_selected_gripper_joint_semantics( gripper_model: str, diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_debug.py b/tests/gen_sim/action_engine/runtime/test_grasp_debug.py new file mode 100644 index 000000000..613f31cdd --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_grasp_debug.py @@ -0,0 +1,133 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from embodichain.gen_sim.action_engine.runtime.grasp_debug import ( + selected_coordinated_grasp_scene, +) +from embodichain.gen_sim.action_engine.runtime.models import ( + ActionOutcome, + GroundedAction, +) +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + HeldObjectState, + ObjectSemantics, +) + + +def _pose(x: float, y: float, z: float) -> torch.Tensor: + pose = torch.eye(4, dtype=torch.float32).unsqueeze(0) + pose[:, :3, 3] = torch.tensor([x, y, z], dtype=torch.float32) + return pose + + +def _outcome(*, success: bool) -> ActionOutcome: + vertices = torch.tensor( + [ + [-0.1, -0.05, -0.02], + [0.1, -0.05, -0.02], + [0.0, 0.05, -0.02], + [0.0, 0.0, 0.02], + ], + dtype=torch.float32, + ) + triangles = torch.tensor( + [[0, 1, 2], [0, 1, 3], [1, 2, 3], [2, 0, 3]], + dtype=torch.int64, + ) + semantics = ObjectSemantics( + affordance=AntipodalAffordance( + object_label="tray", + mesh_vertices=vertices, + mesh_triangles=triangles, + ), + geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + entity_id="tray", + label="tray", + ) + live_object_pose = _pose(0.4, -0.2, 0.75) + left_object_to_eef = _pose(0.0, -0.12, 0.01) + right_object_to_eef = _pose(0.0, 0.13, 0.02) + state = ExecutionState( + last_qpos=torch.zeros(1, 8), + held_objects={ + "physical_left_arm": HeldObjectState( + semantics=semantics, + object_to_eef=left_object_to_eef, + grasp_xpos=torch.bmm(live_object_pose, left_object_to_eef), + ), + "physical_right_arm": HeldObjectState( + semantics=semantics, + object_to_eef=right_object_to_eef, + grasp_xpos=torch.bmm(live_object_pose, right_object_to_eef), + ), + }, + ) + grounded = GroundedAction( + action_class="CoordinatedPickment", + arm="coordinated", + control="coordinated", + target=SimpleNamespace(), + cfg={}, + object_pose=live_object_pose, + object_uid="tray", + ) + return ActionOutcome( + trajectory=torch.zeros(1, 1, 8), + success=torch.tensor([success]), + next_state=state, + grounded=grounded, + ) + + +def test_selected_scene_uses_final_e5_grasps_at_live_object_pose() -> None: + outcome = _outcome(success=True) + + scene = selected_coordinated_grasp_scene( + outcome, + left_control_part="physical_left_arm", + right_control_part="physical_right_arm", + ) + + assert scene is not None + assert scene.object_label == "tray" + assert torch.allclose(scene.object_pose, _pose(0.4, -0.2, 0.75)[0]) + assert torch.allclose( + scene.left_grasp_pose[:3, 3], + torch.tensor([0.4, -0.32, 0.76]), + ) + assert torch.allclose( + scene.right_grasp_pose[:3, 3], + torch.tensor([0.4, -0.07, 0.77]), + ) + + +def test_selected_scene_skips_e5_without_valid_env_zero_grasps() -> None: + assert ( + selected_coordinated_grasp_scene( + _outcome(success=False), + left_control_part="physical_left_arm", + right_control_part="physical_right_arm", + ) + is None + ) diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py new file mode 100644 index 000000000..c743282b8 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py @@ -0,0 +1,124 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 torch + +from embodichain.gen_sim.action_engine.runtime.grasp_diagnostics import ( + _TracingAntipodalGraspPoseGenerator, +) +from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg + + +class _CollisionChecker: + def __init__(self) -> None: + self.calls = 0 + + def query(self, *_args, **_kwargs): + self.calls += 1 + if self.calls == 1: + return torch.tensor([False, True, False]), torch.tensor([0.01, -0.002, 0.0]) + return torch.tensor([True, False]), torch.tensor([-0.004, 0.003]) + + +class _Backend: + device = torch.device("cpu") + _max_deviation_angle = torch.pi / 6 + _approach_direction_samples = 4 + + def __init__(self) -> None: + self.antipodal_pairs = torch.tensor( + [ + [[-0.08, -0.08, 0.0], [0.08, -0.08, 0.0]], + [[-0.08, -0.06, 0.0], [0.08, -0.06, 0.0]], + [[-0.08, 0.06, 0.0], [0.08, 0.06, 0.0]], + [[-0.08, 0.08, 0.0], [0.08, 0.08, 0.0]], + ], + dtype=torch.float32, + ) + self._collision_checker = _CollisionChecker() + + def get_dual_arm_valid_grasp_poses(self, **_kwargs): + pose = torch.eye(4).repeat(3, 1, 1) + left_colliding, _ = self._collision_checker.query(None, pose, torch.ones(3)) + right_colliding, _ = self._collision_checker.query( + None, pose[:2], torch.ones(2) + ) + return { + "left": { + "is_success": True, + "grasp_poses": pose[~left_colliding], + "open_lengths": torch.ones(2), + "total_cost": torch.tensor([0.1, 0.2]), + }, + "right": { + "is_success": True, + "grasp_poses": pose[:2][~right_colliding], + "open_lengths": torch.ones(1), + "total_cost": torch.tensor([0.3]), + }, + } + + +def test_dual_grasp_trace_separates_generation_angle_nms_and_collision( + monkeypatch, +) -> None: + generator = _TracingAntipodalGraspPoseGenerator( + ParallelJawGripperModelCfg(model_id="trace_test") + ) + backend = _Backend() + monkeypatch.setattr(generator, "_backend", lambda *_args: backend) + vertices = torch.tensor( + [ + [-0.1, -0.1, -0.02], + [0.1, -0.1, -0.02], + [0.1, 0.1, 0.02], + [-0.1, 0.1, 0.02], + ] + ) + + result = generator.get_dual_arm_valid_grasp_poses( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2], [0, 2, 3]]), + obj_poses=torch.eye(4).unsqueeze(0), + left_to_right_arm_direction=torch.tensor([0.0, 1.0, 0.0]), + approach_direction=torch.tensor([0.0, 0.0, -1.0]), + middle_empty_ratio=0.4, + ) + + assert result[0] is not None + trace = generator.last_dual_trace + assert trace is not None + assert trace["S1_grasp_pair_generation"]["antipodal_pair_count"] == 4 + assert trace["S2_approach_angle_filtering"] == { + "left_partition_pair_count": 2, + "right_partition_pair_count": 2, + "left_angle_valid_pair_count": 2, + "right_angle_valid_pair_count": 2, + } + assert trace["S3_nms"] == { + "left_candidate_count": 3, + "right_candidate_count": 2, + } + assert trace["S4_collision_filtering"]["left_candidate_count"] == 2 + assert trace["S4_collision_filtering"]["right_candidate_count"] == 1 + assert trace["S5_left_right_pairing"] == { + "left_final_count": 2, + "right_final_count": 1, + "paired": True, + } + assert generator.last_dual_trace is not trace diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index 55af776db..ac6f125a8 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -563,6 +563,23 @@ def test_documented_run_command_arguments_remain_compatible() -> None: assert args.seed == 17 assert args.runtime_backend == "independent" assert args.failure_policy == "stop" + assert args.show_grasp_poses is False + + +def test_run_command_accepts_grasp_pose_visualization() -> None: + args = build_run_parser().parse_args( + [ + "--task_name", + "task4_2", + "--gym_config", + "/tmp/fast_gym_config.json", + "--agent_config", + "/tmp/agent_config.json", + "--show-grasp-poses", + ] + ) + + assert args.show_grasp_poses is True def test_run_command_accepts_continue_failure_policy() -> None: @@ -1693,6 +1710,13 @@ def execute_trajectory( assert executor._object_owners["tray"] == [expected_owner] assert executor._arm_owners["left_arm"] == (["tray"] if expect_held else [None]) assert executor._arm_owners["right_arm"] == (["tray"] if expect_held else [None]) + release = result.planner_traces[0]["physical_release"] + assert release["both_grippers_open"] == [opens] + assert release["released"] == [opens] + assert release["left_gripper_open_error"] == pytest.approx([0.0]) + assert release["right_gripper_open_error"] == pytest.approx( + [0.0 if opens else torch.linalg.vector_norm(env.close_state).item()] + ) @pytest.mark.parametrize( @@ -1701,6 +1725,7 @@ def execute_trajectory( ) def test_coordinated_pickment_commits_only_after_physical_dual_hold( monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, closes_both: bool, expected_failed: bool, expect_held: bool, @@ -1724,7 +1749,13 @@ def test_coordinated_pickment_commits_only_after_physical_dual_hold( ) ) ) - executor = ProgramExecutor(program, env, record_runtime=False) + executor = ProgramExecutor( + program, + env, + record_runtime=False, + show_grasp_poses=True, + ) + executor._runtime_output_dir = tmp_path step = program.semantic_steps[0] edge = program.edges[0] executor._assignments[step.id] = ["coordinated"] @@ -1764,6 +1795,14 @@ def execute_trajectory(*_args: Any, **_kwargs: Any) -> list[torch.Tensor]: return [] monkeypatch.setattr(executor.adapter, "execute_trajectory", execute_trajectory) + visualized: list[tuple[str, ActionOutcome]] = [] + monkeypatch.setattr( + executor, + "_write_selected_coordinated_grasp_pose", + lambda selected_step, selected_outcome: visualized.append( + (selected_step.id, selected_outcome) + ), + ) result = executor._execute_coordinated(edge, step, torch.tensor([False])) @@ -1778,6 +1817,15 @@ def execute_trajectory(*_args: Any, **_kwargs: Any) -> list[torch.Tensor]: assert executor._object_owners["tray"] == [expected_owner] assert executor._arm_owners["left_arm"] == (["tray"] if expect_held else [None]) assert executor._arm_owners["right_arm"] == (["tray"] if expect_held else [None]) + assert visualized == [(step.id, outcome)] + physical = result.planner_traces[0]["physical_execution"] + assert physical["dual_hold_predicate"] == [expect_held] + assert physical["semantic_target_reached"] == [True] + assert physical["accepted"] == [expect_held] + assert physical["object_displacement_distance"] == pytest.approx([0.0]) + assert physical["arms"]["left_arm"]["gripper_closed"] == [True] + assert physical["arms"]["right_arm"]["gripper_closed"] == [closes_both] + assert len(physical["arms"]["left_arm"]["eef_relation_position_error"]) == 1 def _handover_held_state( @@ -4735,6 +4783,25 @@ def test_both_grippers_open_uses_the_live_reset_posture() -> None: assert opened.tolist() == [False] +def test_gripper_predicates_measure_transmission_joint_not_commanded_mimics() -> None: + env = _FakeEnv() + env.agent_gripper_state_joint_indices = { + "left": (0,), + "right": (0,), + } + env.left_arm_init_gripper_state = torch.zeros(1, 2) + env.right_arm_init_gripper_state = torch.zeros(1, 2) + env.robot._qpos[:, env.left_eef_joints] = torch.tensor([0.02, 0.30]) + env.robot._qpos[:, env.right_eef_joints] = torch.tensor([0.03, -0.25]) + + opened = evaluate_predicate( + env, + {"type": "both_grippers_open", "tolerance": 0.08}, + ) + + assert opened.tolist() == [True] + + def test_object_supported_by_requires_overlap_and_vertical_contact() -> None: support_z = 0.75 payload_z = support_z + 0.05 + 0.02 + 0.005 @@ -6977,12 +7044,14 @@ def run(self, **kwargs: Any) -> Any: result = env_module.ActionEngineEnv.create_demo_action_list.__wrapped__( env, failure_policy="continue", + show_grasp_poses=True, runtime_run_id="run", episode_index=3, ) assert result is expected assert captured["failure_policy"] == "continue" + assert captured["show_grasp_poses"] is True assert captured["run"] == {"run_id": "run", "episode_index": 3} diff --git a/tests/gen_sim/action_engine/test_gripper_profiles.py b/tests/gen_sim/action_engine/test_gripper_profiles.py index 672d65bde..6b6c2aa1d 100644 --- a/tests/gen_sim/action_engine/test_gripper_profiles.py +++ b/tests/gen_sim/action_engine/test_gripper_profiles.py @@ -58,7 +58,7 @@ def test_pgi_profile_owns_asset_control_mimic_tcp_and_grasp_geometry() -> None: assert profile.grasp_model.opening_margin == pytest.approx(0.03) -def test_robotiq_profile_preserves_existing_rotation_and_joint_semantics() -> None: +def test_robotiq_profile_separates_commanded_mimics_from_state_joint() -> None: profile = get_gripper_profile("robotiq") assert profile.asset_path == ("Robotiq/robotiq_arg2f_140/robotiq_arg2f_140.urdf") @@ -70,6 +70,31 @@ def test_robotiq_profile_preserves_existing_rotation_and_joint_semantics() -> No "left_right_inner_knuckle_joint", "left_right_inner_finger_joint", ) + assert profile.control_joint_names("right") == ( + "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", + ) + assert profile.mimic_joint_names("left") == ( + "left_inner_knuckle_joint", + "left_inner_finger_joint", + "left_right_outer_knuckle_joint", + "left_right_inner_knuckle_joint", + "left_right_inner_finger_joint", + ) + assert profile.state_joint_names("left") == ("left_finger_joint",) + assert profile.state_joint_names("right") == ("right_finger_joint",) + assert profile.state_joint_indices("left") == (0,) + assert profile.state_joint_indices("right") == (0,) + assert set(profile.state_joint_names("left")).isdisjoint( + profile.mimic_joint_names("left") + ) + assert set(profile.state_joint_names("right")).isdisjoint( + profile.mimic_joint_names("right") + ) assert profile.open_positions == (0.0, 0.0, 0.0, 0.0, 0.0, 0.0) assert profile.close_positions == (0.7, -0.7, 0.7, -0.7, -0.7, 0.7) assert profile.tcp_transform == ( diff --git a/tests/gen_sim/action_engine/test_solver_profiles.py b/tests/gen_sim/action_engine/test_solver_profiles.py new file mode 100644 index 000000000..04c1d5cc6 --- /dev/null +++ b/tests/gen_sim/action_engine/test_solver_profiles.py @@ -0,0 +1,42 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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.solver_profiles import ( + IK_SOLVER_MODES, + resolve_ik_solver_mode, +) + + +def test_auto_solver_preserves_current_robot_family_defaults() -> None: + assert IK_SOLVER_MODES == ("auto", "ur", "pytorch") + assert resolve_ik_solver_mode("auto", "dual_ur10") == "ur" + assert resolve_ik_solver_mode("auto", "dual_ur5") == "ur" + assert resolve_ik_solver_mode("auto", "dual_franka") == "pytorch" + + +def test_explicit_solver_mode_is_strict_and_profile_compatible() -> None: + assert resolve_ik_solver_mode("pytorch", "dual_ur10") == "pytorch" + assert resolve_ik_solver_mode("ur", "dual_ur10") == "ur" + + with pytest.raises(ValueError, match="Franka.*URSolver"): + resolve_ik_solver_mode("ur", "dual_franka") + for invalid in ("", "UR", "torch", None): + with pytest.raises((TypeError, ValueError), match="auto.*ur.*pytorch"): + resolve_ik_solver_mode(invalid, "dual_ur10") # type: ignore[arg-type] diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py index e4ca7d5d5..27b6d029b 100644 --- a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -601,6 +601,7 @@ def generator(_scene, output, **kwargs): assert result.bound assert generator_calls assert generator_calls[0]["gripper_model"] == "pgi" + assert generator_calls[0]["ik_solver"] == "auto" assert not (result.output_dir / ".task_engine_input").exists() grounded = json.loads( (result.output_dir / "grounded_task_plan.json").read_text(encoding="utf-8") @@ -885,6 +886,8 @@ def run(self, request, **kwargs): str(tmp_path / "history"), "--planner-mode", "ik_interp", + "--ik-solver", + "pytorch", ] ) == 0 @@ -892,6 +895,7 @@ def run(self, request, **kwargs): planning_cfg = captured["planning_cfg"] assert planning_cfg.planner == {"mode": "ik_interp"} + assert planning_cfg.ik_solver == "pytorch" assert json.loads(capsys.readouterr().out)["status"] == "prepared" @@ -1031,6 +1035,29 @@ def test_public_cli_exposes_prepare_run_and_run_all_modes() -> None: assert arguments.dataset_saving is True assert arguments.failure_policy == "stop" assert arguments.planner_mode is None + assert arguments.ik_solver is None + assert arguments.show_grasp_poses is False + + +def test_run_all_cli_accepts_grasp_pose_visualization() -> None: + arguments = cli.build_parser().parse_args( + [ + "run-all", + "--mode", + "scene", + "--task-id", + "task", + "--instruction", + "move the tray", + "--scene", + "scene", + "--output-root", + "history", + "--show-grasp-poses", + ] + ) + + assert arguments.show_grasp_poses is True def test_prepare_cli_stops_before_simulator_execution( diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py index d7068c126..3936cc2e9 100644 --- a/tests/gen_sim/task_engine/test_parallel_workflow.py +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -273,9 +273,11 @@ def __init__( successes: list[list[bool]], *, expected_dataset_saving: bool = False, + expected_show_grasp_poses: bool = False, ) -> None: self.successes = successes self.expected_dataset_saving = expected_dataset_saving + self.expected_show_grasp_poses = expected_show_grasp_poses self.calls = 0 def __call__( @@ -287,11 +289,13 @@ def __call__( num_envs: int, dataset_saving: bool = False, failure_policy: str = "stop", + show_grasp_poses: bool = False, ): values = self.successes[min(self.calls, len(self.successes) - 1)] self.calls += 1 assert len(values) == num_envs assert dataset_saving is self.expected_dataset_saving + assert show_grasp_poses is self.expected_show_grasp_poses assert failure_policy == "stop" return { "status": "succeeded" if all(values) else "failed", @@ -303,6 +307,32 @@ def __call__( } +def test_parallel_workflow_propagates_grasp_pose_visualization( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor( + [[True]], + expected_show_grasp_poses=True, + ), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + planning_cfg=TaskEnginePlanningCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=1), + show_grasp_poses=True, + ) + + assert result.succeeded + + @pytest.mark.parametrize("existing", [False, True]) @pytest.mark.parametrize("edit", [False, True]) def test_parallel_workflow_supports_all_four_scene_inputs( @@ -405,12 +435,14 @@ def test_parallel_workflow_accepts_one_success_and_publishes_all_graphs( "candidate_count": 3, "planning_mode": "offline", "gripper_model": "pgi", + "ik_solver": "auto", "max_episodes": 1, "max_episode_steps": 6000, } assert manifest["configuration"]["execution"]["dataset_saving"] is False assert coordinator.kwargs[0]["max_episode_steps"] == 6000 assert coordinator.kwargs[0]["gripper_model"] == "pgi" + assert coordinator.kwargs[0]["ik_solver"] == "auto" assert coordinator.kwargs[0]["final_inspection"]["scene_revision_id"] == "0" * 64 assert ( coordinator.kwargs[0]["unbound_action_plan"]["candidate_id"] == "candidate_01" @@ -502,6 +534,9 @@ def test_subprocess_executor_controls_dataset_saving_and_copies_trajectory( trajectory = tmp_path / "trajectory-source" trajectory.mkdir() (trajectory / "episode.json").write_text("{}\n", encoding="utf-8") + grasp_image = trajectory / "env_0000" / "grasp_poses" / "task_01.png" + grasp_image.parent.mkdir(parents=True) + grasp_image.write_bytes(b"grasp-pose-png") captured = {} provenance = build_execution_provenance(episode_seed=7) @@ -551,6 +586,7 @@ def fake_run(command, log_path): num_envs=4, dataset_saving=dataset_saving, failure_policy="continue", + show_grasp_poses=True, ) assert report["status"] == "succeeded" @@ -563,10 +599,14 @@ def fake_run(command, log_path): assert " prepare" not in " ".join(captured["command"]) assert " workflow" not in " ".join(captured["command"]) assert ("--filter_dataset_saving" in captured["command"]) is expects_filter + assert "--show-grasp-poses" in captured["command"] assert captured["command"][-2:] == ["--failure-policy", "continue"] assert captured["log_path"] == attempt / "action.log" assert (attempt / "action.log").read_text(encoding="utf-8") == "child output\n" assert (attempt / "trajectory" / "episode.json").is_file() + assert ( + attempt / "trajectory" / "env_0000" / "grasp_poses" / "task_01.png" + ).read_bytes() == b"grasp-pose-png" process = json.loads((attempt / "process.json").read_text(encoding="utf-8")) assert process["combined_log"] == "action.log" assert process["stdout"] == "ok" diff --git a/tests/gen_sim/task_engine/test_workflow.py b/tests/gen_sim/task_engine/test_workflow.py index 0ea7ba9db..07a5d41b7 100644 --- a/tests/gen_sim/task_engine/test_workflow.py +++ b/tests/gen_sim/task_engine/test_workflow.py @@ -273,6 +273,7 @@ def test_packaged_workflow_configuration_uses_recovery_defaults() -> None: assert planning.candidate_count == 3 assert planning.planning_mode == "offline" assert planning.gripper_model == "pgi" + assert planning.ik_solver == "auto" assert planning.max_episodes == 1 assert planning.max_episode_steps == 6000 assert execution.num_envs == 1 @@ -298,6 +299,7 @@ def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: candidate_count: 7 planning_mode: offline gripper_model: robotiq + ik_solver: pytorch max_episodes: 2 max_episode_steps: 5000 planner: @@ -317,6 +319,7 @@ def test_workflow_configuration_can_be_tuned_from_yaml(tmp_path: Path) -> None: assert workflow.max_action_attempts == 5 assert planning.candidate_count == 7 assert planning.gripper_model == "robotiq" + assert planning.ik_solver == "pytorch" assert planning.max_episodes == 2 assert planning.max_episode_steps == 5000 assert planning.planner == {"mode": "toppra"} @@ -349,5 +352,7 @@ def test_planning_configuration_rejects_invalid_values() -> None: TaskEnginePlanningCfg(planning_mode="unsupported") with pytest.raises(ValueError, match="pgi.*robotiq"): TaskEnginePlanningCfg(gripper_model="unsupported") + with pytest.raises(ValueError, match="auto.*ur.*pytorch"): + TaskEnginePlanningCfg(ik_solver="unsupported") with pytest.raises(ValueError, match="mode cannot be combined"): TaskEnginePlanningCfg(planner={"mode": "toppra", "dynamic_collision": True}) From 5f5dbce595240341342219c713475a7b558878b5 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Fri, 28 Aug 2026 01:56:21 +0800 Subject: [PATCH 79/85] fix(gen-sim): harden coordinated pickment safety --- .../action_engine/config/defaults.yaml | 7 + .../gen_sim/action_engine/runtime/actions.py | 557 ++++++++++++++++-- .../runtime/coordinated_safety.py | 391 ++++++++++++ .../runtime/grasp_diagnostics.py | 181 +++++- .../action_engine/runtime/test_actions.py | 95 +++ .../runtime/test_coordinated_safety.py | 185 ++++++ .../runtime/test_grasp_diagnostics.py | 53 ++ 7 files changed, 1414 insertions(+), 55 deletions(-) create mode 100644 embodichain/gen_sim/action_engine/runtime/coordinated_safety.py create mode 100644 tests/gen_sim/action_engine/runtime/test_coordinated_safety.py diff --git a/embodichain/gen_sim/action_engine/config/defaults.yaml b/embodichain/gen_sim/action_engine/config/defaults.yaml index 93ed15678..e7c1ff81c 100644 --- a/embodichain/gen_sim/action_engine/config/defaults.yaml +++ b/embodichain/gen_sim/action_engine/config/defaults.yaml @@ -230,6 +230,13 @@ runtime: middle_empty_ratio: 0.4 grasp_opening_margin: 0.02 grasp_seed: 17393 + grasp_pair_candidate_count: 3 + minimum_grasp_separation: 0.08 + minimum_grasp_lateral_gap: 0.05 + maximum_joint_step: 0.25 + maximum_wrist_orientation_error: 0.20 + inter_arm_capsule_radius: 0.04 + minimum_inter_arm_clearance: 0.01 is_filter_ground_collision: false release_sample_interval: 60 release_gripper_tolerance: 0.08 diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index ac035b64f..83d14cf54 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -54,6 +54,8 @@ MotionPolicy, ObjectSemantics, PlanningContext, + PlanningFailure, + PlannerDiagnostics, RecoveryPolicy, RobotObservation, RigidObjectSceneProvider, @@ -80,6 +82,7 @@ from embodichain.utils.math import matrix_from_quat, quat_from_matrix, quat_slerp from .body_grasp import AxisAlignBodyGraspAdapter +from .coordinated_safety import _trajectory_safety_report from .grasp_diagnostics import _TracingAntipodalGraspPoseGenerator from .models import ActionOutcome, GroundedAction from .state import ExecutionState @@ -448,34 +451,53 @@ def plan( if grasp_seed is None else self._isolated_random_seed(int(grasp_seed)) ) - with seed_context, _capture_retreat_warnings(capture_warnings) as warnings: + pair_context = self._coordinated_pair_selection_context( + candidate_engine, + candidate, + context, + ) + with ( + seed_context, + pair_context, + _capture_retreat_warnings(capture_warnings) as warnings, + ): candidate_plan = candidate_engine.plan(candidate_invocation, context) coordinated_trace = candidate.motion_policy.get("coordinated_grasp") if isinstance(coordinated_trace, dict): - grasp_stages = self._latest_coordinated_grasp_trace(candidate_engine) - if grasp_stages is not None: - coordinated_trace["stages"] = grasp_stages - coordinated_search_attempts.append( - { - "candidate_index": coordinated_trace.get("candidate_index"), - "approach_candidate_label": coordinated_trace.get( - "approach_candidate_label" - ), - "approach_direction": coordinated_trace.get( - "approach_direction" - ), - "middle_empty_ratio": coordinated_trace.get( - "selected_middle_empty_ratio" - ), - "plan_success": candidate_plan.plan_success.detach() - .cpu() - .tolist(), - "planner_messages": list( - candidate_plan.diagnostics.messages - ), - "stages": deepcopy(grasp_stages), - } - ) + grasp_stages = ( + self._latest_coordinated_grasp_trace(candidate_engine) or {} + ) + coordinated_trace["stages"] = grasp_stages + raw_plan_success = candidate_plan.plan_success.detach().clone() + candidate_plan, trajectory_audit = self._audit_coordinated_trajectory( + candidate, + candidate_invocation, + candidate_plan, + context, + grasp_stages, + ) + coordinated_trace["trajectory_audit"] = trajectory_audit + coordinated_search_attempts.append( + { + "candidate_index": coordinated_trace.get("candidate_index"), + "approach_candidate_label": coordinated_trace.get( + "approach_candidate_label" + ), + "approach_direction": coordinated_trace.get( + "approach_direction" + ), + "middle_empty_ratio": coordinated_trace.get( + "selected_middle_empty_ratio" + ), + "raw_plan_success": raw_plan_success.cpu().tolist(), + "plan_success": candidate_plan.plan_success.detach() + .cpu() + .tolist(), + "planner_messages": list(candidate_plan.diagnostics.messages), + "stages": deepcopy(grasp_stages), + "trajectory_audit": deepcopy(trajectory_audit), + } + ) if capture_warnings: candidate_search_warnings.extend(warnings) candidate_search_attempts += 1 @@ -1013,6 +1035,11 @@ def add_approach(label: str, direction: torch.Tensor) -> None: grasp_seed = grounded.cfg.get("grasp_seed", _COORDINATED_GRASP_SEED) if type(grasp_seed) is not int or grasp_seed < 0: raise ValueError("grasp_seed must be a non-negative integer.") + pair_candidate_count = grounded.cfg.get("grasp_pair_candidate_count", 3) + if type(pair_candidate_count) is not int: + raise ValueError("grasp_pair_candidate_count must be an integer.") + if not 1 <= pair_candidate_count <= 8: + raise ValueError("grasp_pair_candidate_count must be in [1, 8].") trace = { "strategy": "live_geometry_approach_search", "local_pca_axes": eigenvectors.detach().cpu().tolist(), @@ -1037,6 +1064,7 @@ def add_approach(label: str, direction: torch.Tensor) -> None: ], "candidate_middle_empty_ratios": list(ratios), "grasp_seed": grasp_seed, + "grasp_pair_candidate_count": pair_candidate_count, } candidates: list[GroundedAction] = [] candidate_index = 0 @@ -1044,35 +1072,40 @@ def add_approach(label: str, direction: torch.Tensor) -> None: approach_candidates ): for ratio in ratios: - cfg = { - **grounded.cfg, - "left_to_right_arm_direction": shared_direction.clone(), - "approach_direction": approach.clone(), - "middle_empty_ratio": ratio, - "grasp_seed": grasp_seed, - } - motion_policy = { - **grounded.motion_policy, - "grasp_seed": grasp_seed, - "coordinated_grasp": { - **trace, - "candidate_index": candidate_index, - "approach_candidate_index": approach_index, - "approach_candidate_label": approach_label, - "approach_direction": approach.detach().cpu().tolist(), - "selected_middle_empty_ratio": ratio, - }, - } - candidates.append( - replace( - grounded, - target=replace(target, object_initial_pose=live_pose.clone()), - cfg=cfg, - object_pose=live_pose.clone(), - motion_policy=motion_policy, + for pair_rank in range(pair_candidate_count): + cfg = { + **grounded.cfg, + "left_to_right_arm_direction": shared_direction.clone(), + "approach_direction": approach.clone(), + "middle_empty_ratio": ratio, + "grasp_seed": grasp_seed, + "grasp_pair_rank": pair_rank, + } + motion_policy = { + **grounded.motion_policy, + "grasp_seed": grasp_seed, + "coordinated_grasp": { + **trace, + "candidate_index": candidate_index, + "approach_candidate_index": approach_index, + "approach_candidate_label": approach_label, + "approach_direction": approach.detach().cpu().tolist(), + "selected_middle_empty_ratio": ratio, + "grasp_pair_rank": pair_rank, + }, + } + candidates.append( + replace( + grounded, + target=replace( + target, object_initial_pose=live_pose.clone() + ), + cfg=cfg, + object_pose=live_pose.clone(), + motion_policy=motion_policy, + ) ) - ) - candidate_index += 1 + candidate_index += 1 return tuple(candidates) @contextmanager @@ -1110,6 +1143,424 @@ def _latest_coordinated_grasp_trace( return None return generator.last_dual_trace + def _audit_pose_interpolation( + self, + start: torch.Tensor, + end: torch.Tensor, + waypoint_count: int, + *, + interpolate_orientation: bool, + ) -> torch.Tensor: + """Build the Cartesian reference used only by the GenSim FK audit.""" + if waypoint_count <= 0: + return start[:, None].repeat(1, 0, 1, 1) + weights = torch.linspace( + 0.0, + 1.0, + waypoint_count, + dtype=start.dtype, + device=start.device, + ) + result = start[:, None].repeat(1, waypoint_count, 1, 1) + result[:, :, :3, 3] = torch.lerp( + start[:, None, :3, 3], + end[:, None, :3, 3], + weights[None, :, None], + ) + if not interpolate_orientation: + return result + start_quat = quat_from_matrix(start[:, :3, :3]) + end_quat = quat_from_matrix(end[:, :3, :3]) + end_quat = torch.where( + torch.sum(start_quat * end_quat, dim=1, keepdim=True) < 0.0, + -end_quat, + end_quat, + ) + for waypoint_index, weight in enumerate(weights.tolist()): + interpolated = torch.stack( + [ + quat_slerp(start_quat[row], end_quat[row], tau=float(weight)) + for row in range(start.shape[0]) + ] + ) + result[:, waypoint_index, :3, :3] = matrix_from_quat(interpolated) + return result + + def _coordinated_selected_grasp_poses( + self, + grasp_stages: Mapping[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + rows = grasp_stages.get("environment_rows") + if rows is None: + rows = [grasp_stages] + left: list[torch.Tensor] = [] + right: list[torch.Tensor] = [] + for row in rows: + pair = row.get("pair_selection") + if not isinstance(pair, Mapping) or not bool(pair.get("selected", False)): + raise ValueError( + "Trajectory audit requires one selected grasp pair per row." + ) + left.append( + torch.as_tensor( + pair["selected_left_pose"], + dtype=torch.float32, + device=self.device, + ) + ) + right.append( + torch.as_tensor( + pair["selected_right_pose"], + dtype=torch.float32, + device=self.device, + ) + ) + return torch.stack(left), torch.stack(right) + + def _arm_trajectory_fk( + self, + positions: torch.Tensor, + control_part: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return every TCP pose and every serial-chain link point.""" + joint_ids = list(self.env.robot.get_joint_ids(name=control_part)) + arm_qpos = positions[:, :, joint_ids] + batch_size, waypoint_count, dof = arm_qpos.shape + solver = self.env.robot.get_solver(name=control_part) + flat_qpos = arm_qpos.reshape(batch_size * waypoint_count, dof) + local_eef = solver.get_fk(flat_qpos).reshape( + batch_size, + waypoint_count, + 4, + 4, + ) + base_pose = self.env.robot.get_link_pose( + link_name=solver.root_link_name, + to_matrix=True, + ).to(device=positions.device, dtype=positions.dtype) + eef = torch.matmul(base_pose[:, None], local_eef) + chain = getattr(solver, "pk_serial_chain", None) + if chain is None: + raise ValueError(f"{control_part} has no serial chain for capsule audit.") + link_transforms = chain.forward_kinematics(flat_qpos, end_only=False) + link_points: list[torch.Tensor] = [] + for transform in link_transforms.values(): + local = transform.get_matrix().reshape( + batch_size, + waypoint_count, + 4, + 4, + ) + world = torch.matmul(base_pose[:, None], local) + link_points.append(world[:, :, :3, 3]) + link_points.append(eef[:, :, :3, 3]) + return eef, torch.stack(link_points, dim=2) + + @staticmethod + def _audit_batched_pose( + value: Any, + *, + batch_size: int, + device: torch.device | str, + name: str, + ) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(batch_size, 1, 1) + if pose.shape != (batch_size, 4, 4) or not bool(torch.isfinite(pose).all()): + raise ValueError(f"{name} must have finite shape ({batch_size}, 4, 4).") + return pose + + def _coordinated_audit_references( + self, + candidate: GroundedAction, + invocation: ActionInvocation, + plan: ActionPlan, + actual_left: torch.Tensor, + actual_right: torch.Tensor, + grasp_stages: Mapping[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + positions = plan.joint_trajectory + assert positions is not None + batch_size = positions.batch_size + initial = self._audit_batched_pose( + candidate.target.object_initial_pose, + batch_size=batch_size, + device=self.device, + name="object_initial_pose", + ) + target = self._audit_batched_pose( + candidate.target.object_target_pose, + batch_size=batch_size, + device=self.device, + name="object_target_pose", + ) + left_grasp, right_grasp = self._coordinated_selected_grasp_poses(grasp_stages) + left_relation = torch.bmm(torch.linalg.inv(initial), left_grasp) + right_relation = torch.bmm(torch.linalg.inv(initial), right_grasp) + desired_left = actual_left.clone() + desired_right = actual_right.clone() + segments = {segment.name: segment for segment in plan.segments} + + def assign( + name: str, + left_value: torch.Tensor, + right_value: torch.Tensor, + ) -> None: + segment = segments[name] + desired_left[:, segment.start : segment.stop] = left_value + desired_right[:, segment.start : segment.stop] = right_value + + approach = segments["approach"] + assign( + "approach", + self._audit_pose_interpolation( + actual_left[:, approach.start], + left_grasp, + approach.waypoint_count, + interpolate_orientation=True, + ), + self._audit_pose_interpolation( + actual_right[:, approach.start], + right_grasp, + approach.waypoint_count, + interpolate_orientation=True, + ), + ) + close = segments["close"] + assign( + "close", + left_grasp[:, None].repeat(1, close.waypoint_count, 1, 1), + right_grasp[:, None].repeat(1, close.waypoint_count, 1, 1), + ) + lift_pose = initial.clone() + lift_pose[:, 2, 3] += float(invocation.skill_options.lift_height) + lift = segments["lift"] + lift_object = self._audit_pose_interpolation( + initial, + lift_pose, + lift.waypoint_count, + interpolate_orientation=False, + ) + assign( + "lift", + torch.matmul(lift_object, left_relation[:, None]), + torch.matmul(lift_object, right_relation[:, None]), + ) + move = segments["move"] + move_object = self._audit_pose_interpolation( + lift_pose, + target, + move.waypoint_count, + interpolate_orientation=True, + ) + assign( + "move", + torch.matmul(move_object, left_relation[:, None]), + torch.matmul(move_object, right_relation[:, None]), + ) + hold = segments["hold"] + assign( + "hold", + torch.matmul(target, left_relation)[:, None].repeat( + 1, hold.waypoint_count, 1, 1 + ), + torch.matmul(target, right_relation)[:, None].repeat( + 1, hold.waypoint_count, 1, 1 + ), + ) + return desired_left, desired_right + + def _audit_coordinated_trajectory( + self, + candidate: GroundedAction, + invocation: ActionInvocation, + plan: ActionPlan, + context: PlanningContext, + grasp_stages: Mapping[str, Any], + ) -> tuple[ActionPlan, dict[str, Any]]: + """Reject unsafe E5 joint paths before they can become executable.""" + del context + raw_success = plan.plan_success.to(device=self.device) + if plan.joint_trajectory is None or not bool(raw_success.any()): + return plan, { + "success": raw_success.detach().cpu().tolist(), + "skipped": True, + "reason": "planner_failed_or_missing_trajectory", + } + positions = plan.joint_trajectory.positions.to( + device=self.device, + dtype=torch.float32, + ) + try: + left_arm, _, _ = self._parts("left_arm") + right_arm, _, _ = self._parts("right_arm") + left_eef, left_links = self._arm_trajectory_fk(positions, left_arm) + right_eef, right_links = self._arm_trajectory_fk(positions, right_arm) + desired_left, desired_right = self._coordinated_audit_references( + candidate, + invocation, + plan, + left_eef, + right_eef, + grasp_stages, + ) + direction = torch.as_tensor( + candidate.cfg["left_to_right_arm_direction"], + dtype=torch.float32, + device=self.device, + ) + report = _trajectory_safety_report( + left_qpos=positions[:, :, self.env.robot.get_joint_ids(name=left_arm)], + right_qpos=positions[ + :, :, self.env.robot.get_joint_ids(name=right_arm) + ], + left_eef=left_eef, + right_eef=right_eef, + desired_left_eef=desired_left, + desired_right_eef=desired_right, + left_link_points=left_links, + right_link_points=right_links, + left_to_right_direction=direction, + maximum_joint_step=float(candidate.cfg.get("maximum_joint_step", 0.25)), + maximum_orientation_error=float( + candidate.cfg.get("maximum_wrist_orientation_error", 0.20) + ), + minimum_lateral_gap=float( + candidate.cfg.get("minimum_grasp_lateral_gap", 0.05) + ), + capsule_radius=float( + candidate.cfg.get("inter_arm_capsule_radius", 0.04) + ), + minimum_capsule_clearance=float( + candidate.cfg.get("minimum_inter_arm_clearance", 0.01) + ), + orientation_start_index={ + segment.name: segment.start for segment in plan.segments + }["close"], + ) + audited_success = raw_success & report.success.to(device=self.device) + audit_trace = { + "success": audited_success.detach().cpu().tolist(), + "failed_checks": report.failed_checks, + "metrics": report.metrics, + "orientation_audit_start": "close", + "capsule_model": { + "left_link_segments": left_links.shape[2] - 1, + "right_link_segments": right_links.shape[2] - 1, + }, + "thresholds": { + "maximum_joint_step": float( + candidate.cfg.get("maximum_joint_step", 0.25) + ), + "maximum_wrist_orientation_error": float( + candidate.cfg.get("maximum_wrist_orientation_error", 0.20) + ), + "minimum_lateral_gap": float( + candidate.cfg.get("minimum_grasp_lateral_gap", 0.05) + ), + "inter_arm_capsule_radius": float( + candidate.cfg.get("inter_arm_capsule_radius", 0.04) + ), + "minimum_inter_arm_clearance": float( + candidate.cfg.get("minimum_inter_arm_clearance", 0.01) + ), + }, + } + except Exception as exc: + audited_success = torch.zeros_like(raw_success) + audit_trace = { + "success": audited_success.detach().cpu().tolist(), + "failed_checks": {"audit_error": raw_success.cpu().tolist()}, + "error": f"{type(exc).__name__}: {exc}", + } + if torch.equal(audited_success, raw_success): + return plan, audit_trace + failed_rows = ( + torch.nonzero( + raw_success & ~audited_success, + as_tuple=False, + ) + .flatten() + .tolist() + ) + message = ( + f"GenSim trajectory safety audit rejected environment(s) {failed_rows}." + ) + diagnostics = PlannerDiagnostics( + backend=plan.diagnostics.backend, + messages=(*plan.diagnostics.messages, message), + metadata={ + **dict(plan.diagnostics.metadata), + "gensim_trajectory_audit": audit_trace, + }, + failure=PlanningFailure("gensim_trajectory_safety_failed"), + ) + return ( + replace( + plan, + plan_success=audited_success, + diagnostics=diagnostics, + ), + audit_trace, + ) + + @contextmanager + def _coordinated_pair_selection_context( + self, + engine: AtomicActionEngine, + candidate: GroundedAction, + context: PlanningContext, + ) -> Iterator[None]: + """Bind live wrists and arm bases to one synchronous E5 generator call.""" + trace = candidate.motion_policy.get("coordinated_grasp") + if not isinstance(trace, Mapping): + yield + return + _, left_hand, _ = self._parts("left_arm") + if left_hand is None: + yield + return + generator = engine.grasp_pose_generators.get(left_hand) + if not isinstance(generator, _TracingAntipodalGraspPoseGenerator): + yield + return + left_arm, _, _ = self._parts("left_arm") + right_arm, _, _ = self._parts("right_arm") + qpos = context.robot.qpos.to(device=self.device, dtype=torch.float32) + left_ids = list(self.env.robot.get_joint_ids(name=left_arm)) + right_ids = list(self.env.robot.get_joint_ids(name=right_arm)) + left_eef = self.env.robot.compute_fk( + qpos[:, left_ids], + name=left_arm, + to_matrix=True, + ) + right_eef = self.env.robot.compute_fk( + qpos[:, right_ids], + name=right_arm, + to_matrix=True, + ) + left_base, right_base = self._coordinated_arm_bases() + with generator.dual_arm_selection_context( + left_eef=left_eef, + right_eef=right_eef, + left_base=left_base, + right_base=right_base, + left_to_right_direction=torch.as_tensor( + candidate.cfg["left_to_right_arm_direction"], + dtype=torch.float32, + device=self.device, + ), + pair_rank=int(candidate.cfg.get("grasp_pair_rank", 0)), + minimum_separation=float( + candidate.cfg.get("minimum_grasp_separation", 0.08) + ), + minimum_lateral_gap=float( + candidate.cfg.get("minimum_grasp_lateral_gap", 0.05) + ), + ): + yield + def _search_reachable_retreat( self, *, diff --git a/embodichain/gen_sim/action_engine/runtime/coordinated_safety.py b/embodichain/gen_sim/action_engine/runtime/coordinated_safety.py new file mode 100644 index 000000000..5bd150d82 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/coordinated_safety.py @@ -0,0 +1,391 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Pure dual-arm grasp and trajectory safety checks owned by GenSim.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +import torch + +__all__: list[str] = [] + + +@dataclass(frozen=True, slots=True) +class _CanonicalizedGraspPoses: + poses: torch.Tensor + flipped: torch.Tensor + selected_rotation_radians: torch.Tensor + alternative_rotation_radians: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _RankedGraspPairs: + ranked_pairs: tuple[tuple[int, int], ...] + scores: tuple[float, ...] + rejection_counts: dict[str, int] + + +@dataclass(frozen=True, slots=True) +class _TrajectorySafetyReport: + success: torch.Tensor + failed_checks: dict[str, list[bool]] + metrics: dict[str, list[float]] + + +def _rotation_distance( + actual: torch.Tensor, + expected: torch.Tensor, +) -> torch.Tensor: + relative = torch.matmul(actual.transpose(-1, -2), expected) + cosine = (torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) - 1.0) * 0.5 + return torch.acos(torch.clamp(cosine, -1.0, 1.0)) + + +def _canonicalize_parallel_jaw_poses( + poses: torch.Tensor, + live_eef_pose: torch.Tensor, +) -> _CanonicalizedGraspPoses: + """Choose the local-Z half-turn equivalent nearest one live wrist pose.""" + poses = torch.as_tensor(poses, dtype=torch.float32) + live = torch.as_tensor(live_eef_pose, dtype=poses.dtype, device=poses.device) + if poses.ndim != 3 or poses.shape[1:] != (4, 4): + raise ValueError("poses must have shape (N, 4, 4).") + if live.shape != (4, 4): + raise ValueError("live_eef_pose must have shape (4, 4).") + half_turn = torch.eye(4, dtype=poses.dtype, device=poses.device) + half_turn[0, 0] = -1.0 + half_turn[1, 1] = -1.0 + alternatives = torch.matmul(poses, half_turn) + live_rot = live[:3, :3].unsqueeze(0).expand(poses.shape[0], -1, -1) + original_distance = _rotation_distance(live_rot, poses[:, :3, :3]) + alternative_distance = _rotation_distance(live_rot, alternatives[:, :3, :3]) + flipped = alternative_distance < original_distance + selected = torch.where(flipped[:, None, None], alternatives, poses) + selected_distance = torch.where( + flipped, + alternative_distance, + original_distance, + ) + rejected_distance = torch.where( + flipped, + original_distance, + alternative_distance, + ) + return _CanonicalizedGraspPoses( + poses=selected, + flipped=flipped, + selected_rotation_radians=selected_distance, + alternative_rotation_radians=rejected_distance, + ) + + +def _segments_intersect_2d( + first_start: torch.Tensor, + first_end: torch.Tensor, + second_start: torch.Tensor, + second_end: torch.Tensor, +) -> bool: + epsilon = 1.0e-7 + + def orientation(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> float: + ab = b - a + ac = c - a + return float(ab[0] * ac[1] - ab[1] * ac[0]) + + def on_segment(a: torch.Tensor, b: torch.Tensor, point: torch.Tensor) -> bool: + return bool( + float(torch.min(a[0], b[0])) - epsilon + <= float(point[0]) + <= float(torch.max(a[0], b[0])) + epsilon + and float(torch.min(a[1], b[1])) - epsilon + <= float(point[1]) + <= float(torch.max(a[1], b[1])) + epsilon + ) + + a = first_start[:2] + b = first_end[:2] + c = second_start[:2] + d = second_end[:2] + abc = orientation(a, b, c) + abd = orientation(a, b, d) + cda = orientation(c, d, a) + cdb = orientation(c, d, b) + if abc * abd < 0.0 and cda * cdb < 0.0: + return True + return ( + (abs(abc) <= epsilon and on_segment(a, b, c)) + or (abs(abd) <= epsilon and on_segment(a, b, d)) + or (abs(cda) <= epsilon and on_segment(c, d, a)) + or (abs(cdb) <= epsilon and on_segment(c, d, b)) + ) + + +def _rank_non_crossing_grasp_pairs( + left_poses: torch.Tensor, + right_poses: torch.Tensor, + *, + left_costs: torch.Tensor, + right_costs: torch.Tensor, + left_rotation_costs: torch.Tensor, + right_rotation_costs: torch.Tensor, + left_base: torch.Tensor, + right_base: torch.Tensor, + left_to_right_direction: torch.Tensor, + minimum_separation: float, + minimum_lateral_gap: float, +) -> _RankedGraspPairs: + """Rank only distinct, ordered grasp pairs with non-crossing XY routes.""" + left_poses = torch.as_tensor(left_poses, dtype=torch.float32) + right_poses = torch.as_tensor(right_poses, dtype=torch.float32) + direction = torch.as_tensor( + left_to_right_direction, + dtype=torch.float32, + device=left_poses.device, + ) + direction = direction / torch.linalg.vector_norm(direction).clamp_min(1.0e-8) + if minimum_separation < 0.0 or minimum_lateral_gap < 0.0: + raise ValueError("Pair separation constraints must be non-negative.") + rejection_counts = {"reversed": 0, "too_close": 0, "path_crossing": 0} + ranked: list[tuple[float, int, int]] = [] + left_base_position = torch.as_tensor(left_base, dtype=torch.float32)[:3, 3] + right_base_position = torch.as_tensor(right_base, dtype=torch.float32)[:3, 3] + for left_index, left_pose in enumerate(left_poses): + left_position = left_pose[:3, 3] + left_projection = float(torch.dot(left_position, direction)) + for right_index, right_pose in enumerate(right_poses): + right_position = right_pose[:3, 3] + right_projection = float(torch.dot(right_position, direction)) + reversed_pair = ( + left_projection + float(minimum_lateral_gap) > right_projection + ) + too_close = bool( + torch.linalg.vector_norm(left_position - right_position) + < float(minimum_separation) + ) + crossing = _segments_intersect_2d( + left_base_position, + left_position, + right_base_position, + right_position, + ) + rejection_counts["reversed"] += int(reversed_pair) + rejection_counts["too_close"] += int(too_close) + rejection_counts["path_crossing"] += int(crossing) + if reversed_pair or too_close or crossing: + continue + route_length = torch.linalg.vector_norm( + left_position - left_base_position + ) + torch.linalg.vector_norm(right_position - right_base_position) + score = ( + float(left_costs[left_index]) + + float(right_costs[right_index]) + + float(left_rotation_costs[left_index]) / math.pi + + float(right_rotation_costs[right_index]) / math.pi + + 0.05 * float(route_length) + ) + ranked.append((score, left_index, right_index)) + ranked.sort(key=lambda item: (item[0], item[1], item[2])) + return _RankedGraspPairs( + ranked_pairs=tuple((left, right) for _, left, right in ranked), + scores=tuple(score for score, _, _ in ranked), + rejection_counts=rejection_counts, + ) + + +def _point_segment_distance( + point: torch.Tensor, + start: torch.Tensor, + end: torch.Tensor, +) -> torch.Tensor: + segment = end - start + denominator = torch.sum(segment * segment, dim=-1).clamp_min(1.0e-12) + fraction = torch.sum((point - start) * segment, dim=-1) / denominator + closest = start + torch.clamp(fraction, 0.0, 1.0)[..., None] * segment + return torch.linalg.vector_norm(point - closest, dim=-1) + + +def _segment_distance( + first_start: torch.Tensor, + first_end: torch.Tensor, + second_start: torch.Tensor, + second_end: torch.Tensor, +) -> torch.Tensor: + first = first_end - first_start + second = second_end - second_start + offset = first_start - second_start + a = torch.sum(first * first, dim=-1).clamp_min(1.0e-12) + b = torch.sum(first * second, dim=-1) + c = torch.sum(second * second, dim=-1).clamp_min(1.0e-12) + d = torch.sum(first * offset, dim=-1) + e = torch.sum(second * offset, dim=-1) + denominator = a * c - b * b + first_fraction = (b * e - c * d) / denominator.clamp_min(1.0e-12) + second_fraction = (a * e - b * d) / denominator.clamp_min(1.0e-12) + interior = ( + (denominator > 1.0e-12) + & (first_fraction >= 0.0) + & (first_fraction <= 1.0) + & (second_fraction >= 0.0) + & (second_fraction <= 1.0) + ) + first_closest = first_start + first_fraction[..., None] * first + second_closest = second_start + second_fraction[..., None] * second + interior_distance = torch.linalg.vector_norm( + first_closest - second_closest, + dim=-1, + ) + endpoint_distance = ( + torch.stack( + ( + _point_segment_distance(first_start, second_start, second_end), + _point_segment_distance(first_end, second_start, second_end), + _point_segment_distance(second_start, first_start, first_end), + _point_segment_distance(second_end, first_start, first_end), + ), + dim=-1, + ) + .min(dim=-1) + .values + ) + return torch.where(interior, interior_distance, endpoint_distance) + + +def _minimum_interarm_capsule_clearance( + left_link_points: torch.Tensor, + right_link_points: torch.Tensor, + *, + capsule_radius: float, +) -> torch.Tensor: + """Return minimum surface clearance over every inter-arm link pair.""" + left = torch.as_tensor(left_link_points, dtype=torch.float32) + right = torch.as_tensor(right_link_points, dtype=torch.float32) + if left.ndim != 4 or right.ndim != 4 or left.shape[:2] != right.shape[:2]: + raise ValueError("Link points must have matching shape prefixes (B, T, L, 3).") + if left.shape[-1] != 3 or right.shape[-1] != 3: + raise ValueError("Link points must end in xyz coordinates.") + if left.shape[2] < 2 or right.shape[2] < 2: + raise ValueError("Each arm must provide at least two link points.") + if capsule_radius < 0.0: + raise ValueError("capsule_radius must be non-negative.") + minimum = torch.full( + left.shape[:2], + torch.inf, + dtype=left.dtype, + device=left.device, + ) + for left_index in range(left.shape[2] - 1): + for right_index in range(right.shape[2] - 1): + distance = _segment_distance( + left[:, :, left_index], + left[:, :, left_index + 1], + right[:, :, right_index], + right[:, :, right_index + 1], + ) + minimum = torch.minimum(minimum, distance) + return minimum - 2.0 * float(capsule_radius) + + +def _trajectory_safety_report( + *, + left_qpos: torch.Tensor, + right_qpos: torch.Tensor, + left_eef: torch.Tensor, + right_eef: torch.Tensor, + desired_left_eef: torch.Tensor, + desired_right_eef: torch.Tensor, + left_link_points: torch.Tensor, + right_link_points: torch.Tensor, + left_to_right_direction: torch.Tensor, + maximum_joint_step: float, + maximum_orientation_error: float, + minimum_lateral_gap: float, + capsule_radius: float, + minimum_capsule_clearance: float, + orientation_start_index: int = 0, +) -> _TrajectorySafetyReport: + """Evaluate all hard E5 post-plan continuity and inter-arm constraints.""" + if left_qpos.ndim != 3 or right_qpos.ndim != 3: + raise ValueError("Arm qpos trajectories must have shape (B, T, DOF).") + if left_qpos.shape[:2] != right_qpos.shape[:2] or left_qpos.shape[1] < 2: + raise ValueError("Arm qpos trajectories must share at least two waypoints.") + waypoint_count = left_qpos.shape[1] + if not 0 <= orientation_start_index < waypoint_count: + raise ValueError("orientation_start_index must select a trajectory waypoint.") + left_step = torch.amax(torch.abs(torch.diff(left_qpos, dim=1)), dim=(1, 2)) + right_step = torch.amax(torch.abs(torch.diff(right_qpos, dim=1)), dim=(1, 2)) + joint_step = torch.maximum(left_step, right_step) + left_orientation = torch.amax( + _rotation_distance( + left_eef[:, orientation_start_index:, :3, :3], + desired_left_eef[:, orientation_start_index:, :3, :3], + ), + dim=1, + ) + right_orientation = torch.amax( + _rotation_distance( + right_eef[:, orientation_start_index:, :3, :3], + desired_right_eef[:, orientation_start_index:, :3, :3], + ), + dim=1, + ) + orientation = torch.maximum(left_orientation, right_orientation) + direction = torch.as_tensor( + left_to_right_direction, + dtype=left_eef.dtype, + device=left_eef.device, + ) + direction = direction / torch.linalg.vector_norm(direction).clamp_min(1.0e-8) + lateral_gap = ( + torch.sum( + (right_eef[:, :, :3, 3] - left_eef[:, :, :3, 3]) * direction, + dim=2, + ) + .min(dim=1) + .values + ) + capsule_clearance = ( + _minimum_interarm_capsule_clearance( + left_link_points, + right_link_points, + capsule_radius=capsule_radius, + ) + .min(dim=1) + .values + ) + failures = { + "joint_step": joint_step > float(maximum_joint_step), + "orientation": orientation > float(maximum_orientation_error), + "lateral_order": lateral_gap < float(minimum_lateral_gap), + "capsule_collision": capsule_clearance < float(minimum_capsule_clearance), + } + failed = torch.zeros_like(joint_step, dtype=torch.bool) + for value in failures.values(): + failed |= value + return _TrajectorySafetyReport( + success=~failed, + failed_checks={ + name: value.detach().cpu().tolist() for name, value in failures.items() + }, + metrics={ + "maximum_joint_step": joint_step.detach().cpu().tolist(), + "maximum_orientation_error": orientation.detach().cpu().tolist(), + "minimum_lateral_gap": lateral_gap.detach().cpu().tolist(), + "minimum_capsule_clearance": capsule_clearance.detach().cpu().tolist(), + }, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py index ea273a8f9..3e5ac2da7 100644 --- a/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py +++ b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py @@ -18,29 +18,199 @@ from __future__ import annotations +from contextlib import contextmanager from copy import deepcopy -from typing import Any +from dataclasses import dataclass +from typing import Any, Iterator import torch import torch.nn.functional as F from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator +from .coordinated_safety import ( + _canonicalize_parallel_jaw_poses, + _rank_non_crossing_grasp_pairs, +) + __all__: list[str] = [] +@dataclass(frozen=True, slots=True) +class _DualGraspSelectionContext: + left_eef: torch.Tensor + right_eef: torch.Tensor + left_base: torch.Tensor + right_base: torch.Tensor + left_to_right_direction: torch.Tensor + pair_rank: int + minimum_separation: float + minimum_lateral_gap: float + + class _TracingAntipodalGraspPoseGenerator(AntipodalGraspPoseGenerator): """Retain compact S1-S5 evidence from the concrete GenSim grasp backend.""" def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._last_dual_trace: dict[str, Any] | None = None + self._selection_context: _DualGraspSelectionContext | None = None @property def last_dual_trace(self) -> dict[str, Any] | None: """Return an owned snapshot of the most recent dual-grasp trace.""" return deepcopy(self._last_dual_trace) + @contextmanager + def dual_arm_selection_context( + self, + *, + left_eef: torch.Tensor, + right_eef: torch.Tensor, + left_base: torch.Tensor, + right_base: torch.Tensor, + left_to_right_direction: torch.Tensor, + pair_rank: int, + minimum_separation: float, + minimum_lateral_gap: float, + ) -> Iterator[None]: + """Install one invocation-local arm context for pair-aware selection.""" + if self._selection_context is not None: + raise RuntimeError("Dual grasp selection context cannot be nested.") + if type(pair_rank) is not int or pair_rank < 0: + raise ValueError("pair_rank must be a non-negative integer.") + self._selection_context = _DualGraspSelectionContext( + left_eef=torch.as_tensor(left_eef, dtype=torch.float32).clone(), + right_eef=torch.as_tensor(right_eef, dtype=torch.float32).clone(), + left_base=torch.as_tensor(left_base, dtype=torch.float32).clone(), + right_base=torch.as_tensor(right_base, dtype=torch.float32).clone(), + left_to_right_direction=torch.as_tensor( + left_to_right_direction, dtype=torch.float32 + ).clone(), + pair_rank=pair_rank, + minimum_separation=float(minimum_separation), + minimum_lateral_gap=float(minimum_lateral_gap), + ) + try: + yield + finally: + self._selection_context = None + + @staticmethod + def _failed_arm_result(reference: torch.Tensor) -> dict[str, Any]: + return { + "is_success": False, + "grasp_poses": torch.eye( + 4, + dtype=torch.float32, + device=reference.device, + ), + "open_lengths": 0.0, + "total_cost": torch.zeros(1, device=reference.device), + } + + def _select_pair( + self, + result: dict[str, dict[str, Any]] | None, + *, + row_index: int, + ) -> tuple[dict[str, dict[str, Any]] | None, dict[str, Any] | None]: + context = self._selection_context + if context is None or result is None: + return result, None + left = result["left"] + right = result["right"] + if not left.get("is_success", False) or not right.get("is_success", False): + return result, { + "requested_pair_rank": context.pair_rank, + "valid_pair_count": 0, + "selected": False, + "reason": "one_or_both_arms_have_no_candidates", + } + left_poses = torch.as_tensor(left["grasp_poses"], dtype=torch.float32) + right_poses = torch.as_tensor(right["grasp_poses"], dtype=torch.float32) + if left_poses.ndim == 2: + left_poses = left_poses.unsqueeze(0) + if right_poses.ndim == 2: + right_poses = right_poses.unsqueeze(0) + left_canonical = _canonicalize_parallel_jaw_poses( + left_poses, + context.left_eef[row_index], + ) + right_canonical = _canonicalize_parallel_jaw_poses( + right_poses, + context.right_eef[row_index], + ) + ranking = _rank_non_crossing_grasp_pairs( + left_canonical.poses, + right_canonical.poses, + left_costs=torch.as_tensor(left["total_cost"], dtype=torch.float32), + right_costs=torch.as_tensor(right["total_cost"], dtype=torch.float32), + left_rotation_costs=left_canonical.selected_rotation_radians, + right_rotation_costs=right_canonical.selected_rotation_radians, + left_base=context.left_base[row_index], + right_base=context.right_base[row_index], + left_to_right_direction=context.left_to_right_direction, + minimum_separation=context.minimum_separation, + minimum_lateral_gap=context.minimum_lateral_gap, + ) + trace: dict[str, Any] = { + "requested_pair_rank": context.pair_rank, + "valid_pair_count": len(ranking.ranked_pairs), + "rejection_counts": dict(ranking.rejection_counts), + "left_half_turn_count": int(left_canonical.flipped.sum().item()), + "right_half_turn_count": int(right_canonical.flipped.sum().item()), + "selected": context.pair_rank < len(ranking.ranked_pairs), + } + if context.pair_rank >= len(ranking.ranked_pairs): + trace["reason"] = "requested_pair_rank_unavailable" + return { + "left": self._failed_arm_result(left_poses), + "right": self._failed_arm_result(right_poses), + }, trace + left_index, right_index = ranking.ranked_pairs[context.pair_rank] + left_pose = left_canonical.poses[left_index] + right_pose = right_canonical.poses[right_index] + trace.update( + { + "selected_left_index": left_index, + "selected_right_index": right_index, + "selected_pair_score": ranking.scores[context.pair_rank], + "selected_left_half_turn": bool(left_canonical.flipped[left_index]), + "selected_right_half_turn": bool(right_canonical.flipped[right_index]), + "selected_left_rotation_radians": float( + left_canonical.selected_rotation_radians[left_index] + ), + "selected_right_rotation_radians": float( + right_canonical.selected_rotation_radians[right_index] + ), + "selected_left_pose": left_pose.detach().cpu().tolist(), + "selected_right_pose": right_pose.detach().cpu().tolist(), + "selected_separation": float( + torch.linalg.vector_norm(left_pose[:3, 3] - right_pose[:3, 3]) + ), + } + ) + + def selected_arm( + arm: dict[str, Any], + poses: torch.Tensor, + index: int, + ) -> dict[str, Any]: + open_lengths = torch.as_tensor(arm["open_lengths"]) + costs = torch.as_tensor(arm["total_cost"], dtype=torch.float32) + return { + "is_success": True, + "grasp_poses": poses[index : index + 1], + "open_lengths": open_lengths[index : index + 1], + "total_cost": costs.new_zeros(1), + } + + return { + "left": selected_arm(left, left_canonical.poses, left_index), + "right": selected_arm(right, right_canonical.poses, right_index), + }, trace + @staticmethod def _transform_points(points: torch.Tensor, pose: torch.Tensor) -> torch.Tensor: return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] @@ -134,6 +304,7 @@ def traced_query(*args: Any, **query_kwargs: Any): checker.query = original_query row_traces: list[dict[str, Any]] = [] + selected_results: list[dict[str, dict[str, Any]] | None] = [] for row_index, (object_pose, approach, result) in enumerate( zip(poses, directions, results, strict=True) ): @@ -153,6 +324,11 @@ def traced_query(*args: Any, **query_kwargs: Any): right = {} if result is None else result["right"] left_final = self._candidate_count(left) right_final = self._candidate_count(right) + selected_result, pair_trace = self._select_pair( + result, + row_index=row_index, + ) + selected_results.append(selected_result) row_traces.append( { "environment_index": row_index, @@ -195,9 +371,10 @@ def traced_query(*args: Any, **query_kwargs: Any): "right_final_count": right_final, "paired": left_final > 0 and right_final > 0, }, + "pair_selection": pair_trace, } ) self._last_dual_trace = ( row_traces[0] if len(row_traces) == 1 else {"environment_rows": row_traces} ) - return results + return selected_results diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 33300de93..2c08b95cc 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -18,6 +18,8 @@ from __future__ import annotations +from contextlib import nullcontext +from dataclasses import replace from types import SimpleNamespace from typing import Any @@ -768,6 +770,99 @@ def test_coordinated_pickment_approach_family_follows_live_shared_reach() -> Non "current", "robot_forward", ] + assert [candidate.cfg["grasp_pair_rank"] for candidate in candidates[:3]] == [ + 0, + 1, + 2, + ] + + +def test_coordinated_pickment_rejects_boolean_pair_candidate_count() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + grounded = _coordinated_grounded(torch.eye(3)) + grounded = replace( + grounded, + cfg={**grounded.cfg, "grasp_pair_candidate_count": True}, + ) + + with pytest.raises(ValueError, match="must be an integer"): + adapter._adapt_coordinated_pickment_grasps(grounded, capability) + + +def test_coordinated_pickment_continues_after_pair_trajectory_audit_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + positions = torch.zeros(2, 2, env.robot.dof) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + successful_plan = ActionPlan( + skill_id="coordinated_pickment", + plan_success=torch.ones(2, dtype=torch.bool), + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + ) + engine = _FakeEngine(lambda *_args: successful_plan) + engine.grasp_pose_generators = {} + monkeypatch.setattr(adapter, "_engine_for", lambda *_args: engine) + monkeypatch.setattr( + adapter, + "_coordinated_pair_selection_context", + lambda *_args: nullcontext(), + ) + monkeypatch.setattr( + adapter, + "_latest_coordinated_grasp_trace", + lambda *_args: { + "pair_selection": { + "selected": True, + "selected_left_pose": torch.eye(4).tolist(), + "selected_right_pose": torch.eye(4).tolist(), + } + }, + ) + audited_ranks: list[int] = [] + + def audit(candidate, _invocation, plan, _context, _stages): + rank = int(candidate.cfg["grasp_pair_rank"]) + audited_ranks.append(rank) + success = torch.full((2,), rank == 1, dtype=torch.bool) + diagnostics = ( + plan.diagnostics + if bool(success.all()) + else PlannerDiagnostics( + backend="fake", + failure=PlanningFailure("trajectory_safety_failed"), + ) + ) + return ( + replace(plan, plan_success=success, diagnostics=diagnostics), + {"success": success.tolist()}, + ) + + monkeypatch.setattr(adapter, "_audit_coordinated_trajectory", audit) + grounded = _coordinated_grounded(torch.eye(3)) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, env.robot.dof)), + ) + + assert audited_ranks == [0, 1] + assert outcome.success.tolist() == [True, True] + trace = outcome.planner_trace["coordinated_grasp"] + assert trace["grasp_pair_rank"] == 1 + assert len(trace["search_attempts"]) == 2 def test_coordinated_pickment_geometry_candidates_are_deterministic_for_tray() -> None: diff --git a/tests/gen_sim/action_engine/runtime/test_coordinated_safety.py b/tests/gen_sim/action_engine/runtime/test_coordinated_safety.py new file mode 100644 index 000000000..b9a3eb5a6 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_coordinated_safety.py @@ -0,0 +1,185 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime.coordinated_safety import ( + _canonicalize_parallel_jaw_poses, + _minimum_interarm_capsule_clearance, + _rank_non_crossing_grasp_pairs, + _segments_intersect_2d, + _trajectory_safety_report, +) + + +def _pose(x: float, y: float, z: float = 0.75) -> torch.Tensor: + pose = torch.eye(4, dtype=torch.float32) + pose[:3, 3] = torch.tensor([x, y, z]) + return pose + + +def test_parallel_jaw_pose_chooses_equivalent_half_turn_nearest_live_eef() -> None: + pose = torch.eye(4).unsqueeze(0) + pose[0, 0, 0] = -1.0 + pose[0, 1, 1] = -1.0 + + result = _canonicalize_parallel_jaw_poses(pose, torch.eye(4)) + + torch.testing.assert_close(result.poses[0], torch.eye(4)) + assert result.flipped.tolist() == [True] + assert result.selected_rotation_radians.tolist() == [0.0] + assert result.alternative_rotation_radians.tolist() == pytest.approx([torch.pi]) + + +def test_pair_ranking_rejects_reversal_overlap_and_xy_crossing() -> None: + left = torch.stack( + ( + _pose(0.0, -0.20), + _pose(1.0, 0.20), + _pose(0.0, 0.18), + ) + ) + right = torch.stack( + ( + _pose(0.0, 0.20), + _pose(-1.0, -0.20), + _pose(0.0, 0.19), + ) + ) + result = _rank_non_crossing_grasp_pairs( + left, + right, + left_costs=torch.tensor([0.0, 10.0, 10.0]), + right_costs=torch.tensor([0.0, 10.0, 10.0]), + left_rotation_costs=torch.zeros(3), + right_rotation_costs=torch.zeros(3), + left_base=_pose(-1.0, -0.30), + right_base=_pose(1.0, 0.30), + left_to_right_direction=torch.tensor([0.0, 1.0, 0.0]), + minimum_separation=0.08, + minimum_lateral_gap=0.05, + ) + + assert result.ranked_pairs[0] == (0, 0) + assert (1, 1) not in result.ranked_pairs # XY paths intersect. + assert (2, 2) not in result.ranked_pairs # Distinct poses are too close. + assert result.rejection_counts["reversed"] > 0 + assert result.rejection_counts["too_close"] > 0 + assert result.rejection_counts["path_crossing"] > 0 + + +def test_xy_route_intersection_includes_touching_and_collinear_overlap() -> None: + assert _segments_intersect_2d( + torch.tensor([0.0, 0.0, 0.0]), + torch.tensor([1.0, 0.0, 0.0]), + torch.tensor([0.5, 0.0, 0.0]), + torch.tensor([1.5, 0.0, 0.0]), + ) + assert _segments_intersect_2d( + torch.tensor([0.0, 0.0, 0.0]), + torch.tensor([1.0, 1.0, 0.0]), + torch.tensor([1.0, 1.0, 0.0]), + torch.tensor([2.0, 1.0, 0.0]), + ) + + +def test_capsule_clearance_covers_every_left_right_link_segment() -> None: + left = torch.tensor( + [[[[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]]], + dtype=torch.float32, + ) + crossing = torch.tensor( + [[[[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]]]], + dtype=torch.float32, + ) + clear = crossing.clone() + clear[..., 2] = 1.0 + + crossing_clearance = _minimum_interarm_capsule_clearance( + left, + crossing, + capsule_radius=0.05, + ) + safe_clearance = _minimum_interarm_capsule_clearance( + left, + clear, + capsule_radius=0.05, + ) + + torch.testing.assert_close(crossing_clearance, torch.tensor([[-0.10]])) + torch.testing.assert_close(safe_clearance, torch.tensor([[0.90]])) + + +def test_trajectory_safety_rejects_orientation_jump_order_and_capsule_collision() -> ( + None +): + left_qpos = torch.zeros(1, 3, 2) + right_qpos = torch.zeros(1, 3, 2) + left_qpos[:, 1, 0] = 0.40 + left_eef = torch.eye(4).repeat(1, 3, 1, 1) + right_eef = torch.eye(4).repeat(1, 3, 1, 1) + desired_left = left_eef.clone() + desired_right = right_eef.clone() + left_eef[:, 1, 0, 0] = -1.0 + left_eef[:, 1, 1, 1] = -1.0 + left_eef[:, :, 1, 3] = torch.tensor([-0.2, 0.2, -0.2]) + right_eef[:, :, 1, 3] = torch.tensor([0.2, -0.2, 0.2]) + left_links = torch.tensor( + [ + [ + [[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + ] + ] + ) + right_links = torch.tensor( + [ + [ + [[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]], + [[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]], + [[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]], + ] + ] + ) + + report = _trajectory_safety_report( + left_qpos=left_qpos, + right_qpos=right_qpos, + left_eef=left_eef, + right_eef=right_eef, + desired_left_eef=desired_left, + desired_right_eef=desired_right, + left_link_points=left_links, + right_link_points=right_links, + left_to_right_direction=torch.tensor([0.0, 1.0, 0.0]), + maximum_joint_step=0.25, + maximum_orientation_error=0.20, + minimum_lateral_gap=0.05, + capsule_radius=0.05, + minimum_capsule_clearance=0.0, + ) + + assert report.success.tolist() == [False] + assert report.failed_checks == { + "joint_step": [True], + "orientation": [True], + "lateral_order": [True], + "capsule_collision": [True], + } diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py index c743282b8..00b3704fb 100644 --- a/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py +++ b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py @@ -122,3 +122,56 @@ def test_dual_grasp_trace_separates_generation_angle_nms_and_collision( "paired": True, } assert generator.last_dual_trace is not trace + + +def test_generator_context_selects_non_crossing_pair_and_canonical_half_turn() -> None: + generator = _TracingAntipodalGraspPoseGenerator( + ParallelJawGripperModelCfg(model_id="pair_test") + ) + left_poses = torch.stack((_pose_with_y(-0.2), _pose_with_y(0.2))) + left_poses[0, 0, 0] = -1.0 + left_poses[0, 1, 1] = -1.0 + right_poses = torch.stack((_pose_with_y(0.2), _pose_with_y(-0.2))) + result = { + "left": { + "is_success": True, + "grasp_poses": left_poses, + "open_lengths": torch.ones(2), + "total_cost": torch.tensor([0.1, 0.0]), + }, + "right": { + "is_success": True, + "grasp_poses": right_poses, + "open_lengths": torch.ones(2), + "total_cost": torch.tensor([0.1, 0.0]), + }, + } + + with generator.dual_arm_selection_context( + left_eef=torch.eye(4).unsqueeze(0), + right_eef=torch.eye(4).unsqueeze(0), + left_base=_pose_with_y(-0.3).unsqueeze(0), + right_base=_pose_with_y(0.3).unsqueeze(0), + left_to_right_direction=torch.tensor([0.0, 1.0, 0.0]), + pair_rank=0, + minimum_separation=0.08, + minimum_lateral_gap=0.05, + ): + selected, trace = generator._select_pair(result, row_index=0) + + assert selected is not None + assert trace is not None + assert trace["selected"] is True + assert trace["selected_left_index"] == 0 + assert trace["selected_right_index"] == 0 + assert trace["selected_left_half_turn"] is True + torch.testing.assert_close( + selected["left"]["grasp_poses"][0, :3, :3], + torch.eye(3), + ) + + +def _pose_with_y(y: float) -> torch.Tensor: + pose = torch.eye(4) + pose[1, 3] = y + return pose From 1921e9ec67e9f7bc4aba41f56e426fd3fa090b59 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:46:31 +0800 Subject: [PATCH 80/85] temp --- .../action_engine/capabilities/atomic.py | 100 ++++++++++---- .../action_engine/domain/task_contracts.py | 2 +- .../gen_sim/action_engine/runtime/actions.py | 18 ++- .../action_engine/runtime/atomic_compat.py | 78 +++++++++-- .../gen_sim/action_engine/runtime/executor.py | 17 ++- .../action_engine/runtime/grounding.py | 7 +- .../gen_sim/action_engine/tasks/recipes.py | 50 ++++++- embodichain/gen_sim/task_engine/defaults.yaml | 2 +- .../capabilities/test_atomic_v2.py | 84 +++++++++--- .../action_engine/runtime/test_actions.py | 2 + .../runtime/test_atomic_compat.py | 127 ++++++++++++++---- .../runtime/test_runtime_contracts.py | 45 ++++++- .../action_engine/tasks/test_factory.py | 80 ++++++++--- .../tasks/test_interpretation.py | 2 + 14 files changed, 488 insertions(+), 126 deletions(-) diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 0f483e7b7..200fda04f 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -382,11 +382,10 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: frozenset({"object"}), frozenset({"arm"}), "single_arm_object", - "preserve", + "hold", "axis_align", motion_base="AxisAlign", verifier="postcondition", - verifier_hook=_verify_axis_alignment, failure_classifier="grasp", contract_resolver_hook=_resolve_axis_align_contract, allows_target_contact=True, @@ -572,6 +571,10 @@ def capability_precondition( target_binding: Mapping[str, Any], ) -> dict[str, Any]: """Build the generic live precondition used to authorize a retry.""" + if target_binding.get("single_release", False): + # Opening a gripper is idempotent. A retry remains safe when the first + # attempt physically released the object but failed terminal tracking. + return {} if target_binding.get("coordinated_release_role") is not None: # Opening a gripper is idempotent. A retry must remain legal when one # hand opened on the first attempt and the physical dual-hold predicate @@ -820,7 +823,7 @@ def _resolve_end_effector_contract( def _resolve_axis_align_contract(node: Mapping[str, Any]) -> ResolvedActionContract: - """Keep the object free while enforcing one verified E2 terminal barrier.""" + """Acquire one object and retain it until an explicit release action.""" object_uid = _required_string(node.get("object_uid"), "node.object_uid") actor = node.get("actor", {}) if not isinstance(actor, Mapping): @@ -832,32 +835,21 @@ def _resolve_axis_align_contract(node: Mapping[str, Any]) -> ResolvedActionContr StateAtom("object_free", object_uid=object_uid), ), effects=( - StateEffect("add", StateAtom("arm_free", arm=arm)), - StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + StateEffect("delete", StateAtom("arm_free", arm=arm)), + StateEffect("delete", StateAtom("object_free", object_uid=object_uid)), + StateEffect( + "add", + StateAtom("object_held", object_uid=object_uid, arm=arm), + ), ), claims=( - ResourceClaim(f"arm:{arm}"), - ResourceClaim(f"object:{object_uid}"), + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), ), - completion="terminal_barrier", failure_policy="task_required", ) -def _verify_axis_alignment( - *, - executor: Any, - step: Any, - arm: str, - outcome: Any, - attempted: torch.Tensor, -) -> torch.Tensor: - """Reuse the E2 live predicate before committing AxisAlign completion.""" - del arm, outcome - verified_failed, success, _ = executor._verify_step(step, ~attempted) - return attempted & success & ~verified_failed - - def _resolve_pour_contract(node: Mapping[str, Any]) -> ResolvedActionContract: """Retain one verified holder until the E3 action chain completes.""" object_uid = _required_string(node.get("object_uid"), "node.object_uid") @@ -907,6 +899,35 @@ def _resolve_joints_contract(node: Mapping[str, Any]) -> ResolvedActionContract: binding = node.get("target_binding", {}) if not isinstance(binding, Mapping): raise ValueError("MoveJoints contract requires a target_binding mapping.") + single_release = binding.get("single_release", False) + if not isinstance(single_release, bool): + raise TypeError("joint_state single_release must be a boolean.") + if single_release: + if ( + node.get("control") != "hand" + or binding.get("source") != "gripper_open" + or binding.get("coordinated_release_role") is not None + ): + raise ValueError( + "Single-arm MoveJoints release requires a hand action targeting " + "gripper_open without a coordinated release role." + ) + return ResolvedActionContract( + requires=(StateAtom("object_held", object_uid=object_uid, arm=arm),), + effects=( + StateEffect( + "delete", + StateAtom("object_held", object_uid=object_uid, arm=arm), + ), + StateEffect("add", StateAtom("arm_free", arm=arm)), + StateEffect("add", StateAtom("object_free", object_uid=object_uid)), + ), + claims=( + ResourceClaim(f"arm:{arm}", lifetime="until_release"), + ResourceClaim(f"object:{object_uid}", lifetime="until_release"), + ), + failure_policy="task_required", + ) release_role = binding.get("coordinated_release_role") if release_role is not None: if ( @@ -1060,9 +1081,42 @@ def _verify_required_home( outcome: Any, attempted: torch.Tensor, ) -> torch.Tensor: - """Verify required cleanup against the live arm joint state.""" + """Verify explicit release or required cleanup against live joint state.""" del step policy = outcome.grounded.motion_policy + if bool(policy.get("single_release", False)): + getter = getattr(executor.env, "get_current_gripper_state_agent", None) + if not callable(getter) or arm not in {"left_arm", "right_arm"}: + return torch.zeros_like(attempted) + values = getter() + index = 0 if arm == "left_arm" else 1 + if not isinstance(values, (tuple, list)) or len(values) <= index: + return torch.zeros_like(attempted) + current = torch.as_tensor( + values[index], + dtype=torch.float32, + device=executor.env.device, + ) + if current.ndim == 1: + current = current.unsqueeze(0).repeat(int(executor.env.num_envs), 1) + expected = torch.as_tensor( + executor.env.open_state, + dtype=current.dtype, + device=current.device, + ).flatten() + repeats = (current.shape[-1] + expected.numel() - 1) // expected.numel() + expected = expected.repeat(repeats)[: current.shape[-1]] + tolerance = float( + policy.get( + "release_gripper_tolerance", + executor.runtime_policy.predicate_fallbacks["gripper_state_tolerance"], + ) + ) + opened = ( + torch.linalg.vector_norm(current - expected.unsqueeze(0), dim=1) + <= tolerance + ) + return attempted & opened if not bool(policy.get("verify_required_home", False)): return attempted env = executor.env diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py index 3d1b268aa..da1ff9397 100644 --- a/embodichain/gen_sim/action_engine/domain/task_contracts.py +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -70,7 +70,7 @@ def normalize_placement_relation(value: Any) -> str: _CORE_ACTIONS: Mapping[str, tuple[str, ...]] = MappingProxyType( { "E1": ("PickUp", "MoveHeldObject", "Place"), - "E2": ("AxisAlign",), + "E2": ("AxisAlign", "MoveHeldObject", "MoveJoints"), "E3": ("PickUp", "MoveHeldObject", "Pour", "Place"), "E4": ("PickUp", "MoveHeldObject", "HandOver", "Place"), "E5": ("CoordinatedPickment",), diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 83d14cf54..3a410e934 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -2293,6 +2293,7 @@ def _binding( return ActionBinding(owner_id=engine.binding_owner_id) slot_parts: dict[str, tuple[str, str | None]] = {} + task_state_keys: dict[str, str] | None = None if capability.config_materializer == "handover": transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" @@ -2330,6 +2331,10 @@ def _binding( f"{action.arm} has no configured {action.control} part." ) slot_parts = {"primary": (motion_part, hand_part)} + if action.control == "hand" and bool( + action.cfg.get("single_release", False) + ): + task_state_keys = {"primary": arm_part} endpoints: dict[str, dict[str, str]] = {} for slot in contract.slots: @@ -2356,9 +2361,13 @@ def _binding( f"{requirement.endpoint_id}." ) endpoints[slot.slot_id] = selected + skill_id = str(capability.action_type.skill_id) + if task_state_keys is None: + return engine.bind_control_parts(skill_id, endpoints) return engine.bind_control_parts( - str(capability.action_type.skill_id), + skill_id, endpoints, + task_state_keys=task_state_keys, ) def _build_config( @@ -2419,6 +2428,10 @@ def _build_single_arm_config( from .atomic_compat import ExactTargetMoveHeldObjectOptions config_type = ExactTargetMoveHeldObjectOptions + elif capability.target_materializer == "joint_state": + from .atomic_compat import ActionEngineMoveJointsOptions + + config_type = ActionEngineMoveJointsOptions if capability.target_materializer == "press": press_depth = policy.pop("press_depth", None) if press_depth is not None and "press_distance" not in policy: @@ -2732,7 +2745,7 @@ def _new_engine( ) -> AtomicActionEngine: from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver - from .atomic_compat import ExactTargetMoveHeldObject + from .atomic_compat import ActionEngineMoveJoints, ExactTargetMoveHeldObject engine = AtomicActionEngine( motion_generator, @@ -2743,6 +2756,7 @@ def _new_engine( ), ) engine.register(ExactTargetMoveHeldObject(), replace=True) + engine.register(ActionEngineMoveJoints(), replace=True) engine.register(HeldObjectHandOver(), replace=True) return engine diff --git a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py index 1953360de..6ca39d26d 100644 --- a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py +++ b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py @@ -18,16 +18,26 @@ from __future__ import annotations -from dataclasses import dataclass - -import torch +from dataclasses import dataclass, replace from embodichain.lab.sim.atomic_actions import ( + ActionPlan, + JointPositionGoal, MoveHeldObject, MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + PlanningContext, + ResolvedActionRequest, + StateDelta, ) -__all__ = ["ExactTargetMoveHeldObject", "ExactTargetMoveHeldObjectOptions"] +__all__ = [ + "ActionEngineMoveJoints", + "ActionEngineMoveJointsOptions", + "ExactTargetMoveHeldObject", + "ExactTargetMoveHeldObjectOptions", +] @dataclass(frozen=True, slots=True, eq=False) @@ -36,15 +46,61 @@ class ExactTargetMoveHeldObjectOptions(MoveHeldObjectOptions): class ExactTargetMoveHeldObject(MoveHeldObject): - """Preserve a selected semantic orientation when explicitly requested.""" + """Action Engine marker for mainline exact-target transport.""" OptionsType = ExactTargetMoveHeldObjectOptions binding_contract = MoveHeldObject.binding_contract - def _apply_automatic_transport_rotation( + +@dataclass(frozen=True, slots=True, eq=False) +class ActionEngineMoveJointsOptions(MoveJointsOptions): + """Joint motion with an explicit optional single-arm release effect.""" + + single_release: bool = False + """Whether a successful gripper-open command releases the held object.""" + + def __post_init__(self) -> None: + if type(self.single_release) is not bool: + raise TypeError("single_release must be a boolean.") + + +class ActionEngineMoveJoints(MoveJoints): + """Preserve ordinary joint motion and commit explicit release nodes.""" + + OptionsType = ActionEngineMoveJointsOptions + binding_contract = MoveJoints.binding_contract + + def _plan( self, - move_eef_xpos: torch.Tensor, - end_arm_xpos: torch.Tensor, - ) -> None: - """Keep target shaping in GenSim grounding and free-yaw search.""" - del move_eef_xpos, end_arm_xpos + request: ResolvedActionRequest[ + JointPositionGoal, + ActionEngineMoveJointsOptions, + ], + context: PlanningContext, + ) -> ActionPlan: + endpoint = request.binding.endpoint("primary", "motion") + task_state_key = endpoint.task_state_key + if request.skill_options.single_release: + if not isinstance(task_state_key, str) or not task_state_key: + raise ValueError( + "Single-arm release requires a non-empty task-state key." + ) + if context.task.get_held_object(task_state_key) is None: + return self.failed_plan( + request, + context, + message=( + "Single-arm release requires an object held by task-state " + f"resource {task_state_key!r}." + ), + ) + + plan = super()._plan(request, context) + if not request.skill_options.single_release: + return plan + return replace( + plan, + expected_effects=StateDelta( + held_object_updates={task_state_key: None}, + ), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/executor.py b/embodichain/gen_sim/action_engine/runtime/executor.py index d07ab4aa0..3770bd486 100644 --- a/embodichain/gen_sim/action_engine/runtime/executor.py +++ b/embodichain/gen_sim/action_engine/runtime/executor.py @@ -3017,13 +3017,16 @@ def observe_waypoint(waypoint_index: int) -> None: from_planned_qpos=capability.state_effect == "preserve_hold", ) self._step_states[(step.id, arm)] = committed_state - self._update_ownership( - step, - arm, - action_class, - committed_state, - successful, - ) + if bool(outcome.grounded.motion_policy.get("single_release", False)): + self._release_ownership(step.object_uid, arm, successful) + else: + self._update_ownership( + step, + arm, + action_class, + committed_state, + successful, + ) edge_failed = ( failed | (~failed & ~assigned) diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index b0f8e8c24..35fdc23cc 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -707,6 +707,8 @@ def ground( if kind == "joint_state": joint_defaults = self.runtime_policy.grounding["joint_state"] source = binding.get("source") + if bool(binding.get("single_release", False)): + policy["single_release"] = True if source == "gripper_closed": policy["sample_interval"] = int( joint_defaults["hand_close_sample_interval"] @@ -809,10 +811,7 @@ def ground( phase="final", orientation_reference_pose=orientation_reference_pose, ) - target = AxisAlignGoal( - semantics=semantics, - object_target_pose=target_object_pose, - ) + target = AxisAlignGoal(semantics=semantics) elif capability.target_materializer == "coordinated_pickment": target_object_pose = self._semantic_target( step, diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index 0b22dd10e..d2c4e41a8 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -435,12 +435,48 @@ def _recipe( {"kind": "object", "object": object_uid}, dependencies, role, + {}, + motion_policy(), + ) + descend = _node( + group_id, + 2, + "MoveHeldObject", + task_type, + object_uid, + actor, + "arm", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "final", + }, + [alignment["id"]], + role, + {}, + motion_policy(), + ) + release = _node( + group_id, + 3, + "MoveJoints", + task_type, + object_uid, + actor, + "hand", + { + "kind": "joint_state", + "source": "gripper_open", + "single_release": True, + }, + [descend["id"]], + role, success, motion_policy(), ) lift_clear = _node( group_id, - 2, + 4, "MoveEndEffector", task_type, object_uid, @@ -451,14 +487,14 @@ def _recipe( "source": "release", "operation": "lift_clear", }, - [alignment["id"]], + [release["id"]], "cleanup", {}, motion_policy(), ) reorient = _node( group_id, - 3, + 5, "MoveEndEffector", task_type, object_uid, @@ -476,7 +512,7 @@ def _recipe( ) post_reorient_lift = _node( group_id, - 4, + 6, "MoveEndEffector", task_type, object_uid, @@ -495,7 +531,7 @@ def _recipe( ) retreat = _node( group_id, - 5, + 7, "MoveEndEffector", task_type, object_uid, @@ -513,7 +549,7 @@ def _recipe( ) home = _node( group_id, - 6, + 8, "MoveJoints", task_type, object_uid, @@ -533,6 +569,8 @@ def _recipe( return ( [ alignment, + descend, + release, lift_clear, reorient, post_reorient_lift, diff --git a/embodichain/gen_sim/task_engine/defaults.yaml b/embodichain/gen_sim/task_engine/defaults.yaml index 4c58490d7..7bc281451 100644 --- a/embodichain/gen_sim/task_engine/defaults.yaml +++ b/embodichain/gen_sim/task_engine/defaults.yaml @@ -25,7 +25,7 @@ planning: candidate_count: 3 planning_mode: offline gripper_model: robotiq - ik_solver: pytorch + ik_solver: auto max_episodes: 1 max_episode_steps: 6000 planner: diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py index 276c815c3..f0804579a 100644 --- a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -128,41 +128,83 @@ def test_axis_align_uses_its_tutorial_motion_policy_base() -> None: assert capability.motion_base == "AxisAlign" -def test_axis_align_verifies_the_live_semantic_postcondition() -> None: +def test_axis_align_retains_ownership_until_explicit_release() -> None: capability = build_atomic_capability_registry().get("AxisAlign") - attempted = torch.tensor([True, False]) - calls = [] + contract = capability.resolve_contract( + { + "object_uid": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + } + ) + + assert capability.state_effect == "hold" + assert capability.verifier_hook is None + assert contract.requires == ( + StateAtom("arm_free", arm="left_arm"), + StateAtom("object_free", object_uid="can"), + ) + assert contract.effects[-1].atom == StateAtom( + "object_held", + object_uid="can", + arm="left_arm", + ) - def verify_step(step, failed): - calls.append((step, failed.clone())) - return failed.clone(), torch.tensor([True, False]), torch.zeros(2, 3) +def test_single_arm_release_contract_frees_the_object_and_arm() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + contract = capability.resolve_contract( + { + "object_uid": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "control": "hand", + "target_binding": { + "kind": "joint_state", + "source": "gripper_open", + "single_release": True, + }, + } + ) + + assert contract.requires == ( + StateAtom("object_held", object_uid="can", arm="left_arm"), + ) + assert [(effect.op, effect.atom.predicate) for effect in contract.effects] == [ + ("delete", "object_held"), + ("add", "arm_free"), + ("add", "object_free"), + ] + assert contract.failure_policy == "task_required" + + +def test_single_arm_release_verifies_the_selected_gripper_is_open() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") executor = SimpleNamespace( - _verify_step=verify_step, - _action_execution_observation=lambda _uid: { - "linear_velocity": torch.zeros(2, 3), - "angular_velocity": torch.zeros(2, 3), - }, + env=SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + open_state=(0.0, 0.0), + get_current_gripper_state_agent=lambda: ( + torch.tensor([[0.0, 0.0], [0.1, 0.1]]), + torch.ones(2, 2), + ), + ), runtime_policy=SimpleNamespace( - execution={ - "support_linear_velocity_tolerance": 0.02, - "support_angular_velocity_tolerance": 0.2, - } + predicate_fallbacks={"gripper_state_tolerance": 1.0e-3} ), ) - step = SimpleNamespace(id="orient", object_uid="can") + outcome = SimpleNamespace( + grounded=SimpleNamespace(motion_policy={"single_release": True}) + ) verified = capability.verifier_hook( executor=executor, - step=step, + step=SimpleNamespace(), arm="left_arm", - outcome=SimpleNamespace(), - attempted=attempted, + outcome=outcome, + attempted=torch.tensor([True, True]), ) assert verified.tolist() == [True, False] - assert calls[0][0] is step - assert calls[0][1].tolist() == [False, True] def test_explicit_required_home_is_safety_required_for_any_task_type() -> None: diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 2c08b95cc..59ff624a1 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -29,6 +29,7 @@ from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver from embodichain.gen_sim.action_engine.runtime import actions from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ActionEngineMoveJoints, ExactTargetMoveHeldObject, ) from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter @@ -223,6 +224,7 @@ def register(self, action: Any, *, replace: bool = False) -> None: assert isinstance(engine, Engine) assert registered == [ (ExactTargetMoveHeldObject, True), + (ActionEngineMoveJoints, True), (HeldObjectHandOver, True), ] diff --git a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py index 2859095c2..54cdd37cc 100644 --- a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py +++ b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py @@ -16,46 +16,30 @@ from __future__ import annotations +from dataclasses import dataclass from types import SimpleNamespace -import pytest -import torch - from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ActionEngineMoveJoints, + ActionEngineMoveJointsOptions, ExactTargetMoveHeldObject, ExactTargetMoveHeldObjectOptions, ) -from embodichain.lab.sim.atomic_actions import MoveHeldObject, MoveHeldObjectOptions - - -def test_grounded_target_transport_never_applies_an_implicit_rotation( - monkeypatch: pytest.MonkeyPatch, -) -> None: - applied = [] - result = object() - - def fake_apply(self, move_eef_xpos, end_arm_xpos) -> None: - del self, move_eef_xpos, end_arm_xpos - applied.append(True) +from embodichain.lab.sim.atomic_actions import ( + MoveHeldObject, + MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + StateDelta, +) - def fake_plan(self, request, context): - del request, context - self._apply_automatic_transport_rotation(torch.eye(4), torch.eye(4)) - return result - monkeypatch.setattr( - MoveHeldObject, - "_apply_automatic_transport_rotation", - fake_apply, - ) - monkeypatch.setattr(MoveHeldObject, "_plan", fake_plan) +def test_grounded_target_transport_uses_mainline_exact_target_contract() -> None: action = ExactTargetMoveHeldObject() assert type(action).__dict__["binding_contract"] is MoveHeldObject.binding_contract - request = SimpleNamespace(skill_options=ExactTargetMoveHeldObjectOptions()) - - assert action._plan(request, object()) is result - assert not applied + assert action._plan.__func__ is MoveHeldObject._plan + assert "_apply_automatic_transport_rotation" not in type(action).__dict__ def test_semantic_transport_config_has_no_task_facing_rotation_switch() -> None: @@ -70,3 +54,88 @@ def test_semantic_transport_config_has_no_task_facing_rotation_switch() -> None: assert isinstance(options, ExactTargetMoveHeldObjectOptions) assert not hasattr(options, "allow_automatic_transport_rotation") + + +def test_joint_config_materializes_single_release_only_when_requested() -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + capability = SimpleNamespace( + config_type=MoveJointsOptions, + target_materializer="joint_state", + ) + + release = adapter._build_single_arm_config( + SimpleNamespace(cfg={"single_release": True}), + capability, + ) + ordinary = adapter._build_single_arm_config( + SimpleNamespace(cfg={}), + capability, + ) + + assert isinstance(release, ActionEngineMoveJointsOptions) + assert release.single_release + assert isinstance(ordinary, ActionEngineMoveJointsOptions) + assert not ordinary.single_release + + +def test_single_release_binds_hand_motion_to_the_arm_held_state_key() -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + adapter._parts = lambda _arm: ("physical_left_arm", "physical_left_hand", 2) + captured = {} + + class Engine: + def bind_control_parts(self, skill_id, endpoints, *, task_state_keys=None): + captured.update( + skill_id=skill_id, + endpoints=endpoints, + task_state_keys=task_state_keys, + ) + return object() + + adapter._binding( + SimpleNamespace( + arm="left_arm", + control="hand", + cfg={"single_release": True}, + ), + SimpleNamespace( + action_type=MoveJoints, + config_materializer="single_arm", + ), + engine=Engine(), + ) + + assert captured == { + "skill_id": "move_joints", + "endpoints": {"primary": {"motion": "physical_left_hand"}}, + "task_state_keys": {"primary": "physical_left_arm"}, + } + + +def test_single_release_plan_removes_only_the_bound_arm_attachment(monkeypatch) -> None: + @dataclass(frozen=True) + class Plan: + expected_effects: object + + monkeypatch.setattr( + MoveJoints, + "_plan", + lambda _self, _request, _context: Plan(expected_effects=object()), + ) + action = ActionEngineMoveJoints() + request = SimpleNamespace( + binding=SimpleNamespace( + endpoint=lambda _slot, _endpoint: SimpleNamespace(task_state_key="left_arm") + ), + skill_options=ActionEngineMoveJointsOptions(single_release=True), + ) + context = SimpleNamespace( + task=SimpleNamespace( + get_held_object=lambda key: object() if key == "left_arm" else None + ) + ) + + plan = action._plan(request, context) + + assert isinstance(plan.expected_effects, StateDelta) + assert dict(plan.expected_effects.held_object_updates) == {"left_arm": None} diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index ac6f125a8..a3d8db192 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -3905,6 +3905,18 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - if candidate.id in orient_step.edge_ids and candidate.actions[0]["atomic_action_class"] == "AxisAlign" ) + orient_descend_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["atomic_action_class"] == "MoveHeldObject" + ) + orient_release_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["target_binding"].get("single_release") is True + ) orient_lift_edges = [ candidate for candidate in program.edges @@ -3952,6 +3964,18 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - arm="right_arm", state=ExecutionState(last_qpos=env.robot.get_qpos()), ) + orient_descend = grounder.ground( + orient_descend_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + orient_release = grounder.ground( + orient_release_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) release_pose = _pose(0.05, 0.2, 0.78) orient_lift = grounder.ground( orient_lift_edge.actions[0], @@ -3992,6 +4016,19 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - assert "approach_direction_mode" not in handover_pickup.cfg assert handover_pickup.cfg["pick_object_part"] == "top" assert isinstance(orient_alignment.target, AxisAlignGoal) + assert isinstance(orient_descend.target, HeldObjectPoseGoal) + assert torch.equal( + orient_descend.target.object_target_pose[:, :2, 3], + orient_alignment.target_object_pose[:, :2, 3], + ) + assert orient_descend.motion_policy["surface_clearance"] == pytest.approx(0.005) + assert isinstance(orient_release.target, JointPositionGoal) + assert orient_release.control == "hand" + assert orient_release.cfg["single_release"] is True + assert torch.equal( + orient_release.target.target, + torch.as_tensor(env.open_state, dtype=torch.float32), + ) assert isinstance(orient_lift.target, EndEffectorPoseGoal) assert torch.equal( orient_lift.target.xpos[:, :2, 3], @@ -4044,10 +4081,7 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - ).any() ) assert orient_alignment.target.grasp_xpos is None - assert torch.equal( - orient_alignment.target.object_target_pose, - orient_alignment.target_object_pose, - ) + assert not hasattr(orient_alignment.target, "object_target_pose") assert orient_alignment.target_object_pose is not None assert isinstance( orient_alignment.target.semantics.affordance, @@ -5585,7 +5619,8 @@ def execute( result = executor.run() assert active_by_edge[lift_edge.id] == [True] - for edge_id in program.semantic_steps[0].edge_ids[2:]: + lift_index = program.semantic_steps[0].edge_ids.index(lift_edge.id) + for edge_id in program.semantic_steps[0].edge_ids[lift_index + 1 :]: assert active_by_edge[edge_id] == [False] assert not bool(result.success[0]) diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index c40b894a9..c07ab0f2b 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -141,6 +141,8 @@ def test_historical_task2_1_uses_axis_align_and_explicit_handover_arms() -> None expected_orient_actions = [ "AxisAlign", + "MoveHeldObject", + "MoveJoints", "MoveEndEffector", "MoveEndEffector", "MoveEndEffector", @@ -240,6 +242,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert [node["atomic_action"] for node in orient_nodes] == [ "AxisAlign", + "MoveHeldObject", + "MoveJoints", "MoveEndEffector", "MoveEndEffector", "MoveEndEffector", @@ -247,9 +251,12 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "MoveJoints", ] assert orient["goal"]["upright_local_axis"] == "z" - assert orient_nodes[0]["postcondition"]["local_axis"] == "z" + assert orient_nodes[0]["postcondition"] == {} + assert orient_nodes[2]["postcondition"]["local_axis"] == "z" assert orient_nodes[0]["motion_policy"] == {"modifiers": []} assert [node["role"] for node in orient_nodes] == [ + "primary", + "primary", "primary", "cleanup", "cleanup", @@ -258,38 +265,47 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "cleanup", ] assert orient_nodes[1]["target_binding"] == { + "kind": "semantic_goal", + "semantic_step": "orient", + "phase": "final", + } + assert orient_nodes[1]["motion_policy"] == {"modifiers": []} + assert orient_nodes[2]["target_binding"] == { + "kind": "joint_state", + "source": "gripper_open", + "single_release": True, + } + assert orient_nodes[2]["control"] == "hand" + assert orient_nodes[3]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "lift_clear", } - assert orient_nodes[1]["motion_policy"] == {"modifiers": []} - assert orient_nodes[2]["target_binding"] == { + assert orient_nodes[3]["motion_policy"] == {"modifiers": []} + assert orient_nodes[4]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "reorient_tool_down", } - assert orient_nodes[3]["target_binding"] == { + assert orient_nodes[5]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "lift_clear", "requires_arm_clear": True, } - assert orient_nodes[4]["target_binding"] == { + assert orient_nodes[6]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "retreat_after_lift", } - assert orient_nodes[5]["target_binding"] == { + assert orient_nodes[7]["target_binding"] == { "kind": "joint_state", "source": "initial", "operation": "e2_home", "required_home": True, } - assert orient_nodes[1]["depends_on"] == [orient_nodes[0]["id"]] - assert orient_nodes[2]["depends_on"] == [orient_nodes[1]["id"]] - assert orient_nodes[3]["depends_on"] == [orient_nodes[2]["id"]] - assert orient_nodes[4]["depends_on"] == [orient_nodes[3]["id"]] - assert orient_nodes[5]["depends_on"] == [orient_nodes[4]["id"]] + for previous, current in zip(orient_nodes, orient_nodes[1:]): + assert current["depends_on"] == [previous["id"]] assert [node["atomic_action"] for node in handover_nodes] == [ "PickUp", "MoveHeldObject", @@ -306,23 +322,55 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert handover_nodes[0]["actor"] == {"mode": "required", "arm": "right_arm"} assert orient_nodes[-1]["contract"]["completion"] == "terminal_barrier" assert orient_nodes[0]["contract"]["failure_policy"] == "task_required" - assert orient_nodes[1]["contract"]["failure_policy"] == "safety_required" - assert orient_nodes[2]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[1]["contract"]["failure_policy"] == "task_required" + assert orient_nodes[2]["contract"]["failure_policy"] == "task_required" assert orient_nodes[3]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[4]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[5]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[6]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[-1]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[1]["contract"]["requires"] == [ - {"predicate": "arm_free", "arm": "left_arm"} + { + "predicate": "object_held", + "object_uid": "interact_can", + "arm": "left_arm", + } ] assert orient_nodes[2]["contract"]["requires"] == [ - {"predicate": "arm_clear", "arm": "left_arm"} + { + "predicate": "object_held", + "object_uid": "interact_can", + "arm": "left_arm", + } + ] + assert [ + (effect["op"], effect["atom"]["predicate"]) + for effect in orient_nodes[0]["contract"]["effects"] + ] == [ + ("delete", "arm_free"), + ("delete", "object_free"), + ("add", "object_held"), + ] + assert [ + (effect["op"], effect["atom"]["predicate"]) + for effect in orient_nodes[2]["contract"]["effects"] + ] == [ + ("delete", "object_held"), + ("add", "arm_free"), + ("add", "object_free"), ] assert orient_nodes[3]["contract"]["requires"] == [ - {"predicate": "arm_clear", "arm": "left_arm"} + {"predicate": "arm_free", "arm": "left_arm"} ] assert orient_nodes[4]["contract"]["requires"] == [ {"predicate": "arm_clear", "arm": "left_arm"} ] + assert orient_nodes[5]["contract"]["requires"] == [ + {"predicate": "arm_clear", "arm": "left_arm"} + ] + assert orient_nodes[6]["contract"]["requires"] == [ + {"predicate": "arm_clear", "arm": "left_arm"} + ] assert any( effect["atom"]["predicate"] == "arm_home" for effect in orient["contract"]["exit_effects"] diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index b9051a2c0..9e91b71a9 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -322,6 +322,8 @@ def caller(**kwargs): handover_nodes = [node for node in graph["nodes"] if node["task_type"] == "E4"] assert orient_actions == [ "AxisAlign", + "MoveHeldObject", + "MoveJoints", "MoveEndEffector", "MoveEndEffector", "MoveEndEffector", From 8899abc7b3d92b405c8073e8553d117073ebf612 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:51:25 +0800 Subject: [PATCH 81/85] refactor(action-engine): implement E2 upright placement with pickup and transport --- .../action_engine/capabilities/atomic.py | 170 ++++++++++++++---- .../action_engine/domain/task_contracts.py | 2 +- .../gen_sim/action_engine/domain/v2.py | 23 ++- .../gen_sim/action_engine/gripper_profiles.py | 6 + .../gen_sim/action_engine/runtime/actions.py | 121 ++++++++++++- .../runtime/grasp_diagnostics.py | 155 ++++++++++++++++ .../gen_sim/action_engine/runtime/state.py | 63 +++++++ .../action_engine/tasks/interpretation.py | 2 +- .../gen_sim/action_engine/tasks/recipes.py | 48 +++-- .../capabilities/test_atomic_v2.py | 96 ++++++++++ .../action_engine/runtime/test_actions.py | 32 +++- .../runtime/test_grasp_diagnostics.py | 50 ++++++ .../runtime/test_runtime_contracts.py | 59 +++--- .../action_engine/tasks/test_factory.py | 56 ++++-- .../tasks/test_interpretation.py | 6 +- .../action_engine/test_gripper_profiles.py | 2 + 16 files changed, 787 insertions(+), 104 deletions(-) diff --git a/embodichain/gen_sim/action_engine/capabilities/atomic.py b/embodichain/gen_sim/action_engine/capabilities/atomic.py index 200fda04f..07b2f6d51 100644 --- a/embodichain/gen_sim/action_engine/capabilities/atomic.py +++ b/embodichain/gen_sim/action_engine/capabilities/atomic.py @@ -26,6 +26,8 @@ import torch +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile + __all__ = [ "ACTION_CONTRACT_VERSION", "AtomicCapability", @@ -434,7 +436,7 @@ def build_atomic_capability_registry() -> AtomicCapabilityRegistry: "control_part", "preserve", "joint_state", - verifier_hook=_verify_required_home, + verifier_hook=_verify_move_joints, contract_resolver_hook=_resolve_joints_contract, ), AtomicCapability( @@ -1073,7 +1075,7 @@ def _verify_arm_clearance( return attempted & clear -def _verify_required_home( +def _verify_move_joints( *, executor: Any, step: Any, @@ -1081,42 +1083,144 @@ def _verify_required_home( outcome: Any, attempted: torch.Tensor, ) -> torch.Tensor: - """Verify explicit release or required cleanup against live joint state.""" - del step + """Route joint effects to their dedicated physical verifier.""" policy = outcome.grounded.motion_policy if bool(policy.get("single_release", False)): - getter = getattr(executor.env, "get_current_gripper_state_agent", None) - if not callable(getter) or arm not in {"left_arm", "right_arm"}: - return torch.zeros_like(attempted) - values = getter() - index = 0 if arm == "left_arm" else 1 - if not isinstance(values, (tuple, list)) or len(values) <= index: - return torch.zeros_like(attempted) - current = torch.as_tensor( - values[index], - dtype=torch.float32, - device=executor.env.device, + return _verify_single_release( + executor=executor, + step=step, + arm=arm, + outcome=outcome, + attempted=attempted, ) - if current.ndim == 1: - current = current.unsqueeze(0).repeat(int(executor.env.num_envs), 1) - expected = torch.as_tensor( - executor.env.open_state, - dtype=current.dtype, - device=current.device, - ).flatten() - repeats = (current.shape[-1] + expected.numel() - 1) // expected.numel() - expected = expected.repeat(repeats)[: current.shape[-1]] - tolerance = float( - policy.get( - "release_gripper_tolerance", - executor.runtime_policy.predicate_fallbacks["gripper_state_tolerance"], - ) + return _verify_required_home( + executor=executor, + arm=arm, + outcome=outcome, + attempted=attempted, + ) + + +def _verify_single_release( + *, + executor: Any, + step: Any, + arm: str, + outcome: Any, + attempted: torch.Tensor, +) -> torch.Tensor: + """Verify normalized hand opening plus stable object support.""" + env = executor.env + stable_support = attempted.clone() + support_reference = getattr(executor, "_support_reference_uid", None) + support_stable_for = getattr(executor, "_support_stable_for", None) + if callable(support_reference) and callable(support_stable_for): + support_uid = support_reference(step) + if not isinstance(support_uid, str) or not support_uid: + stable_support &= False + else: + stable_support &= torch.as_tensor( + support_stable_for(step, support_uid, attempted), + dtype=torch.bool, + device=env.device, + ).reshape(-1) + + upright = attempted.clone() + entity_pose = getattr(executor, "_entity_pose", None) + orientation_satisfied = getattr( + executor, + "_placement_orientation_satisfied", + None, + ) + if callable(entity_pose) and callable(orientation_satisfied): + upright &= torch.as_tensor( + orientation_satisfied(step, entity_pose(step.object_uid)), + dtype=torch.bool, + device=env.device, + ).reshape(-1) + + getter = getattr(env, "get_current_gripper_state_agent", None) + if not callable(getter) or arm not in {"left_arm", "right_arm"}: + return torch.zeros_like(attempted) + values = getter() + index = 0 if arm == "left_arm" else 1 + if not isinstance(values, (tuple, list)) or len(values) <= index: + return torch.zeros_like(attempted) + current = torch.as_tensor( + values[index], + dtype=torch.float32, + device=env.device, + ) + if current.ndim == 1: + current = current.unsqueeze(0).repeat(int(env.num_envs), 1) + expected_open = torch.as_tensor( + env.open_state, + dtype=current.dtype, + device=current.device, + ).flatten() + expected_close = torch.as_tensor( + env.close_state, + dtype=current.dtype, + device=current.device, + ).flatten() + configured = getattr(env, "agent_gripper_state_joint_indices", {}) + side = "left" if arm == "left_arm" else "right" + indices = configured.get(side) if isinstance(configured, Mapping) else None + if indices is not None: + indices = list(indices) + current = current[:, indices] + expected_open = expected_open[indices] + expected_close = expected_close[indices] + else: + repeats = ( + current.shape[-1] + expected_open.numel() - 1 + ) // expected_open.numel() + expected_open = expected_open.repeat(repeats)[: current.shape[-1]] + expected_close = expected_close.repeat(repeats)[: current.shape[-1]] + stroke = torch.linalg.vector_norm(expected_close - expected_open) + if not torch.isfinite(stroke) or stroke <= 1.0e-6: + return torch.zeros_like(attempted) + open_error_fraction = ( + torch.linalg.vector_norm( + current - expected_open.unsqueeze(0), + dim=1, ) - opened = ( - torch.linalg.vector_norm(current - expected.unsqueeze(0), dim=1) - <= tolerance + / stroke + ) + gripper_profile = get_gripper_profile(getattr(env, "agent_gripper_model", "pgi")) + tolerance = float( + outcome.grounded.motion_policy.get( + "release_open_fraction_tolerance", + gripper_profile.release_open_fraction_tolerance, ) - return attempted & opened + ) + opened = open_error_fraction <= tolerance + accepted = attempted & opened & stable_support + planner_trace = getattr(outcome, "planner_trace", None) + if isinstance(planner_trace, dict): + planner_trace["release_verification"] = { + "state_joint_indices": None if indices is None else indices, + "current_state": current.detach().cpu().tolist(), + "expected_open_state": expected_open.detach().cpu().tolist(), + "open_error_fraction": open_error_fraction.detach().cpu().tolist(), + "open_fraction_tolerance": tolerance, + "gripper_open": opened.detach().cpu().tolist(), + "support_stable": stable_support.detach().cpu().tolist(), + "upright": upright.detach().cpu().tolist(), + "accepted": accepted.detach().cpu().tolist(), + } + return accepted + + +def _verify_required_home( + *, + executor: Any, + arm: str, + outcome: Any, + attempted: torch.Tensor, +) -> torch.Tensor: + """Verify an explicit required-home effect against live arm joints.""" + policy = outcome.grounded.motion_policy if not bool(policy.get("verify_required_home", False)): return attempted env = executor.env diff --git a/embodichain/gen_sim/action_engine/domain/task_contracts.py b/embodichain/gen_sim/action_engine/domain/task_contracts.py index da1ff9397..dccbb384e 100644 --- a/embodichain/gen_sim/action_engine/domain/task_contracts.py +++ b/embodichain/gen_sim/action_engine/domain/task_contracts.py @@ -70,7 +70,7 @@ def normalize_placement_relation(value: Any) -> str: _CORE_ACTIONS: Mapping[str, tuple[str, ...]] = MappingProxyType( { "E1": ("PickUp", "MoveHeldObject", "Place"), - "E2": ("AxisAlign", "MoveHeldObject", "MoveJoints"), + "E2": ("PickUp", "MoveHeldObject", "MoveJoints"), "E3": ("PickUp", "MoveHeldObject", "Pour", "Place"), "E4": ("PickUp", "MoveHeldObject", "HandOver", "Place"), "E5": ("CoordinatedPickment",), diff --git a/embodichain/gen_sim/action_engine/domain/v2.py b/embodichain/gen_sim/action_engine/domain/v2.py index f14f278a2..8125d8e26 100644 --- a/embodichain/gen_sim/action_engine/domain/v2.py +++ b/embodichain/gen_sim/action_engine/domain/v2.py @@ -673,7 +673,9 @@ def _validate_task_group_semantics( if precondition.get("type") != "object_held": missing = {"PickUp|object_held precondition"} if contract.success_type == "object_upright" and "AxisAlign" not in actions: - missing = {"MoveHeldObject", "Place"} - actions + missing = {"MoveHeldObject"} - actions + if not _has_single_object_release_effect(group_nodes): + missing.add("object release effect") if "PickUp" not in actions: first = group_nodes[0] precondition = first.get("precondition", {}) @@ -701,6 +703,25 @@ def _validate_task_group_semantics( ) +def _has_single_object_release_effect(nodes: Sequence[Mapping[str, Any]]) -> bool: + """Return whether one node releases a single-arm held object by contract.""" + for node in nodes: + effects = node.get("contract", {}).get("effects", ()) + deleted_held = any( + effect.get("op") == "delete" + and effect.get("atom", {}).get("predicate") == "object_held" + for effect in effects + ) + added_free = any( + effect.get("op") == "add" + and effect.get("atom", {}).get("predicate") == "object_free" + for effect in effects + ) + if deleted_held and added_free: + return True + return False + + def _validate_ownership_transitions( nodes: Sequence[Mapping[str, Any]], groups: Sequence[Mapping[str, Any]], diff --git a/embodichain/gen_sim/action_engine/gripper_profiles.py b/embodichain/gen_sim/action_engine/gripper_profiles.py index b9b3e7bc2..e26fa8c5c 100644 --- a/embodichain/gen_sim/action_engine/gripper_profiles.py +++ b/embodichain/gen_sim/action_engine/gripper_profiles.py @@ -100,6 +100,7 @@ class GripperProfile: drive_stiffness: float drive_damping: float drive_max_effort: float + release_open_fraction_tolerance: float grasp_model: GraspModelSpec def __post_init__(self) -> None: @@ -117,6 +118,8 @@ def __post_init__(self) -> None: raise ValueError( "Gripper control states and limits must match control joints." ) + if not 0.0 < self.release_open_fraction_tolerance <= 1.0: + raise ValueError("release_open_fraction_tolerance must be in (0, 1].") mimic_count = len(self.left_mimic_joints) if not ( len(self.right_mimic_joints) @@ -212,6 +215,7 @@ def runtime_manifest( "open_positions": list(self.open_positions), "close_positions": list(self.close_positions), "control_limits": [list(limit) for limit in self.control_limits], + "release_open_fraction_tolerance": (self.release_open_fraction_tolerance), "tcp": { "parent_frames": dict(tcp_parent_frames), "transform_direction": "parent_link_to_tcp", @@ -253,6 +257,7 @@ def _validate_side(side: _Side) -> None: drive_stiffness=1.0e3, drive_damping=1.0e2, drive_max_effort=1.0e4, + release_open_fraction_tolerance=0.03, grasp_model=GraspModelSpec( model_id="dh_pgi_140_80", min_opening_width=0.003, @@ -324,6 +329,7 @@ def _validate_side(side: _Side) -> None: drive_stiffness=50.0, drive_damping=5.0, drive_max_effort=500.0, + release_open_fraction_tolerance=0.03, grasp_model=GraspModelSpec( model_id="robotiq_arg2f_140", min_opening_width=0.01, diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py index 3a410e934..64b8d518b 100644 --- a/embodichain/gen_sim/action_engine/runtime/actions.py +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -85,7 +85,7 @@ from .coordinated_safety import _trajectory_safety_report from .grasp_diagnostics import _TracingAntipodalGraspPoseGenerator from .models import ActionOutcome, GroundedAction -from .state import ExecutionState +from .state import _CollisionOverrideSceneSnapshot, ExecutionState __all__ = ["AtomicActionAdapter"] @@ -456,12 +456,23 @@ def plan( candidate, context, ) + upright_context = self._upright_grasp_selection_context( + candidate_engine, + candidate, + capability, + ) with ( seed_context, pair_context, + upright_context, _capture_retreat_warnings(capture_warnings) as warnings, ): candidate_plan = candidate_engine.plan(candidate_invocation, context) + self._record_selected_upright_grasp( + candidate, + candidate_plan, + context, + ) coordinated_trace = candidate.motion_policy.get("coordinated_grasp") if isinstance(coordinated_trace, dict): grasp_stages = ( @@ -1505,6 +1516,102 @@ def _audit_coordinated_trajectory( audit_trace, ) + @contextmanager + def _upright_grasp_selection_context( + self, + engine: AtomicActionEngine, + candidate: GroundedAction, + capability: AtomicCapability, + ) -> Iterator[None]: + """Apply the E2 side-grasp policy only to upright pickup.""" + local_axis = candidate.cfg.get("obj_upright_direction") + if ( + capability.target_materializer != "object_grasp" + or candidate.cfg.get("rotate_upright") is None + or local_axis is None + ): + yield + return + _, hand_part, _ = self._parts(candidate.arm) + if hand_part is None: + yield + return + generator = engine.grasp_pose_generators.get(hand_part) + if not isinstance(generator, _TracingAntipodalGraspPoseGenerator): + yield + return + with generator.upright_selection_context( + local_axis=torch.as_tensor( + local_axis, + dtype=torch.float32, + device=self.device, + ) + ): + yield + trace = generator.last_upright_trace + if trace is not None: + candidate.motion_policy["upright_grasp"] = trace + + def _record_selected_upright_grasp( + self, + candidate: GroundedAction, + plan: ActionPlan, + context: PlanningContext, + ) -> None: + """Record the grasp actually selected after IK and downstream screening.""" + trace = candidate.motion_policy.get("upright_grasp") + local_axis = candidate.cfg.get("obj_upright_direction") + if not isinstance(trace, dict) or local_axis is None: + return + held = next( + ( + value + for value in plan.expected_effects.held_object_updates.values() + if value is not None + ), + None, + ) + if held is None or held.semantics.entity_id is None: + return + entity = context.scene.entities.get(held.semantics.entity_id) + vertices = held.semantics.geometry.get("mesh_vertices") + if entity is None or vertices is None: + return + grasp_pose = held.grasp_xpos.to(device=self.device, dtype=torch.float32) + object_pose = entity.pose.to(device=self.device, dtype=torch.float32) + if object_pose.shape == (4, 4): + object_pose = object_pose.unsqueeze(0).expand(self.num_envs, -1, -1) + axis = torch.as_tensor( + local_axis, + dtype=torch.float32, + device=self.device, + ) + axis = axis / torch.linalg.vector_norm(axis) + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=self.device, + ) + vertex_positions = torch.matmul(vertices, axis) + axis_min = vertex_positions.min() + axis_extent = vertex_positions.max() - axis_min + world_axis = torch.matmul(object_pose[:, :3, :3], axis) + relative_centers = grasp_pose[:, :3, 3] - object_pose[:, :3, 3] + axis_positions = torch.sum(relative_centers * world_axis, dim=1) + axis_fractions = (axis_positions - axis_min) / axis_extent + closing_axes = torch.nn.functional.normalize( + grasp_pose[:, :3, 0], + dim=1, + ) + axis_alignment = torch.abs(torch.sum(closing_axes * world_axis, dim=1)) + trace.update( + { + "selected_axis_fraction": axis_fractions.detach().cpu().tolist(), + "selected_axis_alignment": axis_alignment.detach().cpu().tolist(), + "selected_grasp_pose": grasp_pose.detach().cpu().tolist(), + } + ) + @contextmanager def _coordinated_pair_selection_context( self, @@ -1973,6 +2080,9 @@ def _planner_trace( body_grasp = grounded.motion_policy.get("body_grasp") if isinstance(body_grasp, Mapping): trace["body_grasp"] = deepcopy(dict(body_grasp)) + upright_grasp = grounded.motion_policy.get("upright_grasp") + if isinstance(upright_grasp, Mapping): + trace["upright_grasp"] = deepcopy(dict(upright_grasp)) coordinated_grasp = grounded.motion_policy.get("coordinated_grasp") if isinstance(coordinated_grasp, Mapping): trace["coordinated_grasp"] = deepcopy(dict(coordinated_grasp)) @@ -2139,6 +2249,7 @@ def _scene_snapshot( return base exclusion_masks = self._collision_exclusion_masks(grounded, state) entities = dict(base.entities) + collision_pose_overrides: dict[str, torch.Tensor] = {} for uid in dynamic_uids: entity_state = entities.get(uid) if entity_state is None: @@ -2155,18 +2266,20 @@ def _scene_snapshot( ) excluded = exclusion_masks.get(uid) if excluded is not None and bool(excluded.any()): - pose = pose.clone() - pose[excluded, 2, 3] += _COLLISION_PARKING_Z_OFFSET + collision_pose = pose.clone() + collision_pose[excluded, 2, 3] += _COLLISION_PARKING_Z_OFFSET + collision_pose_overrides[uid] = collision_pose entities[uid] = EntityState( pose=pose, confidence=entity_state.confidence, ) - return SceneSnapshot( + return _CollisionOverrideSceneSnapshot( timestamp=base.timestamp, version=base.version, entities=entities, collision_world_revision=base.collision_world_revision, collision_entity_ids=dynamic_uids, + collision_pose_overrides=collision_pose_overrides, ) def _collision_exclusion_masks( diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py index 3e5ac2da7..c17dba7f7 100644 --- a/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py +++ b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py @@ -35,6 +35,17 @@ __all__: list[str] = [] +_UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT = 0.65 +_UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION = 0.35 +_UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION = 0.75 +_UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT = 2.0 +_UPRIGHT_SIDE_GRASP_CANDIDATE_LIMIT = 50 + + +@dataclass(frozen=True, slots=True) +class _UprightGraspSelectionContext: + local_axis: torch.Tensor + @dataclass(frozen=True, slots=True) class _DualGraspSelectionContext: @@ -54,13 +65,42 @@ class _TracingAntipodalGraspPoseGenerator(AntipodalGraspPoseGenerator): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._last_dual_trace: dict[str, Any] | None = None + self._last_upright_trace: dict[str, Any] | None = None self._selection_context: _DualGraspSelectionContext | None = None + self._upright_selection_context: _UprightGraspSelectionContext | None = None @property def last_dual_trace(self) -> dict[str, Any] | None: """Return an owned snapshot of the most recent dual-grasp trace.""" return deepcopy(self._last_dual_trace) + @property + def last_upright_trace(self) -> dict[str, Any] | None: + """Return an owned snapshot of the most recent upright-grasp trace.""" + return deepcopy(self._last_upright_trace) + + @contextmanager + def upright_selection_context( + self, + *, + local_axis: torch.Tensor, + ) -> Iterator[None]: + """Install one invocation-local side-grasp selection policy.""" + if self._upright_selection_context is not None: + raise RuntimeError("Upright grasp selection context cannot be nested.") + axis = torch.as_tensor(local_axis, dtype=torch.float32).reshape(-1) + norm = torch.linalg.vector_norm(axis) + if axis.shape != (3,) or not torch.isfinite(axis).all() or norm <= 1.0e-6: + raise ValueError("Upright grasp local_axis must be one finite 3-vector.") + self._last_upright_trace = None + self._upright_selection_context = _UprightGraspSelectionContext( + local_axis=(axis / norm).clone() + ) + try: + yield + finally: + self._upright_selection_context = None + @contextmanager def dual_arm_selection_context( self, @@ -263,6 +303,121 @@ def _candidate_count(result: dict[str, Any]) -> int: return 0 return int(poses.shape[0]) if poses.ndim == 3 else 0 + def get_valid_grasp_poses( + self, + **kwargs: Any, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Apply v14's side and mid-body preference for upright pickup.""" + results = super().get_valid_grasp_poses(**kwargs) + context = self._upright_selection_context + if context is None: + return results + + vertices = torch.as_tensor(kwargs["mesh_vertices"], dtype=torch.float32) + object_poses = torch.as_tensor(kwargs["obj_poses"], dtype=torch.float32) + local_axis = context.local_axis.to( + device=vertices.device, + dtype=vertices.dtype, + ) + vertex_positions = torch.matmul(vertices, local_axis) + axis_min = vertex_positions.min() + axis_extent = vertex_positions.max() - axis_min + if float(axis_extent) <= 1.0e-6: + raise ValueError("Upright grasp axis must span non-zero object geometry.") + + ranked_results: list[tuple[torch.Tensor, torch.Tensor]] = [] + row_traces: list[dict[str, Any]] = [] + for row_index, (result, object_pose) in enumerate( + zip(results, object_poses, strict=True) + ): + grasp_poses, costs = result + grasp_poses = torch.as_tensor(grasp_poses, dtype=torch.float32) + costs = torch.as_tensor( + costs, + device=grasp_poses.device, + dtype=torch.float32, + ) + if grasp_poses.ndim == 2: + grasp_poses = grasp_poses.unsqueeze(0) + object_pose = object_pose.to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + axis = local_axis.to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + world_upright = torch.matmul(object_pose[:3, :3], axis) + closing_axes = F.normalize(grasp_poses[:, :3, 0], dim=1) + axis_alignment = torch.abs( + torch.sum(closing_axes * world_upright[None], dim=1) + ) + side_compatible = axis_alignment <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT + relative_centers = grasp_poses[:, :3, 3] - object_pose[None, :3, 3] + center_axis_positions = torch.sum( + relative_centers * world_upright[None], + dim=1, + ) + center_fractions = ( + center_axis_positions - axis_min.to(grasp_poses.device) + ) / axis_extent.to(grasp_poses.device) + central_band = ( + center_fractions >= _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION + ) & (center_fractions <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION) + interval = ( + _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION + - _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION + ) + height_penalty = ( + torch.clamp( + _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION - center_fractions, + min=0.0, + ) + + torch.clamp( + center_fractions - _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION, + min=0.0, + ) + ) / interval + adjusted_costs = ( + torch.where( + side_compatible, + costs, + torch.full_like(costs, torch.inf), + ) + + _UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT * height_penalty + ) + ranked = torch.argsort(adjusted_costs)[:_UPRIGHT_SIDE_GRASP_CANDIDATE_LIMIT] + ranked_results.append((grasp_poses[ranked], adjusted_costs[ranked])) + finite_ranked = torch.isfinite(adjusted_costs[ranked]) + best_index = int(ranked[0].item()) if bool(finite_ranked.any()) else None + row_traces.append( + { + "environment_index": row_index, + "local_axis": axis.detach().cpu().tolist(), + "candidate_count": int(grasp_poses.shape[0]), + "side_compatible_count": int(side_compatible.sum().item()), + "central_band_count": int(central_band.sum().item()), + "side_and_central_count": int( + (side_compatible & central_band).sum().item() + ), + "retained_count": int(finite_ranked.sum().item()), + "best_candidate_axis_alignment": ( + None + if best_index is None + else float(axis_alignment[best_index].item()) + ), + "best_candidate_axis_fraction": ( + None + if best_index is None + else float(center_fractions[best_index].item()) + ), + } + ) + self._last_upright_trace = ( + row_traces[0] if len(row_traces) == 1 else {"environment_rows": row_traces} + ) + return ranked_results + def get_dual_arm_valid_grasp_poses(self, **kwargs: Any) -> list[dict | None]: """Run the standard generator while observing its NMS/collision boundary.""" vertices = kwargs["mesh_vertices"] diff --git a/embodichain/gen_sim/action_engine/runtime/state.py b/embodichain/gen_sim/action_engine/runtime/state.py index dbafc4748..7fdafebd7 100644 --- a/embodichain/gen_sim/action_engine/runtime/state.py +++ b/embodichain/gen_sim/action_engine/runtime/state.py @@ -19,18 +19,81 @@ from __future__ import annotations from dataclasses import dataclass, field +from types import MappingProxyType from typing import Mapping import torch from embodichain.lab.sim.atomic_actions import ( HeldObjectState, + SceneSnapshot, TaskState, ) __all__ = ["ExecutionState"] +@dataclass(frozen=True, slots=True, eq=False) +class _CollisionOverrideSceneSnapshot(SceneSnapshot): + """Keep semantic entity poses live while overriding collision poses.""" + + collision_pose_overrides: Mapping[str, torch.Tensor] = field(default_factory=dict) + + def __post_init__(self) -> None: + SceneSnapshot.__post_init__(self) + normalized: dict[str, torch.Tensor] = {} + for entity_id, pose in self.collision_pose_overrides.items(): + if entity_id not in self.collision_entity_ids: + raise ValueError( + "Collision pose overrides must reference collision entities." + ) + if ( + not isinstance(pose, torch.Tensor) + or not pose.is_floating_point() + or pose.dim() not in (2, 3) + or pose.shape[-2:] != (4, 4) + or not bool(torch.isfinite(pose).all().item()) + ): + raise ValueError( + "Collision pose overrides must be finite floating tensors " + "with shape (4, 4) or (B, 4, 4)." + ) + normalized[entity_id] = pose.detach().clone() + object.__setattr__( + self, + "collision_pose_overrides", + MappingProxyType(normalized), + ) + + def collision_obstacle_poses( + self, + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> Mapping[str, torch.Tensor]: + """Return planner poses with intentional-contact rows parked.""" + poses = dict( + SceneSnapshot.collision_obstacle_poses( + self, + batch_size=batch_size, + device=device, + dtype=dtype, + ) + ) + for entity_id, override in self.collision_pose_overrides.items(): + pose = override.to(device=device, dtype=dtype) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).expand(batch_size, -1, -1) + elif pose.shape != (batch_size, 4, 4): + raise ValueError( + f"Collision override {entity_id!r} must match planning " + f"batch size {batch_size}." + ) + poses[entity_id] = pose.clone() + return MappingProxyType(poses) + + @dataclass(slots=True, eq=False) class ExecutionState: """Projected task state paired with the next full-robot planning seed. diff --git a/embodichain/gen_sim/action_engine/tasks/interpretation.py b/embodichain/gen_sim/action_engine/tasks/interpretation.py index be136183b..1a5da8467 100644 --- a/embodichain/gen_sim/action_engine/tasks/interpretation.py +++ b/embodichain/gen_sim/action_engine/tasks/interpretation.py @@ -262,7 +262,7 @@ def _emit_step( { "orientation_goal": "upright", "support_role": "table", - "upright_local_axis": "z", + "upright_local_axis": "auto", } ) elif task_type == "E3": diff --git a/embodichain/gen_sim/action_engine/tasks/recipes.py b/embodichain/gen_sim/action_engine/tasks/recipes.py index d2c4e41a8..217d97992 100644 --- a/embodichain/gen_sim/action_engine/tasks/recipes.py +++ b/embodichain/gen_sim/action_engine/tasks/recipes.py @@ -413,7 +413,7 @@ def _recipe( "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", "z")), + "upright_local_axis": str(params.get("upright_local_axis", "auto")), **_orientation_extensions(params), } if terminal_behavior == "hold": @@ -424,10 +424,11 @@ def _recipe( "local_axis": goal["upright_local_axis"], } if incoming_held_arm is None and terminal_behavior == "place": - alignment = _node( + upright_policy = motion_policy(("orientation", "upright")) + pickup = _node( group_id, 1, - "AxisAlign", + "PickUp", task_type, object_uid, actor, @@ -436,9 +437,9 @@ def _recipe( dependencies, role, {}, - motion_policy(), + upright_policy, ) - descend = _node( + staging = _node( group_id, 2, "MoveHeldObject", @@ -446,19 +447,37 @@ def _recipe( object_uid, actor, "arm", + { + "kind": "semantic_goal", + "semantic_step": group_id, + "phase": "staging", + }, + [pickup["id"]], + role, + {}, + upright_policy, + ) + descend = _node( + group_id, + 3, + "MoveHeldObject", + task_type, + object_uid, + actor, + "arm", { "kind": "semantic_goal", "semantic_step": group_id, "phase": "final", }, - [alignment["id"]], + [staging["id"]], role, {}, - motion_policy(), + upright_policy, ) release = _node( group_id, - 3, + 4, "MoveJoints", task_type, object_uid, @@ -476,7 +495,7 @@ def _recipe( ) lift_clear = _node( group_id, - 4, + 5, "MoveEndEffector", task_type, object_uid, @@ -494,7 +513,7 @@ def _recipe( ) reorient = _node( group_id, - 5, + 6, "MoveEndEffector", task_type, object_uid, @@ -512,7 +531,7 @@ def _recipe( ) post_reorient_lift = _node( group_id, - 6, + 7, "MoveEndEffector", task_type, object_uid, @@ -531,7 +550,7 @@ def _recipe( ) retreat = _node( group_id, - 7, + 8, "MoveEndEffector", task_type, object_uid, @@ -549,7 +568,7 @@ def _recipe( ) home = _node( group_id, - 8, + 9, "MoveJoints", task_type, object_uid, @@ -568,7 +587,8 @@ def _recipe( ) return ( [ - alignment, + pickup, + staging, descend, release, lift_clear, diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py index f0804579a..7c27c824f 100644 --- a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -20,6 +20,7 @@ from dataclasses import dataclass from types import SimpleNamespace +import pytest import torch from embodichain.gen_sim.action_engine.capabilities import ( @@ -183,6 +184,7 @@ def test_single_arm_release_verifies_the_selected_gripper_is_open() -> None: device=torch.device("cpu"), num_envs=2, open_state=(0.0, 0.0), + close_state=(1.0, 1.0), get_current_gripper_state_agent=lambda: ( torch.tensor([[0.0, 0.0], [0.1, 0.1]]), torch.ones(2, 2), @@ -207,6 +209,100 @@ def test_single_arm_release_verifies_the_selected_gripper_is_open() -> None: assert verified.tolist() == [True, False] +def test_single_arm_release_ignores_passive_mimic_joint_residuals() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + executor = SimpleNamespace( + env=SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + open_state=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + close_state=(0.7, -0.7, 0.7, -0.7, -0.7, 0.7), + agent_gripper_model="robotiq", + agent_gripper_state_joint_indices={"left": (0,), "right": (0,)}, + get_current_gripper_state_agent=lambda: ( + torch.tensor( + [ + [0.0, 0.01, -0.02, 0.03, -0.01, 0.02], + [0.1, 0.0, 0.0, 0.0, 0.0, 0.0], + ] + ), + torch.zeros(2, 6), + ), + ), + runtime_policy=SimpleNamespace( + predicate_fallbacks={"gripper_state_tolerance": 1.0e-3} + ), + ) + outcome = SimpleNamespace( + grounded=SimpleNamespace(motion_policy={"single_release": True}) + ) + + verified = capability.verifier_hook( + executor=executor, + step=SimpleNamespace(), + arm="left_arm", + outcome=outcome, + attempted=torch.tensor([True, True]), + ) + + assert verified.tolist() == [True, False] + + +def test_single_arm_release_uses_normalized_opening_and_physical_support() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + executor = SimpleNamespace( + env=SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + open_state=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + close_state=(0.7, -0.7, 0.7, -0.7, -0.7, 0.7), + agent_gripper_model="robotiq", + agent_gripper_state_joint_indices={"left": (0,), "right": (0,)}, + get_current_gripper_state_agent=lambda: ( + torch.tensor( + [ + [0.015, 0.2, -0.2, 0.2, -0.2, 0.2], + [0.015, 0.2, -0.2, 0.2, -0.2, 0.2], + ] + ), + torch.zeros(2, 6), + ), + ), + _support_reference_uid=lambda _step: "table", + _support_stable_for=lambda _step, _support, _active: torch.tensor( + [True, False] + ), + _entity_pose=lambda _uid: torch.eye(4).repeat(2, 1, 1), + _placement_orientation_satisfied=lambda _step, _pose: torch.tensor( + [False, False] + ), + runtime_policy=SimpleNamespace( + predicate_fallbacks={"gripper_state_tolerance": 1.0e-3} + ), + ) + planner_trace: dict[str, object] = {} + outcome = SimpleNamespace( + grounded=SimpleNamespace(motion_policy={"single_release": True}), + planner_trace=planner_trace, + ) + + verified = capability.verifier_hook( + executor=executor, + step=SimpleNamespace(object_uid="can"), + arm="left_arm", + outcome=outcome, + attempted=torch.tensor([True, True]), + ) + + assert verified.tolist() == [True, False] + release = planner_trace["release_verification"] + assert release["open_error_fraction"] == pytest.approx([0.015 / 0.7] * 2) + assert release["gripper_open"] == [True, True] + assert release["support_stable"] == [True, False] + assert release["upright"] == [False, False] + assert release["accepted"] == [True, False] + + def test_explicit_required_home_is_safety_required_for_any_task_type() -> None: capability = build_atomic_capability_registry().get("MoveJoints") base = { diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py index 59ff624a1..6799b6298 100644 --- a/tests/gen_sim/action_engine/runtime/test_actions.py +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -1377,7 +1377,7 @@ def fake_motion_generator(*, cfg: Any) -> object: assert planner.world.collision_cache == {"cuboid": 13, "mesh": 2} -def test_dynamic_scene_parks_contact_target_and_held_rows() -> None: +def test_dynamic_scene_separates_live_semantic_and_parked_collision_poses() -> None: actual = torch.eye(4).repeat(2, 1, 1) actual[:, 2, 3] = torch.tensor([0.7, 0.8]) entities = {uid: _PoseEntity(actual.clone()) for uid in ("target", "held", "other")} @@ -1414,16 +1414,24 @@ def test_dynamic_scene_parks_contact_target_and_held_rows() -> None: ) scene = adapter._scene_snapshot(grounded, state) + collision_poses = scene.collision_obstacle_poses( + batch_size=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + assert torch.equal(scene.entities["target"].pose, actual) + assert torch.equal(scene.entities["held"].pose, actual) + assert torch.equal(scene.entities["other"].pose, actual) assert torch.equal( - scene.entities["target"].pose[:, 2, 3], + collision_poses["target"][:, 2, 3], actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET, ) - assert scene.entities["held"].pose[0, 2, 3] == ( + assert collision_poses["held"][0, 2, 3] == ( actual[0, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET ) - assert scene.entities["held"].pose[1, 2, 3] == actual[1, 2, 3] - assert torch.equal(scene.entities["other"].pose, actual) + assert collision_poses["held"][1, 2, 3] == actual[1, 2, 3] + assert torch.equal(collision_poses["other"], actual) def test_released_object_returns_to_live_dynamic_collision_pose() -> None: @@ -1553,7 +1561,7 @@ def start(self, invocations: tuple[Any, ...], context: Any) -> object: assert captured == {"invocations": ("invocation",), "context": "context"} -def test_retreat_parks_intentional_contact_objects() -> None: +def test_retreat_parks_only_collision_poses_for_contact_objects() -> None: actual = torch.eye(4).repeat(2, 1, 1) entities = { uid: _PoseEntity(actual.clone()) for uid in ("released", "container", "other") @@ -1581,11 +1589,19 @@ def test_retreat_parks_intentional_contact_objects() -> None: grounded, ExecutionState(last_qpos=torch.zeros(2, 8)), ) + collision_poses = scene.collision_obstacle_poses( + batch_size=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) parked_z = actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET - assert torch.equal(scene.entities["released"].pose[:, 2, 3], parked_z) - assert torch.equal(scene.entities["container"].pose[:, 2, 3], parked_z) + assert torch.equal(scene.entities["released"].pose, actual) + assert torch.equal(scene.entities["container"].pose, actual) assert torch.equal(scene.entities["other"].pose, actual) + assert torch.equal(collision_poses["released"][:, 2, 3], parked_z) + assert torch.equal(collision_poses["container"][:, 2, 3], parked_z) + assert torch.equal(collision_poses["other"], actual) def test_action_outcome_commits_state_delta_only_for_verified_rows() -> None: diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py index 00b3704fb..76f7949c2 100644 --- a/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py +++ b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py @@ -21,6 +21,7 @@ from embodichain.gen_sim.action_engine.runtime.grasp_diagnostics import ( _TracingAntipodalGraspPoseGenerator, ) +from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg @@ -171,6 +172,55 @@ def test_generator_context_selects_non_crossing_pair_and_canonical_half_turn() - ) +def test_upright_context_rejects_end_clamps_and_ranks_mid_body_grasps( + monkeypatch, +) -> None: + generator = _TracingAntipodalGraspPoseGenerator( + ParallelJawGripperModelCfg(model_id="upright_test") + ) + poses = torch.eye(4).repeat(4, 1, 1) + poses[:, 1, 3] = torch.tensor([0.5, 0.05, 0.5, 0.7]) + poses[0, :3, 0] = torch.tensor([0.0, 1.0, 0.0]) + poses[0, :3, 1] = torch.tensor([-1.0, 0.0, 0.0]) + + monkeypatch.setattr( + AntipodalGraspPoseGenerator, + "get_valid_grasp_poses", + lambda _self, **_kwargs: [(poses, torch.tensor([0.0, 0.01, 0.2, 0.3]))], + ) + vertices = torch.tensor( + [ + [-0.1, 0.0, -0.1], + [0.1, 0.0, 0.1], + [-0.1, 1.0, 0.1], + [0.1, 1.0, -0.1], + ] + ) + + with generator.upright_selection_context(local_axis=torch.tensor([0.0, 1.0, 0.0])): + result = generator.get_valid_grasp_poses( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2], [1, 2, 3]]), + obj_poses=torch.eye(4).unsqueeze(0), + approach_direction=torch.tensor([0.0, 0.0, -1.0]), + ) + + ranked_poses, ranked_costs = result[0] + assert ranked_poses[0, 1, 3] == 0.5 + assert torch.isfinite(ranked_costs[:3]).all() + assert torch.isinf(ranked_costs[-1]) + trace = generator.last_upright_trace + assert trace is not None + assert trace["local_axis"] == [0.0, 1.0, 0.0] + assert trace["candidate_count"] == 4 + assert trace["side_compatible_count"] == 3 + assert trace["central_band_count"] == 3 + assert trace["side_and_central_count"] == 2 + assert trace["retained_count"] == 3 + assert trace["best_candidate_axis_alignment"] == 0.0 + assert trace["best_candidate_axis_fraction"] == 0.5 + + def _pose_with_y(y: float) -> torch.Tensor: pose = torch.eye(4) pose[1, 3] = y diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index a3d8db192..bb4c626bb 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -3849,7 +3849,7 @@ def test_handover_defers_clearance_verification_to_retreat_action() -> None: assert adapter.capabilities.get("MoveEndEffector").verifier_hook is not None -def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() -> None: +def test_e2_then_handover_reacquires_with_a_separate_transfer_policy() -> None: entity = _FakeEntity("can", _pose(0.0, 0.2, 0.75), _rect_vertices(0.10, 0.03, 0.03)) env = _FakeEnv( { @@ -3899,17 +3899,25 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - handover_step = next( candidate for candidate in program.semantic_steps if candidate.id == "task_02" ) - orient_edge = next( + orient_pickup_edge = next( candidate for candidate in program.edges if candidate.id in orient_step.edge_ids - and candidate.actions[0]["atomic_action_class"] == "AxisAlign" + and candidate.actions[0]["atomic_action_class"] == "PickUp" + ) + orient_staging_edge = next( + candidate + for candidate in program.edges + if candidate.id in orient_step.edge_ids + and candidate.actions[0]["atomic_action_class"] == "MoveHeldObject" + and candidate.actions[0]["target_binding"].get("phase") == "staging" ) orient_descend_edge = next( candidate for candidate in program.edges if candidate.id in orient_step.edge_ids and candidate.actions[0]["atomic_action_class"] == "MoveHeldObject" + and candidate.actions[0]["target_binding"].get("phase") == "final" ) orient_release_edge = next( candidate @@ -3958,8 +3966,14 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - ) grounder = ActionGrounder(program, env, lambda _uid: semantics) - orient_alignment = grounder.ground( - orient_edge.actions[0], + orient_pickup = grounder.ground( + orient_pickup_edge.actions[0], + orient_step, + arm="right_arm", + state=ExecutionState(last_qpos=env.robot.get_qpos()), + ) + orient_staging = grounder.ground( + orient_staging_edge.actions[0], orient_step, arm="right_arm", state=ExecutionState(last_qpos=env.robot.get_qpos()), @@ -4012,16 +4026,31 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - state=ExecutionState(last_qpos=env.robot.get_qpos()), ) - assert "approach_direction_mode" not in orient_alignment.cfg + assert "approach_direction_mode" not in orient_pickup.cfg assert "approach_direction_mode" not in handover_pickup.cfg assert handover_pickup.cfg["pick_object_part"] == "top" - assert isinstance(orient_alignment.target, AxisAlignGoal) + assert orient_pickup.target.grasp_xpos is None + assert isinstance( + orient_pickup.target.semantics.affordance, + AntipodalAffordance, + ) + assert torch.equal( + orient_pickup.cfg["obj_upright_direction"], + torch.tensor([1.0, 0.0, 0.0]), + ) + assert isinstance(orient_staging.target, HeldObjectPoseGoal) assert isinstance(orient_descend.target, HeldObjectPoseGoal) assert torch.equal( orient_descend.target.object_target_pose[:, :2, 3], - orient_alignment.target_object_pose[:, :2, 3], + orient_staging.target.object_target_pose[:, :2, 3], + ) + assert bool( + ( + orient_staging.target.object_target_pose[:, 2, 3] + > orient_descend.target.object_target_pose[:, 2, 3] + ).all() ) - assert orient_descend.motion_policy["surface_clearance"] == pytest.approx(0.005) + assert orient_descend.motion_policy["surface_clearance"] == pytest.approx(0.05) assert isinstance(orient_release.target, JointPositionGoal) assert orient_release.control == "hand" assert orient_release.cfg["single_release"] is True @@ -4080,18 +4109,6 @@ def test_axis_align_then_handover_reacquires_with_a_separate_transfer_policy() - != orient_post_reorient_lift.target.xpos[:, :2, 3] ).any() ) - assert orient_alignment.target.grasp_xpos is None - assert not hasattr(orient_alignment.target, "object_target_pose") - assert orient_alignment.target_object_pose is not None - assert isinstance( - orient_alignment.target.semantics.affordance, - AxisAlignAffordance, - ) - assert torch.equal( - orient_alignment.target.semantics.affordance.internal_axis, - torch.tensor([0.0, 0.0, 1.0]), - ) - assert orient_alignment.cfg["target_axis"] == (0.0, 0.0, 1.0) assert handover_pickup.target.grasp_xpos is None assert isinstance( handover_pickup.target.semantics.affordance, diff --git a/tests/gen_sim/action_engine/tasks/test_factory.py b/tests/gen_sim/action_engine/tasks/test_factory.py index c07ab0f2b..7a8e7b8ed 100644 --- a/tests/gen_sim/action_engine/tasks/test_factory.py +++ b/tests/gen_sim/action_engine/tasks/test_factory.py @@ -134,13 +134,14 @@ def test_historical_task2_1_fixture_preserves_ten_step_semantics() -> None: ) -def test_historical_task2_1_uses_axis_align_and_explicit_handover_arms() -> None: +def test_historical_task2_1_uses_split_upright_transport_and_handover_arms() -> None: task = make_task2_1_historical_spec() graph = _historical_task2_1_graph() actions = _actions_by_task_group(graph) expected_orient_actions = [ - "AxisAlign", + "PickUp", + "MoveHeldObject", "MoveHeldObject", "MoveJoints", "MoveEndEffector", @@ -241,7 +242,8 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - orient = next(group for group in graph["task_groups"] if group["id"] == "orient") assert [node["atomic_action"] for node in orient_nodes] == [ - "AxisAlign", + "PickUp", + "MoveHeldObject", "MoveHeldObject", "MoveJoints", "MoveEndEffector", @@ -250,14 +252,16 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - "MoveEndEffector", "MoveJoints", ] - assert orient["goal"]["upright_local_axis"] == "z" + assert orient["goal"]["upright_local_axis"] == "auto" assert orient_nodes[0]["postcondition"] == {} - assert orient_nodes[2]["postcondition"]["local_axis"] == "z" - assert orient_nodes[0]["motion_policy"] == {"modifiers": []} + assert orient_nodes[3]["postcondition"]["local_axis"] == "auto" + upright_policy = {"modifiers": [{"type": "orientation", "mode": "upright"}]} + assert orient_nodes[0]["motion_policy"] == upright_policy assert [node["role"] for node in orient_nodes] == [ "primary", "primary", "primary", + "primary", "cleanup", "cleanup", "cleanup", @@ -267,38 +271,44 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[1]["target_binding"] == { "kind": "semantic_goal", "semantic_step": "orient", - "phase": "final", + "phase": "staging", } - assert orient_nodes[1]["motion_policy"] == {"modifiers": []} + assert orient_nodes[1]["motion_policy"] == upright_policy assert orient_nodes[2]["target_binding"] == { + "kind": "semantic_goal", + "semantic_step": "orient", + "phase": "final", + } + assert orient_nodes[2]["motion_policy"] == upright_policy + assert orient_nodes[3]["target_binding"] == { "kind": "joint_state", "source": "gripper_open", "single_release": True, } - assert orient_nodes[2]["control"] == "hand" - assert orient_nodes[3]["target_binding"] == { + assert orient_nodes[3]["control"] == "hand" + assert orient_nodes[4]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "lift_clear", } - assert orient_nodes[3]["motion_policy"] == {"modifiers": []} - assert orient_nodes[4]["target_binding"] == { + assert orient_nodes[4]["motion_policy"] == {"modifiers": []} + assert orient_nodes[5]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "reorient_tool_down", } - assert orient_nodes[5]["target_binding"] == { + assert orient_nodes[6]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "lift_clear", "requires_arm_clear": True, } - assert orient_nodes[6]["target_binding"] == { + assert orient_nodes[7]["target_binding"] == { "kind": "policy_pose", "source": "release", "operation": "retreat_after_lift", } - assert orient_nodes[7]["target_binding"] == { + assert orient_nodes[8]["target_binding"] == { "kind": "joint_state", "source": "initial", "operation": "e2_home", @@ -324,10 +334,11 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[0]["contract"]["failure_policy"] == "task_required" assert orient_nodes[1]["contract"]["failure_policy"] == "task_required" assert orient_nodes[2]["contract"]["failure_policy"] == "task_required" - assert orient_nodes[3]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[3]["contract"]["failure_policy"] == "task_required" assert orient_nodes[4]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[5]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[6]["contract"]["failure_policy"] == "safety_required" + assert orient_nodes[7]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[-1]["contract"]["failure_policy"] == "safety_required" assert orient_nodes[1]["contract"]["requires"] == [ { @@ -353,17 +364,21 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - ] assert [ (effect["op"], effect["atom"]["predicate"]) - for effect in orient_nodes[2]["contract"]["effects"] + for effect in orient_nodes[3]["contract"]["effects"] ] == [ ("delete", "object_held"), ("add", "arm_free"), ("add", "object_free"), ] assert orient_nodes[3]["contract"]["requires"] == [ - {"predicate": "arm_free", "arm": "left_arm"} + { + "predicate": "object_held", + "object_uid": "interact_can", + "arm": "left_arm", + } ] assert orient_nodes[4]["contract"]["requires"] == [ - {"predicate": "arm_clear", "arm": "left_arm"} + {"predicate": "arm_free", "arm": "left_arm"} ] assert orient_nodes[5]["contract"]["requires"] == [ {"predicate": "arm_clear", "arm": "left_arm"} @@ -371,6 +386,9 @@ def test_orient_then_handover_releases_then_reacquires_with_role_side_pickup() - assert orient_nodes[6]["contract"]["requires"] == [ {"predicate": "arm_clear", "arm": "left_arm"} ] + assert orient_nodes[7]["contract"]["requires"] == [ + {"predicate": "arm_clear", "arm": "left_arm"} + ] assert any( effect["atom"]["predicate"] == "arm_home" for effect in orient["contract"]["exit_effects"] diff --git a/tests/gen_sim/action_engine/tasks/test_interpretation.py b/tests/gen_sim/action_engine/tasks/test_interpretation.py index 9e91b71a9..9932625b9 100644 --- a/tests/gen_sim/action_engine/tasks/test_interpretation.py +++ b/tests/gen_sim/action_engine/tasks/test_interpretation.py @@ -311,7 +311,8 @@ def caller(**kwargs): "object_02": "orange_can", } assert ( - grounded.task_spec["task_instances"][0]["params"]["upright_local_axis"] == "z" + grounded.task_spec["task_instances"][0]["params"]["upright_local_axis"] + == "auto" ) placement_actions = [ node["atomic_action"] for node in graph["nodes"] if node["task_type"] == "E1" @@ -321,7 +322,8 @@ def caller(**kwargs): ] handover_nodes = [node for node in graph["nodes"] if node["task_type"] == "E4"] assert orient_actions == [ - "AxisAlign", + "PickUp", + "MoveHeldObject", "MoveHeldObject", "MoveJoints", "MoveEndEffector", diff --git a/tests/gen_sim/action_engine/test_gripper_profiles.py b/tests/gen_sim/action_engine/test_gripper_profiles.py index 6b6c2aa1d..bb5f9f17a 100644 --- a/tests/gen_sim/action_engine/test_gripper_profiles.py +++ b/tests/gen_sim/action_engine/test_gripper_profiles.py @@ -56,6 +56,7 @@ def test_pgi_profile_owns_asset_control_mimic_tcp_and_grasp_geometry() -> None: assert profile.grasp_model.max_opening_width == pytest.approx(0.100) assert profile.grasp_model.finger_length == pytest.approx(0.10) assert profile.grasp_model.opening_margin == pytest.approx(0.03) + assert profile.release_open_fraction_tolerance == pytest.approx(0.03) def test_robotiq_profile_separates_commanded_mimics_from_state_joint() -> None: @@ -88,6 +89,7 @@ def test_robotiq_profile_separates_commanded_mimics_from_state_joint() -> None: assert profile.state_joint_names("left") == ("left_finger_joint",) assert profile.state_joint_names("right") == ("right_finger_joint",) assert profile.state_joint_indices("left") == (0,) + assert profile.release_open_fraction_tolerance == pytest.approx(0.03) assert profile.state_joint_indices("right") == (0,) assert set(profile.state_joint_names("left")).isdisjoint( profile.mimic_joint_names("left") From e8c86b3fa78510fefdb278200bdbdb4e54dd2711 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:37:39 +0800 Subject: [PATCH 82/85] feat(gen-sim): support articulated objects in scene generation and execution --- .../generation/config_builder.py | 44 ++++++- .../action_engine/runtime/grounding.py | 36 +++++- .../action_engine/runtime/predicates.py | 21 +-- .../scene_engine/core/scene_edit_plan.py | 4 +- .../gen_sim/scene_engine/pipeline/api.py | 44 +++++-- .../gen_sim/task_engine/_bundle_runner.py | 64 +++++++-- .../task_engine/orchestration/artifacts.py | 5 +- .../orchestration/scene_adapter.py | 2 +- .../gen_sim/task_engine/scene_backend.py | 25 ++-- embodichain/gen_sim/task_engine/workflow.py | 29 ++++- .../generation/test_generation.py | 75 +++++++++++ .../runtime/test_runtime_contracts.py | 65 ++++++++++ .../gen_sim/scene_engine/test_pipeline_api.py | 121 +++++++++++++++++- .../orchestration/test_coordinator_cli.py | 82 ++++++++++++ .../task_engine/test_parallel_workflow.py | 53 ++++++++ .../gen_sim/task_engine/test_scene_backend.py | 59 ++++++++- 16 files changed, 662 insertions(+), 67 deletions(-) diff --git a/embodichain/gen_sim/action_engine/generation/config_builder.py b/embodichain/gen_sim/action_engine/generation/config_builder.py index 0b2054eeb..b056f1c5f 100644 --- a/embodichain/gen_sim/action_engine/generation/config_builder.py +++ b/embodichain/gen_sim/action_engine/generation/config_builder.py @@ -69,6 +69,19 @@ _DEFAULT_TABLETOP_Z = float(_GENERATION_DEFAULTS["scene"]["default_tabletop_z"]) _DEFAULT_GRIPPER_MODEL = str(_GENERATION_DEFAULTS["task"]["default_gripper_model"]) _DEFAULT_IK_SOLVER = str(_GENERATION_DEFAULTS["task"]["default_ik_solver"]) +_USD_ARTICULATION_SUFFIXES = frozenset({".usd", ".usda", ".usdc"}) +_ARTICULATION_AUTHORING_KEYS = frozenset( + { + "attributes", + "category", + "description", + "is_articulated", + "name", + "proxy_body_scale", + "proxy_glb_fpath", + "role", + } +) _ARM_SLOTS = { "left": {"arm": "left_arm", "eef": "left_eef"}, @@ -383,11 +396,7 @@ def build_fast_gym_config( } if scene.articulations: config["articulation"] = [ - { - key: deepcopy(value) - for key, value in articulation.items() - if key not in {"attributes", "role"} - } + _runtime_articulation_config(articulation) for articulation in scene.articulations ] validate_fast_gym_config(config) @@ -433,6 +442,17 @@ def validate_fast_gym_config(config: dict[str, Any]) -> None: f"existing absolute file: {path}" ) + for articulation in config.get("articulation", []): + path = Path(str(articulation.get("fpath", ""))) + if ( + path.suffix.lower() in _USD_ARTICULATION_SUFFIXES + and articulation.get("build_pk_chain") is not False + ): + raise ValueError( + f"USD articulation {articulation.get('uid')!r} must set " + "build_pk_chain=false." + ) + action_engine = config.get("env", {}).get("extensions", {}).get("action_engine", {}) if action_engine.get("defaults_schema_version") != ACTION_ENGINE_DEFAULTS_SCHEMA: raise ValueError("Gym config has an unexpected defaults schema version.") @@ -455,7 +475,6 @@ def validate_fast_gym_config(config: dict[str, Any]) -> None: or Path(graph_path).name != EXECUTION_PROGRAM_FILENAME ): raise ValueError("Gym config points to an unexpected SeedGraph artifact.") - planning_mode = action_engine.get("planning_mode", "offline") _validate_planning_mode(planning_mode) if planning_mode == "ab": @@ -483,6 +502,19 @@ def validate_fast_gym_config(config: dict[str, Any]) -> None: raise ValueError("Every rigid object must have one live-pose registry entry.") +def _runtime_articulation_config(value: Mapping[str, Any]) -> dict[str, Any]: + """Reduce one source articulation to the simulator-facing config contract.""" + result = { + key: deepcopy(item) + for key, item in value.items() + if key not in _ARTICULATION_AUTHORING_KEYS + } + path = Path(str(result.get("fpath", ""))) + if path.suffix.lower() in _USD_ARTICULATION_SUFFIXES: + result["build_pk_chain"] = False + return result + + def _make_robot( profile_id: str, profile: dict[str, Any], diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py index 35fdc23cc..36a24518c 100644 --- a/embodichain/gen_sim/action_engine/runtime/grounding.py +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -90,7 +90,9 @@ def _batched_pose(value: Any, env: Any) -> torch.Tensor: def _object(env: Any, uid: str) -> Any: entity = env.sim.get_rigid_object(uid) if entity is None: - raise ValueError(f"Unknown rigid object {uid!r}.") + entity = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") return entity @@ -129,9 +131,35 @@ def _local_vertices(entity: Any, env: Any, env_id: int = 0) -> torch.Tensor: def _world_vertices(entity: Any, env: Any, env_id: int) -> torch.Tensor: - vertices = _local_vertices(entity, env, env_id) - pose = _batched_pose(entity.get_local_pose(to_matrix=True), env)[env_id] - return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + get_vertices = getattr(entity, "get_vertices", None) + if callable(get_vertices): + vertices = _local_vertices(entity, env, env_id) + pose = _batched_pose(entity.get_local_pose(to_matrix=True), env)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + get_link_vertices = getattr(entity, "get_link_vert_face", None) + get_link_pose = getattr(entity, "get_link_pose", None) + link_names = getattr(entity, "link_names", ()) + if not callable(get_link_vertices) or not callable(get_link_pose) or not link_names: + raise ValueError("Scene entity exposes no usable collision geometry.") + world_vertices = [] + for link_name in link_names: + vertices, _ = get_link_vertices(link_name) + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=env.device, + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + continue + pose = _batched_pose( + get_link_pose(link_name, to_matrix=True), + env, + )[env_id] + world_vertices.append(vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3]) + if not world_vertices: + raise ValueError("Scene articulation exposes no usable link geometry.") + return torch.cat(world_vertices, dim=0) @dataclass(frozen=True) diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py index afc86ef71..84006d57d 100644 --- a/embodichain/gen_sim/action_engine/runtime/predicates.py +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -77,10 +77,17 @@ def _constant(env: Any, value: bool) -> torch.Tensor: ) -def _pose(env: Any, uid: str) -> torch.Tensor: +def _entity(env: Any, uid: str) -> Any: entity = env.sim.get_rigid_object(uid) if entity is None: - raise ValueError(f"Unknown rigid object {uid!r}.") + entity = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") + return entity + + +def _pose(env: Any, uid: str) -> torch.Tensor: + entity = _entity(env, uid) pose = torch.as_tensor( entity.get_local_pose(to_matrix=True), dtype=torch.float32, @@ -96,9 +103,7 @@ def _position(env: Any, uid: str) -> torch.Tensor: def _world_vertices(env: Any, uid: str, env_id: int) -> torch.Tensor: - entity = env.sim.get_rigid_object(uid) - if entity is None: - raise ValueError(f"Unknown rigid object {uid!r}.") + entity = _entity(env, uid) value = entity.get_vertices(env_ids=[env_id], scale=True) if isinstance(value, (tuple, list)): value = value[0] @@ -118,7 +123,7 @@ def _projected_center_of_mass( world_vertices: torch.Tensor, ) -> torch.Tensor: """Return the live COM projection, with a geometry-center fallback.""" - entity = env.sim.get_rigid_object(uid) + entity = _entity(env, uid) body_data = None if entity is None else getattr(entity, "body_data", None) com_pose = None if body_data is None else getattr(body_data, "com_pose", None) if callable(com_pose): @@ -244,9 +249,7 @@ def _local_axis_index(env: Any, uid: str, axis: Any) -> int: }.get(name) if geometry_axis is None: raise ValueError(f"Unsupported upright local axis {axis!r}.") - entity = env.sim.get_rigid_object(uid) - if entity is None: - raise ValueError(f"Unknown rigid object {uid!r}.") + entity = _entity(env, uid) vertices = entity.get_vertices(env_ids=[0], scale=True) if isinstance(vertices, (tuple, list)): vertices = vertices[0] diff --git a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py index bf62e2249..03b87d718 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py +++ b/embodichain/gen_sim/scene_engine/core/scene_edit_plan.py @@ -47,8 +47,8 @@ class SceneEditOperation: name: Human-readable name required for an added object. description: Generation prompt and semantic description required for an added object. - orientation_state: Optional standing or lying intent for an added - object. Move operations may only preserve the existing state. + pose_description: Optional free-form pose intent for an added or moved + object. """ op: SceneEditOperationType diff --git a/embodichain/gen_sim/scene_engine/pipeline/api.py b/embodichain/gen_sim/scene_engine/pipeline/api.py index 6f219615c..aaddcad55 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/api.py +++ b/embodichain/gen_sim/scene_engine/pipeline/api.py @@ -25,6 +25,9 @@ from pathlib import Path from typing import Any, Final +from embodichain.gen_sim.scene_engine.clients.articulated_generation import ( + ArticulatedGenerationClient, +) from embodichain.gen_sim.scene_engine.clients.geometry_generation import ( GeometryGenerationClient, ) @@ -73,14 +76,15 @@ "materialize_edit", ] -SCENE_BLUEPRINT_SCHEMA: Final = "embodichain.scene-blueprint/v1" -SCENE_EDIT_BLUEPRINT_SCHEMA: Final = "embodichain.scene-edit-blueprint/v1" +SCENE_BLUEPRINT_SCHEMA: Final = "embodichain.scene-blueprint/v2" +SCENE_EDIT_BLUEPRINT_SCHEMA: Final = "embodichain.scene-edit-blueprint/v2" @dataclass(frozen=True) class SceneBlueprintPackage: """In-process scene semantics plus their persisted audit document.""" + schema_version: str blueprint_id: str image_path: Path output_root: Path @@ -88,11 +92,19 @@ class SceneBlueprintPackage: scene: Scene scene_graph: SceneGraph + def __post_init__(self) -> None: + if self.schema_version != SCENE_BLUEPRINT_SCHEMA: + raise ValueError( + "SceneBlueprintPackage schema_version must be " + f"{SCENE_BLUEPRINT_SCHEMA!r}." + ) + @dataclass(frozen=True) class SceneEditBlueprintPackage: """Validated edit intent before added assets and layout are materialized.""" + schema_version: str blueprint_id: str edit_prompt: str output_root: Path @@ -100,6 +112,13 @@ class SceneEditBlueprintPackage: scene_edit_plan: SceneEditPlan updated_scene_graph: SceneGraph + def __post_init__(self) -> None: + if self.schema_version != SCENE_EDIT_BLUEPRINT_SCHEMA: + raise ValueError( + "SceneEditBlueprintPackage schema_version must be " + f"{SCENE_EDIT_BLUEPRINT_SCHEMA!r}." + ) + @dataclass(frozen=True) class SceneMaterialization: @@ -152,6 +171,7 @@ def analyze_image( manifest_path = resolved_output / "scene_blueprint.json" _write_json(manifest_path, document) return SceneBlueprintPackage( + schema_version=SCENE_BLUEPRINT_SCHEMA, blueprint_id=blueprint_id, image_path=resolved_image, output_root=resolved_output, @@ -166,17 +186,24 @@ def materialize_blueprint( *, vlm_client: OpenAICompatibleVLM | None = None, geometry_generation_client: GeometryGenerationClient | None = None, - seed: int | None = None, + articulated_generation_client: ArticulatedGenerationClient | None = None, ) -> SceneMaterialization: """Generate assets and layout for one image-derived blueprint.""" scene = deepcopy(blueprint.scene) scene_graph = deepcopy(blueprint.scene_graph) effective_vlm = vlm_client or OpenAICompatibleVLM.from_dotenv() geometry = geometry_generation_client or GeometryGenerationClient.from_dotenv() - owns_geometry = geometry_generation_client is None + has_articulated_objects = any(item.is_articulated for item in scene.objects) + articulated = articulated_generation_client if has_articulated_objects else None + owns_articulated = False log_info("Starting Objects + Coarse Layout Generation") try: + if has_articulated_objects and articulated is None: + articulated = ArticulatedGenerationClient.from_dotenv() + owns_articulated = True geometry.check_health() + if articulated is not None: + articulated.check_health() scene = generate_scene_and_refine( image_path=blueprint.image_path, output_root=blueprint.output_root, @@ -184,10 +211,12 @@ def materialize_blueprint( scene_graph=scene_graph, geometry_generation_client=geometry, vlm_client=effective_vlm, - seed=seed, + articulated_generation_client=articulated, ) finally: - if owns_geometry: + if owns_articulated and articulated is not None: + articulated.close() + if geometry_generation_client is None: geometry.close() log_info("Completed Objects + Coarse Layout Generation") return _export_materialization( @@ -230,6 +259,7 @@ def analyze_edit( manifest_path = resolved_output / "scene_edit" / "scene_edit_blueprint.json" _write_json(manifest_path, {**payload, "blueprint_id": blueprint_id}) return SceneEditBlueprintPackage( + schema_version=SCENE_EDIT_BLUEPRINT_SCHEMA, blueprint_id=blueprint_id, edit_prompt=normalized_prompt, output_root=resolved_output, @@ -246,7 +276,6 @@ def materialize_edit( image_generation_client: ImageGenerationClient | None = None, geometry_generation_client: GeometryGenerationClient | None = None, image_segmentation_client: ImageSegmentationClient | None = None, - seed: int | None = None, ) -> SceneMaterialization: """Generate added assets, apply layout edits, and export the new revision.""" scene_edit_plan = deepcopy(blueprint.scene_edit_plan) @@ -271,7 +300,6 @@ def materialize_edit( geometry_generation_client=geometry, image_segmentation_client=segmentation, vlm_client=effective_vlm, - seed=seed, ) finally: for client, owned in owned_clients: diff --git a/embodichain/gen_sim/task_engine/_bundle_runner.py b/embodichain/gen_sim/task_engine/_bundle_runner.py index 5bc80bd9b..665cc4ab4 100644 --- a/embodichain/gen_sim/task_engine/_bundle_runner.py +++ b/embodichain/gen_sim/task_engine/_bundle_runner.py @@ -20,12 +20,17 @@ import argparse from contextlib import contextmanager +from copy import deepcopy import json from pathlib import Path import sys +import tempfile from typing import Any, Iterator, Sequence from embodichain.gen_sim.action_engine.agent import ActionAgent +from embodichain.gen_sim.action_engine.generation.config_builder import ( + _runtime_articulation_config, +) from embodichain.gen_sim.action_engine.protocol import ( AGENT_CONFIG_FILENAME, EXECUTION_PROGRAM_FILENAME, @@ -81,20 +86,21 @@ def execute_bundle( write_execution_report(root, rejection) _print_json(rejection.as_mapping()) return 2 - legacy_argv = [ - "--task_name", - task_id, - "--gym_config", - str(gym_config), - "--agent_config", - str(agent_config), - "--task-engine-report", - *run_args, - ] from embodichain.gen_sim.action_engine.cli import run_agent - with _temporary_argv(["run_agent", *legacy_argv]): - return int(run_agent.cli() or 0) + with _runtime_gym_config(gym_config) as runtime_gym_config: + legacy_argv = [ + "--task_name", + task_id, + "--gym_config", + str(runtime_gym_config), + "--agent_config", + str(agent_config), + "--task-engine-report", + *run_args, + ] + with _temporary_argv(["run_agent", *legacy_argv]): + return int(run_agent.cli() or 0) def _preflight_bundle( @@ -179,6 +185,40 @@ def _read_json(path: Path) -> dict[str, Any]: return value +@contextmanager +def _runtime_gym_config(source: Path) -> Iterator[Path]: + """Yield a simulator-only config without modifying a prepared bundle.""" + original = _read_json(source) + normalized = deepcopy(original) + articulations = normalized.get("articulation") + if isinstance(articulations, list): + normalized["articulation"] = [ + _runtime_articulation_config(item) if isinstance(item, dict) else item + for item in articulations + ] + if normalized == original: + yield source + return + + temporary_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + prefix=".task-engine-runtime-", + suffix=".json", + dir=source.parent, + delete=False, + ) as stream: + json.dump(normalized, stream, ensure_ascii=False, indent=2, allow_nan=False) + stream.write("\n") + temporary_path = Path(stream.name) + yield temporary_path + finally: + if temporary_path is not None: + temporary_path.unlink(missing_ok=True) + + @contextmanager def _temporary_argv(arguments: list[str]) -> Iterator[None]: original = sys.argv diff --git a/embodichain/gen_sim/task_engine/orchestration/artifacts.py b/embodichain/gen_sim/task_engine/orchestration/artifacts.py index b941c86dc..f2bd14237 100644 --- a/embodichain/gen_sim/task_engine/orchestration/artifacts.py +++ b/embodichain/gen_sim/task_engine/orchestration/artifacts.py @@ -300,7 +300,10 @@ def relocate(value: Any) -> Any: if isinstance(value, list): return [relocate(item) for item in value] if isinstance(value, dict): - return {key: relocate(item) for key, item in value.items()} + return { + relocate(key) if isinstance(key, str) else key: relocate(item) + for key, item in value.items() + } return value for path in staging.rglob("*.json"): diff --git a/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py b/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py index 8f101e23b..d5dbdc58c 100644 --- a/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py +++ b/embodichain/gen_sim/task_engine/orchestration/scene_adapter.py @@ -282,7 +282,7 @@ def select_objects( candidate_set: TaskCandidateSet | Sequence[Mapping[str, Any]], scene_objects: Sequence[Mapping[str, Any]], *, - source_format: str = "embodichain.scene-blueprint/v1", + source_format: str = "embodichain.scene-blueprint/v2", robot_profile: str | None = None, grounding_caller: GroundingCaller | None = None, adjudicator: Adjudicator | None = None, diff --git a/embodichain/gen_sim/task_engine/scene_backend.py b/embodichain/gen_sim/task_engine/scene_backend.py index dd4a665c5..630d10353 100644 --- a/embodichain/gen_sim/task_engine/scene_backend.py +++ b/embodichain/gen_sim/task_engine/scene_backend.py @@ -30,6 +30,7 @@ resolve_source_scene, ) from embodichain.gen_sim.scene_engine.pipeline import ( + SCENE_BLUEPRINT_SCHEMA, SceneBlueprintPackage, SceneMaterialization, analyze_edit, @@ -147,6 +148,7 @@ def select( return scene_adapter.select_objects( candidate_set, scene_blueprint_objects(analysis.blueprint), + source_format=analysis.blueprint.schema_version, force_most_likely=force_most_likely, ) adaptation = scene_adapter.adapt( @@ -187,7 +189,7 @@ def materialize( assert analysis.blueprint is not None root.mkdir(parents=True, exist_ok=False) blueprint = replace(analysis.blueprint, output_root=root) - materialization = materialize_blueprint(blueprint, seed=seed) + materialization = materialize_blueprint(blueprint) edit_plan = None if edit_prompt is not None: edit_blueprint = analyze_edit( @@ -195,7 +197,7 @@ def materialize( edit_prompt=str(edit_prompt), ) edit_plan = edit_blueprint.scene_edit_plan.to_dict() - materialization = materialize_edit(edit_blueprint, seed=seed) + materialization = materialize_edit(edit_blueprint) revision = _revision(materialization, seed=seed, edit_plan=edit_plan) _write_revision_audit( root, @@ -229,7 +231,7 @@ def materialize( edit_prompt=str(edit_prompt), ) edit_plan = edit_blueprint.scene_edit_plan.to_dict() - materialization = materialize_edit(edit_blueprint, seed=seed) + materialization = materialize_edit(edit_blueprint) if resolved.source_format == "legacy_gym_config": restore_locked_scene_entities(editable_root) else: @@ -301,16 +303,13 @@ def scene_blueprint_objects(blueprint: SceneBlueprintPackage) -> list[dict[str, Returns: Semantic objects with unknown physical fields represented conservatively. """ - nodes = blueprint.scene_graph.node_by_id() + if blueprint.schema_version != SCENE_BLUEPRINT_SCHEMA: + raise ValueError( + "Unsupported Scene Blueprint schema_version " + f"{blueprint.schema_version!r}; expected {SCENE_BLUEPRINT_SCHEMA!r}." + ) result = [] for item in blueprint.scene.objects: - node = nodes.get(item.id) - orientation = None if node is None else node.orientation_state - initial_state = {} - if orientation == "lying": - initial_state["orientation"] = "fallen" - elif orientation == "standing": - initial_state["orientation"] = "upright" result.append( { "uid": item.id, @@ -322,7 +321,9 @@ def scene_blueprint_objects(blueprint: SceneBlueprintPackage) -> list[dict[str, "color": None, "init_pos": [0.0, 0.0, 0.0], "affordances": [], - "initial_state": initial_state, + # Free-form pose descriptions are authoring intent, not measured + # orientation evidence. Final inspection publishes physical state. + "initial_state": {}, "attributes": {}, } ) diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py index b89ddfcd1..b0ca8552a 100644 --- a/embodichain/gen_sim/task_engine/workflow.py +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -48,7 +48,11 @@ from .contracts import canonical_hash from .orchestration.artifacts import ArtifactTransaction from .orchestration.coordinator import PreparationResult, TaskEngineCoordinator -from .orchestration.scene_adapter import CandidateSelection, SceneAdapter +from .orchestration.scene_adapter import ( + CandidateSelection, + SceneAdapter, + SceneAdapterProtocolError, +) from .orchestration.scene_source import SceneSourceRef from .scene_backend import ( SceneAnalysis, @@ -441,7 +445,7 @@ def run( self.scene_adapter, force_most_likely=True, ) - except Exception as exc: + except SceneAdapterProtocolError as exc: state = fail_stage( state, WorkflowStage.CANDIDATE_SELECTION, @@ -457,9 +461,28 @@ def run( run_metadata, state, attempts, - status="input_conflict", + status="failed", failure_class="candidate_selection", ) + except Exception as exc: + state = fail_stage( + state, + WorkflowStage.CANDIDATE_SELECTION, + reason=str(exc), + ) + return self._publish( + transaction, + staging, + normalized, + workflow_cfg, + planning_cfg, + execution_cfg, + run_metadata, + state, + attempts, + status="failed", + failure_class="internal_error", + ) _write_json( staging / "initial_binding_report.json", selection.binding_report ) diff --git a/tests/gen_sim/action_engine/generation/test_generation.py b/tests/gen_sim/action_engine/generation/test_generation.py index d30f2e42b..0f8b8b292 100644 --- a/tests/gen_sim/action_engine/generation/test_generation.py +++ b/tests/gen_sim/action_engine/generation/test_generation.py @@ -48,6 +48,7 @@ from embodichain.gen_sim.action_engine.generation.config_builder import ( build_agent_config, build_fast_gym_config, + validate_fast_gym_config, ) from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile from embodichain.gen_sim.action_engine.generation.generator import ( @@ -395,6 +396,80 @@ def test_fast_gym_config_has_runnable_franka_contract(gym_export: Path) -> None: ] == [14, 16] +def test_fast_gym_config_normalizes_usdc_articulation_runtime_fields( + gym_export: Path, +) -> None: + usdc_path = gym_export / "microwave.usdc" + usdc_path.write_bytes(b"PXR-USDC") + source_path = gym_export / "gym_config.json" + source = json.loads(source_path.read_text(encoding="utf-8")) + source["articulation"] = [ + { + "uid": "microwave_001", + "category": "microwave", + "name": "silver microwave", + "description": "A countertop microwave.", + "is_articulated": True, + "fpath": usdc_path.name, + "proxy_glb_fpath": "mesh_assets/can.glb", + "proxy_body_scale": [1.0, 1.0, 1.0], + "init_pos": [0.2, 0.0, 0.7], + "init_rot": [0.0, 0.0, 0.0], + "body_scale": [1.0, 1.0, 1.0], + "fix_base": True, + } + ] + source_path.write_text(json.dumps(source), encoding="utf-8") + + config = build_fast_gym_config( + prepare_scene(gym_export), + task_name="microwave_reference", + task_description="Place the can beside the microwave.", + robot_profile="franka", + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=2000, + ) + + articulation = config["articulation"][0] + assert articulation["fpath"] == usdc_path.resolve().as_posix() + assert articulation["build_pk_chain"] is False + assert ( + not { + "category", + "name", + "is_articulated", + "proxy_glb_fpath", + "proxy_body_scale", + } + & articulation.keys() + ) + + +def test_fast_gym_config_rejects_usdc_with_pk_chain(gym_export: Path) -> None: + config = build_fast_gym_config( + prepare_scene(gym_export), + task_name="invalid_usdc", + task_description="Reject an invalid runtime articulation.", + robot_profile="franka", + execution_program_hash="a" * 64, + max_episodes=1, + max_episode_steps=2000, + ) + usdc_path = gym_export / "invalid.usdc" + usdc_path.write_bytes(b"PXR-USDC") + config["articulation"] = [ + { + "uid": "microwave_001", + "fpath": usdc_path.resolve().as_posix(), + "build_pk_chain": True, + } + ] + + with pytest.raises(ValueError, match="must set build_pk_chain=false"): + validate_fast_gym_config(config) + + def test_offline_recording_can_be_disabled_but_ab_keeps_audience_recorder( gym_export: Path, monkeypatch: pytest.MonkeyPatch, diff --git a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py index bb4c626bb..0b6997921 100644 --- a/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py +++ b/tests/gen_sim/action_engine/runtime/test_runtime_contracts.py @@ -4350,6 +4350,71 @@ def test_turn_knob_requires_setting_map_and_reuses_twist() -> None: ) +def test_directional_spacing_accepts_articulation_reference() -> None: + moved = _FakeEntity("bowl", _pose(-0.3, 0.0, 0.75), _box_vertices(0.05)) + reference = _FakeArticulation("microwave", 0.0) + reference.link_names = ["drawer_link"] + reference._pose = _pose(0.0, 0.0, 0.85) + reference._vertices = _rect_vertices(0.20, 0.15, 0.15) + program = load_execution_program( + compile_task_agent_v2( + _task_agent( + { + "id": "place", + "operator": "place_relative", + "object": "bowl", + "actor": {"mode": "required", "arm": "left_arm"}, + "goal": { + "reference_object": "microwave", + "relation": "left_of", + }, + "depends_on": [], + } + ) + ) + ) + env = _FakeEnv( + entities={"bowl": moved}, + articulations={"microwave": reference}, + ) + grounder = ActionGrounder(program, env, lambda _uid: None) + + spacing = grounder._relative_object_spacing( + "bowl", + "microwave", + axis=0, + nominal=0.0, + clearance=0.02, + ) + + assert spacing == pytest.approx(0.27) + + +def test_relative_position_predicate_accepts_articulation_reference() -> None: + moved = _FakeEntity("bowl", _pose(-0.3, 0.0, 0.75), _box_vertices(0.05)) + reference = _FakeArticulation("microwave", 0.0) + reference.link_names = ["drawer_link"] + reference._pose = _pose(0.0, 0.0, 0.85) + reference._vertices = _rect_vertices(0.20, 0.15, 0.15) + env = _FakeEnv( + entities={"bowl": moved}, + articulations={"microwave": reference}, + ) + + result = evaluate_predicate( + env, + { + "type": "object_relative_position", + "object": "bowl", + "reference_object": "microwave", + "relation": "left_of", + "minimum_distance": -1.0, + }, + ) + + assert bool(result[0]) + + def test_handover_clearance_verifier_checks_distance_and_transfer_side() -> None: entities = { "can": _FakeEntity("can", _pose(0.0, -0.2, 1.0), _box_vertices(0.03)), diff --git a/tests/gen_sim/scene_engine/test_pipeline_api.py b/tests/gen_sim/scene_engine/test_pipeline_api.py index 9b10361f7..5bf79bc10 100644 --- a/tests/gen_sim/scene_engine/test_pipeline_api.py +++ b/tests/gen_sim/scene_engine/test_pipeline_api.py @@ -20,6 +20,8 @@ import json from pathlib import Path +import pytest + from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_edit_plan import SceneEditPlan from embodichain.gen_sim.scene_engine.core.scene_graph import ( @@ -38,6 +40,15 @@ def check_health(self) -> None: self.health_checks += 1 +class _OwnedClient(_HealthyClient): + def __init__(self) -> None: + super().__init__() + self.close_calls = 0 + + def close(self) -> None: + self.close_calls += 1 + + def _materialization( *, scene: Scene, @@ -93,6 +104,8 @@ def fake_understand_scene(**kwargs): document = json.loads(package.manifest_path.read_text(encoding="utf-8")) assert segmentation.health_checks == 1 + assert package.schema_version == api.SCENE_BLUEPRINT_SCHEMA + assert document["schema_version"] == "embodichain.scene-blueprint/v2" assert document["blueprint_id"] == package.blueprint_id assert document["scene_graph"] == graph.to_dict() assert document["artifacts"][0]["path"].endswith("table-mask.png") @@ -126,6 +139,8 @@ def import_scene_and_graph(self): ) document = json.loads(package.manifest_path.read_text(encoding="utf-8")) + assert package.schema_version == api.SCENE_EDIT_BLUEPRINT_SCHEMA + assert document["schema_version"] == "embodichain.scene-edit-blueprint/v2" assert document["blueprint_id"] == package.blueprint_id assert document["scene_edit_plan"] == plan.to_dict() assert document["updated_scene_graph"] == graph.to_dict() @@ -139,6 +154,7 @@ def test_materialize_blueprint_does_not_mutate_audited_snapshot( manifest_path = tmp_path / "scene_blueprint.json" manifest_path.write_text("audited blueprint\n", encoding="utf-8") package = api.SceneBlueprintPackage( + schema_version=api.SCENE_BLUEPRINT_SCHEMA, blueprint_id="blueprint", image_path=tmp_path / "input.png", output_root=tmp_path, @@ -150,7 +166,8 @@ def test_materialize_blueprint_does_not_mutate_audited_snapshot( original_graph = deepcopy(graph.to_dict()) def fake_generate_scene_and_refine(**kwargs): - assert kwargs["seed"] == 31 + assert "seed" not in kwargs + assert kwargs["articulated_generation_client"] is None assert kwargs["scene"] is not package.scene assert kwargs["scene_graph"] is not package.scene_graph kwargs["scene"].objects[0].name = "materialized table" @@ -175,7 +192,6 @@ def fake_generate_scene_and_refine(**kwargs): package, vlm_client=object(), geometry_generation_client=_HealthyClient(), - seed=31, ) assert result.scene.objects[0].name == "materialized table" @@ -193,6 +209,7 @@ def test_materialize_edit_does_not_mutate_audited_snapshot( manifest_path = tmp_path / "scene_edit_blueprint.json" manifest_path.write_text("audited edit blueprint\n", encoding="utf-8") package = api.SceneEditBlueprintPackage( + schema_version=api.SCENE_EDIT_BLUEPRINT_SCHEMA, blueprint_id="edit-blueprint", edit_prompt="Keep the scene unchanged.", output_root=tmp_path, @@ -204,7 +221,7 @@ def test_materialize_edit_does_not_mutate_audited_snapshot( original_graph = deepcopy(graph.to_dict()) def fake_prepare_scene_edit_assets(**kwargs): - assert kwargs["seed"] == 32 + assert "seed" not in kwargs return [] monkeypatch.setattr( @@ -235,10 +252,106 @@ def fake_edit_layout(**kwargs): image_generation_client=clients[0], geometry_generation_client=clients[1], image_segmentation_client=clients[2], - seed=32, ) assert result.scene.objects[0].name == "edited table" assert package.scene_edit_plan.to_dict() == original_plan assert package.updated_scene_graph.to_dict() == original_graph assert manifest_path.read_text(encoding="utf-8") == "audited edit blueprint\n" + + +def test_scene_blueprint_package_rejects_v1_schema(tmp_path: Path) -> None: + scene, graph = _table_scene() + + with pytest.raises(ValueError, match="scene-blueprint/v2"): + api.SceneBlueprintPackage( + schema_version="embodichain.scene-blueprint/v1", + blueprint_id="legacy", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=tmp_path / "scene_blueprint.json", + scene=scene, + scene_graph=graph, + ) + + +def test_scene_edit_blueprint_package_rejects_v1_schema(tmp_path: Path) -> None: + scene, graph = _table_scene() + plan = SceneEditPlan(scene=scene, scene_graph=graph, operations=[]) + + with pytest.raises(ValueError, match="scene-edit-blueprint/v2"): + api.SceneEditBlueprintPackage( + schema_version="embodichain.scene-edit-blueprint/v1", + blueprint_id="legacy-edit", + edit_prompt="Keep the scene unchanged.", + output_root=tmp_path, + manifest_path=tmp_path / "scene_edit_blueprint.json", + scene_edit_plan=plan, + updated_scene_graph=graph, + ) + + +def test_materialize_blueprint_owns_articulated_client_lifecycle( + tmp_path: Path, + monkeypatch, +) -> None: + scene, graph = _table_scene() + scene.objects.append( + SceneObject( + id="microwave_001", + kind="asset", + category="microwave", + name="microwave", + description="An articulated microwave.", + is_articulated=True, + ) + ) + graph.nodes.append( + SceneGraphNode( + object_id="microwave_001", + parent_id="table", + parent_relation="on", + pose_description="Stand upright on its base.", + ) + ) + package = api.SceneBlueprintPackage( + schema_version=api.SCENE_BLUEPRINT_SCHEMA, + blueprint_id="articulated", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=tmp_path / "scene_blueprint.json", + scene=scene, + scene_graph=graph, + ) + articulated = _OwnedClient() + monkeypatch.setattr( + api.ArticulatedGenerationClient, + "from_dotenv", + lambda: articulated, + ) + + def fake_generate_scene_and_refine(**kwargs): + assert kwargs["articulated_generation_client"] is articulated + return kwargs["scene"] + + monkeypatch.setattr( + api, "generate_scene_and_refine", fake_generate_scene_and_refine + ) + monkeypatch.setattr( + api, + "_export_materialization", + lambda *, scene, scene_graph, output_root: _materialization( + scene=scene, + scene_graph=scene_graph, + output_root=output_root, + ), + ) + + api.materialize_blueprint( + package, + vlm_client=object(), + geometry_generation_client=_HealthyClient(), + ) + + assert articulated.health_checks == 1 + assert articulated.close_calls == 1 diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py index 27b6d029b..1cfa88de0 100644 --- a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -313,6 +313,34 @@ def test_artifact_transaction_rolls_back_and_preserves_existing_output( assert not (output / "partial.txt").exists() +def test_artifact_transaction_relocates_paths_in_json_mapping_keys( + tmp_path: Path, +) -> None: + output = tmp_path / "bundle" + with ArtifactTransaction(output) as transaction: + assert transaction.staging_dir is not None + staging = transaction.staging_dir.resolve().as_posix() + (transaction.staging_dir / "manifest.json").write_text( + json.dumps( + { + "config_path": f"{staging}/scene/scene_config.json", + "asset_sha256": {f"{staging}/scene/asset.usdc": "hash"}, + } + ), + encoding="utf-8", + ) + transaction.commit() + + manifest = json.loads((output / "manifest.json").read_text(encoding="utf-8")) + assert ( + manifest["config_path"] + == f"{output.resolve().as_posix()}/scene/scene_config.json" + ) + assert list(manifest["asset_sha256"]) == [ + f"{output.resolve().as_posix()}/scene/asset.usdc" + ] + + def test_prepare_rejects_output_overlapping_read_only_source(tmp_path: Path) -> None: source = tmp_path / "gym_project" source.mkdir() @@ -767,6 +795,60 @@ def fake_cli() -> None: assert str(bundle / AGENT_CONFIG_FILENAME) in captured[0] +def test_private_bundle_runner_normalizes_legacy_usdc_config( + tmp_path: Path, +) -> None: + source = tmp_path / FAST_GYM_CONFIG_FILENAME + original = { + "articulation": [ + { + "uid": "microwave", + "category": "microwave", + "is_articulated": True, + "fpath": "/scene/microwave.usdc", + "proxy_glb_fpath": "microwave.glb", + } + ] + } + source.write_text(json.dumps(original), encoding="utf-8") + + with bundle_runner._runtime_gym_config(source) as runtime_path: + runtime = json.loads(runtime_path.read_text(encoding="utf-8")) + assert runtime_path != source + assert runtime["articulation"] == [ + { + "uid": "microwave", + "fpath": "/scene/microwave.usdc", + "build_pk_chain": False, + } + ] + temporary_path = runtime_path + + assert not temporary_path.exists() + assert json.loads(source.read_text(encoding="utf-8")) == original + + +def test_private_bundle_runner_reuses_normalized_gym_config(tmp_path: Path) -> None: + source = tmp_path / FAST_GYM_CONFIG_FILENAME + source.write_text( + json.dumps( + { + "articulation": [ + { + "uid": "microwave", + "fpath": "/scene/microwave.usdc", + "build_pk_chain": False, + } + ] + } + ), + encoding="utf-8", + ) + + with bundle_runner._runtime_gym_config(source) as runtime_path: + assert runtime_path == source + + @pytest.mark.parametrize( ("mode", "image", "scene", "edit"), [ diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py index 3936cc2e9..f073e3685 100644 --- a/tests/gen_sim/task_engine/test_parallel_workflow.py +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -40,6 +40,7 @@ ) from embodichain.gen_sim.task_engine.orchestration.scene_adapter import ( CandidateSelection, + SceneAdapterProtocolError, ) from embodichain.gen_sim.task_engine.orchestration.scene_source import SceneSourceRef from embodichain.gen_sim.task_engine.scene_backend import SceneAnalysis, SceneRevision @@ -150,12 +151,14 @@ def __init__( input_barrier: Barrier | None = None, materialize_barrier: Barrier | None = None, materialize_failures: int = 0, + selection_error: Exception | None = None, ) -> None: self.selection = selection self.input_kind = input_kind self.input_barrier = input_barrier self.materialize_barrier = materialize_barrier self.materialize_failures = materialize_failures + self.selection_error = selection_error self.seeds: list[int] = [] def analyze(self, request, output_root) -> SceneAnalysis: @@ -169,6 +172,8 @@ def analyze(self, request, output_root) -> SceneAnalysis: ) def select(self, *_args, **_kwargs) -> CandidateSelection: + if self.selection_error is not None: + raise self.selection_error return self.selection def materialize( @@ -366,6 +371,54 @@ def test_parallel_workflow_supports_all_four_scene_inputs( assert result.succeeded +def test_candidate_selection_programming_error_is_internal_failure( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _SceneBackend( + _selection(candidates), + selection_error=AttributeError("stale cross-engine field"), + ) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True]]), + ) + + result = workflow.run(_request(tmp_path)) + + assert result.status == "failed" + assert result.failure_class == "internal_error" + state = json.loads( + (result.output_dir / "workflow_state.json").read_text(encoding="utf-8") + ) + assert state["events"][-1]["reason"] == "stale cross-engine field" + + +def test_candidate_selection_protocol_error_is_not_input_conflict( + tmp_path: Path, +) -> None: + candidates = _candidate_set() + scene = _SceneBackend( + _selection(candidates), + selection_error=SceneAdapterProtocolError("invalid grounding response"), + ) + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=scene, + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor([[True]]), + ) + + result = workflow.run(_request(tmp_path)) + + assert result.status == "failed" + assert result.failure_class == "candidate_selection" + + def test_parallel_workflow_preserves_requested_robot_profile(tmp_path: Path) -> None: candidates = _candidate_set() coordinator = _Coordinator(["bound"]) diff --git a/tests/gen_sim/task_engine/test_scene_backend.py b/tests/gen_sim/task_engine/test_scene_backend.py index cc6f85bd9..f93dff965 100644 --- a/tests/gen_sim/task_engine/test_scene_backend.py +++ b/tests/gen_sim/task_engine/test_scene_backend.py @@ -29,11 +29,13 @@ ) from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.api import ( + SCENE_BLUEPRINT_SCHEMA, SceneBlueprintPackage, SceneMaterialization, ) import embodichain.gen_sim.task_engine.scene_backend as scene_backend_module from embodichain.gen_sim.task_engine.scene_backend import ( + SceneAnalysis, SceneEngineBackend, scene_blueprint_objects, ) @@ -88,7 +90,9 @@ def _scene_export(tmp_path: Path) -> Path: return export.parent -def test_blueprint_objects_preserve_semantics_without_geometry(tmp_path: Path) -> None: +def test_blueprint_objects_keep_pose_description_orientation_unknown( + tmp_path: Path, +) -> None: scene = Scene( objects=[ SceneObject("table", "table", "table", "table", "A table."), @@ -98,10 +102,16 @@ def test_blueprint_objects_preserve_semantics_without_geometry(tmp_path: Path) - graph = SceneGraph( nodes=[ SceneGraphNode("table", None), - SceneGraphNode("cup", "table", "on", orientation_state="lying"), + SceneGraphNode( + "cup", + "table", + "on", + pose_description="Lie flat on the support surface.", + ), ] ) package = SceneBlueprintPackage( + schema_version=SCENE_BLUEPRINT_SCHEMA, blueprint_id="blueprint", image_path=tmp_path / "input.png", output_root=tmp_path, @@ -114,11 +124,51 @@ def test_blueprint_objects_preserve_semantics_without_geometry(tmp_path: Path) - cup = next(item for item in objects if item["uid"] == "cup") assert cup["description"] == "A red cup." - assert cup["initial_state"] == {"orientation": "fallen"} + assert cup["initial_state"] == {} assert cup["affordances"] == [] assert cup["init_pos"] == [0.0, 0.0, 0.0] +def test_backend_select_passes_v2_blueprint_contract(tmp_path: Path) -> None: + scene = Scene(objects=[SceneObject("table", "table", "table", "table", "A table.")]) + graph = SceneGraph(nodes=[SceneGraphNode("table", None)]) + package = SceneBlueprintPackage( + schema_version=SCENE_BLUEPRINT_SCHEMA, + blueprint_id="blueprint", + image_path=tmp_path / "input.png", + output_root=tmp_path, + manifest_path=tmp_path / "scene_blueprint.json", + scene=scene, + scene_graph=graph, + ) + analysis = SceneAnalysis( + input_kind="image", + source=package.image_path, + blueprint=package, + source_fingerprint=None, + ) + captured = {} + marker = object() + + class CapturingAdapter: + def select_objects(self, candidate_set, scene_objects, **kwargs): + captured["candidate_set"] = candidate_set + captured["scene_objects"] = scene_objects + captured.update(kwargs) + return marker + + result = SceneEngineBackend().select( + analysis, + {"candidate": "value"}, + CapturingAdapter(), + force_most_likely=True, + ) + + assert result is marker + assert captured["source_format"] == "embodichain.scene-blueprint/v2" + assert captured["scene_objects"][0]["uid"] == "table" + + def test_existing_scene_edit_creates_revision_and_never_writes_source( tmp_path: Path, monkeypatch, @@ -153,8 +203,7 @@ def fake_analyze_edit(*, output_root, edit_prompt): ), ) - def fake_materialize_edit(blueprint, *, seed=None): - assert seed == 7 + def fake_materialize_edit(blueprint): return SceneMaterialization( scene=Scene(), scene_graph=SceneGraph(nodes=[SceneGraphNode("table", None)]), From 2beaf92d666b901e2dfd394c6b27fa2daa954403 Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Tue, 1 Sep 2026 11:06:10 +0800 Subject: [PATCH 83/85] temp --- embodichain/gen_sim/scene_engine/cli/start.py | 54 +++++++++++- .../gen_sim/scene_engine/pipeline/generate.py | 21 ++++- .../pipeline/utils/scene_layout_utils.py | 66 ++++++++++++++ tests/gen_sim/scene_engine/test_config.py | 85 ++++++++++++++++++- .../scene_engine/test_scene_generation.py | 74 ++++++++++++++++ 5 files changed, 294 insertions(+), 6 deletions(-) diff --git a/embodichain/gen_sim/scene_engine/cli/start.py b/embodichain/gen_sim/scene_engine/cli/start.py index 6d46f32dc..5d0c752c6 100644 --- a/embodichain/gen_sim/scene_engine/cli/start.py +++ b/embodichain/gen_sim/scene_engine/cli/start.py @@ -18,6 +18,7 @@ import argparse from collections.abc import Sequence +import math from pathlib import Path from embodichain.gen_sim.scene_engine.pipeline.generate import generate_scene_from_image @@ -31,9 +32,12 @@ def cli_scene_engine( output_root: str | Path, *, edit_prompt: str | None = None, + scene_z_rotation_degrees: float = 0.0, ) -> None: """Generate a scene from an image, edit an export, or do both in sequence.""" resolved_output_root = Path(output_root).expanduser().resolve() + if not math.isfinite(scene_z_rotation_degrees): + raise ValueError("scene_z_rotation_degrees must be finite.") if edit_prompt is not None: edit_prompt = edit_prompt.strip() if not edit_prompt: @@ -60,6 +64,7 @@ def cli_scene_engine( generate_scene_from_image( image_path=resolved_image_path, output_root=resolved_output_root, + scene_z_rotation_degrees=scene_z_rotation_degrees, ) if edit_prompt is not None: edit_scene( @@ -93,9 +98,56 @@ def main(argv: Sequence[str] | None = None) -> None: default=None, help="Text instruction for editing an existing or newly generated output root", ) + parser.add_argument( + "--scene-z-rotation-degrees", + "--scene_z_rotation_degrees", + "--prompt2scene-scene-z-rotation-degrees", + "--prompt2scene_scene_z_rotation_degrees", + dest="scene_z_rotation_degrees", + type=float, + default=0.0, + help=( + "Final counterclockwise world-Z rotation applied to the complete " + "generated scene. Defaults to 0." + ), + ) + parser.add_argument( + "--target-body-scale-mode", + "--target_body_scale_mode", + choices=("preserve",), + default="preserve", + help=( + "Compatibility option for the direct-GLB Scene Engine path; source " + "scale is always preserved." + ), + ) + parser.add_argument( + "--prompt2scene-mesh-x-rotation-degrees", + "--prompt2scene_mesh_x_rotation_degrees", + type=_zero_mesh_x_rotation, + default=0.0, + help=( + "Compatibility option for direct GLB loading. It must remain 0; " + "DexSim performs the GLTF y-up conversion without a baked 90-degree fix." + ), + ) args = parser.parse_args(argv) - cli_scene_engine(args.image, args.output_root, edit_prompt=args.edit_prompt) + cli_scene_engine( + args.image, + args.output_root, + edit_prompt=args.edit_prompt, + scene_z_rotation_degrees=args.scene_z_rotation_degrees, + ) + + +def _zero_mesh_x_rotation(value: str) -> float: + rotation = float(value) + if not math.isfinite(rotation) or rotation != 0.0: + raise argparse.ArgumentTypeError( + "prompt2scene mesh X rotation must be 0 for direct GLB loading." + ) + return rotation if __name__ == "__main__": diff --git a/embodichain/gen_sim/scene_engine/pipeline/generate.py b/embodichain/gen_sim/scene_engine/pipeline/generate.py index 71155a40a..5f65660d7 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/generate.py +++ b/embodichain/gen_sim/scene_engine/pipeline/generate.py @@ -41,13 +41,28 @@ generate_scene_and_refine, ) from embodichain.gen_sim.scene_engine.pipeline.utils.scene_exporter import SceneExporter +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + rotate_scene_z_up_world, +) def generate_scene_from_image( image_path: str | Path, output_root: str | Path, + *, + scene_z_rotation_degrees: float = 0.0, ) -> Scene: - """Generate the initial core scene state from an input image.""" + """Generate the initial core scene state from an input image. + + Args: + image_path: Source tabletop image. + output_root: Directory receiving intermediate and exported artifacts. + scene_z_rotation_degrees: Final counterclockwise world-z rotation applied + rigidly to the complete scene before export. + + Returns: + The final scene in the rotated export frame. + """ resolved_output_root = Path(output_root).expanduser().resolve() resolved_output_root.mkdir(parents=True, exist_ok=True) @@ -100,6 +115,10 @@ def generate_scene_from_image( # 3. Scene Export log_info("Starting Scene Export") + rotate_scene_z_up_world( + scene=scene, + rotation_degrees=scene_z_rotation_degrees, + ) scene_exporter = SceneExporter( scene=scene, scene_graph=scene_graph, diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py index 855342920..4838e0c31 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_layout_utils.py @@ -17,7 +17,9 @@ from __future__ import annotations import numpy as np +from scipy.spatial.transform import Rotation +from embodichain.gen_sim.scene_engine.core.scene import Scene from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject from embodichain.gen_sim.scene_engine.pipeline.utils.scene_generation_utils import ( layout_object_to_transform_matrix, @@ -75,6 +77,70 @@ def translate_scene_object_y_up_by_z_up_delta( ] +def rotate_scene_z_up_world(*, scene: Scene, rotation_degrees: float) -> None: + """Rotate every final scene pose and support point about world z. + + The generated meshes remain unchanged. The same rigid world transform is + applied to positions, orientations, and persisted z-up XY support metadata + so the exported scene stays internally consistent. + + Args: + scene: Final generated scene whose object poses use the internal y-up frame. + rotation_degrees: Counterclockwise world-z rotation in degrees. + + Raises: + ValueError: If the rotation is not finite or a scene pose is incomplete. + """ + angle = float(rotation_degrees) + if not np.isfinite(angle): + raise ValueError("rotation_degrees must be finite.") + if angle % 360.0 == 0.0: + return + + world_rotation = np.eye(4) + world_rotation[:3, :3] = Rotation.from_euler("z", angle, degrees=True).as_matrix() + basis = y_up_to_z_up_matrix() + inverse_basis = np.linalg.inv(basis) + rotation_xy = world_rotation[:2, :2] + + for scene_object in scene.objects: + y_up_layout = scene_object_y_up_layout(scene_object) + z_up_transform = ( + basis @ layout_object_to_transform_matrix(y_up_layout) @ inverse_basis + ) + rotated_y_up_transform = inverse_basis @ world_rotation @ z_up_transform @ basis + rotated_layout = transform_matrix_to_layout_object( + scene_object.id, + rotated_y_up_transform, + ) + scene_object.pos = rotated_layout["pos"] + scene_object.rot = rotated_layout["rot"] + scene_object.scale = rotated_layout["scale"] + + if scene_object.center_xy is not None: + scene_object.center_xy = ( + rotation_xy + @ np.asarray( + two_floats(scene_object.center_xy, field_name="center_xy"), + dtype=float, + ) + ).tolist() + for field_name in ("support_contour_xy", "support_optimization_rect_xy"): + points = getattr(scene_object, field_name) + if points is None: + continue + setattr( + scene_object, + field_name, + [ + (rotation_xy @ np.asarray(two_floats(point, field_name=field_name))) + .astype(float) + .tolist() + for point in points + ], + ) + + def measure_scene_object_z_up_world_aabb( *, scene_object: SceneObject ) -> list[list[float]]: diff --git a/tests/gen_sim/scene_engine/test_config.py b/tests/gen_sim/scene_engine/test_config.py index 3754210d9..96bebf4ad 100644 --- a/tests/gen_sim/scene_engine/test_config.py +++ b/tests/gen_sim/scene_engine/test_config.py @@ -60,6 +60,9 @@ def test_scene_engine_help_exposes_only_runtime_arguments( output = capsys.readouterr().out assert "--image" in output assert "--output_root" in output + assert "--prompt2scene_scene_z_rotation_degrees" in output + assert "--prompt2scene_mesh_x_rotation_degrees" in output + assert "--target_body_scale_mode" in output assert "gen_sim/.env" in output assert "--config" not in output @@ -70,23 +73,92 @@ def test_scene_engine_cli_forwards_validated_paths( ) -> None: image_path = tmp_path / "scene.png" image_path.write_bytes(b"png") - captured: dict[str, Path] = {} + captured: dict[str, object] = {} - def generate_scene(*, image_path: Path, output_root: Path) -> None: + def generate_scene( + *, + image_path: Path, + output_root: Path, + scene_z_rotation_degrees: float, + ) -> None: captured["image_path"] = image_path captured["output_root"] = output_root + captured["scene_z_rotation_degrees"] = scene_z_rotation_degrees monkeypatch.setattr(start, "generate_scene_from_image", generate_scene) output_root = tmp_path / "output" - start.cli_scene_engine(image_path, output_root) + start.cli_scene_engine( + image_path, + output_root, + scene_z_rotation_degrees=180.0, + ) assert captured == { "image_path": image_path.resolve(), "output_root": output_root.resolve(), + "scene_z_rotation_degrees": 180.0, } +def test_scene_engine_main_accepts_legacy_direct_glb_options( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + image_path = tmp_path / "scene.png" + image_path.write_bytes(b"png") + captured: dict[str, object] = {} + + def cli_scene_engine( + image: str, + output_root: str, + *, + edit_prompt: str | None, + scene_z_rotation_degrees: float, + ) -> None: + captured.update( + image=image, + output_root=output_root, + edit_prompt=edit_prompt, + scene_z_rotation_degrees=scene_z_rotation_degrees, + ) + + monkeypatch.setattr(start, "cli_scene_engine", cli_scene_engine) + + start.main( + [ + "--image", + str(image_path), + "--output_root", + str(tmp_path / "output"), + "--target_body_scale_mode", + "preserve", + "--prompt2scene_scene_z_rotation_degrees", + "180", + "--prompt2scene_mesh_x_rotation_degrees", + "0", + ] + ) + + assert captured["scene_z_rotation_degrees"] == 180.0 + + +def test_scene_engine_main_rejects_legacy_mesh_x_rotation( + tmp_path: Path, +) -> None: + with pytest.raises(SystemExit) as exc_info: + start.main( + [ + "--output_root", + str(tmp_path / "output"), + "--prompt2scene_mesh_x_rotation_degrees", + "90", + ] + ) + + assert exc_info.value.code == 2 + + def test_scene_engine_cli_edits_existing_output_without_an_image( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -116,7 +188,12 @@ def test_scene_engine_cli_generates_then_edits_when_both_inputs_exist( image_path.write_bytes(b"png") call_order: list[str] = [] - def generate_scene(*, image_path: Path, output_root: Path) -> None: + def generate_scene( + *, + image_path: Path, + output_root: Path, + scene_z_rotation_degrees: float, + ) -> None: call_order.append("generate") def edit_scene(*, output_root: Path, edit_prompt: str) -> None: diff --git a/tests/gen_sim/scene_engine/test_scene_generation.py b/tests/gen_sim/scene_engine/test_scene_generation.py index b8c633884..584d8d214 100644 --- a/tests/gen_sim/scene_engine/test_scene_generation.py +++ b/tests/gen_sim/scene_engine/test_scene_generation.py @@ -49,6 +49,11 @@ layout_object_to_transform_matrix, transform_matrix_to_layout_object, ) +from embodichain.gen_sim.scene_engine.pipeline.utils.scene_layout_utils import ( + rotate_scene_z_up_world, + scene_object_y_up_layout, + y_up_to_z_up_matrix, +) def _z_up_rotation_from_y_up_layout(layout: dict[str, object]) -> np.ndarray: @@ -223,6 +228,75 @@ def test_visual_yaws_replace_coarse_rotations_but_preserve_positions() -> None: assert np.allclose(yawed_layout["pos"], [0.1, 0.2, 0.3]) +def test_rotate_scene_z_up_world_rotates_complete_scene_and_support_metadata() -> None: + scene = Scene( + objects=[ + SceneObject( + id="table", + kind="table", + category="table", + name="table", + description="table", + rot=[0.0, 0.0, 0.0], + pos=[0.0, 0.0, 0.0], + scale=[1.0, 1.0, 1.0], + center_xy=[1.0, 2.0], + support_contour_xy=[[1.0, 2.0], [-1.0, 2.0]], + support_optimization_rect_xy=[[1.0, 1.0], [-1.0, 1.0]], + ), + SceneObject( + id="book_001", + kind="asset", + category="book", + name="book", + description="book", + rot=[10.0, 20.0, 30.0], + pos=[1.0, 2.0, 3.0], + scale=[1.0, 2.0, 3.0], + center_xy=[3.0, 4.0], + ), + ] + ) + basis = y_up_to_z_up_matrix() + inverse_basis = np.linalg.inv(basis) + original_z_up_transforms = { + scene_object.id: ( + basis + @ layout_object_to_transform_matrix(scene_object_y_up_layout(scene_object)) + @ inverse_basis + ) + for scene_object in scene.objects + } + expected_world_rotation = np.eye(4) + expected_world_rotation[:3, :3] = Rotation.from_euler( + "z", 180.0, degrees=True + ).as_matrix() + + rotate_scene_z_up_world(scene=scene, rotation_degrees=180.0) + + for scene_object in scene.objects: + actual_z_up_transform = ( + basis + @ layout_object_to_transform_matrix(scene_object_y_up_layout(scene_object)) + @ inverse_basis + ) + assert np.allclose( + actual_z_up_transform, + expected_world_rotation @ original_z_up_transforms[scene_object.id], + ) + assert np.allclose(scene.table.center_xy, [-1.0, -2.0]) + assert np.allclose(scene.table.support_contour_xy, [[-1.0, -2.0], [1.0, -2.0]]) + assert np.allclose( + scene.table.support_optimization_rect_xy, + [[-1.0, -1.0], [1.0, -1.0]], + ) + + +def test_rotate_scene_z_up_world_rejects_non_finite_angle() -> None: + with pytest.raises(ValueError, match="rotation_degrees must be finite"): + rotate_scene_z_up_world(scene=Scene(), rotation_degrees=float("nan")) + + def test_articulated_usdcs_use_visible_rgba_in_scene_order( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, From 8eeb6e884d9878987f0267cd62ada5376ecd09fb Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:53:49 +0800 Subject: [PATCH 84/85] fix(graspkit): switch convex collision decomposition to V-HACD --- .../graspkit/pg_grasp/collision_checker.py | 13 +- tests/toolkits/test_batch_convex_collision.py | 7 +- .../toolkits/test_convex_collision_checker.py | 128 ++++++++++++++++++ 3 files changed, 144 insertions(+), 4 deletions(-) create mode 100644 tests/toolkits/test_convex_collision_checker.py diff --git a/embodichain/toolkits/graspkit/pg_grasp/collision_checker.py b/embodichain/toolkits/graspkit/pg_grasp/collision_checker.py index 6a212c64e..7f0b5284f 100644 --- a/embodichain/toolkits/graspkit/pg_grasp/collision_checker.py +++ b/embodichain/toolkits/graspkit/pg_grasp/collision_checker.py @@ -27,7 +27,7 @@ import open3d as o3d from typing import List, Tuple, Union -from dexsim.kit.meshproc import convex_decomposition_coacd +from dexsim.kit.meshproc import convex_decomposition_vhacd from embodichain.utils.warp import convex_signed_distance_kernel from embodichain.utils.device_utils import standardize_device_string @@ -37,6 +37,8 @@ __all__ = ["ConvexCollisionCheckerCfg", "ConvexCollisionChecker"] +_CONVEX_DECOMPOSITION_CACHE_TAG = "vhacd_v1" + @configclass class ConvexCollisionCheckerCfg: @@ -86,7 +88,10 @@ def __init__( self.cache_path = os.path.join( CONVEX_DECOMPOSITION_CACHE_DIR, - f"{mesh_hash}_{max_decomposition_hulls}.pkl", + ( + f"{mesh_hash}_{max_decomposition_hulls}_" + f"{_CONVEX_DECOMPOSITION_CACHE_TAG}.pkl" + ), ) if not os.path.isfile(self.cache_path): @@ -299,9 +304,11 @@ def _compute_plane_equations( mesh = o3d.t.geometry.TriangleMesh() mesh.vertex.positions = o3d.core.Tensor(vertices, dtype=o3d.core.Dtype.Float32) mesh.triangle.indices = o3d.core.Tensor(faces, dtype=o3d.core.Dtype.Int32) - is_success, out_mesh_list = convex_decomposition_coacd( + is_success, out_mesh_list = convex_decomposition_vhacd( mesh, max_convex_hull_num=max_decomposition_hulls ) + if not is_success or not out_mesh_list: + raise RuntimeError("V-HACD convex decomposition failed.") convex_vert_face_list = [] for out_mesh in out_mesh_list: verts = out_mesh.vertex.positions.numpy() diff --git a/tests/toolkits/test_batch_convex_collision.py b/tests/toolkits/test_batch_convex_collision.py index 5e6255f4a..3a70b6778 100644 --- a/tests/toolkits/test_batch_convex_collision.py +++ b/tests/toolkits/test_batch_convex_collision.py @@ -30,6 +30,8 @@ pytestmark = pytest.mark.gpu +_EXPECTED_VHACD_MAX_SURFACE_DISTANCE = 0.5945 + def batch_convex_collision_query(device=torch.device("cuda")): mug_path = get_data_path("ScannedBottle/moliwulong_processed.ply") @@ -73,7 +75,10 @@ def batch_convex_collision_query(device=torch.device("cuda")): is_pose_collide = is_point_collide.any(dim=1) pose_surface_distance = point_surface_distance.min(dim=1).values assert is_pose_collide.sum().item() == 1 - assert abs(pose_surface_distance.max().item() - 0.8492) < 1e-2 + assert ( + abs(pose_surface_distance.max().item() - _EXPECTED_VHACD_MAX_SURFACE_DISTANCE) + < 1e-2 + ) def test_batch_convex_collision_cpu(): diff --git a/tests/toolkits/test_convex_collision_checker.py b/tests/toolkits/test_convex_collision_checker.py new file mode 100644 index 000000000..f0b0f3e3a --- /dev/null +++ b/tests/toolkits/test_convex_collision_checker.py @@ -0,0 +1,128 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import hashlib +from pathlib import Path + +import numpy as np +import pytest +import torch + +from embodichain.toolkits.graspkit.pg_grasp import collision_checker as module +from embodichain.toolkits.graspkit.pg_grasp.collision_checker import ( + ConvexCollisionChecker, +) + + +def _tetrahedron() -> tuple[np.ndarray, np.ndarray]: + vertices = np.array( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + faces = np.array( + [ + [0, 2, 1], + [0, 1, 3], + [0, 3, 2], + [1, 2, 3], + ], + dtype=np.int32, + ) + return vertices, faces + + +def test_plane_equations_use_vhacd(monkeypatch: pytest.MonkeyPatch) -> None: + vertices, faces = _tetrahedron() + calls: list[int] = [] + + def fake_vhacd(mesh, *, max_convex_hull_num: int): + calls.append(max_convex_hull_num) + return True, (mesh,) + + monkeypatch.setattr(module, "convex_decomposition_vhacd", fake_vhacd) + + plane_equations = ConvexCollisionChecker._compute_plane_equations( + vertices, + faces, + max_decomposition_hulls=16, + ) + + assert calls == [16] + assert len(plane_equations) == 1 + + +def test_vhacd_failure_is_reported(monkeypatch: pytest.MonkeyPatch) -> None: + vertices, faces = _tetrahedron() + + monkeypatch.setattr( + module, + "convex_decomposition_vhacd", + lambda *_args, **_kwargs: (False, ()), + ) + + with pytest.raises(RuntimeError, match="V-HACD convex decomposition failed"): + ConvexCollisionChecker._compute_plane_equations( + vertices, + faces, + max_decomposition_hulls=16, + ) + + +def test_vhacd_cache_does_not_reuse_legacy_backend( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + vertices, faces = _tetrahedron() + mesh_hash = hashlib.md5(vertices.tobytes() + faces.tobytes()).hexdigest() + legacy_path = tmp_path / f"{mesh_hash}_16.pkl" + legacy_path.write_bytes(b"legacy CoACD cache") + calls: list[int] = [] + + def fake_plane_equations( + _vertices: np.ndarray, + _faces: np.ndarray, + max_decomposition_hulls: int, + ) -> list[tuple[np.ndarray, np.ndarray]]: + calls.append(max_decomposition_hulls) + return [ + ( + np.array([[1.0, 0.0, 0.0]], dtype=np.float32), + np.array([0.0], dtype=np.float32), + ) + ] + + monkeypatch.setattr(module, "CONVEX_DECOMPOSITION_CACHE_DIR", tmp_path) + monkeypatch.setattr( + ConvexCollisionChecker, + "_compute_plane_equations", + staticmethod(fake_plane_equations), + ) + + checker = ConvexCollisionChecker( + torch.from_numpy(vertices), + torch.from_numpy(faces), + max_decomposition_hulls=16, + ) + + assert calls == [16] + assert checker.cache_path == str(tmp_path / f"{mesh_hash}_16_vhacd_v1.pkl") From 9125f3481fbcbcffe36ad9623320aabbec05f8af Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:47:25 +0800 Subject: [PATCH 85/85] feat(task-engine): add --open-window simulator execution option --- embodichain/gen_sim/task_engine/cli.py | 12 +++++ embodichain/gen_sim/task_engine/workflow.py | 16 +++++- .../orchestration/test_coordinator_cli.py | 53 ++++++++++++++++++- .../task_engine/test_parallel_workflow.py | 38 ++++++++++++- 4 files changed, 115 insertions(+), 4 deletions(-) diff --git a/embodichain/gen_sim/task_engine/cli.py b/embodichain/gen_sim/task_engine/cli.py index 3b7ee42c3..22d20a8f0 100644 --- a/embodichain/gen_sim/task_engine/cli.py +++ b/embodichain/gen_sim/task_engine/cli.py @@ -69,6 +69,7 @@ def build_parser() -> argparse.ArgumentParser: "run-all", help="Prepare and execute one complete workflow." ) _add_workflow_arguments(run_all_parser) + _add_open_window_argument(run_all_parser) run_parser = subparsers.add_parser( "run", help="Execute an already prepared Task Engine bundle." ) @@ -79,6 +80,7 @@ def build_parser() -> argparse.ArgumentParser: run_parser.add_argument("--num-envs", type=int, default=None) run_parser.add_argument("--dataset-saving", action="store_true") run_parser.add_argument("--show-grasp-poses", action="store_true") + _add_open_window_argument(run_parser) _add_failure_policy_argument(run_parser) return parser @@ -145,6 +147,14 @@ def _add_failure_policy_argument(parser: argparse.ArgumentParser) -> None: ) +def _add_open_window_argument(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--open-window", + action="store_true", + help="Open the native DexSim window during simulator execution.", + ) + + def main(argv: Sequence[str] | None = None) -> int: """Dispatch one Task Engine workflow command.""" parser = build_parser() @@ -218,6 +228,7 @@ def _run_workflow( dataset_saving=args.dataset_saving, failure_policy=args.failure_policy, show_grasp_poses=args.show_grasp_poses, + open_window=bool(getattr(args, "open_window", False)), run_id=allocation.run_id, created_at=allocation.created_at, execute=execute, @@ -252,6 +263,7 @@ def _run_prepared_bundle(args: argparse.Namespace) -> int: dataset_saving=bool(args.dataset_saving), failure_policy=args.failure_policy, show_grasp_poses=bool(args.show_grasp_poses), + open_window=bool(args.open_window), ) environments = report.get("environments", ()) successes = [ diff --git a/embodichain/gen_sim/task_engine/workflow.py b/embodichain/gen_sim/task_engine/workflow.py index b0ca8552a..338d5e89a 100644 --- a/embodichain/gen_sim/task_engine/workflow.py +++ b/embodichain/gen_sim/task_engine/workflow.py @@ -112,6 +112,7 @@ def __call__( dataset_saving: bool = False, failure_policy: str = "stop", show_grasp_poses: bool = False, + open_window: bool = False, ) -> Mapping[str, Any]: """Run one simulator attempt and preserve its report and trajectory. @@ -124,6 +125,7 @@ def __call__( failure_policy: Whether failed dependencies stop or permit downstream diagnostic execution. show_grasp_poses: Whether to write the valid E5 grasp pair as a PNG. + open_window: Whether to open the native DexSim execution window. Returns: Validated Action Engine execution report. @@ -134,6 +136,8 @@ def __call__( raise ValueError("failure_policy must be 'stop' or 'continue'.") if not isinstance(show_grasp_poses, bool): raise TypeError("show_grasp_poses must be a boolean.") + if not isinstance(open_window, bool): + raise TypeError("open_window must be a boolean.") attempt_root.mkdir(parents=True, exist_ok=False) command = [ sys.executable, @@ -146,8 +150,9 @@ def __call__( str(num_envs), "--seed", str(seed), - "--headless", ] + if not open_window: + command.append("--headless") if not dataset_saving: command.append("--filter_dataset_saving") if show_grasp_poses: @@ -157,7 +162,8 @@ def __call__( print( "[Task Engine] Starting " f"{attempt_root.name}: seed={seed}, num_envs={num_envs}, " - f"dataset_saving={dataset_saving}, failure_policy={failure_policy}", + f"dataset_saving={dataset_saving}, open_window={open_window}, " + f"failure_policy={failure_policy}", flush=True, ) completed = _run_streaming_process(command, log_path) @@ -307,6 +313,7 @@ def run( dataset_saving: bool = False, failure_policy: str = "stop", show_grasp_poses: bool = False, + open_window: bool = False, run_id: str | None = None, created_at: datetime | None = None, overwrite: bool = False, @@ -327,6 +334,7 @@ def run( failure_policy: Whether failed dependencies stop or permit downstream diagnostic execution. show_grasp_poses: Whether to write the valid E5 grasp pair as a PNG. + open_window: Whether simulator attempts open the native DexSim window. run_id: Optional externally allocated run identifier. created_at: Optional timezone-aware run creation timestamp. overwrite: Whether to atomically replace an existing run directory. @@ -340,6 +348,8 @@ def run( raise TypeError("dataset_saving must be a boolean.") if not isinstance(show_grasp_poses, bool): raise TypeError("show_grasp_poses must be a boolean.") + if not isinstance(open_window, bool): + raise TypeError("open_window must be a boolean.") if failure_policy not in {"stop", "continue"}: raise ValueError("failure_policy must be 'stop' or 'continue'.") if workflow_cfg is None or planning_cfg is None or execution_cfg is None: @@ -902,6 +912,8 @@ def run( } if show_grasp_poses: execution_options["show_grasp_poses"] = True + if open_window: + execution_options["open_window"] = True report = self.action_executor( preparation.output_dir, action_root, diff --git a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py index 1cfa88de0..0b4570aa7 100644 --- a/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py +++ b/tests/gen_sim/task_engine/orchestration/test_coordinator_cli.py @@ -1121,7 +1121,7 @@ def test_public_cli_exposes_prepare_run_and_run_all_modes() -> None: assert arguments.show_grasp_poses is False -def test_run_all_cli_accepts_grasp_pose_visualization() -> None: +def test_run_all_cli_accepts_execution_visualization_options() -> None: arguments = cli.build_parser().parse_args( [ "run-all", @@ -1136,10 +1136,59 @@ def test_run_all_cli_accepts_grasp_pose_visualization() -> None: "--output-root", "history", "--show-grasp-poses", + "--open-window", ] ) assert arguments.show_grasp_poses is True + assert arguments.open_window is True + + +def test_run_all_cli_forwards_open_window( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = {} + + class FakeWorkflow: + def __init__(self, **_kwargs) -> None: + pass + + def run(self, request, **kwargs): + captured.update(kwargs) + output = Path(request["output_dir"]) + return SimpleNamespace( + status="succeeded", + succeeded=True, + failure_class=None, + output_dir=output, + manifest_path=output / "run_manifest.json", + final_bundle=output / "final" / "bundle", + ) + + monkeypatch.setattr(cli, "SceneAdapter", lambda **_kwargs: object()) + monkeypatch.setattr(cli, "TaskEngineWorkflow", FakeWorkflow) + + result = cli.main( + [ + "run-all", + "--mode", + "image", + "--task-id", + "task", + "--instruction", + "place the cup", + "--image", + str(tmp_path / "input.png"), + "--output-root", + str(tmp_path / "history"), + "--open-window", + ] + ) + + assert result == 0 + assert captured["execute"] is True + assert captured["open_window"] is True def test_prepare_cli_stops_before_simulator_execution( @@ -1199,6 +1248,7 @@ def __call__(self, _bundle, output, **kwargs): Path(output).mkdir() assert kwargs["num_envs"] == 2 assert kwargs["failure_policy"] == "continue" + assert kwargs["open_window"] is True return { "status": "failed", "environments": [ @@ -1220,6 +1270,7 @@ def __call__(self, _bundle, output, **kwargs): "2", "--failure-policy", "continue", + "--open-window", ] ) diff --git a/tests/gen_sim/task_engine/test_parallel_workflow.py b/tests/gen_sim/task_engine/test_parallel_workflow.py index f073e3685..80569418a 100644 --- a/tests/gen_sim/task_engine/test_parallel_workflow.py +++ b/tests/gen_sim/task_engine/test_parallel_workflow.py @@ -279,10 +279,12 @@ def __init__( *, expected_dataset_saving: bool = False, expected_show_grasp_poses: bool = False, + expected_open_window: bool = False, ) -> None: self.successes = successes self.expected_dataset_saving = expected_dataset_saving self.expected_show_grasp_poses = expected_show_grasp_poses + self.expected_open_window = expected_open_window self.calls = 0 def __call__( @@ -295,12 +297,14 @@ def __call__( dataset_saving: bool = False, failure_policy: str = "stop", show_grasp_poses: bool = False, + open_window: bool = False, ): values = self.successes[min(self.calls, len(self.successes) - 1)] self.calls += 1 assert len(values) == num_envs assert dataset_saving is self.expected_dataset_saving assert show_grasp_poses is self.expected_show_grasp_poses + assert open_window is self.expected_open_window assert failure_policy == "stop" return { "status": "succeeded" if all(values) else "failed", @@ -338,6 +342,30 @@ def test_parallel_workflow_propagates_grasp_pose_visualization( assert result.succeeded +def test_parallel_workflow_propagates_open_window(tmp_path: Path) -> None: + candidates = _candidate_set() + workflow = TaskEngineWorkflow( + task_agent=_TaskAgent(candidates), + scene_backend=_SceneBackend(_selection(candidates)), + action_agent=_ActionAgent(), + coordinator=_Coordinator(["bound"]), + action_executor=_Executor( + [[True]], + expected_open_window=True, + ), + ) + + result = workflow.run( + _request(tmp_path), + workflow_cfg=TaskEngineWorkflowCfg(), + planning_cfg=TaskEnginePlanningCfg(), + execution_cfg=TaskEngineExecutionCfg(num_envs=1), + open_window=True, + ) + + assert result.succeeded + + @pytest.mark.parametrize("existing", [False, True]) @pytest.mark.parametrize("edit", [False, True]) def test_parallel_workflow_supports_all_four_scene_inputs( @@ -576,11 +604,17 @@ def test_final_candidate_rebinding_updates_attempt_unbound_audit( ("dataset_saving", "expects_filter"), [(False, True), (True, False)], ) -def test_subprocess_executor_controls_dataset_saving_and_copies_trajectory( +@pytest.mark.parametrize( + ("open_window", "expects_headless"), + [(False, True), (True, False)], +) +def test_subprocess_executor_controls_launch_options_and_copies_trajectory( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, dataset_saving: bool, expects_filter: bool, + open_window: bool, + expects_headless: bool, ) -> None: bundle = tmp_path / "bundle" bundle.mkdir() @@ -640,6 +674,7 @@ def fake_run(command, log_path): dataset_saving=dataset_saving, failure_policy="continue", show_grasp_poses=True, + open_window=open_window, ) assert report["status"] == "succeeded" @@ -652,6 +687,7 @@ def fake_run(command, log_path): assert " prepare" not in " ".join(captured["command"]) assert " workflow" not in " ".join(captured["command"]) assert ("--filter_dataset_saving" in captured["command"]) is expects_filter + assert ("--headless" in captured["command"]) is expects_headless assert "--show-grasp-poses" in captured["command"] assert captured["command"][-2:] == ["--failure-policy", "continue"] assert captured["log_path"] == attempt / "action.log"