# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""Episode-level combined expansion contracts.
The contracts in this module deliberately do not depend on Gym, a simulator,
or a renderer. They provide the deterministic schedule that a host runner
uses to combine reference families, affordance branches, trajectory variants,
and one visual profile per physical episode.
"""
from __future__ import annotations
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
import hashlib
import json
import math
from pathlib import Path
from typing import Any
import yaml
from embodichain.utils import configclass
__all__ = [
"CombinedExpansionProfile",
"ReferenceFamilySpec",
"CubeInitialPoseProvider",
"CycleRecipe",
"CallRecipe",
"CandidateRecipe",
"PhysicalSlotPool",
"SlotReservation",
"enumerate_candidate_recipes",
"load_visual_profile_registry",
"schedule_digest",
]
def _mapping(value: object, name: str) -> Mapping[str, Any]:
if not isinstance(value, Mapping):
raise ValueError(f"{name} must be a mapping")
return value
def _keys(data: Mapping[str, Any], allowed: set[str], name: str) -> None:
unknown = set(data) - allowed
if unknown:
raise ValueError(f"unknown fields in {name}: {sorted(unknown, key=str)}")
def _text(value: object, name: str) -> str:
if not isinstance(value, str) or not value.strip() or value != value.strip():
raise ValueError(f"{name} must be a nonempty stripped string")
return value
def _bool(value: object, name: str) -> bool:
if type(value) is not bool:
raise ValueError(f"{name} must be a boolean")
return value
def _int(value: object, name: str, minimum: int = 0) -> int:
if type(value) is not int or value < minimum:
raise ValueError(f"{name} must be an integer >= {minimum}")
return value
def _finite(value: object, name: str, *, positive: bool = False) -> float:
if type(value) not in (int, float) or not math.isfinite(value):
raise ValueError(f"{name} must be finite")
result = float(value)
if positive and result <= 0:
raise ValueError(f"{name} must be positive")
return result
def _strings(value: object, name: str, *, nonempty: bool = True) -> tuple[str, ...]:
if not isinstance(value, (list, tuple)):
raise ValueError(f"{name} must be a sequence")
result = tuple(_text(item, name) for item in value)
if nonempty and not result:
raise ValueError(f"{name} must be nonempty")
if len(set(result)) != len(result):
raise ValueError(f"{name} must contain unique values")
return result
def _validate_visual_operation(operation: Mapping[str, Any], name: str) -> None:
"""Validate the scope and bounds of one visual profile operation."""
scope = operation.get("scope")
if scope == "per_environment":
if operation.get("kind") == "global_sun":
raise ValueError("global_sun must use scope 'global'")
return
if scope != "global" or operation.get("kind") != "global_sun":
raise ValueError(f"{name} requests unsupported global state")
if operation.get("light_uid") != "main_light":
raise ValueError(f"{name}.light_uid must be 'main_light'")
color = operation.get("color")
if not isinstance(color, (list, tuple)) or len(color) != 3:
raise ValueError(f"{name}.color must contain exactly three values")
for channel in color:
value = _finite(channel, f"{name}.color value")
if not 0 <= value <= 1:
raise ValueError(f"{name}.color values must be in [0, 1]")
intensity = _finite(operation.get("intensity"), f"{name}.intensity")
if not 0 <= intensity <= 10:
raise ValueError(f"{name}.intensity must be in [0, 10]")
@configclass
class _CombinedSourceCfg:
kind: str = "task_program"
source_id: str = "repeated_cube_pick_place"
source_revision: str = "config:repeated_pick_place_v1"
unit_scope: str = "episode"
template_id: str = "repeated_pick_place_episode"
phase_permissions: dict[str, tuple[str, ...]] = {}
phase_kinds: dict[str, str] = {}
call_count: int = 6
affordance_call_indices: tuple[int, ...] = (0, 2, 4)
trajectory_call_indices: tuple[int, ...] = (1, 3, 5)
def __post_init__(self) -> None:
if self.kind != "task_program":
raise ValueError("combined source.kind must be task_program")
for name in ("source_id", "source_revision", "template_id"):
_text(getattr(self, name), f"source.{name}")
if self.unit_scope != "episode":
raise ValueError("combined source.unit_scope must be episode")
permissions = {
phase_id: tuple(operators)
for phase_id, operators in self.phase_permissions.items()
}
kinds = dict(self.phase_kinds)
if set(permissions) != set(kinds):
raise ValueError(
"source.phase_permissions and source.phase_kinds must name "
"the same phases"
)
for phase_id, operators in permissions.items():
_text(phase_id, "source phase")
if len(set(operators)) != len(operators):
raise ValueError(f"source.phase_permissions.{phase_id} must be unique")
if any(
operator not in ("joint_residual", "via_points")
for operator in operators
):
raise ValueError(
f"source.phase_permissions.{phase_id} contains an unsupported operator"
)
for phase_id, kind in kinds.items():
_text(phase_id, "source phase")
if kind not in ("free", "contact", "hold"):
raise ValueError(
f"source.phase_kinds.{phase_id} must be free, contact, or hold"
)
self.phase_permissions = permissions
self.phase_kinds = kinds
_int(self.call_count, "source.call_count", 1)
self.affordance_call_indices = tuple(self.affordance_call_indices)
self.trajectory_call_indices = tuple(self.trajectory_call_indices)
for field_name, indices in (
("source.affordance_call_indices", self.affordance_call_indices),
("source.trajectory_call_indices", self.trajectory_call_indices),
):
if len(set(indices)) != len(indices) or any(
type(index) is not int or not 0 <= index < self.call_count
for index in indices
):
raise ValueError(f"{field_name} must contain unique call indices")
@configclass
class _SceneRandomizationCfg:
enabled: bool = True
reference_family_count: int = 4
profile: str = "cube_initial_pose"
profile_file: str = (
"extensions/embodichain_tasks/configs/components/randomization_profiles/"
"cube_initial_pose.yaml"
)
change_scope: str = "cube_initial_pose_only"
keep_drop_targets: bool = True
def __post_init__(self) -> None:
_bool(self.enabled, "scene_randomization.enabled")
_int(
self.reference_family_count, "scene_randomization.reference_family_count", 1
)
_text(self.profile, "scene_randomization.profile")
_text(self.profile_file, "scene_randomization.profile_file")
if self.enabled and self.change_scope != "cube_initial_pose_only":
raise ValueError(
"scene_randomization.change_scope must be cube_initial_pose_only"
)
if self.enabled and not self.keep_drop_targets:
raise ValueError("combined expansion must keep authored drop targets")
@configclass
class _AffordanceCfg:
enabled: bool = True
branches_per_family: int = 4
max_proposals: int = 128
fallback: str = "next_surviving_branch"
def __post_init__(self) -> None:
_bool(self.enabled, "affordance.enabled")
_int(self.branches_per_family, "affordance.branches_per_family", 1)
_int(self.max_proposals, "affordance.max_proposals", 1)
if self.fallback != "next_surviving_branch":
raise ValueError("affordance.fallback must be next_surviving_branch")
@configclass
class _SpatialCfg:
enabled: bool = True
operators: tuple[str, ...] = ("joint_residual", "via_points")
joint_offset_scale: float = 0.05
dedup_tolerance: float | None = None
via_count: int = 1
def __post_init__(self) -> None:
self.operators = _strings(self.operators, "trajectory.spatial.operators")
if any(item not in ("joint_residual", "via_points") for item in self.operators):
raise ValueError(
"trajectory.spatial.operators contains an unsupported operator"
)
_bool(self.enabled, "trajectory.spatial.enabled")
self.joint_offset_scale = _finite(
self.joint_offset_scale, "trajectory.spatial.joint_offset_scale"
)
if not 0 <= self.joint_offset_scale <= 1:
raise ValueError("trajectory.spatial.joint_offset_scale must be in [0, 1]")
if self.dedup_tolerance is not None:
self.dedup_tolerance = _finite(
self.dedup_tolerance,
"trajectory.spatial.dedup_tolerance",
positive=True,
)
# The tolerance decides which proposals count as genuinely new. At
# or above the offset scale every proposal measures as a duplicate
# of the reference, so the configured amplitude would be discarded
# and only the nominal variant would survive.
if self.enabled and self.dedup_tolerance >= self.joint_offset_scale:
raise ValueError(
"trajectory.spatial.dedup_tolerance must stay below "
"joint_offset_scale"
)
_int(self.via_count, "trajectory.spatial.via_count", 1)
@configclass
class _TimingCfg:
enabled: bool = False
def __post_init__(self) -> None:
_bool(self.enabled, "trajectory.timing.enabled")
@configclass
class _TrajectoryCfg:
variants_per_family: int = 4
spatial: _SpatialCfg = _SpatialCfg()
timing: _TimingCfg = _TimingCfg()
def __post_init__(self) -> None:
_int(self.variants_per_family, "trajectory.variants_per_family", 1)
@configclass
class _VisualCfg:
enabled: bool = True
profile_file: str = (
"extensions/embodichain_tasks/configs/components/randomization_profiles/rgb_visual.yaml"
)
profiles: tuple[str, ...] = (
"rgb_canonical",
"rgb_material_01",
"rgb_camera_01",
"rgb_light_01",
)
apply_per_rollout: bool = True
balance: str = "round_robin"
restore_before_next_candidate: bool = True
def __post_init__(self) -> None:
self.profile_file = _text(self.profile_file, "visual.profile_file")
self.profiles = _strings(
self.profiles,
"visual.profiles",
nonempty=self.enabled,
)
_bool(self.enabled, "visual.enabled")
_bool(self.apply_per_rollout, "visual.apply_per_rollout")
_bool(
self.restore_before_next_candidate, "visual.restore_before_next_candidate"
)
if self.enabled and not self.apply_per_rollout:
raise ValueError("enabled visual profiles must apply_per_rollout")
if self.balance != "round_robin":
raise ValueError("visual.balance must be round_robin")
@configclass
class _ObservationCfg:
profiles: tuple[str, ...] = ("rgb_lowres",)
capture_rgb: bool = True
capture_video: bool = False
def __post_init__(self) -> None:
self.profiles = _strings(self.profiles, "observation.profiles", nonempty=False)
_bool(self.capture_rgb, "observation.capture_rgb")
_bool(self.capture_video, "observation.capture_video")
@configclass
class _SchedulingCfg:
policy: str = "round_robin_fifo"
candidate_budget: int = 64
reference_family_budget: int = 4
def __post_init__(self) -> None:
if self.policy != "round_robin_fifo":
raise ValueError("combined scheduling.policy must be round_robin_fifo")
for name in (
"candidate_budget",
"reference_family_budget",
):
_int(getattr(self, name), f"scheduling.{name}", 1)
@configclass
class _ExecutionCfg:
max_inflight: int = 16
compatibility: str = "strict"
def __post_init__(self) -> None:
_int(self.max_inflight, "execution.max_inflight", 1)
if self.compatibility != "strict":
raise ValueError("execution.compatibility must be strict")
@configclass
class _PersistenceCfg:
sink: str = "local_artifact"
output_dir: str = "/tmp/repeated-pick-place-expansion"
def __post_init__(self) -> None:
if self.sink not in {"local_artifact", "lerobot"}:
raise ValueError(
"combined persistence.sink must be local_artifact or lerobot"
)
self.output_dir = _text(self.output_dir, "persistence.output_dir")
def _construct(cls: type, data: object, name: str, allowed: set[str]) -> Any:
mapping = _mapping(data, name)
_keys(mapping, allowed, name)
values = dict(mapping)
return cls(**values)
[docs]
@configclass
class CombinedExpansionProfile:
"""Strict episode-level profile for the combined showcase."""
schema_version: int = 1
source: _CombinedSourceCfg = _CombinedSourceCfg()
scene_randomization: _SceneRandomizationCfg = _SceneRandomizationCfg()
affordance: _AffordanceCfg = _AffordanceCfg()
trajectory: _TrajectoryCfg = _TrajectoryCfg()
visual: _VisualCfg = _VisualCfg()
observation: _ObservationCfg = _ObservationCfg()
scheduling: _SchedulingCfg = _SchedulingCfg()
execution: _ExecutionCfg = _ExecutionCfg()
persistence: _PersistenceCfg = _PersistenceCfg()
def __post_init__(self) -> None:
if self.schema_version != 1:
raise ValueError("combined schema_version must be 1")
m = self.scene_randomization.reference_family_count
a = self.affordance.branches_per_family
t = self.trajectory.variants_per_family
if self.scheduling.candidate_budget < m * a * t:
raise ValueError("scheduling.candidate_budget must cover m × a × t recipes")
if self.scheduling.reference_family_budget < m:
raise ValueError(
"scheduling.reference_family_budget must cover all families"
)
if self.visual.enabled and not self.visual.profiles:
raise ValueError("enabled visual expansion requires profiles")
[docs]
@classmethod
def from_mapping(cls, data: Mapping[str, Any]) -> "CombinedExpansionProfile":
"""Decode the canonical combined schema without evaluating callables."""
allowed = {
"schema_version",
"source",
"scene_randomization",
"affordance",
"trajectory",
"visual",
"observation",
"scheduling",
"execution",
"persistence",
}
_keys(data, allowed, "combined profile")
values = {
"schema_version": _int(data.get("schema_version", 1), "schema_version", 1),
"source": _construct(
_CombinedSourceCfg,
data.get("source", {}),
"source",
{
"kind",
"source_id",
"source_revision",
"unit_scope",
"template_id",
"phase_permissions",
"phase_kinds",
"call_count",
"affordance_call_indices",
"trajectory_call_indices",
},
),
"scene_randomization": _construct(
_SceneRandomizationCfg,
data.get("scene_randomization", {}),
"scene_randomization",
{
"enabled",
"reference_family_count",
"profile",
"profile_file",
"change_scope",
"keep_drop_targets",
},
),
"affordance": _construct(
_AffordanceCfg,
data.get("affordance", {}),
"affordance",
{"enabled", "branches_per_family", "max_proposals", "fallback"},
),
"visual": _construct(
_VisualCfg,
data.get("visual", {}),
"visual",
{
"enabled",
"profile_file",
"profiles",
"apply_per_rollout",
"balance",
"restore_before_next_candidate",
},
),
"observation": _construct(
_ObservationCfg,
data.get("observation", {}),
"observation",
{"profiles", "capture_rgb", "capture_video"},
),
"scheduling": _construct(
_SchedulingCfg,
data.get("scheduling", {}),
"scheduling",
{
"policy",
"candidate_budget",
"reference_family_budget",
},
),
"execution": _construct(
_ExecutionCfg,
data.get("execution", {}),
"execution",
{
"max_inflight",
"compatibility",
},
),
"persistence": _construct(
_PersistenceCfg,
data.get("persistence", {}),
"persistence",
{
"sink",
"output_dir",
},
),
}
trajectory = _mapping(data.get("trajectory", {}), "trajectory")
_keys(trajectory, {"variants_per_family", "spatial", "timing"}, "trajectory")
spatial = _construct(
_SpatialCfg,
trajectory.get("spatial", {}),
"trajectory.spatial",
{
"enabled",
"operators",
"joint_offset_scale",
"dedup_tolerance",
"via_count",
},
)
timing = _construct(
_TimingCfg, trajectory.get("timing", {}), "trajectory.timing", {"enabled"}
)
values["trajectory"] = _TrajectoryCfg(
variants_per_family=_int(
trajectory.get("variants_per_family", 4),
"trajectory.variants_per_family",
1,
),
spatial=spatial,
timing=timing,
)
return cls(**values)
[docs]
def validate_for_num_envs(self, num_envs: int) -> None:
"""Validate the host-owned slot capacity before simulator construction."""
_int(num_envs, "num_envs", 1)
if self.execution.max_inflight != num_envs:
raise ValueError(
"execution.max_inflight must equal the resolved num_envs; "
"expansion executes one vectorized batch"
)
[docs]
def validate_registries(
self,
*,
visual_profile_ids: Sequence[str],
observation_profile_ids: Sequence[str],
) -> None:
"""Validate that host registries contain each requested profile once."""
visual = _strings(visual_profile_ids, "visual_profile_ids")
observations = _strings(observation_profile_ids, "observation_profile_ids")
if self.visual.enabled and set(visual) != set(self.visual.profiles):
raise ValueError(
"visual profile registry does not match the combined profile"
)
if set(observations) != set(self.observation.profiles):
raise ValueError(
"observation profile registry does not match the combined profile"
)
def _read_visual_profile_registry(path: str | Path) -> Mapping[str, Any]:
"""Read and validate the canonical visual profile registry mapping.
The deployment permits one explicitly declared global sun operation. It is
applied to the whole simulator and therefore must be shared by every row
in a batch; other unsupported global renderer mutations are rejected.
"""
with Path(path).open(encoding="utf-8") as stream:
data = yaml.safe_load(stream)
root = _mapping(data, "visual profile registry")
_keys(root, {"profile_id", "revision", "profiles"}, "visual profile registry")
profiles = _mapping(root.get("profiles", {}), "visual profile registry.profiles")
ids = tuple(_text(profile_id, "visual profile ID") for profile_id in profiles)
if len(set(ids)) != len(ids):
raise ValueError("visual profile IDs must be unique")
for profile_id, value in profiles.items():
profile = _mapping(value, f"visual profile {profile_id}")
_keys(profile, {"seed_offset", "operations"}, f"visual profile {profile_id}")
_int(
profile.get("seed_offset", 0), f"visual profile {profile_id}.seed_offset", 0
)
operations = profile.get("operations", [])
if not isinstance(operations, (list, tuple)) or not operations:
raise ValueError(f"visual profile {profile_id} must declare operations")
for operation in operations:
operation_mapping = _mapping(
operation, f"visual profile {profile_id}.operation"
)
_validate_visual_operation(
operation_mapping,
f"visual profile {profile_id}.operation",
)
return root
[docs]
def load_visual_profile_registry(
path: str | Path,
*,
requested_profile_ids: Sequence[str] | None = None,
) -> tuple[str, ...]:
"""Load and validate seedable visual profile IDs.
This compatibility helper exposes only IDs; full profile resolution uses
the same canonical parser through :class:`VisualProfileRegistry`.
"""
root = _read_visual_profile_registry(path)
profiles = _mapping(root["profiles"], "visual profile registry.profiles")
ids = tuple(profiles)
if requested_profile_ids is not None:
requested = _strings(requested_profile_ids, "requested_profile_ids")
if set(requested) != set(ids):
raise ValueError("requested visual profiles do not match the registry")
return ids
[docs]
@dataclass(frozen=True)
class ReferenceFamilySpec:
"""One immutable legal initial cube pose and its authored targets."""
reference_family_id: str
cube_position: tuple[float, float, float]
cube_quaternion_xyzw: tuple[float, float, float, float]
drop_targets: tuple[tuple[str, tuple[float, float, float]], ...]
deterministic_seed: int
profile_revision: str = "cube_initial_pose:v1"
scene_signature: str = "repeated_pick_place_scene_v1"
def __post_init__(self) -> None:
_text(self.reference_family_id, "reference_family_id")
if len(self.cube_position) != 3 or len(self.cube_quaternion_xyzw) != 4:
raise ValueError(
"cube pose must contain position(3) and quaternion_xyzw(4)"
)
for value in (*self.cube_position, *self.cube_quaternion_xyzw):
_finite(value, "cube pose value")
if type(self.deterministic_seed) is not int or self.deterministic_seed < 0:
raise ValueError("deterministic_seed must be non-negative")
_text(self.profile_revision, "profile_revision")
_text(self.scene_signature, "scene_signature")
names = [name for name, _ in self.drop_targets]
if len(names) != len(set(names)):
raise ValueError("drop target IDs must be unique")
[docs]
class CubeInitialPoseProvider:
"""Return deterministic authored cube pose families for a expansion job."""
[docs]
def __init__(
self,
poses: Sequence[tuple[float, float, float]] | None = None,
*,
family_ids: Sequence[str] | None = None,
profile_revision: str = "cube_initial_pose:v1",
drop_targets: Sequence[tuple[str, tuple[float, float, float]]] = (
("drop_pose_00", (-0.40, 0.48, 0.10)),
("drop_pose_01", (-0.42, -0.08, 0.10)),
),
seed: int = 0,
) -> None:
positions = tuple(
poses
or (
(-0.42, -0.08, 0.025),
(-0.38, -0.02, 0.025),
(-0.46, -0.14, 0.025),
(-0.42, 0.04, 0.025),
)
)
if not positions:
raise ValueError("at least one cube pose is required")
if type(seed) is not int or seed < 0:
raise ValueError("seed must be non-negative")
targets = tuple(
(
_text(name, "drop target ID"),
tuple(_finite(value, "drop target coordinate") for value in position),
)
for name, position in drop_targets
)
if any(len(position) != 3 for _, position in targets):
raise ValueError("drop target positions must have three coordinates")
self._positions = positions
if family_ids is None:
ids = tuple(f"cube_pose_{index:02d}" for index in range(len(positions)))
else:
ids = tuple(_text(value, "reference family ID") for value in family_ids)
if len(ids) != len(positions) or len(set(ids)) != len(ids):
raise ValueError("family_ids must align with unique authored poses")
self._family_ids = ids
self._profile_revision = _text(profile_revision, "profile_revision")
self._targets = targets
self._seed = seed
[docs]
@classmethod
def from_yaml(cls, path: str | Path, *, seed: int = 0) -> "CubeInitialPoseProvider":
"""Load authored cube families from a strict pose profile."""
with Path(path).open(encoding="utf-8") as stream:
root = yaml.safe_load(stream)
data = _mapping(root, "initial-pose profile")
_keys(
data,
{"profile_id", "revision", "change_scope", "keep_drop_targets", "poses"},
"initial-pose profile",
)
if data.get("change_scope") != "cube_initial_pose_only":
raise ValueError("initial-pose profile must change only cube_initial_pose")
if data.get("keep_drop_targets") is not True:
raise ValueError("initial-pose profile must keep authored drop targets")
poses = []
family_ids = []
for item in data.get("poses", ()):
value = _mapping(item, "initial-pose profile pose")
_keys(
value,
{"id", "position", "quaternion_xyzw"},
"initial-pose profile pose",
)
position = tuple(
_finite(number, "cube position") for number in value["position"]
)
if len(position) != 3:
raise ValueError("cube position must contain three values")
quaternion = tuple(
_finite(number, "cube quaternion")
for number in value["quaternion_xyzw"]
)
if quaternion != (0.0, 0.0, 0.0, 1.0):
raise ValueError(
"only the authored identity cube quaternion is supported"
)
poses.append(position)
family_ids.append(_text(value["id"], "reference family ID"))
return cls(
poses,
family_ids=family_ids,
profile_revision=_text(data.get("revision"), "profile revision"),
seed=seed,
)
[docs]
def enumerate(self, count: int | None = None) -> tuple[ReferenceFamilySpec, ...]:
"""Enumerate the first ``count`` legal deterministic families."""
limit = len(self._positions) if count is None else _int(count, "count", 1)
if limit > len(self._positions):
raise ValueError("initial-pose provider does not expose enough legal poses")
result = []
for index, position in enumerate(self._positions[:limit]):
if len(position) != 3:
raise ValueError("cube positions must contain three coordinates")
result.append(
ReferenceFamilySpec(
reference_family_id=self._family_ids[index],
cube_position=tuple(
_finite(value, "cube position") for value in position
),
cube_quaternion_xyzw=(0.0, 0.0, 0.0, 1.0),
drop_targets=self._targets,
deterministic_seed=self._seed + index,
profile_revision=self._profile_revision,
)
)
return tuple(result)
[docs]
@dataclass(frozen=True)
class CycleRecipe:
"""Requested affordance and trajectory ordinals for one Pick/Place cycle."""
cycle_index: int
affordance_requested: int
trajectory_requested: int
[docs]
@dataclass(frozen=True)
class CallRecipe:
"""Requested affordance and trajectory ordinals for one workflow call."""
call_index: int
affordance_requested: int
trajectory_requested: int
[docs]
@dataclass(frozen=True)
class CandidateRecipe:
"""Complete episode recipe; identity is independent of simulator row."""
recipe_index: int
candidate_id: str
reference_family_id: str
physical_geometry_family_id: str
cycle_schedule: tuple[CycleRecipe, ...]
call_schedule: tuple[CallRecipe, ...]
visual_profile_id: str
visual_seed: int
source_unit_id: str
source_revision: str
deterministic_seed_digest: str
[docs]
@dataclass(frozen=True)
class SlotReservation:
"""Opaque reservation returned by :class:`PhysicalSlotPool`."""
slot_id: int
candidate_id: str
compatibility_key: str
token: str
[docs]
class PhysicalSlotPool:
"""Generic compatible-slot reservation with exact release semantics.
A slot can advertise a strict compatibility key. Candidates remain
independent of the numeric slot they receive; releasing a stale or foreign
reservation raises rather than freeing another candidate's slot.
"""
[docs]
def __init__(
self,
slot_count: int,
*,
compatibility_keys: Mapping[int, str] | None = None,
) -> None:
self._slot_count = _int(slot_count, "slot_count", 1)
keys = {} if compatibility_keys is None else dict(compatibility_keys)
if any(
type(slot) is not int or not 0 <= slot < self._slot_count for slot in keys
):
raise ValueError("compatibility_keys contains an invalid slot")
self._keys = {
slot: _text(value, "compatibility key") for slot, value in keys.items()
}
self._reservations: dict[int, SlotReservation] = {}
@property
def capacity(self) -> int:
"""Return the number of physical slots."""
return self._slot_count
@property
def available_count(self) -> int:
"""Return the number of currently unreserved compatible rows."""
return self._slot_count - len(self._reservations)
[docs]
def reserve(self, candidate_id: str, compatibility_key: str) -> SlotReservation:
"""Reserve the first free row compatible with ``compatibility_key``."""
_text(candidate_id, "candidate_id")
_text(compatibility_key, "compatibility_key")
for slot_id in range(self._slot_count):
if slot_id in self._reservations:
continue
advertised = self._keys.get(slot_id)
if advertised is not None and advertised != compatibility_key:
continue
token = _seed_digest("slot", slot_id, candidate_id, compatibility_key)
reservation = SlotReservation(
slot_id, candidate_id, compatibility_key, token
)
self._reservations[slot_id] = reservation
return reservation
raise BufferError("no compatible physical slot is available")
[docs]
def release(self, reservation: SlotReservation) -> None:
"""Release exactly the reservation returned by :meth:`reserve`."""
if not isinstance(reservation, SlotReservation):
raise TypeError("reservation must be a SlotReservation")
current = self._reservations.get(reservation.slot_id)
if current != reservation:
raise ValueError("reservation is stale or owned by another candidate")
del self._reservations[reservation.slot_id]
[docs]
def reservation_for(self, slot_id: int) -> SlotReservation | None:
"""Inspect a row without exposing mutable pool state."""
_int(slot_id, "slot_id", 0)
if slot_id >= self._slot_count:
raise ValueError("slot_id is outside the pool")
return self._reservations.get(slot_id)
def _seed_digest(*values: object) -> str:
payload = json.dumps(values, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
[docs]
def enumerate_candidate_recipes(
profile: CombinedExpansionProfile,
*,
job_id: str = "repeated_pick_place_expansion",
families: Sequence[ReferenceFamilySpec] | None = None,
) -> tuple[CandidateRecipe, ...]:
"""Enumerate the deterministic ``m × a × t`` episode recipes."""
if not isinstance(profile, CombinedExpansionProfile):
raise TypeError("profile must be a CombinedExpansionProfile")
_text(job_id, "job_id")
family_count = profile.scene_randomization.reference_family_count
family_values = (
tuple(families)
if families is not None
else CubeInitialPoseProvider().enumerate(family_count)
)
if len(family_values) < family_count:
raise ValueError("families does not contain the requested reference count")
a_count = profile.affordance.branches_per_family
t_count = profile.trajectory.variants_per_family
affordance_calls = set(profile.source.affordance_call_indices)
trajectory_calls = set(profile.source.trajectory_call_indices)
visual_profiles = profile.visual.profiles or ("rgb_canonical",)
result: list[CandidateRecipe] = []
for family_index, family in enumerate(family_values[:family_count]):
for a0 in range(a_count):
for t0 in range(t_count):
recipe_index = family_index * (a_count * t_count) + a0 * t_count + t0
call_schedule = tuple(
CallRecipe(
call_index=call_index,
affordance_requested=(
0
if recipe_index == 0 or call_index not in affordance_calls
else (a0 + call_index // 2 + 2 * family_index) % a_count
),
trajectory_requested=(
0
if recipe_index == 0 or call_index not in trajectory_calls
else (t0 + 2 * (call_index // 2) + family_index) % t_count
),
)
for call_index in range(profile.source.call_count)
)
cycle_count = max(1, (profile.source.call_count + 1) // 2)
cycles = tuple(
CycleRecipe(
cycle_index=cycle,
affordance_requested=call_schedule[
min(2 * cycle, len(call_schedule) - 1)
].affordance_requested,
trajectory_requested=call_schedule[
min(2 * cycle + 1, len(call_schedule) - 1)
].trajectory_requested,
)
for cycle in range(cycle_count)
)
visual_index = (
0 if recipe_index == 0 else recipe_index % len(visual_profiles)
)
digest = _seed_digest(
job_id,
profile.source.source_revision,
family.reference_family_id,
family.cube_position,
family.cube_quaternion_xyzw,
family.drop_targets,
family.profile_revision,
family.scene_signature,
recipe_index,
tuple(
(c.affordance_requested, c.trajectory_requested)
for c in call_schedule
),
tuple(
(c.affordance_requested, c.trajectory_requested) for c in cycles
),
)
candidate_payload = {
"job_id": job_id,
"revision": profile.source.source_revision,
"family": family.reference_family_id,
"cube_position": family.cube_position,
"cube_quaternion_xyzw": family.cube_quaternion_xyzw,
"drop_targets": family.drop_targets,
"profile_revision": family.profile_revision,
"scene_signature": family.scene_signature,
"recipe_index": recipe_index,
"calls": [
(c.affordance_requested, c.trajectory_requested)
for c in call_schedule
],
"cycles": [
(c.affordance_requested, c.trajectory_requested) for c in cycles
],
"visual": visual_profiles[visual_index],
}
candidate_id = hashlib.sha256(
json.dumps(
candidate_payload, sort_keys=True, separators=(",", ":")
).encode()
).hexdigest()
result.append(
CandidateRecipe(
recipe_index=recipe_index,
candidate_id=f"episode_{recipe_index:04d}_{candidate_id[:16]}",
reference_family_id=family.reference_family_id,
physical_geometry_family_id=family.reference_family_id,
cycle_schedule=cycles,
call_schedule=call_schedule,
visual_profile_id=visual_profiles[visual_index],
visual_seed=int(digest[:16], 16),
source_unit_id=profile.source.source_id,
source_revision=profile.source.source_revision,
deterministic_seed_digest=digest,
)
)
return tuple(result)
[docs]
def schedule_digest(recipes: Sequence[CandidateRecipe]) -> str:
"""Hash the complete schedule in order for deterministic rerun checks."""
payload = [
{
"recipe_index": item.recipe_index,
"candidate_id": item.candidate_id,
"family": item.reference_family_id,
"cycles": [
[cycle.affordance_requested, cycle.trajectory_requested]
for cycle in item.cycle_schedule
],
"calls": [
[call.affordance_requested, call.trajectory_requested]
for call in item.call_schedule
],
"visual": item.visual_profile_id,
"seed": item.deterministic_seed_digest,
}
for item in recipes
]
return hashlib.sha256(
json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
).hexdigest()