Source code for embodichain_benchmarks.motion_generation.artifacts

# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------

"""Reproducibility and raw-result artifacts for planner benchmarks."""

from __future__ import annotations

import importlib.metadata
import json
import math
import platform
import subprocess
from collections.abc import Mapping
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any

import torch
import yaml

from .config import SuiteCfg, stable_hash, suite_to_dict
from .contracts import EvaluationDomain
from .models import BenchmarkCase, TrialRecord

if TYPE_CHECKING:
    from .embodiments import EmbodimentResolution

__all__ = [
    "TrialJsonlWriter",
    "create_run_directory",
    "environment_metadata",
    "write_case_manifest",
    "write_json",
    "write_resolved_suite",
]


[docs] def create_run_directory(output_root: str | Path, suite_name: str) -> Path: """Create one timestamped benchmark run directory.""" timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S_%f") run_dir = Path(output_root) / suite_name / timestamp run_dir.mkdir(parents=True, exist_ok=False) return run_dir
def _package_version(package: str) -> str | None: """Return an installed distribution version when available.""" try: return importlib.metadata.version(package) except importlib.metadata.PackageNotFoundError: return None def _git_commit() -> str | None: """Return the framework checkout revision, independent of caller cwd.""" repository = Path(__file__).resolve().parents[3] if not (repository / ".git").exists(): return None try: result = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=repository, check=False, capture_output=True, text=True, ) except FileNotFoundError: return None return result.stdout.strip() or None
[docs] def environment_metadata() -> dict[str, object]: """Collect software and hardware metadata needed to interpret results.""" gpu_name = None if torch.cuda.is_available(): gpu_name = torch.cuda.get_device_name(torch.cuda.current_device()) return { "generated_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "git_commit": _git_commit(), "platform": platform.platform(), "python": platform.python_version(), "processor": platform.processor() or None, "torch": torch.__version__, "cuda_runtime": torch.version.cuda, "cuda_available": torch.cuda.is_available(), "gpu": gpu_name, "curobo": _package_version("nvidia-curobo"), "dexsim_engine": _package_version("dexsim-engine"), "embodichain": _package_version("embodichain"), }
[docs] def write_json(path: str | Path, value: object) -> Path: """Write one UTF-8 JSON artifact with stable formatting.""" output = Path(path) output.parent.mkdir(parents=True, exist_ok=True) output.write_text( json.dumps( _to_json_value(value), indent=2, sort_keys=True, ensure_ascii=False, default=str, allow_nan=False, ) + "\n", encoding="utf-8", ) return output
[docs] def write_resolved_suite(path: str | Path, suite: SuiteCfg) -> Path: """Write the fully resolved suite as YAML.""" output = Path(path) output.parent.mkdir(parents=True, exist_ok=True) output.write_text( yaml.safe_dump(suite_to_dict(suite), sort_keys=False), encoding="utf-8" ) return output
def _case_to_dict(case: BenchmarkCase) -> dict[str, Any]: """Serialize a fixed case without losing tensor numeric values.""" return { "suite_version": case.suite_version, "track": case.track, "scenario_id": case.scenario_id, "case_id": case.case_id, "seed": case.seed, "batch_size": case.batch_size, "num_waypoints": case.num_waypoints, "path_shape": case.path_shape, "start_state_bin": case.start_state_bin, "robot_id": case.robot_id, "skill_id": case.skill_id, "object_id": case.object_id, "task_difficulty": case.task_difficulty, "primary_success": case.primary_success, "start_qpos": case.start_qpos.detach().cpu().tolist(), "full_start_qpos": ( None if case.full_start_qpos is None else case.full_start_qpos.detach().cpu().tolist() ), "target_waypoints": case.target_waypoints.detach().cpu().tolist(), "case_parameters": _to_json_value(case.case_parameters), "domain": _domain_to_dict(case.domain), "track_group": case.track_group, "embodiment_id": case.embodiment_id, "required_capabilities": sorted(case.required_capabilities), "validity_evidence": { "method": ( case.case_parameters.get("reachability_evidence", "reference_qpos_fk") if case.skill_id == "N/A" else ( "joint_limits_and_reference_fk" if case.skill_id == "move_joints" else "independent_sequential_ik" ) ), "reference_qpos": case.reference_qpos.detach().cpu().tolist(), }, } def _domain_to_dict(domain: object) -> dict[str, object] | None: """Serialize a domain identity without leaking set representations.""" if domain is None: return None values = asdict(domain) values["required_capabilities"] = sorted(values.get("required_capabilities", ())) if not values["required_capabilities"]: values.pop("required_capabilities") return values def _to_json_value(value: object) -> object: """Recursively preserve tensors and numeric case configuration values.""" if isinstance(value, float) and not math.isfinite(value): return None if isinstance(value, torch.Tensor): return _to_json_value(value.detach().cpu().tolist()) if isinstance(value, dict): return {str(key): _to_json_value(item) for key, item in value.items()} if isinstance(value, (set, frozenset)): return sorted(_to_json_value(item) for item in value) if isinstance(value, (list, tuple)): return [_to_json_value(item) for item in value] return value
[docs] def write_case_manifest( path: str | Path, cases: list[BenchmarkCase], *, suite: SuiteCfg | None = None, planner_config_hashes: Mapping[str, str] | None = None, embodiment_resolution: "EmbodimentResolution | None" = None, ) -> Path: """Write the algorithm-independent case manifest.""" payload: dict[str, object] = { "case_schema_version": 3, "cases": [_case_to_dict(case) for case in cases], } payload["cases_sha256"] = stable_hash(payload["cases"]) if suite is not None: payload["resolved_suite_version"] = suite.suite_version payload["motion_evaluation_protocol"] = { "version": "start_boundary_v1", "start_state_tolerance_rad": suite.protocol.start_state_tolerance_rad, "collision_validation": "not_evaluated", "path_validation": "ordered_waypoint_arrival", "reference_path_metric": "sampled_translation_to_waypoint_polyline_v1", } if embodiment_resolution is not None: payload["resolved_embodiment"] = embodiment_resolution.to_metadata() elif suite.embodiment is None: payload["resolved_embodiment"] = { "id": suite.robot.id, "provider": suite.robot.provider, } else: payload["resolved_embodiment"] = { "component": suite.embodiment.component, "overrides": _to_json_value(suite.embodiment.overrides), } if planner_config_hashes: payload["planner_config_hashes"] = dict(sorted(planner_config_hashes.items())) return write_json( path, payload, )
[docs] class TrialJsonlWriter: """Append numeric raw trial records to one JSONL artifact."""
[docs] def __init__(self, path: str | Path) -> None: self.path = Path(path) self.path.parent.mkdir(parents=True, exist_ok=True) self.path.write_text("", encoding="utf-8")
[docs] def append(self, record: TrialRecord) -> None: """Append one record and flush it immediately for recoverability.""" with self.path.open("a", encoding="utf-8") as file: file.write( json.dumps( record.to_dict(), ensure_ascii=False, default=str, allow_nan=False, ) + "\n" )
def _read_case_manifest( path: str | Path, device: torch.device, *, expected_embodiment: EmbodimentResolution, ) -> list[BenchmarkCase]: """Validate and reconstruct a frozen algorithm-independent population.""" payload = json.loads(Path(path).read_text(encoding="utf-8")) if payload.get("case_schema_version") != 3: raise ValueError("Only case_schema_version 3 manifests can be replayed.") rows = payload.get("cases") if not isinstance(rows, list) or not rows: raise ValueError("A replay manifest must contain a nonempty cases list.") if payload.get("cases_sha256") != stable_hash(rows): raise ValueError("Frozen case manifest hash mismatch.") frozen_embodiment = payload.get("resolved_embodiment") expected = expected_embodiment.to_metadata() if not isinstance(frozen_embodiment, dict) or any( frozen_embodiment.get(name) != expected[name] for name in ("embodiment_id", "config_hash", "source") ): raise ValueError( "Frozen case embodiment is missing or differs from the resolved suite embodiment." ) cases = [] ids = set() for row in rows: values = dict(row) evidence = values.pop("validity_evidence") values["reference_qpos"] = evidence["reference_qpos"] for field in ( "start_qpos", "full_start_qpos", "target_waypoints", "reference_qpos", ): if values[field] is not None: values[field] = torch.tensor( values[field], device=device, dtype=torch.float32 ) if not torch.isfinite(values[field]).all(): raise ValueError(f"Frozen {field} must be finite.") if values["domain"] is not None: domain = dict(values["domain"]) domain["required_capabilities"] = frozenset( domain.get("required_capabilities", ()) ) values["domain"] = EvaluationDomain(**domain) values["required_capabilities"] = frozenset(values["required_capabilities"]) case = BenchmarkCase(**values) b, k = case.batch_size, case.num_waypoints if ( case.case_id in ids or b < 1 or k < 1 or case.start_qpos.ndim != 2 or case.start_qpos.shape[0] != b or case.target_waypoints.shape != (b, k, 4, 4) or case.reference_qpos.shape != (b, k, case.start_qpos.shape[-1]) ): raise ValueError("Frozen case identity or tensor shape is invalid.") ids.add(case.case_id) cases.append(case) return cases