diff --git a/embodichain/lab/gym/envs/expert_program/__init__.py b/embodichain/lab/gym/envs/expert_program/__init__.py index e15b40619..66a8b55b3 100644 --- a/embodichain/lab/gym/envs/expert_program/__init__.py +++ b/embodichain/lab/gym/envs/expert_program/__init__.py @@ -155,6 +155,10 @@ create_simulation_expert_program_adapter, ) from .simulation_handover import ConfiguredHandOverPoseProvider +from .simulation_parallel_safety import ( + CuroboParallelCommandSafetyValidator, + CuroboParallelSafetyValidatorFactory, +) from .simulation_policies import ( SimulationSegmentPolicyPort, default_simulation_settle_presets, @@ -189,6 +193,8 @@ "ControlPartResourceBinding", "ConfiguredHandOverPoseProvider", "CyclicPoseTargetCfg", + "CuroboParallelCommandSafetyValidator", + "CuroboParallelSafetyValidatorFactory", "DeclarativeCfgValue", "DemoBridgeError", "EXPERT_PROGRAM_SCHEMA_VERSION", diff --git a/embodichain/lab/gym/envs/expert_program/catalog.py b/embodichain/lab/gym/envs/expert_program/catalog.py index 8b468b1bb..9df60f462 100644 --- a/embodichain/lab/gym/envs/expert_program/catalog.py +++ b/embodichain/lab/gym/envs/expert_program/catalog.py @@ -1457,14 +1457,27 @@ def create_parallel_safety_validator( *, simulation: object, robot: object, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, ) -> ParallelCommandSafetyValidator | None: """Create and strictly validate the registration-owned live safety gate.""" self.assert_unchanged() factory = self.parallel_safety_factory if factory is None: return None + 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 robot: + raise ValueError("engine and factory must reference the exact same robot.") with self._parallel_safety_validator_lock: - validator = factory.create(simulation=simulation, robot=robot) + validator = factory.create( + simulation=simulation, + robot=robot, + scene_registry=scene_registry, + engine=engine, + ) if not isinstance(validator, ParallelCommandSafetyValidator): raise TypeError( "parallel_safety_factory.create() must return a " diff --git a/embodichain/lab/gym/envs/expert_program/extensions.py b/embodichain/lab/gym/envs/expert_program/extensions.py index a0d549f02..f5b019d8b 100644 --- a/embodichain/lab/gym/envs/expert_program/extensions.py +++ b/embodichain/lab/gym/envs/expert_program/extensions.py @@ -28,7 +28,7 @@ from dataclasses import dataclass, fields, is_dataclass from enum import Enum from types import MappingProxyType -from typing import ClassVar, Protocol, runtime_checkable +from typing import ClassVar, Protocol, runtime_checkable, TYPE_CHECKING import torch @@ -57,6 +57,10 @@ RuntimeTransportActionEncoder, ) +if TYPE_CHECKING: + from embodichain.lab.sim.atomic_actions import AtomicActionEngine + from embodichain.lab.sim.skills import SceneRegistry + VersionedKey = tuple[str, str] """Exact ``(provider_or_projector_id, revision)`` registry key.""" @@ -450,8 +454,10 @@ def create( *, simulation: object, robot: object, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, ) -> ParallelCommandSafetyValidator: - """Create one live validator bound to the exact simulation and robot.""" + """Create one live gate bound to the exact assembled runtime.""" @dataclass(frozen=True, slots=True) diff --git a/embodichain/lab/gym/envs/expert_program/simulation_environment.py b/embodichain/lab/gym/envs/expert_program/simulation_environment.py index 37718121d..63b848cb8 100644 --- a/embodichain/lab/gym/envs/expert_program/simulation_environment.py +++ b/embodichain/lab/gym/envs/expert_program/simulation_environment.py @@ -1060,6 +1060,8 @@ def create_parallel_command_safety_validator( validator = self._registration.create_parallel_safety_validator( simulation=self._simulation, robot=self._robot, + scene_registry=scene_registry, + engine=engine, ) if not isinstance(validator, ParallelCommandSafetyValidator): raise TypeError( diff --git a/embodichain/lab/gym/envs/expert_program/simulation_parallel_safety.py b/embodichain/lab/gym/envs/expert_program/simulation_parallel_safety.py new file mode 100644 index 000000000..97317e3f6 --- /dev/null +++ b/embodichain/lab/gym/envs/expert_program/simulation_parallel_safety.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. +# ---------------------------------------------------------------------------- + +"""cuRobo-backed physical safety gate for synchronized simulation commands.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +import math +from typing import ClassVar + +import torch + +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + JointPositionPayload, + JointPositionTarget, + RuntimeCommandFrame, +) +from embodichain.lab.sim.planners import CuroboPlanner, MotionGenerator +from embodichain.lab.sim.skills import ( + ParallelSafetyError, + RegistrySceneProvider, + SceneRegistry, +) + + +def _identifier(value: object, *, field_name: str) -> str: + """Validate one exact identifier.""" + if type(value) is not str or not value or value != value.strip(): + raise ValueError( + f"{field_name} must be a non-empty string without outer whitespace." + ) + return value + + +@dataclass(frozen=True, slots=True) +class CuroboParallelSafetyValidatorFactory: + """Create exact-sample collision gates for one aggregate control part. + + ``validation_control_part`` must contain every joint that any parallel + branch can command. A common dual-arm example is ``"dual_arm"``. The + cuRobo model for that part remains the authoritative bounds, self-collision, + and world-collision model. + + Args: + validation_control_part: Aggregate robot control part containing every + joint that a parallel lane may command. + max_joint_step: Maximum absolute joint displacement between collision + samples in radians or the joint's native linear unit. + max_interpolation_samples: Fail-closed upper bound on samples per frame. + """ + + validator_id: ClassVar[str] = "builtin.simulation.curobo_parallel_safety" + revision: ClassVar[str] = "1" + supported_transport_ids: ClassVar[frozenset[str]] = frozenset( + {JointPositionTarget.TRANSPORT_ID} + ) + + validation_control_part: str + max_joint_step: float = 0.025 + max_interpolation_samples: int = 256 + + def __post_init__(self) -> None: + _identifier( + self.validation_control_part, + field_name="validation_control_part", + ) + if isinstance(self.max_joint_step, bool) or not isinstance( + self.max_joint_step, + (int, float), + ): + raise TypeError("max_joint_step must be a real number.") + normalized_step = float(self.max_joint_step) + if not math.isfinite(normalized_step) or normalized_step <= 0.0: + raise ValueError("max_joint_step must be finite and positive.") + object.__setattr__(self, "max_joint_step", normalized_step) + if ( + type(self.max_interpolation_samples) is not int + or self.max_interpolation_samples < 2 + or self.max_interpolation_samples > 4096 + ): + raise ValueError( + "max_interpolation_samples must be an integer in [2, 4096]." + ) + + def create( + self, + *, + simulation: object, + robot: object, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + ) -> CuroboParallelCommandSafetyValidator: + """Create one fresh validator bound to the assembled live runtime.""" + del simulation + 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 robot: + raise ValueError("engine and factory must reference the exact same robot.") + return CuroboParallelCommandSafetyValidator( + robot=robot, + motion_generator=engine.motion_generator, + scene_registry=scene_registry, + validation_control_part=self.validation_control_part, + max_joint_step=self.max_joint_step, + max_interpolation_samples=self.max_interpolation_samples, + ) + + +class CuroboParallelCommandSafetyValidator: + """Validate the exact synchronized joint segment before transport dispatch. + + Args: + robot: Live robot supplying measured joint state and control-part IDs. + motion_generator: Runtime motion generator backed by exact cuRobo. + scene_registry: Authoritative live collision-scene registry. + validation_control_part: Aggregate control part for merged commands. + max_joint_step: Maximum displacement between collision samples. + max_interpolation_samples: Fail-closed sample-count upper bound. + """ + + def __init__( + self, + *, + robot: object, + motion_generator: MotionGenerator, + scene_registry: SceneRegistry, + validation_control_part: str, + max_joint_step: float, + max_interpolation_samples: int, + ) -> None: + if type(scene_registry) is not SceneRegistry: + raise TypeError("scene_registry must be exactly SceneRegistry.") + if not isinstance(motion_generator, MotionGenerator): + raise TypeError("motion_generator must be a MotionGenerator.") + if type(motion_generator.planner) is not CuroboPlanner: + raise TypeError( + "CuroboParallelCommandSafetyValidator requires the active " + "CuroboPlanner backend." + ) + if not motion_generator.supports_joint_trajectory_validation: + raise ValueError( + "The active motion generator does not validate exact joint " + "trajectories." + ) + get_joint_ids = getattr(robot, "get_joint_ids", None) + if not callable(get_joint_ids): + raise TypeError("robot must provide get_joint_ids().") + joint_ids = tuple(get_joint_ids(name=validation_control_part)) + if not joint_ids or not all( + type(joint_id) is int and joint_id >= 0 for joint_id in joint_ids + ): + raise ValueError( + "The validation control part must resolve non-negative joint IDs." + ) + if len(set(joint_ids)) != len(joint_ids): + raise ValueError("The validation control part joint IDs must be unique.") + self._robot = robot + self._motion_generator = motion_generator + self._scene_registry = scene_registry + self._validation_control_part = validation_control_part + self._validation_joint_ids = joint_ids + self._local_joint_columns = { + joint_id: index for index, joint_id in enumerate(joint_ids) + } + self._max_joint_step = max_joint_step + self._max_interpolation_samples = max_interpolation_samples + self._scene_provider: RegistrySceneProvider | None = None + self._scene_timestamp = 0.0 + + def validate( + self, + *, + branch_frames: Mapping[str, RuntimeCommandFrame], + merged_frame: RuntimeCommandFrame, + ) -> None: + """Reject a merged command whose exact interpolated segment collides.""" + if not isinstance(branch_frames, Mapping) or len(branch_frames) < 2: + raise TypeError("branch_frames must contain at least two branch frames.") + if type(merged_frame) is not RuntimeCommandFrame: + raise TypeError("merged_frame must be exactly RuntimeCommandFrame.") + for branch_id, frame in branch_frames.items(): + _identifier(branch_id, field_name="parallel branch IDs") + if type(frame) is not RuntimeCommandFrame: + raise TypeError( + "branch_frames values must be exact RuntimeCommandFrame values." + ) + if not torch.equal(frame.env_ids, merged_frame.env_ids): + raise ValueError("Parallel branch and merged env_ids must match.") + + active = merged_frame.active_mask + if not bool(active.any().item()): + return + current = self._current_control_part_qpos(merged_frame.env_ids) + target = current.clone() + commanded_joint_ids: set[int] = set() + for command in merged_frame.commands: + if ( + type(command.target) is not JointPositionTarget + or type(command.payload) is not JointPositionPayload + ): + raise ParallelSafetyError( + "cuRobo parallel safety accepts only exact joint-position " + "targets and payloads." + ) + missing = sorted( + set(command.target.joint_ids).difference(self._local_joint_columns) + ) + if missing: + raise ParallelSafetyError( + f"Parallel target {command.target.target_id!r} commands joints " + f"{missing} outside validation control part " + f"{self._validation_control_part!r}." + ) + for payload_column, joint_id in enumerate(command.target.joint_ids): + if joint_id in commanded_joint_ids: + raise ParallelSafetyError( + f"Merged parallel commands overlap on joint {joint_id}." + ) + commanded_joint_ids.add(joint_id) + target[:, self._local_joint_columns[joint_id]] = ( + command.payload.positions[:, payload_column] + ) + target = torch.where(active[:, None], target, current) + trajectory = self._interpolate(current, target) + obstacle_poses = self._obstacle_poses( + env_ids=merged_frame.env_ids, + device=trajectory.device, + dtype=trajectory.dtype, + ) + validity = self._motion_generator.validate_joint_trajectory( + trajectory, + control_part=self._validation_control_part, + obstacle_poses=obstacle_poses, + ) + row_valid = validity.all(dim=1) + failed = active & ~row_valid + if not bool(failed.any().item()): + return + failed_rows = failed.nonzero(as_tuple=False).flatten() + failed_env_ids = merged_frame.env_ids.index_select(0, failed_rows) + first_invalid_samples = tuple( + int((~validity[row]).nonzero(as_tuple=False)[0, 0].item()) + for row in failed_rows.detach().cpu().tolist() + ) + raise ParallelSafetyError( + "Merged parallel joint segment is not collision-free for env IDs " + f"{tuple(failed_env_ids.detach().cpu().tolist())}; first invalid " + f"samples={first_invalid_samples}." + ) + + def _current_control_part_qpos(self, env_ids: torch.Tensor) -> torch.Tensor: + """Read current full robot state and select the validator joint order.""" + getter = getattr(self._robot, "get_qpos", None) + if not callable(getter): + raise TypeError("robot must provide get_qpos().") + full = getter(target=False) + if ( + not isinstance(full, torch.Tensor) + or not full.is_floating_point() + or full.dim() != 2 + or not bool(torch.isfinite(full).all().item()) + ): + raise ValueError("robot.get_qpos() must return finite floating (B, D).") + if env_ids.device != full.device: + raise ValueError("Parallel env_ids and robot qpos must share a device.") + if ( + bool((env_ids < 0).any().item()) + or int(env_ids.max().item()) >= full.shape[0] + ): + raise ValueError("Parallel env_ids do not address robot qpos rows.") + if max(self._validation_joint_ids) >= full.shape[1]: + raise ValueError( + "Validation control-part joint IDs exceed robot qpos width." + ) + rows = full.index_select(0, env_ids) + columns = torch.tensor( + self._validation_joint_ids, + dtype=torch.long, + device=full.device, + ) + return rows.index_select(1, columns).clone() + + def _interpolate( + self, + current: torch.Tensor, + target: torch.Tensor, + ) -> torch.Tensor: + """Densify the exact controller segment under a bounded joint step.""" + max_delta = float((target - current).abs().max().item()) + sample_count = max(2, math.ceil(max_delta / self._max_joint_step) + 1) + if sample_count > self._max_interpolation_samples: + raise ParallelSafetyError( + "Merged parallel joint segment needs " + f"{sample_count} collision samples at max_joint_step=" + f"{self._max_joint_step}, exceeding configured limit " + f"{self._max_interpolation_samples}." + ) + alpha = torch.linspace( + 0.0, + 1.0, + sample_count, + device=current.device, + dtype=current.dtype, + ) + return ( + current[:, None, :] + alpha[None, :, None] * (target - current)[:, None, :] + ) + + def _obstacle_poses( + self, + *, + env_ids: torch.Tensor, + device: torch.device, + dtype: torch.dtype, + ) -> Mapping[str, torch.Tensor] | None: + """Observe the exact dynamic collision world for this safety decision.""" + if not self._scene_registry.dynamic_collision_entity_ids: + return None + if self._scene_provider is None: + self._scene_provider = self._scene_registry.make_scene_provider( + batch_size=int(env_ids.numel()) + ) + snapshot = self._scene_provider.snapshot( + timestamp=self._scene_timestamp, + env_ids=env_ids, + ) + self._scene_timestamp += 1.0 + return snapshot.collision_obstacle_poses( + batch_size=int(env_ids.numel()), + device=device, + dtype=dtype, + ) + + +__all__ = [ + "CuroboParallelCommandSafetyValidator", + "CuroboParallelSafetyValidatorFactory", +] diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index d58f7b019..d12a347ed 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -179,6 +179,9 @@ def __init__(self, cfg: BasePlannerCfg): supports_collision_world_updates: bool = False """Whether per-plan dynamic obstacle poses can update the collision world.""" + supports_joint_trajectory_validation: bool = False + """Whether exact joint samples can be checked against bounds/collisions.""" + @property def dynamic_collision_entity_ids(self) -> tuple[str, ...]: """Return canonical entity IDs accepted for dynamic pose updates.""" @@ -253,6 +256,35 @@ def with_collision_world( """ return options + def validate_joint_trajectory( + self, + trajectory: torch.Tensor, + *, + control_part: str, + obstacle_poses: Mapping[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Validate exact joint samples without replacing their path. + + Backends that implement this contract must evaluate every supplied + sample against joint bounds, self-collision, and their configured world + collision model. They return a boolean mask with shape ``(B, T)``. + + Args: + trajectory: Simulator-order joint samples with shape ``(B, T, D)``. + control_part: Robot control part whose ordered joints form ``D``. + obstacle_poses: Optional current dynamic-obstacle world poses. + + Returns: + Per-environment, per-sample validity mask. + + Raises: + NotImplementedError: Always for the base planner. + """ + del trajectory, control_part, obstacle_poses + raise NotImplementedError( + f"{type(self).__name__} does not validate exact joint trajectories." + ) + @validate_plan_options @abstractmethod def plan( diff --git a/embodichain/lab/sim/planners/curobo/curobo_planner.py b/embodichain/lab/sim/planners/curobo/curobo_planner.py index 6259134de..6f73c356f 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo/curobo_planner.py @@ -726,6 +726,7 @@ def _require_curobo(log_level: str = "error") -> "Any": try: planner_mod = importlib.import_module("curobo.motion_planner") batch_mod = importlib.import_module("curobo.batch_motion_planner") + collision_mod = importlib.import_module("curobo.collision_checking") types_mod = importlib.import_module("curobo.types") except ModuleNotFoundError as exc: raise ImportError( @@ -740,6 +741,8 @@ def _require_curobo(log_level: str = "error") -> "Any": MotionPlanner=planner_mod.MotionPlanner, MotionPlannerCfg=planner_mod.MotionPlannerCfg, BatchMotionPlanner=batch_mod.BatchMotionPlanner, + RobotCollisionChecker=collision_mod.RobotCollisionChecker, + RobotCollisionCheckerCfg=collision_mod.RobotCollisionCheckerCfg, JointState=types_mod.JointState, Pose=types_mod.Pose, GoalToolPose=types_mod.GoalToolPose, @@ -827,6 +830,7 @@ class CuroboPlanner(BasePlanner): supported_move_types = frozenset({MoveType.EEF_MOVE, MoveType.JOINT_MOVE}) supports_collision_world_updates = True + supports_joint_trajectory_validation = True @property def preserve_plan_samples(self) -> bool: @@ -982,6 +986,105 @@ def prepare_backend( "use_cuda_graph": backend.use_cuda_graph, } + def validate_joint_trajectory( + self, + trajectory: torch.Tensor, + *, + control_part: str, + obstacle_poses: Mapping[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Validate every supplied joint sample with cuRobo collision models. + + The samples are not replanned or replaced. They are mapped from the + simulator's control-part order into the exact cuRobo model, then checked + against joint bounds, self-collision, and the live world collision + checker. Calling the configuration validator once per horizon sample + works around cuRobo 0.8's configuration-only ``validate`` contract while + retaining batched environments. + """ + if ( + not isinstance(trajectory, torch.Tensor) + or not trajectory.is_floating_point() + or trajectory.dim() != 3 + or 0 in trajectory.shape + or not bool(torch.isfinite(trajectory).all().item()) + ): + raise ValueError("trajectory must be finite floating shape (B, T, D).") + batch_size, horizon, dof = trajectory.shape + backend = self._get_backend( + control_part, + batch_size, + MoveType.JOINT_MOVE, + ) + if dof != len(backend.sim_joint_names): + raise ValueError( + f"Trajectory for {control_part!r} has {dof} joints, expected " + f"{len(backend.sim_joint_names)}." + ) + + dynamic_ids = tuple(self.cfg.world.dynamic_obstacle_names) + if dynamic_ids and obstacle_poses is None: + raise ValueError( + "Live dynamic-obstacle poses are required for exact trajectory " + "validation." + ) + poses = None if obstacle_poses is None else dict(obstacle_poses) + if poses is not None: + _validate_dynamic_obstacles(poses, list(dynamic_ids)) + missing = sorted(set(dynamic_ids).difference(poses)) + extra = sorted(set(poses).difference(dynamic_ids)) + if missing or extra: + raise ValueError( + "Dynamic collision obstacle IDs do not match the planner " + f"configuration; missing={missing}, extra={extra}." + ) + if poses: + base_pose_inv = pose_inv(self._get_sim_base_pose(backend, batch_size)) + self.update_dynamic_obstacles( + poses, + backend, + base_pose_inv, + ) + + if backend.collision_checker is None: + checker_cfg = self._bindings.RobotCollisionCheckerCfg.load_from_config( + robot_config=backend.profile.robot_config_path, + scene_collision_checker=(backend.planner.scene_collision_checker), + device_cfg=self._bindings.DeviceCfg(device=self._curobo_device), + num_envs=batch_size, + collision_activation_distance=self.cfg.collision_activation_distance, + ) + backend.collision_checker = self._bindings.RobotCollisionChecker( + checker_cfg + ) + + self._to_curobo_joint_state(trajectory[:, 0], backend) + assert backend.sim_to_curobo_col_idx is not None + curobo_trajectory = trajectory.to( + device=self._curobo_device, + dtype=torch.float32, + ).index_select(-1, backend.sim_to_curobo_col_idx) + env_query_idx = ( + torch.arange(batch_size, device=self._curobo_device, dtype=torch.int32) + if self.cfg.world.multi_env + else None + ) + samples: list[torch.Tensor] = [] + device_context = ( + torch.cuda.device(self._curobo_device) + if self._curobo_device.type == "cuda" + else nullcontext() + ) + with device_context: + for sample_index in range(horizon): + sample = curobo_trajectory[:, sample_index : sample_index + 1] + valid = backend.collision_checker.validate( + sample, + env_query_idx=env_query_idx, + ) + samples.append(valid[:, 0].to(torch.bool)) + return torch.stack(samples, dim=1).to(trajectory.device) + def with_collision_world( self, options: PlanOptions, @@ -2441,6 +2544,7 @@ class _CuroboBackend: batch_size: int use_cuda_graph: bool planning_mode: MoveType + collision_checker: "Any | None" = None # Lazily-built device-tensor caches for the shared post-processing. The # cuRobo joint order and the profile's fixed transforms are stable for a # planner's life, so these are built once on first use and reused across diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index 0c3e9a4a4..678301241 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -168,6 +168,18 @@ def supports_dynamic_collision_world(self) -> bool: """ return getattr(self.planner, "supports_collision_world_updates", False) is True + @property + def supports_joint_trajectory_validation(self) -> bool: + """Whether the backend checks exact joint samples for collisions.""" + return ( + getattr( + self.planner, + "supports_joint_trajectory_validation", + False, + ) + is True + ) + @property def dynamic_collision_entity_ids(self) -> tuple[str, ...]: """Return canonical dynamic-obstacle IDs declared by the planner.""" @@ -321,6 +333,56 @@ def bind_collision_world( ) return bound + def validate_joint_trajectory( + self, + trajectory: torch.Tensor, + *, + control_part: str, + obstacle_poses: Mapping[str, torch.Tensor] | None = None, + ) -> torch.Tensor: + """Check exact joint samples through the selected planner backend. + + Args: + trajectory: Simulator-order joint samples with shape ``(B, T, D)``. + control_part: Robot control part whose ordered joints form ``D``. + obstacle_poses: Optional live dynamic-obstacle poses. + + Returns: + Boolean validity mask with shape ``(B, T)`` on the trajectory device. + """ + if not self.supports_joint_trajectory_validation: + raise ValueError( + f"Planner {type(self.planner).__name__} does not support exact " + "joint-trajectory collision validation." + ) + if not isinstance(trajectory, torch.Tensor): + raise TypeError("trajectory must be a torch.Tensor.") + if ( + not trajectory.is_floating_point() + or trajectory.dim() != 3 + or 0 in trajectory.shape + or not bool(torch.isfinite(trajectory).all().item()) + ): + raise ValueError( + "trajectory must be finite floating shape (B, T, D) with " + "non-zero dimensions." + ) + if type(control_part) is not str or not control_part: + raise ValueError("control_part must be a non-empty string.") + validity = self.planner.validate_joint_trajectory( + trajectory, + control_part=control_part, + obstacle_poses=obstacle_poses, + ) + if not isinstance(validity, torch.Tensor): + raise TypeError("Planner.validate_joint_trajectory() must return a tensor.") + if validity.dtype != torch.bool or validity.shape != trajectory.shape[:2]: + raise ValueError( + "Planner.validate_joint_trajectory() must return bool shape " + f"{tuple(trajectory.shape[:2])}." + ) + return validity.to(trajectory.device).clone() + def resolve_plan_options( self, plan_opts: PlanOptions | None, diff --git a/tests/gym/envs/expert_program/test_catalog.py b/tests/gym/envs/expert_program/test_catalog.py index 9374d7932..776bbad69 100644 --- a/tests/gym/envs/expert_program/test_catalog.py +++ b/tests/gym/envs/expert_program/test_catalog.py @@ -21,6 +21,7 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, replace from threading import Event, Lock +from types import SimpleNamespace from typing import ClassVar import pytest @@ -37,7 +38,11 @@ decode_expert_program, ) from embodichain.lab.gym.utils.registration import EnvSpec -from embodichain.lab.sim.atomic_actions import Affordance, PlanningContext +from embodichain.lab.sim.atomic_actions import ( + Affordance, + AtomicActionEngine, + PlanningContext, +) from embodichain.lab.sim.skills import ( PLACEMENT_TARGET_AFFORDANCE_REVISION, PLACE_ON_AFFORDANCE_CAPABILITY, @@ -54,6 +59,7 @@ SceneEntityManifest, SceneManifest, SceneObjectRef, + SceneRegistry, SemanticRelationTarget, RegisteredSemanticCall, SemanticCallDescriptor, @@ -248,9 +254,11 @@ def create( *, simulation: object, robot: object, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, ) -> ParallelCommandSafetyValidator: """Return one independent protocol-compatible safety gate.""" - del simulation, robot + del simulation, robot, scene_registry, engine return _AcceptParallelSafety() @@ -285,9 +293,11 @@ def create( *, simulation: object, robot: object, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, ) -> ParallelCommandSafetyValidator: """Block the first call so a second call can attempt registration entry.""" - del simulation, robot + del simulation, robot, scene_registry, engine with self._state_lock: call_index = self._calls type(self)._calls += 1 @@ -344,6 +354,14 @@ def _registration() -> SimulationExpertProgramRegistration: ) +def _parallel_live_inputs() -> tuple[object, SceneRegistry, AtomicActionEngine]: + """Build minimal identity-consistent inputs for factory lifecycle tests.""" + robot = object() + engine = AtomicActionEngine.__new__(AtomicActionEngine) + engine._planning_services = SimpleNamespace(robot=robot) # type: ignore[attr-defined] + return robot, SceneRegistry(), engine + + def _operate_articulation_payload( *, target: str, @@ -547,9 +565,12 @@ def test_parallel_preflight_accepts_exact_registration_owned_safety_factory() -> ) compiled = registration.catalog.preflight(program) + robot, scene_registry, engine = _parallel_live_inputs() validator = registration.create_parallel_safety_validator( simulation=object(), - robot=object(), + robot=robot, + scene_registry=scene_registry, + engine=engine, ) assert tuple(compiled.iter_segments())[0].parallel_block is not None @@ -566,8 +587,15 @@ class InvalidParallelSafetyFactory: {"robot.joint_position"} ) - def create(self, *, simulation: object, robot: object) -> object: - del simulation, robot + def create( + self, + *, + simulation: object, + robot: object, + scene_registry: SceneRegistry, + engine: AtomicActionEngine, + ) -> object: + del simulation, robot, scene_registry, engine return object() registration = SimulationExpertProgramRegistration( @@ -576,10 +604,13 @@ def create(self, *, simulation: object, robot: object) -> object: parallel_safety_factory=InvalidParallelSafetyFactory(), ) + robot, scene_registry, engine = _parallel_live_inputs() with pytest.raises(TypeError, match="must return a ParallelCommandSafetyValidator"): registration.create_parallel_safety_validator( simulation=object(), - robot=object(), + robot=robot, + scene_registry=scene_registry, + engine=engine, ) @@ -592,11 +623,14 @@ def test_parallel_safety_creation_and_history_are_one_registration_lock_scope() robot_profile_binding=create_cube_robot_profile_binding(), parallel_safety_factory=factory_type(), ) + robot, scene_registry, engine = _parallel_live_inputs() def create_validator() -> ParallelCommandSafetyValidator | None: return registration.create_parallel_safety_validator( simulation=object(), - robot=object(), + robot=robot, + scene_registry=scene_registry, + engine=engine, ) with ThreadPoolExecutor(max_workers=2) as executor: diff --git a/tests/gym/envs/expert_program/test_extensions.py b/tests/gym/envs/expert_program/test_extensions.py index e51f8ef8d..602013fc5 100644 --- a/tests/gym/envs/expert_program/test_extensions.py +++ b/tests/gym/envs/expert_program/test_extensions.py @@ -250,8 +250,15 @@ class _MobileSafetyFactory: {_MobileTarget.TRANSPORT_ID} ) - def create(self, *, simulation: object, robot: object) -> _SafetyValidator: - del simulation, robot + def create( + self, + *, + simulation: object, + robot: object, + scene_registry: object, + engine: object, + ) -> _SafetyValidator: + del simulation, robot, scene_registry, engine return _SafetyValidator() diff --git a/tests/gym/envs/expert_program/test_simulation_environment.py b/tests/gym/envs/expert_program/test_simulation_environment.py index 0a2eea1ed..37df143c8 100644 --- a/tests/gym/envs/expert_program/test_simulation_environment.py +++ b/tests/gym/envs/expert_program/test_simulation_environment.py @@ -1064,8 +1064,17 @@ class _RegisteredParallelSafetyFactory: {JointPositionTarget.TRANSPORT_ID} ) - def create(self, *, simulation: object, robot: object) -> _RegisteredParallelSafety: + def create( + self, + *, + simulation: object, + robot: object, + scene_registry: object, + engine: object, + ) -> _RegisteredParallelSafety: assert getattr(simulation, "get_robot")(getattr(robot, "uid")) is robot + assert scene_registry is not None + assert getattr(engine, "robot") is robot return _RegisteredParallelSafety() @@ -1075,8 +1084,15 @@ class _ReusedParallelSafetyFactory(_RegisteredParallelSafetyFactory): validator_id: ClassVar[str] = "test.reused_parallel_safety" _validator: ClassVar[_RegisteredParallelSafety] = _RegisteredParallelSafety() - def create(self, *, simulation: object, robot: object) -> _RegisteredParallelSafety: - del simulation, robot + def create( + self, + *, + simulation: object, + robot: object, + scene_registry: object, + engine: object, + ) -> _RegisteredParallelSafety: + del simulation, robot, scene_registry, engine return self._validator @@ -1090,8 +1106,15 @@ class _AlternatingReusedParallelSafetyFactory(_RegisteredParallelSafetyFactory): ) _next_index: ClassVar[int] = 0 - def create(self, *, simulation: object, robot: object) -> _RegisteredParallelSafety: - del simulation, robot + def create( + self, + *, + simulation: object, + robot: object, + scene_registry: object, + engine: object, + ) -> _RegisteredParallelSafety: + del simulation, robot, scene_registry, engine validator = self._validators[self._next_index % len(self._validators)] type(self)._next_index += 1 return validator diff --git a/tests/gym/envs/expert_program/test_simulation_parallel_safety.py b/tests/gym/envs/expert_program/test_simulation_parallel_safety.py new file mode 100644 index 000000000..ce037af0c --- /dev/null +++ b/tests/gym/envs/expert_program/test_simulation_parallel_safety.py @@ -0,0 +1,238 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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 the production cuRobo parallel-command safety gate.""" + +from __future__ import annotations + +from types import MethodType, SimpleNamespace + +import pytest +import torch + +from embodichain.lab.gym.envs.expert_program import ( + CuroboParallelCommandSafetyValidator, + CuroboParallelSafetyValidatorFactory, +) +from embodichain.lab.sim.atomic_actions import ( + AtomicActionEngine, + EndpointCommand, + JointPositionPayload, + JointPositionTarget, + RuntimeCommandFrame, +) +from embodichain.lab.sim.planners import CuroboPlanner, MotionGenerator +from embodichain.lab.sim.skills import ParallelSafetyError, SceneRegistry + + +class _Robot: + """Expose current full joint state and one aggregate validation part.""" + + def __init__(self) -> None: + self.qpos = torch.zeros(2, 3) + + def get_qpos(self, *, target: bool = False) -> torch.Tensor: + assert target is False + return self.qpos.clone() + + def get_joint_ids(self, *, name: str) -> list[int]: + assert name == "dual_arm" + return [0, 1] + + +def _motion_generator( + *, + reject_env: int | None = None, +) -> tuple[MotionGenerator, list[torch.Tensor]]: + """Build a no-CUDA shell around the exact CuroboPlanner type.""" + observed: list[torch.Tensor] = [] + planner = CuroboPlanner.__new__(CuroboPlanner) + + def validate_joint_trajectory( + self: CuroboPlanner, + trajectory: torch.Tensor, + *, + control_part: str, + obstacle_poses: object, + ) -> torch.Tensor: + del self, obstacle_poses + assert control_part == "dual_arm" + observed.append(trajectory.clone()) + validity = torch.ones(trajectory.shape[:2], dtype=torch.bool) + if reject_env is not None: + validity[reject_env, trajectory.shape[1] // 2] = False + return validity + + planner.validate_joint_trajectory = MethodType( # type: ignore[method-assign] + validate_joint_trajectory, + planner, + ) + generator = MotionGenerator.__new__(MotionGenerator) + generator.planner = planner + return generator, observed + + +def _frame( + *commands: EndpointCommand, + active: tuple[bool, bool] = (True, True), +) -> RuntimeCommandFrame: + """Build one two-row runtime frame.""" + return RuntimeCommandFrame( + commands=commands, + active_mask=torch.tensor(active, dtype=torch.bool), + env_ids=torch.tensor((0, 1), dtype=torch.long), + hold_duration=torch.full((2,), 0.05), + ) + + +def _command( + target_id: str, + joint_id: int, + positions: tuple[float, float], +) -> EndpointCommand: + """Build one single-joint batched command.""" + return EndpointCommand( + target=JointPositionTarget(target_id, (joint_id,)), + payload=JointPositionPayload( + positions=torch.tensor(positions, dtype=torch.float32).unsqueeze(1) + ), + ) + + +def _validator( + *, + reject_env: int | None = None, + max_joint_step: float = 0.05, + max_interpolation_samples: int = 16, +) -> tuple[CuroboParallelCommandSafetyValidator, list[torch.Tensor]]: + """Create one validator with a deterministic collision backend shell.""" + generator, observed = _motion_generator(reject_env=reject_env) + return ( + CuroboParallelCommandSafetyValidator( + robot=_Robot(), + motion_generator=generator, + scene_registry=SceneRegistry(), + validation_control_part="dual_arm", + max_joint_step=max_joint_step, + max_interpolation_samples=max_interpolation_samples, + ), + observed, + ) + + +def test_curobo_parallel_safety_checks_exact_dense_merged_segment() -> None: + """Disjoint lane targets become one densely sampled aggregate trajectory.""" + validator, observed = _validator() + left = _command("left_arm", 0, (0.1, 0.0)) + right = _command("right_arm", 1, (0.2, -0.1)) + + validator.validate( + branch_frames={"left": _frame(left), "right": _frame(right)}, + merged_frame=_frame(left, right), + ) + + assert len(observed) == 1 + trajectory = observed[0] + # float32 represents 0.2 just above the mathematical value, so the strict + # 0.05 maximum step requires five intervals rather than rounding down. + assert trajectory.shape == (2, 6, 2) + torch.testing.assert_close(trajectory[:, 0], torch.zeros(2, 2)) + torch.testing.assert_close( + trajectory[:, -1], + torch.tensor(((0.1, 0.2), (0.0, -0.1))), + ) + + +def test_curobo_parallel_safety_reports_row_local_collision() -> None: + """One invalid environment rejects dispatch with its stable env ID.""" + validator, _ = _validator(reject_env=1) + left = _command("left_arm", 0, (0.1, 0.1)) + right = _command("right_arm", 1, (0.2, 0.2)) + + with pytest.raises(ParallelSafetyError, match=r"env IDs \(1,\)"): + validator.validate( + branch_frames={"left": _frame(left), "right": _frame(right)}, + merged_frame=_frame(left, right), + ) + + +def test_curobo_parallel_safety_rejects_uncovered_joint() -> None: + """Every outgoing joint must belong to the aggregate collision model.""" + validator, _ = _validator() + left = _command("left_arm", 0, (0.1, 0.1)) + hand = _command("left_hand", 2, (0.2, 0.2)) + + with pytest.raises(ParallelSafetyError, match="outside validation control part"): + validator.validate( + branch_frames={"left": _frame(left), "hand": _frame(hand)}, + merged_frame=_frame(left, hand), + ) + + +def test_curobo_parallel_safety_fails_instead_of_under_sampling() -> None: + """The configured memory bound cannot silently enlarge the joint step.""" + validator, _ = _validator( + max_joint_step=0.01, + max_interpolation_samples=4, + ) + left = _command("left_arm", 0, (0.1, 0.1)) + right = _command("right_arm", 1, (0.0, 0.0)) + + with pytest.raises(ParallelSafetyError, match="exceeding configured limit"): + validator.validate( + branch_frames={"left": _frame(left), "right": _frame(right)}, + merged_frame=_frame(left, right), + ) + + +@pytest.mark.parametrize( + "kwargs", + ( + {"validation_control_part": ""}, + {"validation_control_part": "dual_arm", "max_joint_step": 0.0}, + { + "validation_control_part": "dual_arm", + "max_interpolation_samples": 1, + }, + ), +) +def test_curobo_parallel_safety_factory_validates_configuration( + kwargs: dict[str, object], +) -> None: + """Invalid safety declarations fail during task registration.""" + with pytest.raises((TypeError, ValueError)): + CuroboParallelSafetyValidatorFactory(**kwargs) # type: ignore[arg-type] + + +def test_curobo_parallel_safety_factory_binds_exact_runtime_components() -> None: + """The production factory consumes the assembled engine and registry.""" + robot = _Robot() + motion_generator, _ = _motion_generator() + engine = AtomicActionEngine.__new__(AtomicActionEngine) + engine._planning_services = SimpleNamespace( # type: ignore[attr-defined] + robot=robot, + motion_generator=motion_generator, + ) + factory = CuroboParallelSafetyValidatorFactory(validation_control_part="dual_arm") + + validator = factory.create( + simulation=object(), + robot=robot, + scene_registry=SceneRegistry(), + engine=engine, + ) + + assert type(validator) is CuroboParallelCommandSafetyValidator diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index 8cfd6e2cd..b76a9f7cb 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -698,6 +698,74 @@ def test_dynamic_update_uses_registry_id_in_curobo_backend(): assert [(name, env_idx) for name, _, env_idx in updates] == [("registry_cube", 0)] +def test_validate_joint_trajectory_checks_every_exact_sample_in_curobo_order(): + """The collision gate preserves samples and maps simulator joint order.""" + planner = object.__new__(CuroboPlanner) + planner.cfg = CuroboPlannerCfg( + robot_uid="robot", + world=CuroboWorldCfg(multi_env=True), + ) + planner._curobo_device = torch.device("cpu") + joint_states = [] + collision_queries = [] + + def from_position(position, *, joint_names): + joint_states.append((position.clone(), tuple(joint_names))) + return SimpleNamespace(position=position) + + def validate(sample, *, env_query_idx): + collision_queries.append((sample.clone(), env_query_idx.clone())) + validity = torch.ones(sample.shape[:2], dtype=torch.bool) + if len(collision_queries) == 2: + validity[1, 0] = False + return validity + + planner._bindings = SimpleNamespace( + JointState=SimpleNamespace(from_position=from_position), + ) + backend = SimpleNamespace( + sim_joint_names=["sim_left", "sim_right"], + sim_to_curobo_col_idx=None, + collision_checker=SimpleNamespace(validate=validate), + profile=SimpleNamespace( + sim_to_curobo_joint_names={ + "sim_left": "curobo_left", + "sim_right": "curobo_right", + }, + ), + planner=SimpleNamespace( + joint_names=["curobo_right", "curobo_left"], + ), + ) + planner._get_backend = lambda control_part, batch_size, move_type: backend + trajectory = torch.tensor( + ( + ((0.0, 1.0), (0.1, 1.1), (0.2, 1.2)), + ((2.0, 3.0), (2.1, 3.1), (2.2, 3.2)), + ), + dtype=torch.float32, + ) + + validity = planner.validate_joint_trajectory( + trajectory, + control_part="dual_arm", + ) + + assert torch.equal( + validity, + torch.tensor(((True, True, True), (True, False, True))), + ) + assert len(collision_queries) == trajectory.shape[1] + for sample_index, (sample, env_query_idx) in enumerate(collision_queries): + torch.testing.assert_close( + sample[:, 0], + trajectory[:, sample_index].flip(dims=(-1,)), + ) + assert torch.equal(env_query_idx, torch.tensor((0, 1), dtype=torch.int32)) + torch.testing.assert_close(joint_states[0][0], trajectory[:, 0].flip(dims=(-1,))) + assert joint_states[0][1] == ("curobo_right", "curobo_left") + + def test_generate_mesh_world_yaml_assembles_schema(tmp_path): rigid_object = _FakeRigidObject( "demo_block",