# ----------------------------------------------------------------------------
# Copyright (c) 2021-2026 DexForce Technology Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ----------------------------------------------------------------------------
"""Production simulation assembly for Gym-backed Task Programs.
This module owns the reusable live wiring between declarative simulation
bindings and :class:`TaskProgramEnvironmentAdapter`. A task supplies one
immutable registration; this factory constructs the motion generator,
atomic-action engine, planning observation port, effect-evidence providers,
and segment-policy port without task-local motion code.
The resulting runtime is intentionally Gym-only. Its buffered command sink
must remain attached to :class:`TaskProgramDemoBridge`, which advances the shared
clock only after an ordinary ``env.step()`` consumes a yielded command. It is
therefore an internal Task Program executor, not a third user-facing action
entry point.
"""
from __future__ import annotations
from collections.abc import Callable, Iterable, Mapping
from copy import deepcopy
from dataclasses import dataclass, fields
import math
from types import MappingProxyType
from typing import Protocol, TYPE_CHECKING
import torch
from embodichain.lab.gym.envs.expert_trajectory import JointCommandMode
from embodichain.lab.sim.atomic_actions import (
AtomicActionEngine,
EntityState,
ObservedArticulationJointState,
PlanningContext,
RobotObservation,
SceneProvider,
SceneSnapshot,
TaskState,
)
from embodichain.lab.sim.motion.motion_generator import MotionGenCfg, MotionGenerator
from embodichain.lab.sim.motion.planners import (
BasePlannerCfg,
CuroboAutoGenCfg,
CuroboPlannerCfg,
CuroboWorldCfg,
NeuralPlannerCfg,
ToppraPlannerCfg,
TrapezoidalPlannerCfg,
)
from embodichain.lab.task_program.compiler.lowering import (
RegisteredSemanticLowerer,
)
from embodichain.lab.task_program.semantics.evidence import (
ControlPartRobotEvidenceSource,
ControlPartSimulationEvidenceProvider,
EffectEvidenceProvider,
SceneArticulationEvidenceProvider,
)
from embodichain.lab.task_program.runtime.parallel_executor import (
ParallelCommandSafetyValidator,
)
from embodichain.lab.task_program.semantics.profiles import RobotSkillProfile
from embodichain.lab.task_program.semantics.scene import (
SceneCollisionRole,
SceneCollisionWorldMode,
RegistrySceneProvider,
SceneRegistry,
)
from embodichain.lab.gym.envs.task_program.bridge import (
EnvironmentStepClock,
GymPlanningObservationProvider,
)
from ..catalog import SimulationTaskProgramRegistration
from ..environment import (
TaskProgramEnvironmentAdapter,
TaskProgramEnvironmentFactory,
PlanningObservationPort,
)
from .policies import SimulationSegmentPolicyPort
if TYPE_CHECKING:
from embodichain.lab.sim.objects import Robot
from embodichain.lab.sim.sim_manager import SimulationManager
from embodichain.toolkits.graspkit import GraspPoseGenerator
MotionGeneratorFactory = Callable[[], MotionGenerator]
"""Zero-argument factory that must return one fresh motion generator."""
_PLANNER_CFG_TYPES: Mapping[str, type[BasePlannerCfg]] = {
"curobo": CuroboPlannerCfg,
"neural": NeuralPlannerCfg,
"toppra": ToppraPlannerCfg,
"trapezoidal": TrapezoidalPlannerCfg,
}
def _planner_cfg_from_config(
value: Mapping[str, object] | None,
*,
robot_uid: str,
) -> BasePlannerCfg | None:
"""Build one typed planner config from an execution-policy declaration.
The robot UID and simulation-manager instance are runtime-owned values and
are deliberately excluded from the serialized policy. cuRobo's live
collision objects are likewise attached by the simulation factory after
the selected environment has been initialized.
"""
if value is None:
return None
if not isinstance(value, Mapping):
raise TypeError("planner_config must be a mapping or None.")
unknown = sorted(set(value).difference({"type", "config"}))
if unknown:
raise ValueError(f"planner_config contains unsupported fields: {unknown}.")
planner_type = value.get("type")
if type(planner_type) is not str or not planner_type:
raise ValueError("planner_config.type must be a non-empty string.")
cfg_type = _PLANNER_CFG_TYPES.get(planner_type)
if cfg_type is None:
raise ValueError(
f"planner_config.type {planner_type!r} is unsupported; "
f"supported types are {sorted(_PLANNER_CFG_TYPES)}."
)
options = value.get("config", {})
if not isinstance(options, Mapping):
raise TypeError("planner_config.config must be a mapping.")
options = deepcopy(dict(options))
reserved = sorted(
{"planner_type", "robot_uid", "sim_instance_id"}.intersection(options)
)
if reserved:
raise ValueError(
"planner_config.config cannot override runtime-owned fields: "
f"{reserved}."
)
if planner_type == "curobo":
if "world" in options:
world = options["world"]
if not isinstance(world, Mapping):
raise TypeError("planner_config.config.world must be a mapping.")
world = deepcopy(dict(world))
if "rigid_objects" in world:
raise ValueError(
"planner_config.config.world.rigid_objects is runtime-owned; "
"declare scene collision roles instead."
)
if "dynamic_obstacle_names" in world:
raise ValueError(
"planner_config.config.world.dynamic_obstacle_names is "
"runtime-owned; declare dynamic scene collision roles instead."
)
options["world"] = CuroboWorldCfg(**world)
if "auto_gen" in options:
auto_gen = options["auto_gen"]
if not isinstance(auto_gen, Mapping):
raise TypeError("planner_config.config.auto_gen must be a mapping.")
options["auto_gen"] = CuroboAutoGenCfg(**deepcopy(dict(auto_gen)))
try:
planner_cfg = cfg_type(robot_uid=robot_uid, **options)
except (TypeError, ValueError) as exc:
raise ValueError(
f"Invalid planner_config for planner type {planner_type!r}: {exc}"
) from exc
planner_cfg.validate()
return planner_cfg
class SimulationTaskProgramEnvironment(Protocol):
"""Minimal Gym environment surface used by the simulation factory."""
sim: SimulationManager
robot: Robot
@property
def step_dt(self) -> float:
"""Return the authoritative Gym control cadence in seconds."""
def _positive_finite(value: float, *, field_name: str) -> float:
"""Validate one positive finite real number."""
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError(f"{field_name} must be a real number.")
normalized = float(value)
if not math.isfinite(normalized) or normalized <= 0.0:
raise ValueError(f"{field_name} must be finite and positive.")
return normalized
def _non_negative_finite(value: float, *, field_name: str) -> float:
"""Validate one non-negative finite real number."""
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError(f"{field_name} must be a real number.")
normalized = float(value)
if not math.isfinite(normalized) or normalized < 0.0:
raise ValueError(f"{field_name} must be finite and non-negative.")
return normalized
def _robot_uid(robot: Robot) -> str:
"""Return one strict live robot UID."""
uid = getattr(robot, "uid", None)
if type(uid) is not str or not uid or uid != uid.strip():
raise ValueError(
"robot.uid must be a non-empty string without outer whitespace."
)
return uid
def _full_robot_tensor(
robot: Robot,
getter_name: str,
*,
required: bool,
reference: torch.Tensor | None = None,
) -> torch.Tensor | None:
"""Read and validate one full-robot floating state tensor."""
getter = getattr(robot, getter_name, None)
if not callable(getter):
if required:
raise TypeError(f"robot must provide {getter_name}().")
return None
value = getter()
if not isinstance(value, torch.Tensor):
raise TypeError(f"robot.{getter_name}() must return a torch.Tensor.")
if not value.is_floating_point() or value.dim() != 2:
raise ValueError(
f"robot.{getter_name}() must return floating shape (B, robot_dof)."
)
if value.shape[0] == 0 or value.shape[1] == 0:
raise ValueError(f"robot.{getter_name}() dimensions must be non-zero.")
if reference is not None and (
value.shape != reference.shape or value.device != reference.device
):
raise ValueError(
f"robot.{getter_name}() must match robot.get_qpos() shape and device."
)
if not bool(torch.isfinite(value).all().item()):
raise ValueError(f"robot.{getter_name}() must contain only finite values.")
return value.clone()
class SharedTickSceneProvider(SceneProvider):
"""Share one immutable scene snapshot across consumers in the same tick.
``RegistrySceneProvider`` is stateful: every call observes native entities
and updates material-change baselines. Planning observations and multiple
evidence providers can legitimately request the same timestamp. This
wrapper always delegates one full-batch request per tick, then returns the
exact snapshot or an owned ordered-row projection to later consumers.
"""
def __init__(
self,
delegate: RegistrySceneProvider,
full_env_ids: torch.Tensor,
) -> None:
if type(delegate) is not RegistrySceneProvider:
raise TypeError("delegate must be exactly RegistrySceneProvider.")
if (
not isinstance(full_env_ids, torch.Tensor)
or full_env_ids.dtype != torch.long
or full_env_ids.dim() != 1
or full_env_ids.numel() == 0
):
raise ValueError("full_env_ids must be a non-empty 1D int64 tensor.")
if torch.unique(full_env_ids).numel() != full_env_ids.numel():
raise ValueError("full_env_ids must be unique.")
self._delegate = delegate
self._full_env_ids = full_env_ids.clone()
self._row_by_env_id = {
int(env_id): row
for row, env_id in enumerate(full_env_ids.detach().cpu().tolist())
}
self._timestamp: float | None = None
self._snapshot: SceneSnapshot | None = None
@property
def collision_entity_ids(self) -> tuple[str, ...]:
"""Return canonical dynamic collision IDs from the delegate."""
return self._delegate.collision_entity_ids
def snapshot(
self,
*,
timestamp: float,
env_ids: torch.Tensor,
) -> SceneSnapshot:
"""Return the single shared snapshot for ``timestamp`` and ``env_ids``."""
if isinstance(timestamp, bool) or not isinstance(timestamp, (int, float)):
raise TypeError("timestamp must be a real number.")
normalized_timestamp = float(timestamp)
if not math.isfinite(normalized_timestamp) or normalized_timestamp < 0.0:
raise ValueError("timestamp must be finite and non-negative.")
if (
not isinstance(env_ids, torch.Tensor)
or env_ids.dtype != torch.long
or env_ids.dim() != 1
or env_ids.numel() == 0
):
raise ValueError("env_ids must be a non-empty 1D int64 tensor.")
if torch.unique(env_ids).numel() != env_ids.numel():
raise ValueError("env_ids must be unique.")
if env_ids.device != self._full_env_ids.device:
raise ValueError("env_ids must share the full simulation batch device.")
try:
rows = tuple(
self._row_by_env_id[int(env_id)]
for env_id in env_ids.detach().cpu().tolist()
)
except KeyError as exc:
raise ValueError(
f"Environment ID {int(exc.args[0])} is absent from full_env_ids."
) from exc
if self._timestamp is not None:
if normalized_timestamp < self._timestamp:
raise ValueError("Shared scene snapshot timestamps must be monotonic.")
if normalized_timestamp == self._timestamp:
assert self._snapshot is not None
return self._select_rows(self._snapshot, rows)
snapshot = self._delegate.snapshot(
timestamp=normalized_timestamp,
env_ids=self._full_env_ids.clone(),
)
if not isinstance(snapshot, SceneSnapshot):
raise TypeError(
"RegistrySceneProvider.snapshot() must return SceneSnapshot."
)
if snapshot.timestamp != normalized_timestamp:
raise ValueError("Scene snapshot timestamp must match the requested tick.")
self._timestamp = normalized_timestamp
self._snapshot = snapshot
return self._select_rows(snapshot, rows)
def _select_rows(
self,
snapshot: SceneSnapshot,
rows: tuple[int, ...],
) -> SceneSnapshot:
"""Project one cached full-batch snapshot to an ordered row subset."""
full_size = int(self._full_env_ids.numel())
if rows == tuple(range(full_size)):
return snapshot
entities: dict[str, EntityState] = {}
for entity_id, state in snapshot.entities.items():
pose = state.pose
if pose.dim() == 3:
if pose.shape[0] != full_size:
raise ValueError(
f"Scene entity {entity_id!r} batch does not match "
"full_env_ids."
)
index = torch.tensor(rows, dtype=torch.long, device=pose.device)
pose = pose.index_select(0, index)
entities[entity_id] = EntityState(pose, confidence=state.confidence)
articulation_joints: dict[tuple[str, str], ObservedArticulationJointState] = {}
for address, state in snapshot.articulation_joints.items():
position = state.position
valid = state.valid_mask
if position.dim() == 2:
if position.shape[0] != full_size:
raise ValueError(
f"Scene articulation joint {address!r} batch does not "
"match full_env_ids."
)
index = torch.tensor(rows, dtype=torch.long, device=position.device)
position = position.index_select(0, index)
if valid is not None:
valid = valid.index_select(0, index.to(valid.device))
articulation_joints[address] = ObservedArticulationJointState(
position,
valid,
)
revisions = snapshot.collision_world_revisions(full_size)
return SceneSnapshot(
timestamp=snapshot.timestamp,
version=snapshot.version,
entities=entities,
collision_world_revision=tuple(revisions[row] for row in rows),
collision_entity_ids=snapshot.collision_entity_ids,
articulation_joints=articulation_joints,
)
class SimulationPlanningObservationProvider(GymPlanningObservationProvider):
"""Gym planning observations backed by live robot and shared scene state."""
def __init__(
self,
robot: Robot,
scene_provider: SharedTickSceneProvider,
clock: EnvironmentStepClock,
env_ids: torch.Tensor,
*,
owner_token: object,
) -> None:
if type(scene_provider) is not SharedTickSceneProvider:
raise TypeError("scene_provider must be exactly SharedTickSceneProvider.")
if type(clock) is not EnvironmentStepClock:
raise TypeError("clock must be exactly EnvironmentStepClock.")
qpos = _full_robot_tensor(robot, "get_qpos", required=True)
assert qpos is not None
if (
not isinstance(env_ids, torch.Tensor)
or env_ids.dtype != torch.long
or env_ids.shape != (qpos.shape[0],)
):
raise ValueError("env_ids must be int64 with one ID per robot row.")
if env_ids.device != qpos.device:
raise ValueError("env_ids and robot qpos must share a device.")
if torch.unique(env_ids).numel() != env_ids.numel():
raise ValueError("env_ids must be unique.")
self._robot = robot
self._scene_provider = scene_provider
self._clock = clock
self._env_ids = env_ids.clone()
self._owner_token = owner_token
super().__init__(self._capture)
@property
def scene_provider(self) -> SharedTickSceneProvider:
"""Return the snapshot-sharing scene provider used by evidence ports."""
return self._scene_provider
@property
def env_ids(self) -> torch.Tensor:
"""Return stable ordered simulation row IDs."""
return self._env_ids.clone()
def is_owned_by(self, owner_token: object) -> bool:
"""Return whether this provider belongs to one factory instance."""
return self._owner_token is owner_token
def _capture(self, task_state: TaskState) -> PlanningContext:
"""Capture one synchronized robot and scene observation."""
qpos = _full_robot_tensor(self._robot, "get_qpos", required=True)
assert qpos is not None
if (
qpos.shape[0] != self._env_ids.numel()
or qpos.device != self._env_ids.device
):
raise ValueError("Robot batch shape or device changed after assembly.")
qvel = _full_robot_tensor(
self._robot,
"get_qvel",
required=False,
reference=qpos,
)
if qvel is None:
qvel = torch.zeros_like(qpos)
qeffort = _full_robot_tensor(
self._robot,
"get_qf",
required=False,
reference=qpos,
)
timestamp = self._clock.now()
scene = self._scene_provider.snapshot(
timestamp=timestamp,
env_ids=self._env_ids.clone(),
)
return PlanningContext(
robot=RobotObservation(
timestamp=timestamp,
qpos=qpos,
qvel=qvel,
qeffort=qeffort,
),
task=task_state,
scene=scene,
env_ids=self._env_ids,
control_dt=self._clock.step_dt,
)
[docs]
class SimulationTaskProgramFactory(TaskProgramEnvironmentFactory):
"""Build every live Task Program component from one task registration.
Args:
simulation: Exact live simulation that owns ``robot`` and scene UIDs.
robot: Exact robot selected for planning and evidence acquisition.
registration: Frozen scene, profile, and extension declarations.
step_dt: Authoritative Gym control cadence.
joint_command_mode: Environment-owned expert joint command mode.
planner_cfg: Explicit planner configuration. ``None`` selects TOPPRA
for ``robot.uid``.
planner_config: Optional executable-free planner declaration from an
execution policy. It is converted to ``planner_cfg`` after the
selected runtime robot is known.
motion_generator_factory: Optional fresh-generator factory. It is
mutually exclusive with ``planner_cfg`` and intended for custom
planners and isolated tests.
grasp_pose_generators: Standalone grasp-pose services keyed by grasp
endpoint target ID, normally a robot control-part name.
translation_threshold: Material scene translation threshold.
rotation_threshold: Material scene rotation threshold.
The Gym cadence is attached to each live :class:`PlanningContext`. Motion
policy presets describe behavior; they do not own the environment's
command period.
"""
[docs]
def __init__(
self,
simulation: SimulationManager,
robot: Robot,
registration: SimulationTaskProgramRegistration,
*,
step_dt: float,
joint_command_mode: JointCommandMode = "position",
planner_cfg: BasePlannerCfg | None = None,
motion_generator_factory: MotionGeneratorFactory | None = None,
grasp_pose_generators: Mapping[str, GraspPoseGenerator] | None = None,
translation_threshold: float = 1.0e-4,
rotation_threshold: float = 1.0e-3,
) -> None:
if type(registration) is not SimulationTaskProgramRegistration:
raise TypeError(
"registration must be exactly SimulationTaskProgramRegistration."
)
registration.assert_unchanged()
selected_scene_binding = registration.scene_binding
selected_profile_binding = registration.robot_profile_binding
selected_endpoint_adapters = dict(registration.endpoint_adapter_map)
selected_settle_presets = registration.settle_presets
if planner_cfg is not None and motion_generator_factory is not None:
raise ValueError(
"planner_cfg and motion_generator_factory are mutually exclusive."
)
if planner_cfg is not None and not isinstance(planner_cfg, BasePlannerCfg):
raise TypeError("planner_cfg must be a BasePlannerCfg or None.")
if motion_generator_factory is not None and not callable(
motion_generator_factory
):
raise TypeError("motion_generator_factory must be callable or None.")
if grasp_pose_generators is not None and not isinstance(
grasp_pose_generators, Mapping
):
raise TypeError("grasp_pose_generators must be a mapping or None.")
if joint_command_mode not in ("position", "position_velocity"):
raise ValueError(
"joint_command_mode must be 'position' or 'position_velocity'."
)
robot_uid = _robot_uid(robot)
get_robot = getattr(simulation, "get_robot", None)
if not callable(get_robot):
raise TypeError("simulation must provide get_robot().")
if get_robot(robot_uid) is not robot:
raise ValueError(
f"simulation.get_robot({robot_uid!r}) must return the exact "
"selected robot."
)
selected_planner_cfg = deepcopy(planner_cfg)
if (
selected_planner_cfg is not None
and selected_planner_cfg.robot_uid != robot_uid
):
raise ValueError(
f"planner_cfg.robot_uid must equal selected robot UID {robot_uid!r}."
)
self._simulation = simulation
self._robot = robot
self._registration = registration
self._scene_binding = selected_scene_binding
self._robot_profile_binding = selected_profile_binding
self._step_dt = _positive_finite(step_dt, field_name="step_dt")
self._joint_command_mode = joint_command_mode
self._planner_cfg = selected_planner_cfg
self._motion_generator_factory = motion_generator_factory
self._grasp_pose_generators = (
{} if grasp_pose_generators is None else dict(grasp_pose_generators)
)
self._endpoint_adapters = selected_endpoint_adapters
self._translation_threshold = _non_negative_finite(
translation_threshold,
field_name="translation_threshold",
)
self._rotation_threshold = _non_negative_finite(
rotation_threshold,
field_name="rotation_threshold",
)
self._owner_token = object()
qpos = _full_robot_tensor(robot, "get_qpos", required=True)
assert qpos is not None
self._env_ids = torch.arange(
qpos.shape[0],
dtype=torch.long,
device=qpos.device,
)
self._segment_policy_port = SimulationSegmentPolicyPort(
simulation,
robot,
selected_scene_binding,
settle_presets=selected_settle_presets,
env_ids=self._env_ids,
)
[docs]
@classmethod
def from_environment(
cls,
environment: SimulationTaskProgramEnvironment,
*,
registration: SimulationTaskProgramRegistration,
planner_cfg: BasePlannerCfg | None = None,
planner_config: Mapping[str, object] | None = None,
motion_generator_factory: MotionGeneratorFactory | None = None,
grasp_pose_generators: Mapping[str, GraspPoseGenerator] | None = None,
translation_threshold: float = 1.0e-4,
rotation_threshold: float = 1.0e-3,
) -> SimulationTaskProgramFactory:
"""Create a registered factory from the standard Gym environment surface."""
simulation = getattr(environment, "sim", None)
robot = getattr(environment, "robot", None)
try:
step_dt = environment.step_dt
except AttributeError as exc:
raise TypeError("environment must expose step_dt.") from exc
if simulation is None or robot is None:
raise TypeError("environment must expose non-None sim and robot values.")
if planner_cfg is not None and planner_config is not None:
raise ValueError("planner_cfg and planner_config are mutually exclusive.")
if planner_config is not None:
planner_cfg = _planner_cfg_from_config(
planner_config,
robot_uid=_robot_uid(robot),
)
expert_trajectory_cfg = getattr(
getattr(environment, "cfg", None),
"expert_trajectory",
None,
)
joint_command_mode = getattr(
expert_trajectory_cfg,
"joint_command_mode",
"position",
)
return cls(
simulation,
robot,
registration,
step_dt=step_dt,
joint_command_mode=joint_command_mode,
planner_cfg=planner_cfg,
motion_generator_factory=motion_generator_factory,
grasp_pose_generators=grasp_pose_generators,
translation_threshold=translation_threshold,
rotation_threshold=rotation_threshold,
)
@property
def scene_registry_id(self) -> str:
"""Return the exact configured scene-registry ID."""
return self._scene_binding.registry_id
@property
def robot_profile_id(self) -> str:
"""Return the exact configured robot-profile ID."""
return self._robot_profile_binding.profile_id
@property
def step_dt(self) -> float:
"""Return the authoritative Gym control cadence."""
return self._step_dt
@property
def segment_policy_port(self) -> SimulationSegmentPolicyPort:
"""Return the shared simulation post-policy and validator port."""
return self._segment_policy_port
@property
def task_program_registration(self) -> SimulationTaskProgramRegistration:
"""Return the exact standard registration owned by this factory."""
return self._registration
[docs]
def registration_owned_segment_policy_ports(
self,
) -> tuple[SimulationSegmentPolicyPort, SimulationSegmentPolicyPort]:
"""Return registration-owned post-policy and validator ports."""
return self._segment_policy_port, self._segment_policy_port
[docs]
def create_registered_semantic_lowerers(
self,
*,
scene_registry: SceneRegistry,
engine: AtomicActionEngine,
) -> tuple[RegisteredSemanticLowerer, ...]:
"""Create fresh registration-owned lowerers for this runtime assembly."""
return self._registration.create_registered_semantic_lowerers(
simulation=self._simulation,
robot=self._robot,
scene_registry=scene_registry,
engine=engine,
)
[docs]
def create_scene_registry(self) -> SceneRegistry:
"""Build one fresh authoritative registry from the task registration."""
registry = self._scene_binding.build(self._simulation)
self._registration.validate_scene_registry(registry)
return registry
[docs]
def create_robot_skill_profile(self) -> RobotSkillProfile:
"""Build the robot profile used to lower calls to Atomic Skills."""
profile = self._robot_profile_binding.build(self._robot)
self._registration.validate_robot_profile(profile)
return profile
[docs]
def create_atomic_action_engine(
self,
profile: RobotSkillProfile,
) -> AtomicActionEngine:
"""Create a fresh engine around the selected planner and exact profile."""
if not isinstance(profile, RobotSkillProfile):
raise TypeError("profile must be a RobotSkillProfile.")
if profile.profile_id != self.robot_profile_id:
raise ValueError(
f"profile ID must be {self.robot_profile_id!r}, got "
f"{profile.profile_id!r}."
)
motion_generator = self._create_motion_generator()
if motion_generator.robot is not self._robot:
raise ValueError(
"Motion generator must own the exact robot selected by the factory."
)
engine = AtomicActionEngine(
motion_generator,
control_profiles=profile.action_control_profiles(),
grasp_pose_generators=self._grasp_pose_generators,
)
self._registration.validate_engine(engine)
return engine
[docs]
def create_planning_observation_provider(
self,
*,
scene_registry: SceneRegistry,
engine: AtomicActionEngine,
clock: EnvironmentStepClock,
) -> PlanningObservationPort:
"""Create one planning port and planner-validated shared scene provider."""
if type(scene_registry) is not SceneRegistry:
raise TypeError("scene_registry must be exactly SceneRegistry.")
if not isinstance(engine, AtomicActionEngine):
raise TypeError("engine must be an AtomicActionEngine.")
if engine.robot is not self._robot:
raise ValueError("engine must own the exact factory robot.")
if type(clock) is not EnvironmentStepClock:
raise TypeError("clock must be exactly EnvironmentStepClock.")
if clock.step_dt != self._step_dt:
raise ValueError("clock.step_dt must equal the factory Gym cadence.")
provider = scene_registry.make_planning_scene_provider(
engine.motion_generator,
batch_size=int(self._env_ids.numel()),
translation_threshold=self._translation_threshold,
rotation_threshold=self._rotation_threshold,
)
shared = SharedTickSceneProvider(provider, self._env_ids)
return SimulationPlanningObservationProvider(
self._robot,
shared,
clock,
self._env_ids,
owner_token=self._owner_token,
)
[docs]
def create_effect_evidence_providers(
self,
*,
scene_registry: SceneRegistry,
engine: AtomicActionEngine,
observation_provider: PlanningObservationPort,
) -> Iterable[EffectEvidenceProvider]:
"""Create built-in control-part and articulation evidence providers."""
if type(scene_registry) is not SceneRegistry:
raise TypeError("scene_registry must be exactly SceneRegistry.")
if not isinstance(engine, AtomicActionEngine):
raise TypeError("engine must be an AtomicActionEngine.")
if engine.robot is not self._robot:
raise ValueError("engine must own the exact factory robot.")
if type(observation_provider) is not SimulationPlanningObservationProvider:
raise TypeError(
"observation_provider must be exactly "
"SimulationPlanningObservationProvider."
)
if not observation_provider.is_owned_by(self._owner_token):
raise ValueError("observation_provider belongs to another factory.")
scene_provider = observation_provider.scene_provider
providers: list[EffectEvidenceProvider] = []
if isinstance(self._robot, ControlPartRobotEvidenceSource):
registered_provider = (
self._registration.create_control_part_evidence_provider(
simulation=self._simulation,
robot=self._robot,
scene_registry=scene_registry,
engine=engine,
scene_provider=scene_provider,
)
)
providers.append(
registered_provider
if registered_provider is not None
else ControlPartSimulationEvidenceProvider(
self._robot,
scene_provider=scene_provider,
)
)
providers.append(
SceneArticulationEvidenceProvider(scene_provider=scene_provider)
)
return tuple(providers)
[docs]
def create_parallel_command_safety_validator(
self,
*,
scene_registry: SceneRegistry,
engine: AtomicActionEngine,
observation_provider: PlanningObservationPort,
) -> ParallelCommandSafetyValidator:
"""Create the registration-owned validator for this runtime assembly."""
if type(scene_registry) is not SceneRegistry:
raise TypeError("scene_registry must be exactly SceneRegistry.")
if (
not isinstance(engine, AtomicActionEngine)
or engine.robot is not self._robot
):
raise ValueError("engine must own the exact factory robot.")
if type(observation_provider) is not SimulationPlanningObservationProvider:
raise TypeError(
"observation_provider must be exactly "
"SimulationPlanningObservationProvider."
)
if not observation_provider.is_owned_by(self._owner_token):
raise ValueError("observation_provider belongs to another factory.")
validator = self._registration.create_parallel_safety_validator(
simulation=self._simulation,
robot=self._robot,
scene_registry=scene_registry,
engine=engine,
)
if validator is None:
raise RuntimeError(
"The task registration does not declare a parallel safety factory."
)
return validator
[docs]
def create_adapter(self) -> TaskProgramEnvironmentAdapter:
"""Create the exact Gym adapter with shared simulation policy ports."""
return TaskProgramEnvironmentAdapter(
self,
step_dt=self._step_dt,
joint_command_mode=self._joint_command_mode,
registration=self._registration,
)
def _create_motion_generator(self) -> MotionGenerator:
"""Create and validate one exact motion generator."""
if self._motion_generator_factory is not None:
generator = self._motion_generator_factory()
else:
planner_cfg = (
ToppraPlannerCfg(robot_uid=_robot_uid(self._robot))
if self._planner_cfg is None
else deepcopy(self._planner_cfg)
)
planner_cfg.sim_instance_id = self._simulation.instance_id
if isinstance(planner_cfg, CuroboPlannerCfg):
self._bind_curobo_scene(planner_cfg)
generator = MotionGenerator(MotionGenCfg(planner_cfg=planner_cfg))
if not isinstance(generator, MotionGenerator):
raise TypeError(
"motion_generator_factory must return a MotionGenerator instance."
)
return generator
def _bind_curobo_scene(self, planner_cfg: CuroboPlannerCfg) -> None:
"""Attach live scene collision objects to a configured cuRobo planner."""
if planner_cfg.world.rigid_objects is not None:
return
collision_bindings = tuple(
binding
for binding in self._scene_binding.rigid_objects
if binding.collision_role is not SceneCollisionRole.NONE
)
unsupported_bindings = tuple(
binding.entity_id
for bindings in (
self._scene_binding.rigidized_articulations,
self._scene_binding.articulations,
self._scene_binding.links,
)
for binding in bindings
if binding.collision_role is not SceneCollisionRole.NONE
)
if unsupported_bindings:
raise ValueError(
"Configured cuRobo collision worlds support only rigid-object "
"scene bindings; unsupported collision entities: "
f"{sorted(unsupported_bindings)}."
)
objects: dict[str, object] = {}
get_rigid_object = getattr(self._simulation, "get_rigid_object", None)
if collision_bindings and not callable(get_rigid_object):
raise TypeError("simulation must provide get_rigid_object().")
for binding in collision_bindings:
entity = get_rigid_object(binding.simulation_uid)
if entity is None:
raise KeyError(
f"Simulation UID {binding.simulation_uid!r} selected for "
f"collision entity {binding.entity_id!r} was not found."
)
objects[binding.entity_id] = entity
world = planner_cfg.world
dynamic_obstacle_names = [
binding.entity_id
for binding in collision_bindings
if binding.collision_role is SceneCollisionRole.DYNAMIC
]
world_values = {
field.name: getattr(world, field.name) for field in fields(CuroboWorldCfg)
}
world_values["rigid_objects"] = objects or None
world_values["dynamic_obstacle_names"] = dynamic_obstacle_names
collision_mode = self._scene_binding.collision_world_mode
if collision_mode is not None:
world_values["multi_env"] = (
collision_mode is SceneCollisionWorldMode.PER_ENV
)
planner_cfg.world = CuroboWorldCfg(**world_values)
[docs]
@dataclass(frozen=True, slots=True, init=False)
class SimulationTaskProgramAdapterFactory:
"""Create one standard simulation adapter after an environment is initialized.
Args:
registration: Immutable provider-free task integration.
grasp_pose_generator_factories: Zero-argument factories keyed by runtime
grasp endpoint target ID. Each environment receives fresh services.
planner_config: Optional execution-policy planner declaration. The live
factory supplies the selected robot and simulation instance values.
"""
_registration: SimulationTaskProgramRegistration
_grasp_pose_generator_factories: Mapping[
str,
Callable[[], GraspPoseGenerator],
]
_planner_config: Mapping[str, object] | None
[docs]
def __init__(
self,
registration: SimulationTaskProgramRegistration,
*,
grasp_pose_generator_factories: (
Mapping[str, Callable[[], GraspPoseGenerator]] | None
) = None,
planner_config: Mapping[str, object] | None = None,
) -> None:
if type(registration) is not SimulationTaskProgramRegistration:
raise TypeError(
"registration must be exactly SimulationTaskProgramRegistration."
)
if grasp_pose_generator_factories is not None and not isinstance(
grasp_pose_generator_factories,
Mapping,
):
raise TypeError("grasp_pose_generator_factories must be a mapping or None.")
if planner_config is not None and not isinstance(planner_config, Mapping):
raise TypeError("planner_config must be a mapping or None.")
factories: dict[str, Callable[[], GraspPoseGenerator]] = {}
for target_id, factory in (grasp_pose_generator_factories or {}).items():
if (
type(target_id) is not str
or not target_id
or target_id != target_id.strip()
):
raise ValueError(
"grasp_pose_generator_factories keys must be non-empty strings "
"without outer whitespace."
)
if not callable(factory):
raise TypeError(
"grasp_pose_generator_factories values must be callable."
)
factories[target_id] = factory
object.__setattr__(self, "_registration", registration)
object.__setattr__(
self,
"_grasp_pose_generator_factories",
MappingProxyType(factories),
)
object.__setattr__(
self,
"_planner_config",
(
None
if planner_config is None
else MappingProxyType(deepcopy(dict(planner_config)))
),
)
@property
def registration(self) -> SimulationTaskProgramRegistration:
"""Return the exact static registration owned by this factory."""
return self._registration
[docs]
def create_adapter(
self,
environment: object,
) -> TaskProgramEnvironmentAdapter:
"""Create fresh grasp services and bind the initialized environment."""
from embodichain.toolkits.graspkit import GraspPoseGenerator
self._registration.assert_unchanged()
generators: dict[str, GraspPoseGenerator] = {}
for target_id, factory in self._grasp_pose_generator_factories.items():
generator = factory()
if not isinstance(generator, GraspPoseGenerator):
raise TypeError(
"A grasp-pose generator factory must return a "
"GraspPoseGenerator instance."
)
generators[target_id] = generator
kwargs: dict[str, object] = {
"registration": self._registration,
"grasp_pose_generators": generators,
}
if self._planner_config is not None:
kwargs["planner_config"] = self._planner_config
return create_simulation_task_program_adapter(environment, **kwargs)
[docs]
def create_simulation_task_program_adapter(
environment: SimulationTaskProgramEnvironment,
*,
registration: SimulationTaskProgramRegistration,
planner_cfg: BasePlannerCfg | None = None,
planner_config: Mapping[str, object] | None = None,
motion_generator_factory: MotionGeneratorFactory | None = None,
grasp_pose_generators: Mapping[str, GraspPoseGenerator] | None = None,
translation_threshold: float = 1.0e-4,
rotation_threshold: float = 1.0e-3,
) -> TaskProgramEnvironmentAdapter:
"""Create a complete production adapter from one standard Gym environment.
This is the intended task-side one-line integration. Standard tasks pass
one immutable ``registration`` that owns their pre-simulation catalog and
every runtime extension.
Args:
environment: Standard Gym simulation environment exposing ``sim``,
``robot``, and ``step_dt``.
registration: Required immutable task registration.
planner_cfg: Optional planner configuration owned by the factory.
planner_config: Optional executable-free planner declaration from an
execution policy.
motion_generator_factory: Optional factory for one fresh motion generator.
grasp_pose_generators: Standalone grasp-pose services keyed by grasp
endpoint target ID.
translation_threshold: Scene translation revision threshold.
rotation_threshold: Scene rotation revision threshold.
Returns:
Complete production Task Program environment adapter.
"""
factory = SimulationTaskProgramFactory.from_environment(
environment,
registration=registration,
planner_cfg=planner_cfg,
planner_config=planner_config,
motion_generator_factory=motion_generator_factory,
grasp_pose_generators=grasp_pose_generators,
translation_threshold=translation_threshold,
rotation_threshold=rotation_threshold,
)
return factory.create_adapter()
__all__: list[str] = []