# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""Case construction contract for Atomic Skill benchmark integrations."""
from __future__ import annotations
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import TYPE_CHECKING
if TYPE_CHECKING:
import torch
from embodichain.lab.sim.atomic_actions import (
ActionInvocation,
SceneSnapshot,
TaskState,
)
from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory
from .config import SuiteCfg, TrackCfg
from .contracts import (
ExecutionObservation,
PhysicalStageEvaluator,
ReplayObservationWindow,
)
from .models import BenchmarkCase, CaseOutcome
from .planners.base import PlannerAdapter
from .scenarios.atomic_skill import AtomicSkillScenario
__all__ = [
"AtomicSkillCaseProvider",
"register_atomic_skill_provider",
"create_atomic_skill_provider",
"atomic_skill_provider_names",
]
[docs]
class AtomicSkillCaseProvider(ABC):
"""Generate and ground one Atomic Action without planner-specific logic."""
skill_id: str
requires_gripper = False
samples_grasps = False
"""Whether the skill samples grasps itself and needs a grasp generator."""
[docs]
@abstractmethod
def generate_case(
self,
scenario: "AtomicSkillScenario",
suite: SuiteCfg,
track: TrackCfg,
config: Mapping[str, object],
*,
seed: int,
batch_size: int,
) -> BenchmarkCase:
"""Generate one frozen case and independent IK validity evidence."""
[docs]
@abstractmethod
def build_invocation(
self,
scenario: "AtomicSkillScenario",
case: BenchmarkCase,
adapter: "PlannerAdapter",
) -> ActionInvocation:
"""Ground the case into one planner-independent action invocation."""
[docs]
def object_id(self, case: BenchmarkCase) -> str | None:
"""Return the manipulated object identifier when one exists."""
return case.object_id
[docs]
def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None:
"""Return the first lift waypoint that should release object dynamics."""
return None
[docs]
def articulation_effect_segment(self, case: BenchmarkCase) -> str | None:
"""Return the segment that starts the articulation-displacement diagnostic."""
del case
return None
[docs]
def initial_task_state(
self, scenario: "AtomicSkillScenario", case: BenchmarkCase
) -> TaskState | None:
"""Return an optional symbolic precondition for isolated action testing."""
del scenario, case
return None
[docs]
def prefix_invocations(
self, scenario: "AtomicSkillScenario", case: BenchmarkCase
) -> tuple[ActionInvocation, ...]:
"""Return actions compiled before this skill to establish its precondition.
A skill that needs a held object, such as Pour, returns the PickUp that
acquires it. The prefix is replayed and scored with the skill.
Returns:
Invocations compiled in order before :meth:`build_invocation`.
"""
del scenario, case
return ()
[docs]
def initial_scene(
self, scenario: "AtomicSkillScenario", case: BenchmarkCase
) -> SceneSnapshot | None:
"""Return explicit scene observations the skill plans against.
Returns:
A snapshot that replaces the engine's rigid-object observations, or
``None`` to keep them.
"""
del scenario, case
return None
[docs]
def frozen_grasps(
self, scenario: "AtomicSkillScenario", case: BenchmarkCase
) -> Mapping[str, tuple[torch.Tensor, torch.Tensor]]:
"""Return the grasps a self-sampling skill must use, by gripper part.
Skills that sample their own grasps, such as OpenDoor or HandOver,
receive these manifest grasps instead, so every planner plans the same
contacts.
Returns:
``{gripper_part: (target_pose, grasp_pose)}`` batches; empty when
the skill accepts explicit grasps or needs none.
"""
del scenario, case
return {}
[docs]
def object_rotation_segment(
self, compiled: CompiledTrajectory
) -> tuple[int, int] | None:
"""Return the waypoint window whose object rotation is measured.
Returns:
``(start, stop)`` trajectory indices, or ``None`` when the skill is
not scored on how far it rotates the held object.
"""
del compiled
return None
[docs]
def replay_observation_windows(
self, compiled: CompiledTrajectory, case: BenchmarkCase
) -> tuple[ReplayObservationWindow, ...]:
"""Declare raw simulator traces required by the physical evaluator.
Providers resolve semantic action segments here. The replay adapter
stores only measured tensors, so the evaluator remains independent of
compiled planning state and the live simulator.
"""
del compiled, case
return ()
[docs]
def physical_evaluator(self) -> PhysicalStageEvaluator | None:
"""Return a measured-stage evaluator, or delegate to ``task_result``.
Returns:
An evaluator that receives no compiled planning state.
"""
return None
[docs]
def task_result(
self,
scenario: "AtomicSkillScenario",
case: BenchmarkCase,
compiled: CompiledTrajectory,
observation: ExecutionObservation,
motion_outcomes: tuple[CaseOutcome, ...],
) -> tuple[torch.Tensor, str]:
"""Evaluate a custom provider that does not supply a stage evaluator."""
raise NotImplementedError("Provide physical_evaluator() or task_result().")
AtomicSkillProviderType = type[AtomicSkillCaseProvider]
_ATOMIC_SKILL_PROVIDERS: dict[str, AtomicSkillProviderType] = {}
[docs]
def register_atomic_skill_provider(
skill_id: str, provider_type: AtomicSkillProviderType
) -> None:
"""Register one Atomic Action case provider."""
if not skill_id:
raise ValueError("Atomic skill id must not be empty.")
previous = _ATOMIC_SKILL_PROVIDERS.get(skill_id)
if previous is not None and previous is not provider_type:
raise ValueError(f"Atomic skill provider {skill_id!r} is already registered.")
_ATOMIC_SKILL_PROVIDERS[skill_id] = provider_type
[docs]
def atomic_skill_provider_names() -> tuple[str, ...]:
"""Return registered Atomic Action case providers."""
return tuple(sorted(_ATOMIC_SKILL_PROVIDERS))
[docs]
def create_atomic_skill_provider(skill_id: str) -> AtomicSkillCaseProvider:
"""Construct one registered Atomic Action case provider."""
try:
provider_type = _ATOMIC_SKILL_PROVIDERS[skill_id]
except KeyError as exc:
raise ValueError(
f"Unknown atomic skill {skill_id!r}; registered skills: "
f"{atomic_skill_provider_names()}."
) from exc
return provider_type()