From 34d2353faf5cf0f223bfddc3d20709261b05b09d Mon Sep 17 00:00:00 2001 From: skywhite1024 <129768272+skywhite1024@users.noreply.github.com> Date: Thu, 27 Aug 2026 14:50:47 +0800 Subject: [PATCH] feat(action-engine): add runtime grounding and capability adapters --- .../embodichain.gen_sim.action_engine.rst | 33 + .../gen_sim/action_engine/runtime/actions.py | 3038 ++++++++++++++++ .../action_engine/runtime/atomic_compat.py | 106 + .../action_engine/runtime/body_grasp.py | 297 ++ .../runtime/coordinated_safety.py | 391 ++ .../gen_sim/action_engine/runtime/frames.py | 165 + .../action_engine/runtime/geometry_axes.py | 102 + .../runtime/grasp_diagnostics.py | 535 +++ .../action_engine/runtime/grounding.py | 3157 +++++++++++++++++ .../action_engine/runtime/predicates.py | 891 +++++ .../action_engine/runtime/robot_parts.py | 34 + .../action_engine/runtime/solver_compat.py | 234 ++ .../gen_sim/action_engine/runtime/state.py | 63 + .../action_engine/capabilities/__init__.py | 19 + .../capabilities/test_atomic_v2.py | 469 +++ .../capabilities/test_held_hand_over.py | 311 ++ .../gen_sim/action_engine/runtime/__init__.py | 19 + .../action_engine/runtime/test_actions.py | 1834 ++++++++++ .../runtime/test_atomic_compat.py | 141 + .../action_engine/runtime/test_body_grasp.py | 154 + .../runtime/test_coordinated_safety.py | 185 + .../runtime/test_grasp_diagnostics.py | 227 ++ .../action_engine/tasks/test_e3_pour.py | 132 + 23 files changed, 12537 insertions(+) create mode 100644 embodichain/gen_sim/action_engine/runtime/actions.py create mode 100644 embodichain/gen_sim/action_engine/runtime/atomic_compat.py create mode 100644 embodichain/gen_sim/action_engine/runtime/body_grasp.py create mode 100644 embodichain/gen_sim/action_engine/runtime/coordinated_safety.py create mode 100644 embodichain/gen_sim/action_engine/runtime/frames.py create mode 100644 embodichain/gen_sim/action_engine/runtime/geometry_axes.py create mode 100644 embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py create mode 100644 embodichain/gen_sim/action_engine/runtime/grounding.py create mode 100644 embodichain/gen_sim/action_engine/runtime/predicates.py create mode 100644 embodichain/gen_sim/action_engine/runtime/robot_parts.py create mode 100644 embodichain/gen_sim/action_engine/runtime/solver_compat.py create mode 100644 tests/gen_sim/action_engine/capabilities/__init__.py create mode 100644 tests/gen_sim/action_engine/capabilities/test_atomic_v2.py create mode 100644 tests/gen_sim/action_engine/capabilities/test_held_hand_over.py create mode 100644 tests/gen_sim/action_engine/runtime/__init__.py create mode 100644 tests/gen_sim/action_engine/runtime/test_actions.py create mode 100644 tests/gen_sim/action_engine/runtime/test_atomic_compat.py create mode 100644 tests/gen_sim/action_engine/runtime/test_body_grasp.py create mode 100644 tests/gen_sim/action_engine/runtime/test_coordinated_safety.py create mode 100644 tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py create mode 100644 tests/gen_sim/action_engine/tasks/test_e3_pour.py diff --git a/docs/source/api_reference/embodichain/embodichain.gen_sim.action_engine.rst b/docs/source/api_reference/embodichain/embodichain.gen_sim.action_engine.rst index a7ccfedde..6549bb634 100644 --- a/docs/source/api_reference/embodichain/embodichain.gen_sim.action_engine.rst +++ b/docs/source/api_reference/embodichain/embodichain.gen_sim.action_engine.rst @@ -159,6 +159,30 @@ Supporting planning utilities .. automodule:: embodichain.gen_sim.action_engine.runtime :members: +.. automodule:: embodichain.gen_sim.action_engine.runtime.actions + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.atomic_compat + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.body_grasp + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.coordinated_safety + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.frames + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.geometry_axes + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.grounding + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.grasp_diagnostics + :members: + .. automodule:: embodichain.gen_sim.action_engine.runtime.loader :members: @@ -168,5 +192,14 @@ Supporting planning utilities .. automodule:: embodichain.gen_sim.action_engine.runtime.motion_policy :members: +.. automodule:: embodichain.gen_sim.action_engine.runtime.predicates + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.robot_parts + :members: + +.. automodule:: embodichain.gen_sim.action_engine.runtime.solver_compat + :members: + .. automodule:: embodichain.gen_sim.action_engine.runtime.state :members: diff --git a/embodichain/gen_sim/action_engine/runtime/actions.py b/embodichain/gen_sim/action_engine/runtime/actions.py new file mode 100644 index 000000000..64b8d518b --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/actions.py @@ -0,0 +1,3038 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Adapt Action Engine requests to the shared typed atomic-action planner.""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager, nullcontext +from copy import deepcopy +from dataclasses import replace +import logging +import math +from threading import RLock +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapability, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.config import default_runtime_policy +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile +from embodichain.gen_sim.action_engine.solver_profiles import ( + expected_ik_solver_class, +) +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + ActionPlan, + AntipodalAffordance, + AtomicActionEngine, + AxisAlignGoal, + ControlPartCommandProfile, + CoordinatedPickGoal, + DynamicCollisionMode, + EndEffectorPoseGoal, + EntityState, + ExecutionSession, + MotionPolicy, + ObjectSemantics, + PlanningContext, + PlanningFailure, + PlannerDiagnostics, + RecoveryPolicy, + RobotObservation, + RigidObjectSceneProvider, + SceneProvider, + SceneSnapshot, + StateDelta, +) +from embodichain.lab.sim.planners import ( + CuroboPlannerCfg, + CuroboWorldCfg, + MotionGenCfg, + MotionGenerator, + ToppraPlannerCfg, +) +from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg +from embodichain.toolkits.graspkit.pg_grasp import ( + AntipodalGraspPoseGenerator, + AntipodalGraspPoseGeneratorCfg, + GraspAnnotationCfg, + ParallelJawGraspCollisionCfg, +) +from embodichain.utils import logger as project_logger +from embodichain.utils.logger import log_info, log_warning +from embodichain.utils.math import matrix_from_quat, quat_from_matrix, quat_slerp + +from .body_grasp import AxisAlignBodyGraspAdapter +from .coordinated_safety import _trajectory_safety_report +from .grasp_diagnostics import _TracingAntipodalGraspPoseGenerator +from .models import ActionOutcome, GroundedAction +from .state import _CollisionOverrideSceneSnapshot, ExecutionState + +__all__ = ["AtomicActionAdapter"] + + +_DEFAULT_PLANNER_POLICY: dict[str, Any] = { + "backend": "curobo", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "fallback_strategy": "ik_interp", + "allow_fallback": True, + "dynamic_collision": False, + "static_obstacle_uids": [], + "dynamic_obstacle_uids": [], + "curobo": { + "log_level": "error", + "obstacle_representation": "cuboid", + "multi_env": False, + "use_cuda_graph": True, + "preserve_plan_samples": False, + "max_attempts": 5, + "collision_activation_distance": 0.01, + }, +} + +# Preserve cuRobo's fixed world shape while disabling intentional-contact objects. +_COLLISION_PARKING_Z_OFFSET = -100.0 +_BODY_GRASP_CANDIDATE_LIMIT = 500 +_BODY_GRASP_SEED = 17_392 +_COORDINATED_GRASP_SEED = 17_393 +_FREE_YAW_SAMPLE_COUNT = 8 +_RETREAT_LOG_LOCK = RLock() + + +@contextmanager +def _capture_retreat_warnings(enabled: bool) -> Iterator[list[str]]: + """Capture candidate-level planner warnings during bounded retreat search.""" + messages: list[str] = [] + if not enabled: + yield messages + return + collector = logging.Handler(level=logging.WARNING) + collector.emit = lambda record: messages.append(record.getMessage()) + logger = project_logger.logger + with _RETREAT_LOG_LOCK: + handlers = list(logger.handlers) + propagate = logger.propagate + try: + logger.handlers[:] = [collector] + logger.propagate = False + yield messages + finally: + logger.handlers[:] = handlers + logger.propagate = propagate + + +def _collision_cache_for_world( + representation: str, obstacle_count: int +) -> dict[str, int]: + """Size cuRobo's fixed collision cache for the generated scene.""" + cache = {"cuboid": 8, "mesh": 2} + if representation in cache: + cache[representation] = max(cache[representation], obstacle_count) + return cache + + +def _supported_kwargs(config_type: type, values: Mapping[str, Any]) -> dict[str, Any]: + names: set[str] = set() + for cls in reversed(config_type.__mro__): + names.update(getattr(cls, "__annotations__", {})) + return {key: value for key, value in values.items() if key in names} + + +def _as_hand_qpos(value: Any, dof: int, device: Any) -> torch.Tensor: + if dof == 0: + return torch.empty(0, dtype=torch.float32, device=device) + result = torch.as_tensor(value, dtype=torch.float32, device=device).flatten() + if result.numel() == 0: + return torch.zeros(dof, dtype=torch.float32, device=device) + if result.numel() == 1: + return result.repeat(dof) + if result.numel() >= dof: + return result[:dof] + repeats = (dof + result.numel() - 1) // result.numel() + return result.repeat(repeats)[:dof] + + +def _diagonal_approach_direction( + horizontal: torch.Tensor, + *, + vertical: float = -1.0, +) -> torch.Tensor: + """Combine one normalized horizontal role direction with a vertical component.""" + horizontal = horizontal.to(dtype=torch.float32) + norm = torch.linalg.vector_norm(horizontal) + if float(norm) <= 1.0e-6: + raise ValueError("Handover role direction must be non-zero.") + horizontal = horizontal / norm + direction = torch.stack( + (horizontal[0], horizontal[1], horizontal.new_tensor(float(vertical))) + ) + return direction / torch.linalg.vector_norm(direction) + + +class AtomicActionAdapter: + """Own the shared atomic engine and preserve Action Engine runtime contracts.""" + + def __init__( + self, + env: Any, + *, + grasp_policy: Mapping[str, Any] | None = None, + planner_policy: Mapping[str, Any] | None = None, + capability_registry: Any | None = None, + scene_provider: SceneProvider | None = None, + ) -> None: + self.env = env + self.num_envs = int(env.num_envs) + self.device = env.device + self.gripper_profile = get_gripper_profile( + getattr(env, "agent_gripper_model", "pgi") + ) + self.ik_solver, self.ik_solver_classes = self._resolve_runtime_ik_solver() + if grasp_policy is None: + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + grasp_policy = default_runtime_policy(profile).grasp + grasp_policy = { + **grasp_policy, + **(getattr(env, "agent_grasp_runtime_defaults", {}) or {}), + } + self.grasp_policy = deepcopy(dict(grasp_policy)) + self.planner_policy = deepcopy(_DEFAULT_PLANNER_POLICY) + if planner_policy is not None: + self._merge_planner_policy(self.planner_policy, planner_policy) + if not self.planner_policy.get("static_obstacle_uids"): + configured = getattr(env, "agent_static_obstacle_uids", ()) or () + if configured: + self.planner_policy["static_obstacle_uids"] = [ + str(uid) for uid in configured + ] + else: + get_rigid_object = getattr(env.sim, "get_rigid_object", None) + if callable(get_rigid_object) and get_rigid_object("table") is not None: + self.planner_policy["static_obstacle_uids"] = ["table"] + self.capabilities = capability_registry or build_atomic_capability_registry() + self._motion_generator: MotionGenerator | None = None + self._atomic_engine: AtomicActionEngine | None = None + self._coordinated_engines: dict[tuple[bool, float], AtomicActionEngine] = {} + self._semantics: dict[str, ObjectSemantics] = {} + self._scene_time = 0.0 + if scene_provider is not None and not isinstance(scene_provider, SceneProvider): + raise TypeError("scene_provider must implement SceneProvider.") + self.scene_provider = scene_provider or self._build_scene_provider() + + def _resolve_runtime_ik_solver(self) -> tuple[str, dict[str, str]]: + """Validate a declared bundle solver against the initialized robot.""" + declared = getattr(self.env, "agent_ik_solver", None) + robot = getattr(self.env, "robot", None) + get_solver = getattr(robot, "get_solver", None) + classes: dict[str, str] = {} + if callable(get_solver): + for arm in ("left_arm", "right_arm"): + part = self.env.get_agent_arm_control_part(arm == "left_arm") + classes[arm] = type(get_solver(name=part)).__name__ + if declared is None: + inferred = { + "URSolver": "ur", + "PytorchSolver": "pytorch", + } + modes = {inferred[name] for name in classes.values() if name in inferred} + return (modes.pop() if len(modes) == 1 else "unknown"), classes + declared = str(declared) + expected = expected_ik_solver_class(declared) + for arm, actual in classes.items(): + if actual != expected: + raise ValueError( + f"Runtime {arm} must use {expected} for " + f"agent_ik_solver={declared!r}, got {actual!r}." + ) + return declared, classes + + @staticmethod + def _merge_planner_policy( + target: dict[str, Any], + update: Mapping[str, Any], + ) -> None: + for key, value in update.items(): + if isinstance(value, Mapping) and isinstance(target.get(key), dict): + AtomicActionAdapter._merge_planner_policy(target[key], value) + else: + target[key] = deepcopy(value) + + def initial_state(self) -> ExecutionState: + """Capture the initial full-robot planning seed.""" + return ExecutionState(last_qpos=self.env.robot.get_qpos().clone()) + + def start_session( + self, + grounded: GroundedAction, + state: ExecutionState | None = None, + ) -> ExecutionSession: + """Start one closed-loop AtomicAction session from live scene state. + + ProgramExecutor may continue using its compatibility scheduler for + compound and per-arm merged trajectories. New callers can use this + boundary to adopt feedback-driven execution without constructing + private planning contexts. + """ + capability = self.capabilities.require_executable(grounded.action_class) + state = state or self.initial_state() + grounded = self._select_transport_yaw(grounded, state) + grounded = self._adapt_coordinated_pickment_grasps( + grounded, + capability, + )[0] + context = self._planning_context(state, grounded) + engine = self._engine_for(grounded, capability) + invocation = self._invocation(grounded, capability, engine=engine) + return engine.start((invocation,), context) + + def _build_scene_provider(self) -> SceneProvider | None: + """Create the shared live rigid-object provider when entities are available.""" + sim = getattr(self.env, "sim", None) + if sim is None: + return None + dynamic_uids = tuple( + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ) + list_uids = getattr(sim, "get_rigid_object_uid_list", None) + uids = tuple(str(uid) for uid in list_uids()) if callable(list_uids) else () + if not uids: + uids = dynamic_uids + get_rigid_object = getattr(sim, "get_rigid_object", None) + if not callable(get_rigid_object): + return None + entities = { + uid: entity for uid in uids if (entity := get_rigid_object(uid)) is not None + } + if not entities: + return None + collision_uids = ( + dynamic_uids + if bool(self.planner_policy.get("dynamic_collision", False)) + else () + ) + return RigidObjectSceneProvider( + entities, + collision_entity_ids=collision_uids, + ) + + def semantics(self, uid: str) -> ObjectSemantics: + """Build object semantics once for the stable scene entity ID.""" + cached = self._semantics.get(uid) + if cached is not None: + return cached + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr( + self.env.sim, + "get_articulation", + lambda _uid: None, + )(uid) + if entity is None: + raise ValueError(f"Unknown grasp target {uid!r}.") + active_joint_ids = list(getattr(entity, "active_joint_ids", ())) + if len(active_joint_ids) != 1: + raise ValueError( + "Articulation semantics require exactly one active joint." + ) + backend_entities = getattr( + entity, + "_entities", + getattr(entity, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + joint_name = str(entity.joint_names[active_joint_ids[0]]) + joint_info = backend_entities[0].get_joint_info(joint_name) + child_link = str(getattr(joint_info, "child_link_name", "")) + vertices, triangles = entity.get_link_vert_face(child_link) + else: + vertices = entity.get_vertices(env_ids=[0], scale=True) + triangles = entity.get_triangles(env_ids=[0]) + if isinstance(vertices, (tuple, list)): + vertices = vertices[0] + if isinstance(triangles, (tuple, list)): + triangles = triangles[0] + vertices = torch.as_tensor(vertices, dtype=torch.float32) + triangles = torch.as_tensor(triangles, dtype=torch.int64) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if triangles.ndim == 3 and triangles.shape[0] == 1: + triangles = triangles[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Object {uid!r} has invalid mesh vertices.") + if triangles.ndim != 2 or triangles.shape[-1] != 3 or triangles.numel() == 0: + raise ValueError(f"Object {uid!r} has invalid mesh triangles.") + + semantics = ObjectSemantics( + label=uid, + entity_id=uid, + geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + affordance=AntipodalAffordance( + object_label=uid, + mesh_vertices=vertices, + mesh_triangles=triangles, + ), + ) + self._semantics[uid] = semantics + return semantics + + def plan( + self, + grounded: GroundedAction, + state: ExecutionState | None = None, + ) -> ActionOutcome: + """Plan one grounded primitive through the mainline typed contract.""" + capability = self.capabilities.require_executable(grounded.action_class) + state = state or self.initial_state() + grounded = self._select_transport_yaw(grounded, state) + context = self._planning_context(state, grounded) + coordinated_candidates = self._adapt_coordinated_pickment_grasps( + grounded, + capability, + ) + grounded_candidates = tuple( + reorientation + for coordinated in coordinated_candidates + for candidate in self._adapt_axis_align_body_grasps( + coordinated, + context, + capability, + ) + for reorientation in self._adapt_tool_down_candidates(candidate) + ) + selected: ( + tuple[ + GroundedAction, + ActionInvocation, + ActionPlan, + AtomicActionEngine, + ] + | None + ) = None + selected_warnings: tuple[str, ...] = () + candidate_search_warnings: list[str] = [] + candidate_search_attempts = 0 + coordinated_search_attempts: list[dict[str, Any]] = [] + best_failure_count = self.num_envs + 1 + for candidate in grounded_candidates: + candidate_engine = self._engine_for(candidate, capability) + candidate_invocation = self._invocation( + candidate, + capability, + engine=candidate_engine, + ) + capture_warnings = bool( + candidate.motion_policy.get("retreat_reachability_search", False) + or candidate.motion_policy.get("reorient_tool_down", False) + ) + grasp_seed = candidate.motion_policy.get("grasp_seed") + seed_context = ( + nullcontext() + if grasp_seed is None + else self._isolated_random_seed(int(grasp_seed)) + ) + pair_context = self._coordinated_pair_selection_context( + candidate_engine, + candidate, + context, + ) + upright_context = self._upright_grasp_selection_context( + candidate_engine, + candidate, + capability, + ) + with ( + seed_context, + pair_context, + upright_context, + _capture_retreat_warnings(capture_warnings) as warnings, + ): + candidate_plan = candidate_engine.plan(candidate_invocation, context) + self._record_selected_upright_grasp( + candidate, + candidate_plan, + context, + ) + coordinated_trace = candidate.motion_policy.get("coordinated_grasp") + if isinstance(coordinated_trace, dict): + grasp_stages = ( + self._latest_coordinated_grasp_trace(candidate_engine) or {} + ) + coordinated_trace["stages"] = grasp_stages + raw_plan_success = candidate_plan.plan_success.detach().clone() + candidate_plan, trajectory_audit = self._audit_coordinated_trajectory( + candidate, + candidate_invocation, + candidate_plan, + context, + grasp_stages, + ) + coordinated_trace["trajectory_audit"] = trajectory_audit + coordinated_search_attempts.append( + { + "candidate_index": coordinated_trace.get("candidate_index"), + "approach_candidate_label": coordinated_trace.get( + "approach_candidate_label" + ), + "approach_direction": coordinated_trace.get( + "approach_direction" + ), + "middle_empty_ratio": coordinated_trace.get( + "selected_middle_empty_ratio" + ), + "raw_plan_success": raw_plan_success.cpu().tolist(), + "plan_success": candidate_plan.plan_success.detach() + .cpu() + .tolist(), + "planner_messages": list(candidate_plan.diagnostics.messages), + "stages": deepcopy(grasp_stages), + "trajectory_audit": deepcopy(trajectory_audit), + } + ) + if capture_warnings: + candidate_search_warnings.extend(warnings) + candidate_search_attempts += 1 + failure_count = int((~candidate_plan.plan_success).sum().item()) + if selected is None or failure_count < best_failure_count: + selected = ( + candidate, + candidate_invocation, + candidate_plan, + candidate_engine, + ) + selected_warnings = tuple(warnings) + best_failure_count = failure_count + if failure_count == 0: + break + if selected is None: + raise RuntimeError("Atomic action adaptation produced no plan candidate.") + grounded, invocation, plan, selected_engine = selected + selected_coordinated_trace = grounded.motion_policy.get("coordinated_grasp") + if isinstance(selected_coordinated_trace, dict): + selected_coordinated_trace["search_attempts"] = deepcopy( + coordinated_search_attempts + ) + if bool(grounded.motion_policy.get("reorient_tool_down", False)): + summary = ( + "Tool-down reorientation search: " + f"resolved={int(plan.plan_success.sum())}/{plan.plan_success.numel()}, " + f"attempts={candidate_search_attempts}, " + f"selected_yaw_degrees=" + f"{grounded.motion_policy.get('reorient_selected_yaw_degrees')}, " + f"suppressed_warnings={len(candidate_search_warnings)}." + ) + if bool(plan.plan_success.all()): + log_info(summary) + else: + log_warning(summary) + selected_positions = self._positions_with_agent_holds( + plan, + grounded, + capability, + ) + primary_success = plan.plan_success.to(self.device) + reachability_search = None + if bool(grounded.motion_policy.get("retreat_reachability_search", False)): + ( + grounded, + selected_positions, + primary_success, + reachability_search, + ) = self._search_reachable_retreat( + grounded=grounded, + capability=capability, + state=state, + context=context, + invocation=invocation, + initial_positions=selected_positions, + initial_success=primary_success, + initial_warnings=selected_warnings, + ) + invocation = replace(invocation, goal=grounded.target) + combined_success = primary_success.clone() + fallback_plan: ActionPlan | None = None + use_fallback = torch.zeros_like(combined_success) + fallback_attempted = torch.zeros_like(combined_success) + fallback_success = torch.zeros_like(combined_success) + + fallback_strategy = self.planner_policy.get("fallback_strategy") + collision_safety = str(grounded.motion_policy.get("collision_safety", "auto")) + fallback_allowed = bool(self.planner_policy.get("allow_fallback", True)) and ( + collision_safety != "required" + ) + if ( + fallback_allowed + and invocation.motion_policy.strategy == "motion_gen" + and fallback_strategy in {"ik_interp"} + and not bool(combined_success.all()) + ): + fallback_attempted = ~primary_success + fallback_policy = replace( + invocation.motion_policy, + strategy=str(fallback_strategy), + dynamic_collision_mode=DynamicCollisionMode.OFF, + plan_opts=None, + ) + fallback_plan = selected_engine.plan( + replace(invocation, motion_policy=fallback_policy), + context, + ) + fallback_positions = self._positions_with_agent_holds( + fallback_plan, + grounded, + capability, + ) + fallback_success = fallback_plan.plan_success.to(self.device) + use_fallback = fallback_attempted & fallback_success + selected_positions = self._merge_plan_rows( + selected_positions, + fallback_positions, + use_fallback, + state.last_qpos, + ) + combined_success |= fallback_plan.plan_success.to(self.device) + + options = invocation.skill_options + if capability.config_materializer == "handover": + combined_success &= self._handover_receiver_hold_mask( + selected_positions, + grounded, + options, + tolerance=float( + grounded.motion_policy.get( + "receiver_hold_joint_tolerance", + 2.0e-3, + ) + ), + ) + + terminal_qpos = ( + selected_positions[:, -1] + if selected_positions.shape[1] + else state.last_qpos + ) + primary_rows = combined_success & primary_success + projected_task = plan.expected_effects.apply( + context.task, + primary_rows, + ) + held_keys = set(plan.expected_effects.held_object_updates) + if fallback_plan is not None: + fallback_rows = combined_success & use_fallback + projected_task = fallback_plan.expected_effects.apply( + projected_task, + fallback_rows, + ) + held_keys.update(fallback_plan.expected_effects.held_object_updates) + committed_effects = StateDelta( + held_object_updates={ + key: projected_task.held_objects.get(key) for key in held_keys + }, + ) + next_state = ExecutionState.from_task_state( + projected_task, + last_qpos=torch.where( + combined_success[:, None], terminal_qpos, state.last_qpos + ), + ) + return ActionOutcome( + trajectory=selected_positions, + success=combined_success, + next_state=next_state, + grounded=grounded, + prior_state=state, + expected_effects=committed_effects, + planner_trace={ + **self._planner_trace( + grounded=grounded, + invocation=invocation, + context=context, + state=state, + primary_success=primary_success, + primary_diagnostics=plan.diagnostics, + fallback_allowed=fallback_allowed, + fallback_strategy=( + str(fallback_strategy) + if invocation.motion_policy.strategy == "motion_gen" + and fallback_strategy in {"ik_interp"} + else None + ), + fallback_attempted=fallback_attempted, + fallback_success=fallback_success, + fallback_used=use_fallback, + reachability_search=reachability_search, + ), + # Auditability takes precedence over compactness here: every + # selected planner route retains its complete joint path. + "planned_trajectory": selected_positions.detach().clone(), + "primary_action_diagnostics": deepcopy(dict(plan.diagnostics.metadata)), + "fallback_action_diagnostics": ( + None + if fallback_plan is None + else deepcopy(dict(fallback_plan.diagnostics.metadata)) + ), + "action_segments": { + segment.name: { + "start": int(segment.start), + "stop": int(segment.stop), + } + for segment in plan.segments + }, + }, + ) + + def _adapt_axis_align_body_grasps( + self, + grounded: GroundedAction, + context: PlanningContext, + capability: AtomicCapability, + ) -> tuple[GroundedAction, ...]: + if capability.target_materializer != "axis_align": + return (grounded,) + goal = grounded.target + if not isinstance(goal, AxisAlignGoal) or goal.grasp_xpos is not None: + return (grounded,) + object_pose = grounded.object_pose + if ( + not isinstance(object_pose, torch.Tensor) + or object_pose.shape != (self.num_envs, 4, 4) + or not torch.isfinite(object_pose).all() + ): + raise ValueError( + "AxisAlign body grasp requires a finite grounded live object pose " + f"with shape ({self.num_envs}, 4, 4)." + ) + _, hand_part, _ = self._parts(grounded.arm) + if hand_part is None: + raise ValueError("AxisAlign body grasp requires a configured hand part.") + options = self._build_config(grounded, capability) + selected_approach = options.approach_direction + adaptation = AxisAlignBodyGraspAdapter().adapt( + goal, + object_pose=object_pose, + grasp_generator=self._engine().grasp_pose_generators[hand_part], + approach_direction=selected_approach, + target_axis=options.target_axis, + seed=_BODY_GRASP_SEED, + ) + cfg = dict(grounded.cfg) + cfg["approach_direction"] = selected_approach + candidates: list[GroundedAction] = [] + for adaptation_index, candidate_goal in enumerate(adaptation.alternative_goals): + rank = adaptation.alternative_rank_indices[adaptation_index] + policy = dict(grounded.motion_policy) + policy["body_grasp"] = { + "long_axis_index": adaptation.axes.long_axis_index, + "short_axis_index": adaptation.axes.short_axis_index, + "elongation_ratio": adaptation.axes.elongation_ratio, + "candidate_indices": ( + adaptation.selection.ranked_candidate_indices[:, rank].tolist() + ), + "candidate_counts": ( + adaptation.selection.body_candidate_counts.tolist() + ), + "candidate_rank": rank, + "approach_direction": selected_approach.detach().cpu().tolist(), + } + candidates.append( + replace( + grounded, + target=candidate_goal, + cfg=cfg, + motion_policy=policy, + ) + ) + return tuple(candidates) + + def _adapt_tool_down_candidates( + self, + grounded: GroundedAction, + ) -> tuple[GroundedAction, ...]: + """Build fixed-position, downward-TCP yaw candidates after E2 lift-clear.""" + if not bool(grounded.motion_policy.get("reorient_tool_down", False)): + return (grounded,) + reference = grounded.motion_policy.get("reorient_reference_pose") + if not isinstance(reference, torch.Tensor): + raise ValueError("Tool-down reorientation requires a reference TCP pose.") + reference = reference.to(device=self.device, dtype=torch.float32) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + if reference.shape != (self.num_envs, 4, 4): + raise ValueError( + "Tool-down reorientation reference must have shape " + f"({self.num_envs}, 4, 4)." + ) + waypoint_count = int(grounded.cfg.get("reorient_waypoint_count", 5)) + if not 2 <= waypoint_count <= 16: + raise ValueError("reorient_waypoint_count must be in [2, 16].") + raw_yaws = grounded.cfg.get( + "reorient_yaw_degrees", (0.0, 45.0, -45.0, 90.0, -90.0, 180.0) + ) + if not isinstance(raw_yaws, Sequence) or isinstance(raw_yaws, (str, bytes)): + raise TypeError("reorient_yaw_degrees must be a sequence.") + yaw_degrees = [float(value) for value in raw_yaws] + if not yaw_degrees or any(not math.isfinite(value) for value in yaw_degrees): + raise ValueError("reorient_yaw_degrees must contain finite values.") + + base_rotation = self._downward_tcp_rotation(reference[:, :3, :3]) + candidates = [] + for yaw_degrees_value in yaw_degrees: + yaw = math.radians(yaw_degrees_value) + yaw_rotation = reference.new_tensor( + [ + [math.cos(yaw), -math.sin(yaw), 0.0], + [math.sin(yaw), math.cos(yaw), 0.0], + [0.0, 0.0, 1.0], + ] + ) + target_rotation = torch.matmul(yaw_rotation, base_rotation) + waypoints = self._rotation_waypoints( + reference, + target_rotation, + waypoint_count=waypoint_count, + ) + policy = { + **grounded.motion_policy, + "reorient_selected_yaw_degrees": yaw_degrees_value, + } + candidates.append( + replace( + grounded, + target=EndEffectorPoseGoal(xpos=waypoints), + cfg={**grounded.cfg, **policy}, + motion_policy=policy, + ) + ) + return tuple(candidates) + + def _downward_tcp_rotation(self, rotation: torch.Tensor) -> torch.Tensor: + """Project the current TCP heading while aligning local +Z with world -Z.""" + x_axis = rotation[:, :3, 0].clone() + x_axis[:, 2] = 0.0 + norm = torch.linalg.vector_norm(x_axis, dim=1, keepdim=True) + fallback = rotation[:, :3, 1].clone() + fallback[:, 2] = 0.0 + fallback_norm = torch.linalg.vector_norm(fallback, dim=1, keepdim=True) + world_x = rotation.new_tensor([1.0, 0.0, 0.0]).expand_as(x_axis) + fallback = torch.where( + fallback_norm > 1.0e-6, + fallback / fallback_norm.clamp_min(1.0e-6), + world_x, + ) + x_axis = torch.where( + norm > 1.0e-6, + x_axis / norm.clamp_min(1.0e-6), + fallback, + ) + z_axis = rotation.new_tensor([0.0, 0.0, -1.0]).expand_as(x_axis) + y_axis = torch.linalg.cross(z_axis, x_axis, dim=1) + return torch.stack((x_axis, y_axis, z_axis), dim=2) + + def _rotation_waypoints( + self, + reference: torch.Tensor, + target_rotation: torch.Tensor, + *, + waypoint_count: int, + ) -> torch.Tensor: + """Interpolate fixed-position TCP rotations, excluding the observed start.""" + start_quat = quat_from_matrix(reference[:, :3, :3]) + end_quat = quat_from_matrix(target_rotation) + end_quat = torch.where( + torch.sum(start_quat * end_quat, dim=1, keepdim=True) < 0.0, + -end_quat, + end_quat, + ) + poses = reference[:, None].repeat(1, waypoint_count, 1, 1) + for index in range(waypoint_count): + fraction = float(index + 1) / float(waypoint_count) + interpolated = torch.stack( + [ + quat_slerp(start_quat[env_id], end_quat[env_id], tau=fraction) + for env_id in range(self.num_envs) + ] + ) + poses[:, index, :3, :3] = matrix_from_quat(interpolated) + return poses + + def _adapt_coordinated_pickment_grasps( + self, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> tuple[GroundedAction, ...]: + """Build deterministic live-geometry approach and partition candidates.""" + if capability.config_materializer != "coordinated_pickment": + return (grounded,) + target = self._validate_coordinated_pickment_goal(grounded) + live_pose = grounded.object_pose + expected_shape = (self.num_envs, 4, 4) + if ( + not isinstance(live_pose, torch.Tensor) + or live_pose.shape != expected_shape + or not bool(torch.isfinite(live_pose).all()) + ): + raise ValueError( + "CoordinatedPickment requires a finite grounded live object pose " + f"with shape {expected_shape}." + ) + live_pose = live_pose.to(device=self.device, dtype=torch.float32).clone() + affordance = target.semantics.affordance + assert isinstance(affordance, AntipodalAffordance) + vertices = torch.as_tensor( + affordance.mesh_vertices, + dtype=torch.float32, + device=self.device, + ) + if ( + vertices.ndim != 2 + or vertices.shape[1] != 3 + or vertices.shape[0] < 3 + or not bool(torch.isfinite(vertices).all()) + ): + raise ValueError( + "CoordinatedPickment mesh vertices must be finite with shape (N, 3)." + ) + + left_base, right_base = self._coordinated_arm_bases() + arm_directions = right_base[:, :3, 3] - left_base[:, :3, 3] + arm_norms = torch.linalg.vector_norm(arm_directions, dim=1, keepdim=True) + if not bool(torch.isfinite(arm_directions).all()) or bool( + (arm_norms <= 1.0e-6).any() + ): + raise ValueError( + "Coordinated pickup requires distinct finite left/right arm bases." + ) + arm_directions = arm_directions / arm_norms + shared_direction = arm_directions[0] + if bool( + (torch.matmul(arm_directions, shared_direction).abs() < 1.0 - 1.0e-4).any() + ): + raise ValueError( + "CoordinatedPickment requires one shared base-to-base direction " + "across vectorized environments." + ) + + centered = vertices - vertices.mean(dim=0, keepdim=True) + covariance = centered.transpose(0, 1) @ centered / float(vertices.shape[0]) + eigenvalues, eigenvectors = torch.linalg.eigh(covariance) + principal_local = eigenvectors[:, -1] + world_axes = torch.matmul(live_pose[:, :3, :3], eigenvectors) + principal_world = world_axes[:, :, -1] + principal_world = principal_world / torch.linalg.vector_norm( + principal_world, + dim=1, + keepdim=True, + ).clamp_min(1.0e-6) + arm_alignment = torch.abs((principal_world * arm_directions).sum(dim=1)) + elongation_ratio = torch.sqrt( + eigenvalues[-1].clamp_min(1.0e-12) / eigenvalues[-2].clamp_min(1.0e-12) + ) + elongation_confidence = torch.clamp( + (elongation_ratio - 1.0) / 1.5, + min=0.0, + max=1.0, + ) + base_ratio = float(grounded.cfg.get("middle_empty_ratio", 0.4)) + if not math.isfinite(base_ratio) or not 0.0 <= base_ratio < 1.0: + raise ValueError("middle_empty_ratio must be finite and in [0, 1).") + geometric_ratio = 0.25 + 0.45 * float(arm_alignment.mean()) + confidence = float(elongation_confidence) + preferred_ratio = (1.0 - confidence) * base_ratio + confidence * geometric_ratio + raw_ratios = ( + preferred_ratio, + base_ratio, + preferred_ratio - 0.15, + preferred_ratio + 0.15, + ) + ratios: list[float] = [] + for raw_ratio in raw_ratios: + ratio = min(0.90, max(0.05, float(raw_ratio))) + if not any(abs(ratio - existing) <= 1.0e-6 for existing in ratios): + ratios.append(ratio) + + requested_approach = grounded.cfg.get("approach_direction", (0.0, 0.0, -1.0)) + requested_approach = torch.as_tensor( + requested_approach, + dtype=torch.float32, + device=self.device, + ) + if requested_approach.shape != (3,) or not bool( + torch.isfinite(requested_approach).all() + ): + raise ValueError("approach_direction must be a finite vector shaped (3,).") + approach_norm = torch.linalg.vector_norm(requested_approach) + if float(approach_norm) <= 1.0e-6: + raise ValueError("approach_direction must be non-zero.") + requested_approach = requested_approach / approach_norm + + horizontal_arm = shared_direction.clone() + horizontal_arm[2] = 0.0 + horizontal_arm_norm = torch.linalg.vector_norm(horizontal_arm) + if float(horizontal_arm_norm) <= 1.0e-6: + raise ValueError( + "CoordinatedPickment arm bases must be separated in the world XY plane." + ) + horizontal_arm = horizontal_arm / horizontal_arm_norm + robot_forward = torch.stack( + (-horizontal_arm[1], horizontal_arm[0], horizontal_arm.new_tensor(0.0)) + ) + base_midpoint = 0.5 * (left_base[:, :3, 3] + right_base[:, :3, 3]) + object_from_bases = live_pose[:, :3, 3] - base_midpoint + reach_alignment = torch.sum(object_from_bases[:, :2] * robot_forward[:2], dim=1) + if float(reach_alignment.mean()) < 0.0: + robot_forward = -robot_forward + reach_alignment = -reach_alignment + + down = requested_approach.new_tensor([0.0, 0.0, -1.0]) + approach_candidates: list[tuple[str, torch.Tensor]] = [] + + def add_approach(label: str, direction: torch.Tensor) -> None: + direction = direction.to(device=self.device, dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if not bool(torch.isfinite(direction).all()) or float(norm) <= 1.0e-6: + return + direction = direction / norm + if any( + float(torch.dot(direction, existing)) >= 1.0 - 1.0e-5 + for _, existing in approach_candidates + ): + return + approach_candidates.append((label, direction)) + + add_approach("robot_forward_down", robot_forward + down) + add_approach("current", requested_approach) + add_approach("world_down", down) + add_approach("robot_forward", robot_forward) + + horizontal_axis_index: int | None = None + for axis_index in torch.argsort(eigenvalues).tolist(): + axes = world_axes[:, :, axis_index] + if float(torch.mean(torch.abs(axes[:, 2]))) > 0.75: + continue + reference = axes[0] + consistency = torch.abs(torch.matmul(axes, reference)) + if bool((consistency < 0.90).any()): + continue + horizontal_axis_index = int(axis_index) + horizontal_axis = reference.clone() + if float(torch.dot(horizontal_axis[:2], robot_forward[:2])) < 0.0: + horizontal_axis = -horizontal_axis + add_approach("short_axis_forward_down", horizontal_axis + down) + add_approach("short_axis_forward", horizontal_axis) + add_approach("short_axis_reverse_down", -horizontal_axis + down) + add_approach("short_axis_reverse", -horizontal_axis) + break + + grasp_seed = grounded.cfg.get("grasp_seed", _COORDINATED_GRASP_SEED) + if type(grasp_seed) is not int or grasp_seed < 0: + raise ValueError("grasp_seed must be a non-negative integer.") + pair_candidate_count = grounded.cfg.get("grasp_pair_candidate_count", 3) + if type(pair_candidate_count) is not int: + raise ValueError("grasp_pair_candidate_count must be an integer.") + if not 1 <= pair_candidate_count <= 8: + raise ValueError("grasp_pair_candidate_count must be in [1, 8].") + trace = { + "strategy": "live_geometry_approach_search", + "local_pca_axes": eigenvectors.detach().cpu().tolist(), + "world_pca_axes": world_axes.detach().cpu().tolist(), + "pca_eigenvalues": eigenvalues.detach().cpu().tolist(), + "local_principal_axis": principal_local.detach().cpu().tolist(), + "world_principal_axes": principal_world.detach().cpu().tolist(), + "elongation_ratio": float(elongation_ratio), + "elongation_confidence": confidence, + "arm_axis_alignment": arm_alignment.detach().cpu().tolist(), + "left_to_right_arm_direction": shared_direction.detach().cpu().tolist(), + "robot_forward_direction": robot_forward.detach().cpu().tolist(), + "shared_reach_alignment": reach_alignment.detach().cpu().tolist(), + "requested_approach_direction": requested_approach.detach().cpu().tolist(), + "selected_horizontal_axis_index": horizontal_axis_index, + "approach_candidates": [ + { + "label": label, + "direction": direction.detach().cpu().tolist(), + } + for label, direction in approach_candidates + ], + "candidate_middle_empty_ratios": list(ratios), + "grasp_seed": grasp_seed, + "grasp_pair_candidate_count": pair_candidate_count, + } + candidates: list[GroundedAction] = [] + candidate_index = 0 + for approach_index, (approach_label, approach) in enumerate( + approach_candidates + ): + for ratio in ratios: + for pair_rank in range(pair_candidate_count): + cfg = { + **grounded.cfg, + "left_to_right_arm_direction": shared_direction.clone(), + "approach_direction": approach.clone(), + "middle_empty_ratio": ratio, + "grasp_seed": grasp_seed, + "grasp_pair_rank": pair_rank, + } + motion_policy = { + **grounded.motion_policy, + "grasp_seed": grasp_seed, + "coordinated_grasp": { + **trace, + "candidate_index": candidate_index, + "approach_candidate_index": approach_index, + "approach_candidate_label": approach_label, + "approach_direction": approach.detach().cpu().tolist(), + "selected_middle_empty_ratio": ratio, + "grasp_pair_rank": pair_rank, + }, + } + candidates.append( + replace( + grounded, + target=replace( + target, object_initial_pose=live_pose.clone() + ), + cfg=cfg, + object_pose=live_pose.clone(), + motion_policy=motion_policy, + ) + ) + candidate_index += 1 + return tuple(candidates) + + @contextmanager + def _isolated_random_seed(self, seed: int) -> Iterator[None]: + """Run one stochastic grasp attempt without perturbing global RNG state.""" + if type(seed) is not int or seed < 0: + raise ValueError("seed must be a non-negative integer.") + device = torch.device(self.device) + cuda_devices: list[int] = [] + if device.type == "cuda": + cuda_devices.append( + torch.cuda.current_device() if device.index is None else device.index + ) + with torch.random.fork_rng(devices=cuda_devices): + torch.manual_seed(seed) + if cuda_devices: + torch.cuda.manual_seed_all(seed) + yield + + def _coordinated_arm_bases(self) -> tuple[torch.Tensor, torch.Tensor]: + from .frames import arm_base_poses + + return arm_base_poses(self.env) + + def _latest_coordinated_grasp_trace( + self, + engine: AtomicActionEngine, + ) -> dict[str, Any] | None: + """Return the S1-S5 trace emitted by the shared E5 grasp service.""" + _, hand_part, _ = self._parts("left_arm") + if hand_part is None: + return None + generator = engine.grasp_pose_generators.get(hand_part) + if not isinstance(generator, _TracingAntipodalGraspPoseGenerator): + return None + return generator.last_dual_trace + + def _audit_pose_interpolation( + self, + start: torch.Tensor, + end: torch.Tensor, + waypoint_count: int, + *, + interpolate_orientation: bool, + ) -> torch.Tensor: + """Build the Cartesian reference used only by the GenSim FK audit.""" + if waypoint_count <= 0: + return start[:, None].repeat(1, 0, 1, 1) + weights = torch.linspace( + 0.0, + 1.0, + waypoint_count, + dtype=start.dtype, + device=start.device, + ) + result = start[:, None].repeat(1, waypoint_count, 1, 1) + result[:, :, :3, 3] = torch.lerp( + start[:, None, :3, 3], + end[:, None, :3, 3], + weights[None, :, None], + ) + if not interpolate_orientation: + return result + start_quat = quat_from_matrix(start[:, :3, :3]) + end_quat = quat_from_matrix(end[:, :3, :3]) + end_quat = torch.where( + torch.sum(start_quat * end_quat, dim=1, keepdim=True) < 0.0, + -end_quat, + end_quat, + ) + for waypoint_index, weight in enumerate(weights.tolist()): + interpolated = torch.stack( + [ + quat_slerp(start_quat[row], end_quat[row], tau=float(weight)) + for row in range(start.shape[0]) + ] + ) + result[:, waypoint_index, :3, :3] = matrix_from_quat(interpolated) + return result + + def _coordinated_selected_grasp_poses( + self, + grasp_stages: Mapping[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + rows = grasp_stages.get("environment_rows") + if rows is None: + rows = [grasp_stages] + left: list[torch.Tensor] = [] + right: list[torch.Tensor] = [] + for row in rows: + pair = row.get("pair_selection") + if not isinstance(pair, Mapping) or not bool(pair.get("selected", False)): + raise ValueError( + "Trajectory audit requires one selected grasp pair per row." + ) + left.append( + torch.as_tensor( + pair["selected_left_pose"], + dtype=torch.float32, + device=self.device, + ) + ) + right.append( + torch.as_tensor( + pair["selected_right_pose"], + dtype=torch.float32, + device=self.device, + ) + ) + return torch.stack(left), torch.stack(right) + + def _arm_trajectory_fk( + self, + positions: torch.Tensor, + control_part: str, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return every TCP pose and every serial-chain link point.""" + joint_ids = list(self.env.robot.get_joint_ids(name=control_part)) + arm_qpos = positions[:, :, joint_ids] + batch_size, waypoint_count, dof = arm_qpos.shape + solver = self.env.robot.get_solver(name=control_part) + flat_qpos = arm_qpos.reshape(batch_size * waypoint_count, dof) + local_eef = solver.get_fk(flat_qpos).reshape( + batch_size, + waypoint_count, + 4, + 4, + ) + base_pose = self.env.robot.get_link_pose( + link_name=solver.root_link_name, + to_matrix=True, + ).to(device=positions.device, dtype=positions.dtype) + eef = torch.matmul(base_pose[:, None], local_eef) + chain = getattr(solver, "pk_serial_chain", None) + if chain is None: + raise ValueError(f"{control_part} has no serial chain for capsule audit.") + link_transforms = chain.forward_kinematics(flat_qpos, end_only=False) + link_points: list[torch.Tensor] = [] + for transform in link_transforms.values(): + local = transform.get_matrix().reshape( + batch_size, + waypoint_count, + 4, + 4, + ) + world = torch.matmul(base_pose[:, None], local) + link_points.append(world[:, :, :3, 3]) + link_points.append(eef[:, :, :3, 3]) + return eef, torch.stack(link_points, dim=2) + + @staticmethod + def _audit_batched_pose( + value: Any, + *, + batch_size: int, + device: torch.device | str, + name: str, + ) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(batch_size, 1, 1) + if pose.shape != (batch_size, 4, 4) or not bool(torch.isfinite(pose).all()): + raise ValueError(f"{name} must have finite shape ({batch_size}, 4, 4).") + return pose + + def _coordinated_audit_references( + self, + candidate: GroundedAction, + invocation: ActionInvocation, + plan: ActionPlan, + actual_left: torch.Tensor, + actual_right: torch.Tensor, + grasp_stages: Mapping[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + positions = plan.joint_trajectory + assert positions is not None + batch_size = positions.batch_size + initial = self._audit_batched_pose( + candidate.target.object_initial_pose, + batch_size=batch_size, + device=self.device, + name="object_initial_pose", + ) + target = self._audit_batched_pose( + candidate.target.object_target_pose, + batch_size=batch_size, + device=self.device, + name="object_target_pose", + ) + left_grasp, right_grasp = self._coordinated_selected_grasp_poses(grasp_stages) + left_relation = torch.bmm(torch.linalg.inv(initial), left_grasp) + right_relation = torch.bmm(torch.linalg.inv(initial), right_grasp) + desired_left = actual_left.clone() + desired_right = actual_right.clone() + segments = {segment.name: segment for segment in plan.segments} + + def assign( + name: str, + left_value: torch.Tensor, + right_value: torch.Tensor, + ) -> None: + segment = segments[name] + desired_left[:, segment.start : segment.stop] = left_value + desired_right[:, segment.start : segment.stop] = right_value + + approach = segments["approach"] + assign( + "approach", + self._audit_pose_interpolation( + actual_left[:, approach.start], + left_grasp, + approach.waypoint_count, + interpolate_orientation=True, + ), + self._audit_pose_interpolation( + actual_right[:, approach.start], + right_grasp, + approach.waypoint_count, + interpolate_orientation=True, + ), + ) + close = segments["close"] + assign( + "close", + left_grasp[:, None].repeat(1, close.waypoint_count, 1, 1), + right_grasp[:, None].repeat(1, close.waypoint_count, 1, 1), + ) + lift_pose = initial.clone() + lift_pose[:, 2, 3] += float(invocation.skill_options.lift_height) + lift = segments["lift"] + lift_object = self._audit_pose_interpolation( + initial, + lift_pose, + lift.waypoint_count, + interpolate_orientation=False, + ) + assign( + "lift", + torch.matmul(lift_object, left_relation[:, None]), + torch.matmul(lift_object, right_relation[:, None]), + ) + move = segments["move"] + move_object = self._audit_pose_interpolation( + lift_pose, + target, + move.waypoint_count, + interpolate_orientation=True, + ) + assign( + "move", + torch.matmul(move_object, left_relation[:, None]), + torch.matmul(move_object, right_relation[:, None]), + ) + hold = segments["hold"] + assign( + "hold", + torch.matmul(target, left_relation)[:, None].repeat( + 1, hold.waypoint_count, 1, 1 + ), + torch.matmul(target, right_relation)[:, None].repeat( + 1, hold.waypoint_count, 1, 1 + ), + ) + return desired_left, desired_right + + def _audit_coordinated_trajectory( + self, + candidate: GroundedAction, + invocation: ActionInvocation, + plan: ActionPlan, + context: PlanningContext, + grasp_stages: Mapping[str, Any], + ) -> tuple[ActionPlan, dict[str, Any]]: + """Reject unsafe E5 joint paths before they can become executable.""" + del context + raw_success = plan.plan_success.to(device=self.device) + if plan.joint_trajectory is None or not bool(raw_success.any()): + return plan, { + "success": raw_success.detach().cpu().tolist(), + "skipped": True, + "reason": "planner_failed_or_missing_trajectory", + } + positions = plan.joint_trajectory.positions.to( + device=self.device, + dtype=torch.float32, + ) + try: + left_arm, _, _ = self._parts("left_arm") + right_arm, _, _ = self._parts("right_arm") + left_eef, left_links = self._arm_trajectory_fk(positions, left_arm) + right_eef, right_links = self._arm_trajectory_fk(positions, right_arm) + desired_left, desired_right = self._coordinated_audit_references( + candidate, + invocation, + plan, + left_eef, + right_eef, + grasp_stages, + ) + direction = torch.as_tensor( + candidate.cfg["left_to_right_arm_direction"], + dtype=torch.float32, + device=self.device, + ) + report = _trajectory_safety_report( + left_qpos=positions[:, :, self.env.robot.get_joint_ids(name=left_arm)], + right_qpos=positions[ + :, :, self.env.robot.get_joint_ids(name=right_arm) + ], + left_eef=left_eef, + right_eef=right_eef, + desired_left_eef=desired_left, + desired_right_eef=desired_right, + left_link_points=left_links, + right_link_points=right_links, + left_to_right_direction=direction, + maximum_joint_step=float(candidate.cfg.get("maximum_joint_step", 0.25)), + maximum_orientation_error=float( + candidate.cfg.get("maximum_wrist_orientation_error", 0.20) + ), + minimum_lateral_gap=float( + candidate.cfg.get("minimum_grasp_lateral_gap", 0.05) + ), + capsule_radius=float( + candidate.cfg.get("inter_arm_capsule_radius", 0.04) + ), + minimum_capsule_clearance=float( + candidate.cfg.get("minimum_inter_arm_clearance", 0.01) + ), + orientation_start_index={ + segment.name: segment.start for segment in plan.segments + }["close"], + ) + audited_success = raw_success & report.success.to(device=self.device) + audit_trace = { + "success": audited_success.detach().cpu().tolist(), + "failed_checks": report.failed_checks, + "metrics": report.metrics, + "orientation_audit_start": "close", + "capsule_model": { + "left_link_segments": left_links.shape[2] - 1, + "right_link_segments": right_links.shape[2] - 1, + }, + "thresholds": { + "maximum_joint_step": float( + candidate.cfg.get("maximum_joint_step", 0.25) + ), + "maximum_wrist_orientation_error": float( + candidate.cfg.get("maximum_wrist_orientation_error", 0.20) + ), + "minimum_lateral_gap": float( + candidate.cfg.get("minimum_grasp_lateral_gap", 0.05) + ), + "inter_arm_capsule_radius": float( + candidate.cfg.get("inter_arm_capsule_radius", 0.04) + ), + "minimum_inter_arm_clearance": float( + candidate.cfg.get("minimum_inter_arm_clearance", 0.01) + ), + }, + } + except Exception as exc: + audited_success = torch.zeros_like(raw_success) + audit_trace = { + "success": audited_success.detach().cpu().tolist(), + "failed_checks": {"audit_error": raw_success.cpu().tolist()}, + "error": f"{type(exc).__name__}: {exc}", + } + if torch.equal(audited_success, raw_success): + return plan, audit_trace + failed_rows = ( + torch.nonzero( + raw_success & ~audited_success, + as_tuple=False, + ) + .flatten() + .tolist() + ) + message = ( + f"GenSim trajectory safety audit rejected environment(s) {failed_rows}." + ) + diagnostics = PlannerDiagnostics( + backend=plan.diagnostics.backend, + messages=(*plan.diagnostics.messages, message), + metadata={ + **dict(plan.diagnostics.metadata), + "gensim_trajectory_audit": audit_trace, + }, + failure=PlanningFailure("gensim_trajectory_safety_failed"), + ) + return ( + replace( + plan, + plan_success=audited_success, + diagnostics=diagnostics, + ), + audit_trace, + ) + + @contextmanager + def _upright_grasp_selection_context( + self, + engine: AtomicActionEngine, + candidate: GroundedAction, + capability: AtomicCapability, + ) -> Iterator[None]: + """Apply the E2 side-grasp policy only to upright pickup.""" + local_axis = candidate.cfg.get("obj_upright_direction") + if ( + capability.target_materializer != "object_grasp" + or candidate.cfg.get("rotate_upright") is None + or local_axis is None + ): + yield + return + _, hand_part, _ = self._parts(candidate.arm) + if hand_part is None: + yield + return + generator = engine.grasp_pose_generators.get(hand_part) + if not isinstance(generator, _TracingAntipodalGraspPoseGenerator): + yield + return + with generator.upright_selection_context( + local_axis=torch.as_tensor( + local_axis, + dtype=torch.float32, + device=self.device, + ) + ): + yield + trace = generator.last_upright_trace + if trace is not None: + candidate.motion_policy["upright_grasp"] = trace + + def _record_selected_upright_grasp( + self, + candidate: GroundedAction, + plan: ActionPlan, + context: PlanningContext, + ) -> None: + """Record the grasp actually selected after IK and downstream screening.""" + trace = candidate.motion_policy.get("upright_grasp") + local_axis = candidate.cfg.get("obj_upright_direction") + if not isinstance(trace, dict) or local_axis is None: + return + held = next( + ( + value + for value in plan.expected_effects.held_object_updates.values() + if value is not None + ), + None, + ) + if held is None or held.semantics.entity_id is None: + return + entity = context.scene.entities.get(held.semantics.entity_id) + vertices = held.semantics.geometry.get("mesh_vertices") + if entity is None or vertices is None: + return + grasp_pose = held.grasp_xpos.to(device=self.device, dtype=torch.float32) + object_pose = entity.pose.to(device=self.device, dtype=torch.float32) + if object_pose.shape == (4, 4): + object_pose = object_pose.unsqueeze(0).expand(self.num_envs, -1, -1) + axis = torch.as_tensor( + local_axis, + dtype=torch.float32, + device=self.device, + ) + axis = axis / torch.linalg.vector_norm(axis) + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=self.device, + ) + vertex_positions = torch.matmul(vertices, axis) + axis_min = vertex_positions.min() + axis_extent = vertex_positions.max() - axis_min + world_axis = torch.matmul(object_pose[:, :3, :3], axis) + relative_centers = grasp_pose[:, :3, 3] - object_pose[:, :3, 3] + axis_positions = torch.sum(relative_centers * world_axis, dim=1) + axis_fractions = (axis_positions - axis_min) / axis_extent + closing_axes = torch.nn.functional.normalize( + grasp_pose[:, :3, 0], + dim=1, + ) + axis_alignment = torch.abs(torch.sum(closing_axes * world_axis, dim=1)) + trace.update( + { + "selected_axis_fraction": axis_fractions.detach().cpu().tolist(), + "selected_axis_alignment": axis_alignment.detach().cpu().tolist(), + "selected_grasp_pose": grasp_pose.detach().cpu().tolist(), + } + ) + + @contextmanager + def _coordinated_pair_selection_context( + self, + engine: AtomicActionEngine, + candidate: GroundedAction, + context: PlanningContext, + ) -> Iterator[None]: + """Bind live wrists and arm bases to one synchronous E5 generator call.""" + trace = candidate.motion_policy.get("coordinated_grasp") + if not isinstance(trace, Mapping): + yield + return + _, left_hand, _ = self._parts("left_arm") + if left_hand is None: + yield + return + generator = engine.grasp_pose_generators.get(left_hand) + if not isinstance(generator, _TracingAntipodalGraspPoseGenerator): + yield + return + left_arm, _, _ = self._parts("left_arm") + right_arm, _, _ = self._parts("right_arm") + qpos = context.robot.qpos.to(device=self.device, dtype=torch.float32) + left_ids = list(self.env.robot.get_joint_ids(name=left_arm)) + right_ids = list(self.env.robot.get_joint_ids(name=right_arm)) + left_eef = self.env.robot.compute_fk( + qpos[:, left_ids], + name=left_arm, + to_matrix=True, + ) + right_eef = self.env.robot.compute_fk( + qpos[:, right_ids], + name=right_arm, + to_matrix=True, + ) + left_base, right_base = self._coordinated_arm_bases() + with generator.dual_arm_selection_context( + left_eef=left_eef, + right_eef=right_eef, + left_base=left_base, + right_base=right_base, + left_to_right_direction=torch.as_tensor( + candidate.cfg["left_to_right_arm_direction"], + dtype=torch.float32, + device=self.device, + ), + pair_rank=int(candidate.cfg.get("grasp_pair_rank", 0)), + minimum_separation=float( + candidate.cfg.get("minimum_grasp_separation", 0.08) + ), + minimum_lateral_gap=float( + candidate.cfg.get("minimum_grasp_lateral_gap", 0.05) + ), + ): + yield + + def _search_reachable_retreat( + self, + *, + grounded: GroundedAction, + capability: AtomicCapability, + state: ExecutionState, + context: PlanningContext, + invocation: ActionInvocation, + initial_positions: torch.Tensor, + initial_success: torch.Tensor, + initial_warnings: Sequence[str] = (), + ) -> tuple[GroundedAction, torch.Tensor, torch.Tensor, dict[str, Any]]: + """Select a row-local retreat candidate accepted by the live planner.""" + candidates = self._retreat_search_targets(grounded) + target = getattr(grounded.target, "xpos", None) + if not isinstance(target, torch.Tensor) or len(candidates) <= 1: + return ( + grounded, + initial_positions, + initial_success, + { + "strategy": "bounded_motion_planner", + "attempts": [], + "selected_target_z": ( + None + if not isinstance(target, torch.Tensor) + else target[:, 2, 3] + ), + }, + ) + + selected_target = candidates[0][1].clone() + selected_positions = initial_positions + success = initial_success.clone() + suppressed_warnings = list(initial_warnings) + reference = grounded.motion_policy["retreat_reference_pose"].to( + device=selected_target.device, + dtype=selected_target.dtype, + ) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + selected_candidates = [ + candidates[0][0] if bool(success[env_id]) else "unresolved" + for env_id in range(self.num_envs) + ] + attempts: list[dict[str, Any]] = [ + { + "candidate": candidates[0][0], + "target_z": candidates[0][1][:, 2, 3].detach().clone(), + "target_distance": torch.linalg.vector_norm( + candidates[0][1][:, :2, 3] - reference[:, :2, 3], + dim=1, + ) + .detach() + .clone(), + "success": initial_success.detach().clone(), + } + ] + for label, candidate_target in candidates[1:]: + unresolved = ~success + if not bool(unresolved.any()): + break + row_target = torch.where( + unresolved[:, None, None], + candidate_target, + selected_target, + ) + candidate_grounded = replace( + grounded, + target=EndEffectorPoseGoal(xpos=row_target), + ) + candidate_invocation = replace( + invocation, + goal=candidate_grounded.target, + ) + with _capture_retreat_warnings(True) as warnings: + candidate_plan = self._engine().plan(candidate_invocation, context) + suppressed_warnings.extend(warnings) + candidate_positions = self._positions_with_agent_holds( + candidate_plan, + candidate_grounded, + capability, + ) + candidate_success = candidate_plan.plan_success.to(self.device) + selected_rows = unresolved & candidate_success + selected_positions = self._merge_plan_rows( + selected_positions, + candidate_positions, + selected_rows, + state.last_qpos, + ) + selected_target = torch.where( + selected_rows[:, None, None], + candidate_target, + selected_target, + ) + for env_id in ( + torch.nonzero(selected_rows, as_tuple=False).flatten().tolist() + ): + selected_candidates[env_id] = label + success |= candidate_success + attempts.append( + { + "candidate": label, + "target_z": candidate_target[:, 2, 3].detach().clone(), + "target_distance": torch.linalg.vector_norm( + candidate_target[:, :2, 3] - reference[:, :2, 3], + dim=1, + ) + .detach() + .clone(), + "success": candidate_success.detach().clone(), + } + ) + + selected_distance = torch.linalg.vector_norm( + selected_target[:, :2, 3] - reference[:, :2, 3], dim=1 + ) + metadata = { + "retreat_selected_target_z": selected_target[:, 2, 3].detach().clone(), + "retreat_selected_target_distance": selected_distance.detach().clone(), + "retreat_reachability_found": success.detach().clone(), + } + summary = ( + "Retreat reachability search: " + f"mode={grounded.motion_policy.get('retreat_search_mode', 'vertical_then_baseward')}, " + f"resolved={int(success.sum())}/{success.numel()}, " + f"attempts={len(attempts)}, " + f"selected={selected_candidates}, " + f"suppressed_warnings={len(suppressed_warnings)}." + ) + if bool(success.all()): + log_info(summary) + else: + log_warning(summary) + selected_grounded = replace( + grounded, + target=EndEffectorPoseGoal(xpos=selected_target), + cfg={**grounded.cfg, **metadata}, + motion_policy={**grounded.motion_policy, **metadata}, + ) + return ( + selected_grounded, + selected_positions, + success, + { + "strategy": "bounded_motion_planner", + "attempts": attempts, + "selected_target_z": selected_target[:, 2, 3].detach().clone(), + "selected_target_distance": selected_distance.detach().clone(), + "selected_candidates": selected_candidates, + "suppressed_warnings": len(suppressed_warnings), + }, + ) + + def _retreat_search_targets( + self, + grounded: GroundedAction, + ) -> list[tuple[str, torch.Tensor]]: + """Build bounded height and baseward retreat candidates from live poses.""" + target = getattr(grounded.target, "xpos", None) + reference = grounded.motion_policy.get("retreat_reference_pose") + if not isinstance(target, torch.Tensor) or not isinstance( + reference, torch.Tensor + ): + return [] + target = target.to(device=self.device, dtype=torch.float32) + reference = reference.to(device=self.device, dtype=torch.float32) + if target.shape == (4, 4): + target = target.unsqueeze(0).repeat(self.num_envs, 1, 1) + if reference.shape == (4, 4): + reference = reference.unsqueeze(0).repeat(self.num_envs, 1, 1) + expected = (self.num_envs, 4, 4) + if target.shape != expected or reference.shape != expected: + return [] + + sample_count = int(grounded.cfg.get("retreat_search_samples", 6)) + if not 2 <= sample_count <= 16: + raise ValueError("retreat_search_samples must be in [2, 16].") + search_mode = str( + grounded.motion_policy.get("retreat_search_mode", "vertical_then_baseward") + ) + allowed_modes = {"horizontal_only", "vertical_only", "vertical_then_baseward"} + if search_mode not in allowed_modes: + raise ValueError( + "retreat_search_mode must be 'horizontal_only', 'vertical_only', " + "or 'vertical_then_baseward'." + ) + if search_mode == "horizontal_only": + direction = target[:, :2, 3] - reference[:, :2, 3] + requested_distance = torch.linalg.vector_norm( + direction, dim=1, keepdim=True + ) + if bool((requested_distance <= 1.0e-6).any()): + return [("requested", target.clone())] + direction = direction / requested_distance + minimum_distance = float(grounded.cfg.get("minimum_retreat_distance", 0.05)) + if not math.isfinite(minimum_distance) or minimum_distance < 0.0: + raise ValueError( + "minimum_retreat_distance must be finite and non-negative." + ) + minimum = torch.minimum( + requested_distance[:, 0], + torch.full_like(requested_distance[:, 0], minimum_distance), + ) + fractions = torch.linspace( + 1.0, + 0.0, + sample_count, + dtype=target.dtype, + device=target.device, + ) + distances = ( + minimum[:, None] + + (requested_distance[:, 0] - minimum)[:, None] * fractions[None] + ) + candidates = [("requested", target.clone())] + for index in range(1, sample_count): + candidate = target.clone() + candidate[:, :2, 3] = ( + reference[:, :2, 3] + direction * distances[:, index, None] + ) + candidates.append((f"distance_{index}", candidate)) + return candidates + + minimum_height = float(grounded.cfg.get("minimum_retreat_height", 0.05)) + if not math.isfinite(minimum_height) or minimum_height < 0.0: + raise ValueError("minimum_retreat_height must be finite and non-negative.") + desired_height = torch.clamp( + target[:, 2, 3] - reference[:, 2, 3], + min=0.0, + ) + minimum = torch.minimum( + desired_height, + torch.full_like(desired_height, minimum_height), + ) + fractions = torch.linspace( + 1.0, + 0.0, + sample_count, + dtype=target.dtype, + device=target.device, + ) + heights = ( + minimum[:, None] + (desired_height - minimum)[:, None] * fractions[None] + ) + candidates: list[tuple[str, torch.Tensor]] = [("requested", target.clone())] + for index in range(1, sample_count): + candidate = target.clone() + candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] + candidates.append((f"height_{index}", candidate)) + + if search_mode == "vertical_only": + return candidates + + from .frames import arm_base_poses + + left_base, right_base = arm_base_poses(self.env) + base = left_base if grounded.arm == "left_arm" else right_base + direction = base[:, :2, 3] - reference[:, :2, 3] + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + direction = torch.where( + norm > 1.0e-6, + direction / torch.clamp(norm, min=1.0e-6), + torch.zeros_like(direction), + ) + distance = float(grounded.cfg.get("retreat_distance", 0.10)) + if not math.isfinite(distance) or distance < 0.0: + raise ValueError("retreat_distance must be finite and non-negative.") + for index in range(sample_count): + candidate = target.clone() + candidate[:, :2, 3] = reference[:, :2, 3] + direction * distance + candidate[:, 2, 3] = reference[:, 2, 3] + heights[:, index] + candidates.append((f"baseward_{index}", candidate)) + return candidates + + def _planner_trace( + self, + *, + grounded: GroundedAction, + invocation: ActionInvocation, + context: PlanningContext, + state: ExecutionState, + primary_success: torch.Tensor, + primary_diagnostics: Any, + fallback_allowed: bool, + fallback_strategy: str | None, + fallback_attempted: torch.Tensor, + fallback_success: torch.Tensor, + fallback_used: torch.Tensor, + reachability_search: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """Build compact per-row evidence for the planner route actually used.""" + exclusions = self._collision_exclusion_masks(grounded, state) + obstacle_positions = { + uid: context.scene.entities[uid].pose[:, :3, 3].detach().clone() + for uid in context.scene.collision_entity_ids + } + revisions = torch.as_tensor( + context.scene.collision_world_revisions(self.num_envs), + dtype=torch.int64, + device=self.device, + ) + requested_backend = str(self.planner_policy["backend"]) + primary_strategy = invocation.motion_policy.strategy + diagnostic_backend = str(primary_diagnostics.backend) + if ( + primary_strategy == "motion_gen" + and diagnostic_backend in {"curobo", "toppra"} + and diagnostic_backend != requested_backend + ): + raise RuntimeError( + "Planner backend mismatch: runtime policy requested " + f"{requested_backend!r}, but the action plan reported " + f"{diagnostic_backend!r}." + ) + primary_effective_backend = ( + diagnostic_backend if primary_strategy == "motion_gen" else "ik_interp" + ) + if bool(fallback_used.all()) and bool(fallback_used.any()): + effective_backend = "ik_interp" + effective_strategy = "ik_interp" + elif bool(fallback_used.any()): + effective_backend = "mixed" + effective_strategy = "mixed" + else: + effective_backend = primary_effective_backend + effective_strategy = primary_strategy + planner_failure_reason = None + if not bool(primary_success.all()): + if primary_diagnostics.messages: + planner_failure_reason = "; ".join(primary_diagnostics.messages) + else: + metadata = primary_diagnostics.metadata + for key in ("failure_reason", "reason", "error"): + if metadata.get(key) is not None: + planner_failure_reason = str(metadata[key]) + break + if planner_failure_reason is None: + planner_failure_reason = "planner_reported_failure" + collision_planning_capable = ( + effective_backend == "curobo" and effective_strategy == "motion_gen" + ) + search_budget: dict[str, Any] = { + "requested_backend": requested_backend, + "fallback_enabled": bool(fallback_allowed), + } + if requested_backend == "curobo": + search_budget["primary_max_attempts"] = int( + self.planner_policy.get("curobo", {}).get("max_attempts", 1) + ) + trace = { + "action_class": grounded.action_class, + "arm": grounded.arm, + "gripper_model": self.gripper_profile.model.value, + "ik_solver": self.ik_solver, + "left_solver_class": self.ik_solver_classes.get("left_arm"), + "right_solver_class": self.ik_solver_classes.get("right_arm"), + "planner": requested_backend, + "requested_backend": requested_backend, + "effective_backend": effective_backend, + "effective_strategy": effective_strategy, + "primary_effective_backend": primary_effective_backend, + "primary_strategy": primary_strategy, + "dynamic_collision_mode": invocation.motion_policy.dynamic_collision_mode.value, + "collision_planning_capable": collision_planning_capable, + "collision_check_scope": ( + "static_and_dynamic" + if collision_planning_capable + and invocation.motion_policy.dynamic_collision_mode.value != "off" + else ("static" if collision_planning_capable else "not_supported") + ), + "primary_success": primary_success.detach().clone(), + "planner_failure_reason": planner_failure_reason, + "fallback_allowed": fallback_allowed, + "fallback_strategy": fallback_strategy, + "fallback_attempted": fallback_attempted.detach().clone(), + "fallback_success": fallback_success.detach().clone(), + "fallback_used": fallback_used.detach().clone(), + "fallback_occurred": fallback_attempted.detach().clone(), + "search_budget": search_budget, + "collision_world_revision": revisions, + "collision_obstacle_positions": obstacle_positions, + "collision_exclusions": { + uid: mask.detach().clone() for uid, mask in exclusions.items() + }, + } + if reachability_search is not None: + trace["reachability_search"] = deepcopy(dict(reachability_search)) + options = invocation.skill_options + object_part = getattr(options, "pick_object_part", None) + approach_direction = getattr(options, "approach_direction", None) + if object_part is None: + object_part = getattr(options, "receive_pick_object_part", None) + approach_direction = getattr( + options, + "receive_approach_direction", + approach_direction, + ) + if object_part is not None: + grasp_policy: dict[str, Any] = {"object_part": str(object_part)} + if isinstance(approach_direction, torch.Tensor): + direction = approach_direction.to(dtype=torch.float32) + norm = torch.linalg.vector_norm(direction) + if bool(torch.isfinite(norm)) and float(norm) > 0.0: + grasp_policy["approach_direction"] = ( + (direction / norm).detach().cpu().tolist() + ) + trace["grasp_policy"] = grasp_policy + body_grasp = grounded.motion_policy.get("body_grasp") + if isinstance(body_grasp, Mapping): + trace["body_grasp"] = deepcopy(dict(body_grasp)) + upright_grasp = grounded.motion_policy.get("upright_grasp") + if isinstance(upright_grasp, Mapping): + trace["upright_grasp"] = deepcopy(dict(upright_grasp)) + coordinated_grasp = grounded.motion_policy.get("coordinated_grasp") + if isinstance(coordinated_grasp, Mapping): + trace["coordinated_grasp"] = deepcopy(dict(coordinated_grasp)) + return trace + + def _select_transport_yaw( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> GroundedAction: + """Choose the closest IK-feasible yaw when task semantics leave it free.""" + sample_count = _FREE_YAW_SAMPLE_COUNT if grounded.allow_yaw_search else 1 + capability = self.capabilities.get(grounded.action_class) + if ( + capability.target_materializer != "semantic_held_object" + or sample_count <= 1 + ): + return grounded + target_pose = getattr(grounded.target, "object_target_pose", None) + if not isinstance(target_pose, torch.Tensor): + return grounded + target_pose = target_pose.to(device=self.device, dtype=torch.float32) + if target_pose.shape == (4, 4): + target_pose = target_pose.unsqueeze(0).repeat(self.num_envs, 1, 1) + if target_pose.shape != (self.num_envs, 4, 4): + raise ValueError("Transport target must have shape (4, 4) or (N, 4, 4).") + + arm_part, _, _ = self._parts(grounded.arm) + held = state.get_held_object(arm_part) + if held is None: + return grounded + object_to_eef = held.object_to_eef.to( + device=self.device, + dtype=target_pose.dtype, + ) + if object_to_eef.shape == (4, 4): + object_to_eef = object_to_eef.unsqueeze(0).repeat(self.num_envs, 1, 1) + variants = self._yaw_variants(target_pose, sample_count) + eef_variants = torch.matmul(variants, object_to_eef[:, None]) + joint_ids = list(self.env.robot.get_joint_ids(name=arm_part)) + start_qpos = state.last_qpos[:, joint_ids] + seeds = start_qpos[:, None].expand(-1, sample_count, -1) + success, qpos = self.env.robot.compute_batch_ik( + pose=eef_variants, + name=arm_part, + joint_seed=seeds, + ) + success = torch.as_tensor( + success, + dtype=torch.bool, + device=self.device, + ).reshape(self.num_envs, sample_count) + qpos = torch.as_tensor(qpos, dtype=torch.float32, device=self.device) + success &= torch.isfinite(qpos).all(dim=-1) + distance = torch.linalg.vector_norm(qpos - seeds, dim=-1) + distance = torch.where( + success, + distance, + torch.full_like(distance, torch.inf), + ) + yaw_offsets = torch.matmul( + variants[:, :, :3, :3], + target_pose[:, None, :3, :3].transpose(-1, -2), + ) + yaw_distance = torch.atan2( + yaw_offsets[:, :, 1, 0], + yaw_offsets[:, :, 0, 0], + ).abs() + minimum_yaw = torch.where( + success, + yaw_distance, + torch.full_like(yaw_distance, torch.inf), + ).amin(dim=1) + minimum_rotation = success & torch.isclose( + yaw_distance, + minimum_yaw[:, None], + atol=1.0e-6, + rtol=0.0, + ) + best = torch.where( + minimum_rotation, + distance, + torch.full_like(distance, torch.inf), + ).argmin(dim=1) + env_ids = torch.arange(self.num_envs, device=self.device) + selected = variants[env_ids, best] + selected = torch.where( + success.any(dim=1)[:, None, None], + selected, + target_pose, + ) + return replace( + grounded, + target=replace(grounded.target, object_target_pose=selected), + target_object_pose=selected, + ) + + @staticmethod + def _yaw_variants( + target_pose: torch.Tensor, + sample_count: int, + ) -> torch.Tensor: + signed_steps = [0] + for step in range(1, (sample_count + 1) // 2): + signed_steps.extend((step, -step)) + if sample_count % 2 == 0: + signed_steps.append(sample_count // 2) + angles = target_pose.new_tensor(signed_steps) * (2.0 * math.pi / sample_count) + yaw = target_pose.new_zeros((sample_count, 3, 3)) + yaw[:, 0, 0] = torch.cos(angles) + yaw[:, 0, 1] = -torch.sin(angles) + yaw[:, 1, 0] = torch.sin(angles) + yaw[:, 1, 1] = torch.cos(angles) + yaw[:, 2, 2] = 1.0 + variants = target_pose[:, None].repeat(1, sample_count, 1, 1) + variants[:, :, :3, :3] = torch.matmul(yaw[None], target_pose[:, None, :3, :3]) + return variants + + def _planning_context( + self, + state: ExecutionState, + grounded: GroundedAction, + ) -> PlanningContext: + qpos = state.last_qpos.to(device=self.device, dtype=torch.float32) + get_qvel = getattr(self.env.robot, "get_qvel", None) + qvel = get_qvel() if callable(get_qvel) else None + if not isinstance(qvel, torch.Tensor) or qvel.shape != qpos.shape: + qvel = torch.zeros_like(qpos) + else: + qvel = qvel.to(device=self.device, dtype=qpos.dtype) + return PlanningContext( + robot=RobotObservation(timestamp=self._scene_time, qpos=qpos, qvel=qvel), + task=state.to_task_state(), + scene=self._scene_snapshot(grounded, state), + env_ids=torch.arange( + self.num_envs, + dtype=torch.long, + device=self.device, + ), + control_dt=float(getattr(self.env, "step_dt", 1.0 / 60.0)), + ) + + def _scene_snapshot( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> SceneSnapshot: + dynamic_uids = tuple( + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ) + env_ids = torch.arange( + self.num_envs, + dtype=torch.long, + device=self.device, + ) + if self.scene_provider is None: + base = SceneSnapshot(timestamp=self._scene_time, version=0) + else: + base = self.scene_provider.snapshot( + timestamp=self._scene_time, + env_ids=env_ids, + ) + if not bool(self.planner_policy.get("dynamic_collision", False)): + return base + exclusion_masks = self._collision_exclusion_masks(grounded, state) + entities = dict(base.entities) + collision_pose_overrides: dict[str, torch.Tensor] = {} + for uid in dynamic_uids: + entity_state = entities.get(uid) + if entity_state is None: + raise ValueError( + f"SceneProvider omitted cuRobo dynamic obstacle {uid!r}." + ) + pose = entity_state.pose.to(dtype=torch.float32, device=self.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(self.num_envs, 1, 1) + if pose.shape != (self.num_envs, 4, 4): + raise ValueError( + f"Dynamic obstacle {uid!r} pose must have shape (4, 4) or " + f"({self.num_envs}, 4, 4), got {tuple(pose.shape)}." + ) + excluded = exclusion_masks.get(uid) + if excluded is not None and bool(excluded.any()): + collision_pose = pose.clone() + collision_pose[excluded, 2, 3] += _COLLISION_PARKING_Z_OFFSET + collision_pose_overrides[uid] = collision_pose + entities[uid] = EntityState( + pose=pose, + confidence=entity_state.confidence, + ) + return _CollisionOverrideSceneSnapshot( + timestamp=base.timestamp, + version=base.version, + entities=entities, + collision_world_revision=base.collision_world_revision, + collision_entity_ids=dynamic_uids, + collision_pose_overrides=collision_pose_overrides, + ) + + def _collision_exclusion_masks( + self, + grounded: GroundedAction, + state: ExecutionState, + ) -> dict[str, torch.Tensor]: + """Return per-environment masks for obstacles intentionally in contact.""" + dynamic_uids = { + str(uid) for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + } + masks: dict[str, torch.Tensor] = {} + + def include(uid: str | None, env_mask: torch.Tensor | None = None) -> None: + if uid is None or uid not in dynamic_uids: + return + mask = ( + torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + if env_mask is None + else torch.as_tensor( + env_mask, + dtype=torch.bool, + device=self.device, + ).reshape(-1) + ) + if mask.shape != (self.num_envs,): + raise ValueError( + f"Collision exclusion mask for {uid!r} must have shape " + f"({self.num_envs},), got {tuple(mask.shape)}." + ) + masks[uid] = masks.get(uid, torch.zeros_like(mask)) | mask + + if self.capabilities.get(grounded.action_class).allows_target_contact: + target_uid = grounded.object_uid + if target_uid is None: + target_uid = getattr( + getattr(grounded.target, "semantics", None), + "label", + None, + ) + include(target_uid) + + for held in state.held_objects.values(): + include(held.semantics.entity_id, held.env_mask) + collision_exclusion_uids = grounded.motion_policy.get( + "collision_exclusion_uids", () + ) + if isinstance(collision_exclusion_uids, str): + collision_exclusion_uids = (collision_exclusion_uids,) + for uid in collision_exclusion_uids: + include(str(uid)) + return masks + + def _invocation( + self, + grounded: GroundedAction, + capability: AtomicCapability, + *, + engine: AtomicActionEngine | None = None, + ) -> ActionInvocation: + if capability.resource_mode == "coordinated_object": + strategy = str(self.planner_policy["coordinated_strategy"]) + elif grounded.control == "hand": + strategy = "ik_interp" + else: + strategy = str(self.planner_policy["single_arm_strategy"]) + sample_count = max(2, int(grounded.cfg.get("sample_interval", 50))) + dynamic_collision = bool(self.planner_policy.get("dynamic_collision", False)) + collision_required = ( + grounded.motion_policy.get("collision_safety") == "required" + ) + if dynamic_collision and strategy == "motion_gen": + dynamic_mode = ( + DynamicCollisionMode.REQUIRED + if collision_required + else DynamicCollisionMode.AUTO + ) + else: + dynamic_mode = DynamicCollisionMode.OFF + goal = grounded.target + if capability.config_materializer == "coordinated_pickment": + self._validate_coordinated_pickment_goal(grounded) + return ActionInvocation( + skill_id=str(capability.action_type.skill_id), + goal=goal, + binding=self._binding(grounded, capability, engine=engine), + motion_policy=MotionPolicy( + strategy=strategy, + sample_count=sample_count, + dynamic_collision_mode=dynamic_mode, + ), + recovery_policy=RecoveryPolicy(), + skill_options=self._build_config(grounded, capability), + ) + + @staticmethod + def _validate_coordinated_pickment_goal( + grounded: GroundedAction, + ) -> CoordinatedPickGoal: + """Validate the coordinated goal against the engine-scoped generator.""" + target = grounded.target + if not isinstance(target, CoordinatedPickGoal): + raise TypeError("CoordinatedPickment requires a CoordinatedPickGoal.") + requested = grounded.cfg.get("is_filter_ground_collision") + if requested is not None and not isinstance(requested, bool): + raise TypeError("is_filter_ground_collision must be a boolean.") + if not isinstance(target.semantics.affordance, AntipodalAffordance): + raise TypeError( + "CoordinatedPickment requires an AntipodalAffordance for GenSim " + "grasp filtering." + ) + return target + + def _binding( + self, + action: GroundedAction, + capability: AtomicCapability, + *, + engine: AtomicActionEngine | None = None, + ) -> ActionBinding: + engine = self._engine() if engine is None else engine + contract = getattr(capability.action_type, "binding_contract", None) + if contract is None: + return ActionBinding(owner_id=engine.binding_owner_id) + + slot_parts: dict[str, tuple[str, str | None]] = {} + task_state_keys: dict[str, str] | None = None + if capability.config_materializer == "handover": + transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + transfer_arm, transfer_hand, _ = self._parts(transfer_side) + receive_arm, receive_hand, _ = self._parts(receive_side) + if transfer_hand is None or receive_hand is None: + raise ValueError("HandOver requires two configured end effectors.") + slot_parts = { + "source": (transfer_arm, transfer_hand), + "destination": (receive_arm, receive_hand), + } + elif capability.config_materializer == "coordinated_pickment": + left_arm, left_hand, _ = self._parts("left_arm") + right_arm, right_hand, _ = self._parts("right_arm") + if left_hand is None or right_hand is None: + raise ValueError("Coordinated pickup requires two end effectors.") + slot_parts = { + "left": (left_arm, left_hand), + "right": (right_arm, right_hand), + } + elif capability.config_materializer == "coordinated_placement": + placing_arm, placing_hand, _ = self._parts("left_arm") + support_arm, support_hand, _ = self._parts("right_arm") + if placing_hand is None or support_hand is None: + raise ValueError("Coordinated placement requires two end effectors.") + slot_parts = { + "placing": (placing_arm, placing_hand), + "support": (support_arm, support_hand), + } + else: + arm_part, hand_part, _ = self._parts(action.arm) + motion_part = hand_part if action.control == "hand" else arm_part + if motion_part is None: + raise ValueError( + f"{action.arm} has no configured {action.control} part." + ) + slot_parts = {"primary": (motion_part, hand_part)} + if action.control == "hand" and bool( + action.cfg.get("single_release", False) + ): + task_state_keys = {"primary": arm_part} + + endpoints: dict[str, dict[str, str]] = {} + for slot in contract.slots: + try: + motion_part, hand_part = slot_parts[slot.slot_id] + except KeyError as exc: + raise ValueError( + f"No GenSim binding is available for slot {slot.slot_id!r}." + ) from exc + selected: dict[str, str] = {} + for requirement in slot.endpoints: + if requirement.endpoint_id == "motion": + selected["motion"] = motion_part + elif requirement.endpoint_id == "grasp": + if hand_part is None: + raise ValueError( + f"{capability.name} requires a grasp endpoint for " + f"slot {slot.slot_id!r}." + ) + selected["grasp"] = hand_part + else: + raise ValueError( + f"Unsupported GenSim endpoint {slot.slot_id}." + f"{requirement.endpoint_id}." + ) + endpoints[slot.slot_id] = selected + skill_id = str(capability.action_type.skill_id) + if task_state_keys is None: + return engine.bind_control_parts(skill_id, endpoints) + return engine.bind_control_parts( + skill_id, + endpoints, + task_state_keys=task_state_keys, + ) + + def _build_config( + self, + action: GroundedAction, + capability: AtomicCapability | type, + ) -> Any: + """Build the mainline immutable ``ActionOptions`` value. + + The method name is retained as a narrow compatibility hook for existing + Action Engine tests and extensions; it no longer constructs legacy + hardware-bound ``ActionCfg`` objects. + """ + if isinstance(capability, type): + registered = self.capabilities.require_executable(action.action_class) + if registered.config_type is not capability: + raise ValueError( + f"Options type {capability.__name__!r} does not match " + f"AtomicAction {action.action_class!r}." + ) + capability = registered + if capability.config_materializer_hook is not None: + return capability.config_materializer_hook( + adapter=self, + action=action, + capability=capability, + ) + builder = getattr( + self, + f"_build_{capability.config_materializer}_config", + self._build_single_arm_config, + ) + return builder(action, capability) + + def _config_policy(self, action: GroundedAction) -> dict[str, Any]: + policy = dict(action.cfg) + for key in ( + "postcondition_tolerance", + "relation_distance", + "hover_height", + "staging_lift_height", + "transport_clearance", + "surface_clearance", + "receiver_hold_joint_tolerance", + "post_hold_steps", + ): + policy.pop(key, None) + return policy + + def _build_single_arm_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + config_type = capability.config_type + if capability.target_materializer == "semantic_held_object": + from .atomic_compat import ExactTargetMoveHeldObjectOptions + + config_type = ExactTargetMoveHeldObjectOptions + elif capability.target_materializer == "joint_state": + from .atomic_compat import ActionEngineMoveJointsOptions + + config_type = ActionEngineMoveJointsOptions + if capability.target_materializer == "press": + press_depth = policy.pop("press_depth", None) + if press_depth is not None and "press_distance" not in policy: + policy["press_distance"] = press_depth + approach_mode = policy.pop("approach_direction_mode", None) + if approach_mode == "handover_transfer": + from .frames import robot_frame_axes + + _, lateral = robot_frame_axes(self.env) + outward = lateral[0] if action.arm == "left_arm" else -lateral[0] + policy["approach_direction"] = _diagonal_approach_direction( + -outward.to(device=self.device) + ) + elif approach_mode is not None: + raise ValueError(f"Unknown approach_direction_mode {approach_mode!r}.") + for name in ("approach_direction", "obj_upright_direction", "target_axis"): + if name in policy and not isinstance(policy[name], torch.Tensor): + policy[name] = torch.as_tensor( + policy[name], dtype=torch.float32, device=self.device + ) + return config_type(**_supported_kwargs(config_type, policy)) + + def _build_coordinated_pickment_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + left_base, right_base = self._coordinated_arm_bases() + direction = right_base[0, :3, 3] - left_base[0, :3, 3] + norm = torch.linalg.vector_norm(direction) + if not torch.isfinite(direction).all() or norm <= 1.0e-6: + raise ValueError( + "Coordinated pickup requires distinct finite left/right arm bases." + ) + policy.setdefault("left_to_right_arm_direction", direction / norm) + for name in ("approach_direction", "left_to_right_arm_direction"): + if name in policy and not isinstance(policy[name], torch.Tensor): + policy[name] = torch.as_tensor( + policy[name], dtype=torch.float32, device=self.device + ) + return capability.config_type( + **_supported_kwargs(capability.config_type, policy) + ) + + def _build_coordinated_placement_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + return self._build_single_arm_config(action, capability) + + def _build_handover_config( + self, + action: GroundedAction, + capability: AtomicCapability, + ) -> Any: + policy = self._config_policy(action) + middle = action.cfg.get("middle_object_pose") + final = action.cfg.get("final_object_pose") + if middle is None or final is None: + raise ValueError("HandOver grounding must provide middle and final poses.") + transfer_side = str(action.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + from .frames import robot_frame_axes + + _, lateral = robot_frame_axes(self.env) + receiver_outward = ( + lateral[0] if receive_side == "left_arm" else -lateral[0] + ).to(device=self.device) + receiver_inward_approach = -receiver_outward + policy.update( + { + "middle_object_pose": middle, + # Delivery is represented by a following MoveHeldObject node. + # Keep the receiver fixed while the source retreats here. + "final_object_pose": middle, + "receive_approach_direction": _diagonal_approach_direction( + receiver_inward_approach + ), + } + ) + return capability.config_type( + **_supported_kwargs(capability.config_type, policy) + ) + + def _positions_with_agent_holds( + self, + plan: ActionPlan, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> torch.Tensor: + trajectory = plan.joint_trajectory + if trajectory is None: + raise ValueError( + f"AtomicAction {plan.skill_id!r} did not retain a joint trajectory." + ) + positions = trajectory.positions.to( + device=self.device, + dtype=torch.float32, + ) + hold_steps = int(grounded.cfg.get("post_hold_steps", 0)) + if capability.state_effect != "release" or hold_steps <= 0: + return positions + release = next((item for item in plan.segments if item.name == "release"), None) + if release is None or release.stop <= 0 or release.stop > positions.shape[1]: + return positions + hold = positions[:, release.stop - 1 : release.stop].repeat(1, hold_steps, 1) + return torch.cat( + (positions[:, : release.stop], hold, positions[:, release.stop :]), + dim=1, + ) + + @staticmethod + def _merge_plan_rows( + primary: torch.Tensor, + fallback: torch.Tensor, + use_fallback: torch.Tensor, + hold_qpos: torch.Tensor, + ) -> torch.Tensor: + steps = max(primary.shape[1], fallback.shape[1], 1) + + def padded(value: torch.Tensor) -> torch.Tensor: + if value.shape[1] == 0: + return hold_qpos[:, None].repeat(1, steps, 1) + if value.shape[1] < steps: + value = torch.cat( + (value, value[:, -1:].repeat(1, steps - value.shape[1], 1)), + dim=1, + ) + return value + + primary = padded(primary) + fallback = padded(fallback) + return torch.where(use_fallback[:, None, None], fallback, primary) + + def _handover_receiver_hold_mask( + self, + trajectory: torch.Tensor, + grounded: GroundedAction, + options: Any, + *, + tolerance: float, + ) -> torch.Tensor: + if tolerance < 0.0: + raise ValueError("receiver_hold_joint_tolerance must be non-negative.") + retreat_steps = max(2, int(options.retreat_steps)) + if trajectory.shape[1] < retreat_steps: + return torch.zeros( + self.num_envs, dtype=torch.bool, device=trajectory.device + ) + transfer_side = str(grounded.cfg.get("transfer_arm", "left_arm")) + receive_side = "right_arm" if transfer_side == "left_arm" else "left_arm" + receive_arm, _, _ = self._parts(receive_side) + receiver_ids = self.env.robot.get_joint_ids(name=receive_arm) + receiver = trajectory[:, -retreat_steps:, receiver_ids] + drift = torch.amax(torch.abs(receiver - receiver[:, :1]), dim=(1, 2)) + return torch.isfinite(drift) & (drift <= tolerance) + + def execute_trajectory( + self, + trajectory: torch.Tensor, + *, + active: torch.Tensor, + waypoint_observer: Callable[[int], None] | None = None, + ) -> list[torch.Tensor]: + """Advance the environment while holding inactive vectorized rows.""" + if trajectory.ndim != 3 or trajectory.shape[0] != self.num_envs: + raise ValueError("Execution trajectory must have shape (N, T, robot_dof).") + active = active.to(device=trajectory.device, dtype=torch.bool) + current = self.env.robot.get_qpos().to( + device=trajectory.device, + dtype=trajectory.dtype, + ) + commands: list[torch.Tensor] = [] + for waypoint_index, waypoint in enumerate(trajectory.unbind(dim=1)): + command = torch.where(active[:, None], waypoint, current) + self.env.step(command) + self._scene_time += self._scene_step_duration() + update = getattr(self.env, "update_obj_info", None) + if callable(update): + update() + if waypoint_observer is not None: + waypoint_observer(waypoint_index) + commands.append(command.detach()) + current = command + sync = getattr(self.env, "sync_agent_state_from_qpos", None) + if callable(sync) and commands: + sync(commands[-1]) + return commands + + def _scene_step_duration(self) -> float: + """Return one positive logical waypoint duration for scene timestamps.""" + sim_config = getattr(getattr(self.env, "sim", None), "sim_config", None) + candidates = ( + getattr(self.env, "physics_dt", None), + getattr(sim_config, "physics_dt", None), + ) + for value in candidates: + if isinstance(value, (int, float)) and not isinstance(value, bool): + duration = float(value) + if math.isfinite(duration) and duration > 0.0: + return duration + return 1.0 + + def combine( + self, + outcomes: Mapping[str, ActionOutcome | None], + masks: Mapping[str, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Merge independently planned arm paths into one synchronized stream.""" + present = [item for item in outcomes.values() if item is not None] + if not present: + raise ValueError("At least one arm outcome is required.") + steps = max(int(item.trajectory.shape[1]) for item in present) + current = self.env.robot.get_qpos().to(self.device, dtype=torch.float32) + merged = current[:, None, :].repeat(1, max(steps, 1), 1) + success = torch.ones( + self.num_envs, + dtype=torch.bool, + device=self.device, + ) + for arm, outcome in outcomes.items(): + if outcome is None: + continue + mask = masks[arm].to(self.device, dtype=torch.bool) + success &= ~mask | outcome.success + trajectory = outcome.trajectory + if trajectory.shape[1] == 0: + continue + if trajectory.shape[1] < steps: + padding = trajectory[:, -1:].repeat(1, steps - trajectory.shape[1], 1) + trajectory = torch.cat((trajectory, padding), dim=1) + joint_ids = self.joint_ids(arm, include_hand=True) + if not joint_ids: + continue + selected = merged[:, :, joint_ids] + merged[:, :, joint_ids] = torch.where( + mask[:, None, None], trajectory[:, :, joint_ids], selected + ) + return merged, success + + def joint_ids(self, arm: str, *, include_hand: bool) -> list[int]: + if arm == "coordinated": + return list(range(int(self.env.robot.dof))) + side = "left" if arm == "left_arm" else "right" + result = list(getattr(self.env, f"{side}_arm_joints", ())) + if include_hand: + result.extend(getattr(self.env, f"{side}_eef_joints", ())) + return result + + def _engine(self) -> AtomicActionEngine: + if self._atomic_engine is None: + self._atomic_engine = self._new_engine( + self._generator(), + filter_ground_collision=True, + ) + return self._atomic_engine + + def _engine_for( + self, + grounded: GroundedAction, + capability: AtomicCapability, + ) -> AtomicActionEngine: + if capability.config_materializer != "coordinated_pickment": + return self._engine() + filter_ground_collision = grounded.cfg.get( + "is_filter_ground_collision", + True, + ) + if not isinstance(filter_ground_collision, bool): + raise TypeError("is_filter_ground_collision must be a boolean.") + opening_margin = grounded.cfg.get( + "grasp_opening_margin", + self.gripper_profile.grasp_model.opening_margin, + ) + if isinstance(opening_margin, bool) or not isinstance( + opening_margin, (int, float) + ): + raise TypeError("grasp_opening_margin must be a real number.") + opening_margin = float(opening_margin) + if ( + not math.isfinite(opening_margin) + or opening_margin < 0.0 + or opening_margin >= self.gripper_profile.grasp_model.max_opening_width + ): + raise ValueError( + "grasp_opening_margin must be finite, non-negative, and smaller " + "than the selected gripper's maximum opening width." + ) + profile_margin = self.gripper_profile.grasp_model.opening_margin + if filter_ground_collision and opening_margin == profile_margin: + return self._engine() + cache_key = (filter_ground_collision, opening_margin) + cached = self._coordinated_engines.get(cache_key) + if cached is None: + cached = self._new_engine( + MotionGenerator(cfg=self._motion_generator_cfg()), + filter_ground_collision=filter_ground_collision, + opening_margin=opening_margin, + ) + self._coordinated_engines[cache_key] = cached + return cached + + def _new_engine( + self, + motion_generator: MotionGenerator, + *, + filter_ground_collision: bool, + opening_margin: float | None = None, + ) -> AtomicActionEngine: + from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver + + from .atomic_compat import ActionEngineMoveJoints, ExactTargetMoveHeldObject + + engine = AtomicActionEngine( + motion_generator, + control_profiles=self._control_profiles(), + grasp_pose_generators=self._grasp_pose_generators( + filter_ground_collision=filter_ground_collision, + opening_margin=opening_margin, + ), + ) + engine.register(ExactTargetMoveHeldObject(), replace=True) + engine.register(ActionEngineMoveJoints(), replace=True) + engine.register(HeldObjectHandOver(), replace=True) + return engine + + def _generator(self) -> MotionGenerator: + if self._motion_generator is None: + self._motion_generator = MotionGenerator(cfg=self._motion_generator_cfg()) + return self._motion_generator + + def _motion_generator_cfg(self) -> MotionGenCfg: + backend = str(self.planner_policy.get("backend", "curobo")) + if backend == "curobo": + options = dict(self.planner_policy.get("curobo", {})) + obstacle_uids = tuple( + dict.fromkeys( + [ + *self.planner_policy.get("static_obstacle_uids", ()), + *self.planner_policy.get("dynamic_obstacle_uids", ()), + ] + ) + ) + rigid_objects: dict[str, Any] = {} + for uid in obstacle_uids: + obstacle_uid = str(uid) + entity = self.env.sim.get_rigid_object(obstacle_uid) + if entity is None: + raise ValueError(f"Unknown cuRobo obstacle {uid!r}.") + rigid_objects[obstacle_uid] = entity + obstacle_representation = str( + options.get("obstacle_representation", "cuboid") + ) + world = CuroboWorldCfg( + rigid_objects=rigid_objects or None, + obstacle_representation=obstacle_representation, + collision_cache=_collision_cache_for_world( + obstacle_representation, + len(rigid_objects), + ), + dynamic_obstacle_names=[ + str(uid) + for uid in self.planner_policy.get("dynamic_obstacle_uids", ()) + ], + multi_env=bool(options.get("multi_env", False)), + ) + planner_cfg = CuroboPlannerCfg( + robot_uid=self.env.robot.uid, + log_level=str(options.get("log_level", "error")), + world=world, + use_cuda_graph=bool(options.get("use_cuda_graph", True)), + preserve_plan_samples=bool(options.get("preserve_plan_samples", False)), + max_attempts=int(options.get("max_attempts", 5)), + collision_activation_distance=float( + options.get("collision_activation_distance", 0.01) + ), + ) + elif backend == "toppra": + planner_cfg = ToppraPlannerCfg(robot_uid=self.env.robot.uid) + else: + raise ValueError(f"Unsupported Action Engine planner backend {backend!r}.") + return MotionGenCfg(planner_cfg=planner_cfg) + + def _control_profiles(self) -> dict[str, ControlPartCommandProfile]: + profiles: dict[str, ControlPartCommandProfile] = {} + for side in ("left_arm", "right_arm"): + try: + _, hand_part, hand_dof = self._parts(side) + except ValueError: + continue + if hand_part is None or hand_dof == 0 or hand_part in profiles: + continue + profiles[hand_part] = ControlPartCommandProfile.joint_positions( + open=_as_hand_qpos(self.env.open_state, hand_dof, self.device), + grasp=_as_hand_qpos(self.env.close_state, hand_dof, self.device), + ) + return profiles + + def _grasp_pose_generators( + self, + *, + filter_ground_collision: bool = True, + opening_margin: float | None = None, + ) -> dict[str, AntipodalGraspPoseGenerator]: + """Build one mainline grasp service for each runtime hand endpoint.""" + if not isinstance(filter_ground_collision, bool): + raise TypeError("filter_ground_collision must be a boolean.") + options = self.grasp_policy + geometry = self.gripper_profile.grasp_model + if opening_margin is None: + opening_margin = geometry.opening_margin + if isinstance(opening_margin, bool) or not isinstance( + opening_margin, (int, float) + ): + raise TypeError("opening_margin must be a real number or None.") + opening_margin = float(opening_margin) + if ( + not math.isfinite(opening_margin) + or opening_margin < 0.0 + or opening_margin >= geometry.max_opening_width + ): + raise ValueError( + "opening_margin must be finite, non-negative, and smaller than " + "the selected gripper's maximum opening width." + ) + model = ParallelJawGripperModelCfg( + model_id=geometry.model_id, + min_opening_width=geometry.min_opening_width, + max_opening_width=geometry.max_opening_width, + finger_length=geometry.finger_length, + finger_width=geometry.finger_width, + finger_thickness=geometry.finger_thickness, + palm_depth=geometry.palm_depth, + ) + algorithm = AntipodalGraspPoseGeneratorCfg( + sample_count=int(options["antipodal_n_sample"]), + ray_deviation_angle=float(options["antipodal_max_angle"]), + approach_deviation_angle=float(options["max_deviation_angle"]), + approach_direction_samples=int(options["n_deviated_approach_directions"]), + max_candidates=_BODY_GRASP_CANDIDATE_LIMIT, + ) + collision = ParallelJawGraspCollisionCfg( + point_sample_density=float(options["point_sample_dense"]), + max_decomposition_hulls=int(options["max_decomposition_hulls"]), + opening_margin=opening_margin, + filter_ground_collision=filter_ground_collision, + ) + annotation = GraspAnnotationCfg( + selection_mode="whole_mesh", + viser_port=int(options["viser_port"]), + force_refresh=bool(options["force_grasp_reannotate"]), + ) + shared_generator = _TracingAntipodalGraspPoseGenerator( + model, + algorithm_cfg=algorithm, + collision_cfg=collision, + annotation_cfg=annotation, + ) + generators: dict[str, AntipodalGraspPoseGenerator] = {} + for arm in ("left_arm", "right_arm"): + try: + _, hand_part, _ = self._parts(arm) + except ValueError: + continue + if hand_part is None or hand_part in generators: + continue + generators[hand_part] = shared_generator + return generators + + def _parts(self, arm: str) -> tuple[str, str | None, int]: + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a physical arm, got {arm!r}.") + is_left = arm == "left_arm" + if hasattr(self.env, "get_agent_arm_control_part"): + arm_part = self.env.get_agent_arm_control_part(is_left) + hand_part = self.env.get_agent_eef_control_part(is_left) + else: + arm_part = arm + hand_part = "left_eef" if is_left else "right_eef" + hand_ids = ( + [] + if hand_part is None + else list(self.env.robot.get_joint_ids(name=hand_part)) + ) + return ( + str(arm_part), + None if hand_part is None else str(hand_part), + len(hand_ids), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/atomic_compat.py b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py new file mode 100644 index 000000000..6ca39d26d --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/atomic_compat.py @@ -0,0 +1,106 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Action Engine-specific adapters for mainline atomic actions.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +from embodichain.lab.sim.atomic_actions import ( + ActionPlan, + JointPositionGoal, + MoveHeldObject, + MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + PlanningContext, + ResolvedActionRequest, + StateDelta, +) + +__all__ = [ + "ActionEngineMoveJoints", + "ActionEngineMoveJointsOptions", + "ExactTargetMoveHeldObject", + "ExactTargetMoveHeldObjectOptions", +] + + +@dataclass(frozen=True, slots=True, eq=False) +class ExactTargetMoveHeldObjectOptions(MoveHeldObjectOptions): + """Action Engine transport options for a grounded object target.""" + + +class ExactTargetMoveHeldObject(MoveHeldObject): + """Action Engine marker for mainline exact-target transport.""" + + OptionsType = ExactTargetMoveHeldObjectOptions + binding_contract = MoveHeldObject.binding_contract + + +@dataclass(frozen=True, slots=True, eq=False) +class ActionEngineMoveJointsOptions(MoveJointsOptions): + """Joint motion with an explicit optional single-arm release effect.""" + + single_release: bool = False + """Whether a successful gripper-open command releases the held object.""" + + def __post_init__(self) -> None: + if type(self.single_release) is not bool: + raise TypeError("single_release must be a boolean.") + + +class ActionEngineMoveJoints(MoveJoints): + """Preserve ordinary joint motion and commit explicit release nodes.""" + + OptionsType = ActionEngineMoveJointsOptions + binding_contract = MoveJoints.binding_contract + + def _plan( + self, + request: ResolvedActionRequest[ + JointPositionGoal, + ActionEngineMoveJointsOptions, + ], + context: PlanningContext, + ) -> ActionPlan: + endpoint = request.binding.endpoint("primary", "motion") + task_state_key = endpoint.task_state_key + if request.skill_options.single_release: + if not isinstance(task_state_key, str) or not task_state_key: + raise ValueError( + "Single-arm release requires a non-empty task-state key." + ) + if context.task.get_held_object(task_state_key) is None: + return self.failed_plan( + request, + context, + message=( + "Single-arm release requires an object held by task-state " + f"resource {task_state_key!r}." + ), + ) + + plan = super()._plan(request, context) + if not request.skill_options.single_release: + return plan + return replace( + plan, + expected_effects=StateDelta( + held_object_updates={task_state_key: None}, + ), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/body_grasp.py b/embodichain/gen_sim/action_engine/runtime/body_grasp.py new file mode 100644 index 000000000..560e34b5b --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/body_grasp.py @@ -0,0 +1,297 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pure body-grasp candidate filtering for elongated rigid objects.""" + +from __future__ import annotations + +from dataclasses import dataclass +from dataclasses import replace +from collections.abc import Callable + +import torch + +from embodichain.lab.sim.atomic_actions import ( + AxisAlignAffordance, + AxisAlignGoal, +) +from .geometry_axes import LocalGeometryAxes +from .geometry_axes import analyze_local_geometry_axes + +__all__ = [ + "AxisAlignBodyGraspAdapter", + "BodyGraspAdaptation", + "BodyGraspSelection", + "select_body_grasp_candidates", +] + + +@dataclass(frozen=True, slots=True) +class BodyGraspSelection: + """One selected body grasp per environment row.""" + + success: torch.Tensor + grasp_xpos: torch.Tensor + candidate_indices: torch.Tensor + body_candidate_counts: torch.Tensor + central_candidate_counts: torch.Tensor + radial_candidate_counts: torch.Tensor + minimum_normalized_axial_offset: torch.Tensor + minimum_long_axis_opening_cosine: torch.Tensor + reachable_candidate_counts: torch.Tensor + ranked_grasp_xpos: torch.Tensor + ranked_candidate_indices: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class BodyGraspAdaptation: + """AxisAlign goal plus auditable body-grasp selection metadata.""" + + goal: AxisAlignGoal + alternative_goals: tuple[AxisAlignGoal, ...] + alternative_rank_indices: tuple[int, ...] + axes: LocalGeometryAxes + selection: BodyGraspSelection + + +class AxisAlignBodyGraspAdapter: + """Lower elongated-object semantics into an explicit mainline grasp goal.""" + + def __init__( + self, + *, + body_band_fraction: float = 0.80, + maximum_long_axis_opening_cosine: float = 0.50, + ) -> None: + self.body_band_fraction = body_band_fraction + self.maximum_long_axis_opening_cosine = maximum_long_axis_opening_cosine + + def adapt( + self, + goal: AxisAlignGoal, + *, + object_pose: torch.Tensor, + grasp_generator: object, + approach_direction: torch.Tensor, + target_axis: torch.Tensor, + seed: int, + candidate_feasibility: Callable[[torch.Tensor], torch.Tensor] | None = None, + maximum_adaptations: int = 12, + ) -> BodyGraspAdaptation: + affordance = goal.semantics.affordance + if not isinstance(affordance, AxisAlignAffordance): + raise TypeError("AxisAlign body grasp requires AxisAlignAffordance.") + if affordance.mesh_vertices is None or affordance.mesh_triangles is None: + raise ValueError("AxisAlign body grasp requires indexed mesh geometry.") + axes = analyze_local_geometry_axes(affordance.mesh_vertices) + sampled = self._sample( + grasp_generator, + seed=seed, + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + obj_poses=object_pose, + approach_direction=approach_direction, + obj_longest_axis=None, + is_positive_part=True, + ) + candidates, costs = self._pack(sampled, object_pose) + feasible = ( + None if candidate_feasibility is None else candidate_feasibility(candidates) + ) + selection = select_body_grasp_candidates( + candidates, + costs, + object_pose, + axes, + body_band_fraction=self.body_band_fraction, + maximum_long_axis_opening_cosine=(self.maximum_long_axis_opening_cosine), + feasible=feasible, + ) + if not bool(selection.success.all().item()): + failed = ( + torch.nonzero(~selection.success, as_tuple=False).flatten().tolist() + ) + raise ValueError( + "No central radial body grasp is available for rows " + f"{failed}; central_counts=" + f"{selection.central_candidate_counts.tolist()}, radial_counts=" + f"{selection.radial_candidate_counts.tolist()}, min_axial=" + f"{selection.minimum_normalized_axial_offset.tolist()}, " + "min_opening_cos=" + f"{selection.minimum_long_axis_opening_cosine.tolist()}, " + "reachable_counts=" + f"{selection.reachable_candidate_counts.tolist()}." + ) + adaptation_count = min( + maximum_adaptations, + selection.ranked_grasp_xpos.shape[1], + ) + alternative_ranks = tuple( + int(value) + for value in torch.linspace( + 0, + selection.ranked_grasp_xpos.shape[1] - 1, + adaptation_count, + ) + .round() + .to(torch.int64) + .tolist() + ) + goals = tuple( + replace(goal, grasp_xpos=selection.ranked_grasp_xpos[:, rank]) + for rank in alternative_ranks + ) + del target_axis + return BodyGraspAdaptation( + goal=goals[0], + alternative_goals=goals, + alternative_rank_indices=alternative_ranks, + axes=axes, + selection=selection, + ) + + @staticmethod + def _sample( + generator: object, + *, + seed: int, + **kwargs: object, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + poses = kwargs.get("obj_poses") + if not isinstance(poses, torch.Tensor): + raise TypeError("obj_poses must be a torch.Tensor.") + devices: list[int] = [] + if poses.device.type == "cuda": + devices.append( + torch.cuda.current_device() + if poses.device.index is None + else poses.device.index + ) + with torch.random.fork_rng(devices=devices): + torch.manual_seed(seed) + return generator.get_valid_grasp_poses( # type: ignore[attr-defined] + **kwargs + ) + + @staticmethod + def _pack( + sampled: list[tuple[torch.Tensor, torch.Tensor]], + object_pose: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if len(sampled) != object_pose.shape[0]: + raise ValueError("Grasp generator must return one result per object row.") + count = max((poses.shape[0] for poses, _ in sampled), default=0) + if count == 0: + raise ValueError("Grasp generator returned no candidates.") + candidates = torch.eye( + 4, + dtype=torch.float32, + device=object_pose.device, + ).repeat(object_pose.shape[0], count, 1, 1) + costs = torch.full( + (object_pose.shape[0], count), + torch.inf, + dtype=torch.float32, + device=object_pose.device, + ) + for env_index, (poses, values) in enumerate(sampled): + row_count = poses.shape[0] + if row_count == 0: + continue + candidates[env_index, :row_count] = poses.to( + device=object_pose.device, + dtype=torch.float32, + ) + costs[env_index, :row_count] = values.to( + device=object_pose.device, + dtype=torch.float32, + ) + return candidates, costs + + +def select_body_grasp_candidates( + candidates: torch.Tensor, + costs: torch.Tensor, + object_pose: torch.Tensor, + axes: LocalGeometryAxes, + *, + body_band_fraction: float = 0.80, + maximum_long_axis_opening_cosine: float = 0.50, + feasible: torch.Tensor | None = None, +) -> BodyGraspSelection: + """Keep central radial grasps and reject cap/end grasps.""" + if candidates.ndim != 4 or candidates.shape[-2:] != (4, 4): + raise ValueError("candidates must have shape (B, N, 4, 4).") + if costs.shape != candidates.shape[:2]: + raise ValueError("costs must have shape (B, N).") + if object_pose.shape != (candidates.shape[0], 4, 4): + raise ValueError("object_pose must have shape (B, 4, 4).") + if not 0.0 < body_band_fraction <= 1.0: + raise ValueError("body_band_fraction must be in (0, 1].") + if not 0.0 <= maximum_long_axis_opening_cosine < 1.0: + raise ValueError("maximum_long_axis_opening_cosine must be in [0, 1).") + if feasible is None: + feasible = torch.ones_like(costs, dtype=torch.bool) + if feasible.dtype != torch.bool or feasible.shape != costs.shape: + raise ValueError("feasible must be a bool tensor shaped (B, N).") + + rotation = object_pose[:, :3, :3] + translation = object_pose[:, :3, 3] + local_centers = torch.matmul( + candidates[..., :3, 3] - translation[:, None], + rotation, + ) + center = axes.bounds_center.to( + device=candidates.device, + dtype=candidates.dtype, + ) + long_axis = axes.long_axis.to( + device=candidates.device, + dtype=candidates.dtype, + ) + axial_offset = torch.abs(torch.sum((local_centers - center) * long_axis, dim=-1)) + normalized_axial = axial_offset / max(axes.long_half_extent, 1.0e-8) + within_body = normalized_axial <= body_band_fraction + + world_opening = torch.nn.functional.normalize(candidates[..., :3, 0], dim=-1) + local_opening = torch.matmul(world_opening, rotation) + long_axis_opening = torch.abs(torch.sum(local_opening * long_axis, dim=-1)) + radial = long_axis_opening <= maximum_long_axis_opening_cosine + valid = within_body & radial & feasible & torch.isfinite(costs) + + ranked = torch.where(valid, costs, torch.inf) + best_cost, best_index = ranked.min(dim=1) + env_index = torch.arange(candidates.shape[0], device=candidates.device) + valid_counts = valid.sum(dim=1) + rank_count = int(valid_counts.min().item()) + ranked_indices = torch.argsort(ranked, dim=1)[:, :rank_count] + ranked_grasps = candidates[ + env_index[:, None], + ranked_indices, + ].clone() + return BodyGraspSelection( + success=torch.isfinite(best_cost), + grasp_xpos=candidates[env_index, best_index].clone(), + candidate_indices=best_index.clone(), + body_candidate_counts=valid.sum(dim=1), + central_candidate_counts=within_body.sum(dim=1), + radial_candidate_counts=radial.sum(dim=1), + minimum_normalized_axial_offset=normalized_axial.min(dim=1).values, + minimum_long_axis_opening_cosine=long_axis_opening.min(dim=1).values, + reachable_candidate_counts=feasible.sum(dim=1), + ranked_grasp_xpos=ranked_grasps, + ranked_candidate_indices=ranked_indices.clone(), + ) diff --git a/embodichain/gen_sim/action_engine/runtime/coordinated_safety.py b/embodichain/gen_sim/action_engine/runtime/coordinated_safety.py new file mode 100644 index 000000000..5bd150d82 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/coordinated_safety.py @@ -0,0 +1,391 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Pure dual-arm grasp and trajectory safety checks owned by GenSim.""" + +from __future__ import annotations + +from dataclasses import dataclass +import math + +import torch + +__all__: list[str] = [] + + +@dataclass(frozen=True, slots=True) +class _CanonicalizedGraspPoses: + poses: torch.Tensor + flipped: torch.Tensor + selected_rotation_radians: torch.Tensor + alternative_rotation_radians: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _RankedGraspPairs: + ranked_pairs: tuple[tuple[int, int], ...] + scores: tuple[float, ...] + rejection_counts: dict[str, int] + + +@dataclass(frozen=True, slots=True) +class _TrajectorySafetyReport: + success: torch.Tensor + failed_checks: dict[str, list[bool]] + metrics: dict[str, list[float]] + + +def _rotation_distance( + actual: torch.Tensor, + expected: torch.Tensor, +) -> torch.Tensor: + relative = torch.matmul(actual.transpose(-1, -2), expected) + cosine = (torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) - 1.0) * 0.5 + return torch.acos(torch.clamp(cosine, -1.0, 1.0)) + + +def _canonicalize_parallel_jaw_poses( + poses: torch.Tensor, + live_eef_pose: torch.Tensor, +) -> _CanonicalizedGraspPoses: + """Choose the local-Z half-turn equivalent nearest one live wrist pose.""" + poses = torch.as_tensor(poses, dtype=torch.float32) + live = torch.as_tensor(live_eef_pose, dtype=poses.dtype, device=poses.device) + if poses.ndim != 3 or poses.shape[1:] != (4, 4): + raise ValueError("poses must have shape (N, 4, 4).") + if live.shape != (4, 4): + raise ValueError("live_eef_pose must have shape (4, 4).") + half_turn = torch.eye(4, dtype=poses.dtype, device=poses.device) + half_turn[0, 0] = -1.0 + half_turn[1, 1] = -1.0 + alternatives = torch.matmul(poses, half_turn) + live_rot = live[:3, :3].unsqueeze(0).expand(poses.shape[0], -1, -1) + original_distance = _rotation_distance(live_rot, poses[:, :3, :3]) + alternative_distance = _rotation_distance(live_rot, alternatives[:, :3, :3]) + flipped = alternative_distance < original_distance + selected = torch.where(flipped[:, None, None], alternatives, poses) + selected_distance = torch.where( + flipped, + alternative_distance, + original_distance, + ) + rejected_distance = torch.where( + flipped, + original_distance, + alternative_distance, + ) + return _CanonicalizedGraspPoses( + poses=selected, + flipped=flipped, + selected_rotation_radians=selected_distance, + alternative_rotation_radians=rejected_distance, + ) + + +def _segments_intersect_2d( + first_start: torch.Tensor, + first_end: torch.Tensor, + second_start: torch.Tensor, + second_end: torch.Tensor, +) -> bool: + epsilon = 1.0e-7 + + def orientation(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> float: + ab = b - a + ac = c - a + return float(ab[0] * ac[1] - ab[1] * ac[0]) + + def on_segment(a: torch.Tensor, b: torch.Tensor, point: torch.Tensor) -> bool: + return bool( + float(torch.min(a[0], b[0])) - epsilon + <= float(point[0]) + <= float(torch.max(a[0], b[0])) + epsilon + and float(torch.min(a[1], b[1])) - epsilon + <= float(point[1]) + <= float(torch.max(a[1], b[1])) + epsilon + ) + + a = first_start[:2] + b = first_end[:2] + c = second_start[:2] + d = second_end[:2] + abc = orientation(a, b, c) + abd = orientation(a, b, d) + cda = orientation(c, d, a) + cdb = orientation(c, d, b) + if abc * abd < 0.0 and cda * cdb < 0.0: + return True + return ( + (abs(abc) <= epsilon and on_segment(a, b, c)) + or (abs(abd) <= epsilon and on_segment(a, b, d)) + or (abs(cda) <= epsilon and on_segment(c, d, a)) + or (abs(cdb) <= epsilon and on_segment(c, d, b)) + ) + + +def _rank_non_crossing_grasp_pairs( + left_poses: torch.Tensor, + right_poses: torch.Tensor, + *, + left_costs: torch.Tensor, + right_costs: torch.Tensor, + left_rotation_costs: torch.Tensor, + right_rotation_costs: torch.Tensor, + left_base: torch.Tensor, + right_base: torch.Tensor, + left_to_right_direction: torch.Tensor, + minimum_separation: float, + minimum_lateral_gap: float, +) -> _RankedGraspPairs: + """Rank only distinct, ordered grasp pairs with non-crossing XY routes.""" + left_poses = torch.as_tensor(left_poses, dtype=torch.float32) + right_poses = torch.as_tensor(right_poses, dtype=torch.float32) + direction = torch.as_tensor( + left_to_right_direction, + dtype=torch.float32, + device=left_poses.device, + ) + direction = direction / torch.linalg.vector_norm(direction).clamp_min(1.0e-8) + if minimum_separation < 0.0 or minimum_lateral_gap < 0.0: + raise ValueError("Pair separation constraints must be non-negative.") + rejection_counts = {"reversed": 0, "too_close": 0, "path_crossing": 0} + ranked: list[tuple[float, int, int]] = [] + left_base_position = torch.as_tensor(left_base, dtype=torch.float32)[:3, 3] + right_base_position = torch.as_tensor(right_base, dtype=torch.float32)[:3, 3] + for left_index, left_pose in enumerate(left_poses): + left_position = left_pose[:3, 3] + left_projection = float(torch.dot(left_position, direction)) + for right_index, right_pose in enumerate(right_poses): + right_position = right_pose[:3, 3] + right_projection = float(torch.dot(right_position, direction)) + reversed_pair = ( + left_projection + float(minimum_lateral_gap) > right_projection + ) + too_close = bool( + torch.linalg.vector_norm(left_position - right_position) + < float(minimum_separation) + ) + crossing = _segments_intersect_2d( + left_base_position, + left_position, + right_base_position, + right_position, + ) + rejection_counts["reversed"] += int(reversed_pair) + rejection_counts["too_close"] += int(too_close) + rejection_counts["path_crossing"] += int(crossing) + if reversed_pair or too_close or crossing: + continue + route_length = torch.linalg.vector_norm( + left_position - left_base_position + ) + torch.linalg.vector_norm(right_position - right_base_position) + score = ( + float(left_costs[left_index]) + + float(right_costs[right_index]) + + float(left_rotation_costs[left_index]) / math.pi + + float(right_rotation_costs[right_index]) / math.pi + + 0.05 * float(route_length) + ) + ranked.append((score, left_index, right_index)) + ranked.sort(key=lambda item: (item[0], item[1], item[2])) + return _RankedGraspPairs( + ranked_pairs=tuple((left, right) for _, left, right in ranked), + scores=tuple(score for score, _, _ in ranked), + rejection_counts=rejection_counts, + ) + + +def _point_segment_distance( + point: torch.Tensor, + start: torch.Tensor, + end: torch.Tensor, +) -> torch.Tensor: + segment = end - start + denominator = torch.sum(segment * segment, dim=-1).clamp_min(1.0e-12) + fraction = torch.sum((point - start) * segment, dim=-1) / denominator + closest = start + torch.clamp(fraction, 0.0, 1.0)[..., None] * segment + return torch.linalg.vector_norm(point - closest, dim=-1) + + +def _segment_distance( + first_start: torch.Tensor, + first_end: torch.Tensor, + second_start: torch.Tensor, + second_end: torch.Tensor, +) -> torch.Tensor: + first = first_end - first_start + second = second_end - second_start + offset = first_start - second_start + a = torch.sum(first * first, dim=-1).clamp_min(1.0e-12) + b = torch.sum(first * second, dim=-1) + c = torch.sum(second * second, dim=-1).clamp_min(1.0e-12) + d = torch.sum(first * offset, dim=-1) + e = torch.sum(second * offset, dim=-1) + denominator = a * c - b * b + first_fraction = (b * e - c * d) / denominator.clamp_min(1.0e-12) + second_fraction = (a * e - b * d) / denominator.clamp_min(1.0e-12) + interior = ( + (denominator > 1.0e-12) + & (first_fraction >= 0.0) + & (first_fraction <= 1.0) + & (second_fraction >= 0.0) + & (second_fraction <= 1.0) + ) + first_closest = first_start + first_fraction[..., None] * first + second_closest = second_start + second_fraction[..., None] * second + interior_distance = torch.linalg.vector_norm( + first_closest - second_closest, + dim=-1, + ) + endpoint_distance = ( + torch.stack( + ( + _point_segment_distance(first_start, second_start, second_end), + _point_segment_distance(first_end, second_start, second_end), + _point_segment_distance(second_start, first_start, first_end), + _point_segment_distance(second_end, first_start, first_end), + ), + dim=-1, + ) + .min(dim=-1) + .values + ) + return torch.where(interior, interior_distance, endpoint_distance) + + +def _minimum_interarm_capsule_clearance( + left_link_points: torch.Tensor, + right_link_points: torch.Tensor, + *, + capsule_radius: float, +) -> torch.Tensor: + """Return minimum surface clearance over every inter-arm link pair.""" + left = torch.as_tensor(left_link_points, dtype=torch.float32) + right = torch.as_tensor(right_link_points, dtype=torch.float32) + if left.ndim != 4 or right.ndim != 4 or left.shape[:2] != right.shape[:2]: + raise ValueError("Link points must have matching shape prefixes (B, T, L, 3).") + if left.shape[-1] != 3 or right.shape[-1] != 3: + raise ValueError("Link points must end in xyz coordinates.") + if left.shape[2] < 2 or right.shape[2] < 2: + raise ValueError("Each arm must provide at least two link points.") + if capsule_radius < 0.0: + raise ValueError("capsule_radius must be non-negative.") + minimum = torch.full( + left.shape[:2], + torch.inf, + dtype=left.dtype, + device=left.device, + ) + for left_index in range(left.shape[2] - 1): + for right_index in range(right.shape[2] - 1): + distance = _segment_distance( + left[:, :, left_index], + left[:, :, left_index + 1], + right[:, :, right_index], + right[:, :, right_index + 1], + ) + minimum = torch.minimum(minimum, distance) + return minimum - 2.0 * float(capsule_radius) + + +def _trajectory_safety_report( + *, + left_qpos: torch.Tensor, + right_qpos: torch.Tensor, + left_eef: torch.Tensor, + right_eef: torch.Tensor, + desired_left_eef: torch.Tensor, + desired_right_eef: torch.Tensor, + left_link_points: torch.Tensor, + right_link_points: torch.Tensor, + left_to_right_direction: torch.Tensor, + maximum_joint_step: float, + maximum_orientation_error: float, + minimum_lateral_gap: float, + capsule_radius: float, + minimum_capsule_clearance: float, + orientation_start_index: int = 0, +) -> _TrajectorySafetyReport: + """Evaluate all hard E5 post-plan continuity and inter-arm constraints.""" + if left_qpos.ndim != 3 or right_qpos.ndim != 3: + raise ValueError("Arm qpos trajectories must have shape (B, T, DOF).") + if left_qpos.shape[:2] != right_qpos.shape[:2] or left_qpos.shape[1] < 2: + raise ValueError("Arm qpos trajectories must share at least two waypoints.") + waypoint_count = left_qpos.shape[1] + if not 0 <= orientation_start_index < waypoint_count: + raise ValueError("orientation_start_index must select a trajectory waypoint.") + left_step = torch.amax(torch.abs(torch.diff(left_qpos, dim=1)), dim=(1, 2)) + right_step = torch.amax(torch.abs(torch.diff(right_qpos, dim=1)), dim=(1, 2)) + joint_step = torch.maximum(left_step, right_step) + left_orientation = torch.amax( + _rotation_distance( + left_eef[:, orientation_start_index:, :3, :3], + desired_left_eef[:, orientation_start_index:, :3, :3], + ), + dim=1, + ) + right_orientation = torch.amax( + _rotation_distance( + right_eef[:, orientation_start_index:, :3, :3], + desired_right_eef[:, orientation_start_index:, :3, :3], + ), + dim=1, + ) + orientation = torch.maximum(left_orientation, right_orientation) + direction = torch.as_tensor( + left_to_right_direction, + dtype=left_eef.dtype, + device=left_eef.device, + ) + direction = direction / torch.linalg.vector_norm(direction).clamp_min(1.0e-8) + lateral_gap = ( + torch.sum( + (right_eef[:, :, :3, 3] - left_eef[:, :, :3, 3]) * direction, + dim=2, + ) + .min(dim=1) + .values + ) + capsule_clearance = ( + _minimum_interarm_capsule_clearance( + left_link_points, + right_link_points, + capsule_radius=capsule_radius, + ) + .min(dim=1) + .values + ) + failures = { + "joint_step": joint_step > float(maximum_joint_step), + "orientation": orientation > float(maximum_orientation_error), + "lateral_order": lateral_gap < float(minimum_lateral_gap), + "capsule_collision": capsule_clearance < float(minimum_capsule_clearance), + } + failed = torch.zeros_like(joint_step, dtype=torch.bool) + for value in failures.values(): + failed |= value + return _TrajectorySafetyReport( + success=~failed, + failed_checks={ + name: value.detach().cpu().tolist() for name, value in failures.items() + }, + metrics={ + "maximum_joint_step": joint_step.detach().cpu().tolist(), + "maximum_orientation_error": orientation.detach().cpu().tolist(), + "minimum_lateral_gap": lateral_gap.detach().cpu().tolist(), + "minimum_capsule_clearance": capsule_clearance.detach().cpu().tolist(), + }, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/frames.py b/embodichain/gen_sim/action_engine/runtime/frames.py new file mode 100644 index 000000000..513308484 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/frames.py @@ -0,0 +1,165 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve directional relations in live robot and world frames.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from .robot_parts import arm_control_part + +__all__ = [ + "DIRECTIONAL_RELATIONS", + "arm_base_poses", + "relation_axes", + "relation_offset", + "robot_frame_axes", +] + + +_RELATION_COMPONENTS = { + "left": ("left",), + "left_of": ("left",), + "right": ("right",), + "right_of": ("right",), + "front": ("front",), + "front_of": ("front",), + "in_front_of": ("front",), + "behind": ("back",), + "back": ("back",), + "front_left": ("front", "left"), + "front_left_of": ("front", "left"), + "front_right": ("front", "right"), + "front_right_of": ("front", "right"), + "back_left": ("back", "left"), + "back_left_of": ("back", "left"), + "back_right": ("back", "right"), + "back_right_of": ("back", "right"), +} +DIRECTIONAL_RELATIONS = frozenset(_RELATION_COMPONENTS) + + +def arm_base_poses(env: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Return live world poses of the left and right arm bases.""" + left_part = arm_control_part(env, "left_arm") + right_part = arm_control_part(env, "right_arm") + robot = env.robot + if hasattr(robot, "get_solver") and hasattr(robot, "get_link_pose"): + left_solver = robot.get_solver(name=left_part) + right_solver = robot.get_solver(name=right_part) + left_root = getattr(left_solver, "root_link_name", None) + right_root = getattr(right_solver, "root_link_name", None) + if left_root is None or right_root is None: + raise ValueError("Directional grounding requires both arm root links.") + left = robot.get_link_pose(link_name=left_root, to_matrix=True) + right = robot.get_link_pose(link_name=right_root, to_matrix=True) + elif hasattr(robot, "get_control_part_base_pose"): + left = robot.get_control_part_base_pose(name=left_part, to_matrix=True) + right = robot.get_control_part_base_pose(name=right_part, to_matrix=True) + elif hasattr(env, "get_current_xpos_agent"): + left, right = env.get_current_xpos_agent() + else: + raise ValueError( + "Directional grounding requires live left/right arm-base or TCP poses." + ) + + left = _batched_pose(left, env) + right = _batched_pose(right, env) + return left, right + + +def robot_frame_axes(env: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Return normalized world-space forward and left axes for a dual-arm robot.""" + left, right = arm_base_poses(env) + lateral = left[:, :2, 3] - right[:, :2, 3] + norm = torch.linalg.vector_norm(lateral, dim=1, keepdim=True) + if bool((norm <= 1.0e-6).any()): + raise ValueError("Left and right arm bases must have distinct XY positions.") + lateral = lateral / norm + forward = torch.stack((lateral[:, 1], -lateral[:, 0]), dim=1) + return forward, lateral + + +def relation_axes( + env: Any, + relation: str, + *, + frame: str, +) -> tuple[torch.Tensor, ...]: + """Return signed world-space axes whose projections define a relation.""" + relation = str(relation) + if relation not in DIRECTIONAL_RELATIONS: + return () + if frame == "robot": + forward, lateral = robot_frame_axes(env) + elif frame == "world": + count = int(env.num_envs) + forward = torch.tensor( + [1.0, 0.0], dtype=torch.float32, device=env.device + ).repeat(count, 1) + lateral = torch.tensor( + [0.0, 1.0], dtype=torch.float32, device=env.device + ).repeat(count, 1) + else: + raise ValueError(f"Unsupported directional relation frame {frame!r}.") + + component_axes = { + "front": forward, + "back": -forward, + "left": lateral, + "right": -lateral, + } + return tuple(component_axes[item] for item in _RELATION_COMPONENTS[relation]) + + +def relation_offset( + env: Any, + relation: str, + *, + frame: str, + forward_distance: float, + lateral_distance: float, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor | None: + """Resolve one directional relation into a batched world-space offset.""" + axes = relation_axes(env, relation, frame=frame) + if not axes: + return None + offset = torch.zeros((int(env.num_envs), 3), dtype=dtype, device=device) + components = _RELATION_COMPONENTS[relation] + for component, axis in zip(components, axes): + axis = axis.to(dtype=dtype, device=device) + distance = ( + forward_distance if component in {"front", "back"} else lateral_distance + ) + offset[:, :2] += axis * float(distance) + return offset + + +def _batched_pose(value: Any, env: Any) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + if pose.shape != (int(env.num_envs), 4, 4): + raise ValueError( + "Frame pose must have shape (4, 4) or " + f"({int(env.num_envs)}, 4, 4), got {tuple(pose.shape)}." + ) + return pose diff --git a/embodichain/gen_sim/action_engine/runtime/geometry_axes.py b/embodichain/gen_sim/action_engine/runtime/geometry_axes.py new file mode 100644 index 000000000..33564ab1e --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/geometry_axes.py @@ -0,0 +1,102 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Object-local axis analysis shared by GenSim grounding and grasp lowering.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +__all__ = ["LocalGeometryAxes", "analyze_local_geometry_axes"] + + +@dataclass(frozen=True, slots=True) +class LocalGeometryAxes: + """Validated local AABB axes with a PCA alignment cross-check.""" + + bounds_center: torch.Tensor + extents: torch.Tensor + ordered_axis_indices: tuple[int, int, int] + long_axis_index: int + short_axis_index: int + long_axis: torch.Tensor + long_half_extent: float + elongation_ratio: float + principal_alignment: float + + +def analyze_local_geometry_axes( + vertices: torch.Tensor, + *, + minimum_elongation_ratio: float = 1.10, + minimum_principal_alignment: float = 0.90, +) -> LocalGeometryAxes: + """Resolve stable local long/short axes or fail on ambiguous geometry.""" + if ( + not isinstance(vertices, torch.Tensor) + or not vertices.is_floating_point() + or vertices.ndim != 2 + or vertices.shape[1] != 3 + or vertices.shape[0] < 3 + or not torch.isfinite(vertices).all() + ): + raise ValueError("vertices must be a finite floating tensor shaped (N, 3).") + if minimum_elongation_ratio <= 1.0: + raise ValueError("minimum_elongation_ratio must be greater than one.") + if not 0.0 < minimum_principal_alignment <= 1.0: + raise ValueError("minimum_principal_alignment must be in (0, 1].") + + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + extents = upper - lower + if torch.any(extents <= 1.0e-8): + raise ValueError("Object geometry must have non-zero extent on every axis.") + ordered = torch.argsort(extents, descending=True) + long_index = int(ordered[0].item()) + short_index = int(ordered[-1].item()) + elongation = float((extents[long_index] / extents[int(ordered[1])]).item()) + if elongation < minimum_elongation_ratio: + raise ValueError( + "Object long axis is ambiguous; provide an explicit upright_local_axis." + ) + + centered = vertices - vertices.mean(dim=0, keepdim=True) + covariance = centered.transpose(0, 1) @ centered + _, eigenvectors = torch.linalg.eigh(covariance) + principal = eigenvectors[:, -1] + principal_index = int(torch.argmax(torch.abs(principal)).item()) + alignment = float(torch.abs(principal[principal_index]).item()) + if principal_index != long_index or alignment < minimum_principal_alignment: + raise ValueError( + "Object principal axis is not aligned with its local AABB; provide an " + "explicit local axis instead of inferring long_axis." + ) + + long_axis = torch.zeros(3, dtype=vertices.dtype, device=vertices.device) + long_axis[long_index] = 1.0 + return LocalGeometryAxes( + bounds_center=((lower + upper) * 0.5).clone(), + extents=extents.clone(), + ordered_axis_indices=tuple(int(index) for index in ordered.tolist()), + long_axis_index=long_index, + short_axis_index=short_index, + long_axis=long_axis, + long_half_extent=float((extents[long_index] * 0.5).item()), + elongation_ratio=elongation, + principal_alignment=alignment, + ) diff --git a/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py new file mode 100644 index 000000000..c17dba7f7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grasp_diagnostics.py @@ -0,0 +1,535 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Stage-separated diagnostics for GenSim antipodal grasp generation.""" + +from __future__ import annotations + +from contextlib import contextmanager +from copy import deepcopy +from dataclasses import dataclass +from typing import Any, Iterator + +import torch +import torch.nn.functional as F + +from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator + +from .coordinated_safety import ( + _canonicalize_parallel_jaw_poses, + _rank_non_crossing_grasp_pairs, +) + +__all__: list[str] = [] + +_UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT = 0.65 +_UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION = 0.35 +_UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION = 0.75 +_UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT = 2.0 +_UPRIGHT_SIDE_GRASP_CANDIDATE_LIMIT = 50 + + +@dataclass(frozen=True, slots=True) +class _UprightGraspSelectionContext: + local_axis: torch.Tensor + + +@dataclass(frozen=True, slots=True) +class _DualGraspSelectionContext: + left_eef: torch.Tensor + right_eef: torch.Tensor + left_base: torch.Tensor + right_base: torch.Tensor + left_to_right_direction: torch.Tensor + pair_rank: int + minimum_separation: float + minimum_lateral_gap: float + + +class _TracingAntipodalGraspPoseGenerator(AntipodalGraspPoseGenerator): + """Retain compact S1-S5 evidence from the concrete GenSim grasp backend.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._last_dual_trace: dict[str, Any] | None = None + self._last_upright_trace: dict[str, Any] | None = None + self._selection_context: _DualGraspSelectionContext | None = None + self._upright_selection_context: _UprightGraspSelectionContext | None = None + + @property + def last_dual_trace(self) -> dict[str, Any] | None: + """Return an owned snapshot of the most recent dual-grasp trace.""" + return deepcopy(self._last_dual_trace) + + @property + def last_upright_trace(self) -> dict[str, Any] | None: + """Return an owned snapshot of the most recent upright-grasp trace.""" + return deepcopy(self._last_upright_trace) + + @contextmanager + def upright_selection_context( + self, + *, + local_axis: torch.Tensor, + ) -> Iterator[None]: + """Install one invocation-local side-grasp selection policy.""" + if self._upright_selection_context is not None: + raise RuntimeError("Upright grasp selection context cannot be nested.") + axis = torch.as_tensor(local_axis, dtype=torch.float32).reshape(-1) + norm = torch.linalg.vector_norm(axis) + if axis.shape != (3,) or not torch.isfinite(axis).all() or norm <= 1.0e-6: + raise ValueError("Upright grasp local_axis must be one finite 3-vector.") + self._last_upright_trace = None + self._upright_selection_context = _UprightGraspSelectionContext( + local_axis=(axis / norm).clone() + ) + try: + yield + finally: + self._upright_selection_context = None + + @contextmanager + def dual_arm_selection_context( + self, + *, + left_eef: torch.Tensor, + right_eef: torch.Tensor, + left_base: torch.Tensor, + right_base: torch.Tensor, + left_to_right_direction: torch.Tensor, + pair_rank: int, + minimum_separation: float, + minimum_lateral_gap: float, + ) -> Iterator[None]: + """Install one invocation-local arm context for pair-aware selection.""" + if self._selection_context is not None: + raise RuntimeError("Dual grasp selection context cannot be nested.") + if type(pair_rank) is not int or pair_rank < 0: + raise ValueError("pair_rank must be a non-negative integer.") + self._selection_context = _DualGraspSelectionContext( + left_eef=torch.as_tensor(left_eef, dtype=torch.float32).clone(), + right_eef=torch.as_tensor(right_eef, dtype=torch.float32).clone(), + left_base=torch.as_tensor(left_base, dtype=torch.float32).clone(), + right_base=torch.as_tensor(right_base, dtype=torch.float32).clone(), + left_to_right_direction=torch.as_tensor( + left_to_right_direction, dtype=torch.float32 + ).clone(), + pair_rank=pair_rank, + minimum_separation=float(minimum_separation), + minimum_lateral_gap=float(minimum_lateral_gap), + ) + try: + yield + finally: + self._selection_context = None + + @staticmethod + def _failed_arm_result(reference: torch.Tensor) -> dict[str, Any]: + return { + "is_success": False, + "grasp_poses": torch.eye( + 4, + dtype=torch.float32, + device=reference.device, + ), + "open_lengths": 0.0, + "total_cost": torch.zeros(1, device=reference.device), + } + + def _select_pair( + self, + result: dict[str, dict[str, Any]] | None, + *, + row_index: int, + ) -> tuple[dict[str, dict[str, Any]] | None, dict[str, Any] | None]: + context = self._selection_context + if context is None or result is None: + return result, None + left = result["left"] + right = result["right"] + if not left.get("is_success", False) or not right.get("is_success", False): + return result, { + "requested_pair_rank": context.pair_rank, + "valid_pair_count": 0, + "selected": False, + "reason": "one_or_both_arms_have_no_candidates", + } + left_poses = torch.as_tensor(left["grasp_poses"], dtype=torch.float32) + right_poses = torch.as_tensor(right["grasp_poses"], dtype=torch.float32) + if left_poses.ndim == 2: + left_poses = left_poses.unsqueeze(0) + if right_poses.ndim == 2: + right_poses = right_poses.unsqueeze(0) + left_canonical = _canonicalize_parallel_jaw_poses( + left_poses, + context.left_eef[row_index], + ) + right_canonical = _canonicalize_parallel_jaw_poses( + right_poses, + context.right_eef[row_index], + ) + ranking = _rank_non_crossing_grasp_pairs( + left_canonical.poses, + right_canonical.poses, + left_costs=torch.as_tensor(left["total_cost"], dtype=torch.float32), + right_costs=torch.as_tensor(right["total_cost"], dtype=torch.float32), + left_rotation_costs=left_canonical.selected_rotation_radians, + right_rotation_costs=right_canonical.selected_rotation_radians, + left_base=context.left_base[row_index], + right_base=context.right_base[row_index], + left_to_right_direction=context.left_to_right_direction, + minimum_separation=context.minimum_separation, + minimum_lateral_gap=context.minimum_lateral_gap, + ) + trace: dict[str, Any] = { + "requested_pair_rank": context.pair_rank, + "valid_pair_count": len(ranking.ranked_pairs), + "rejection_counts": dict(ranking.rejection_counts), + "left_half_turn_count": int(left_canonical.flipped.sum().item()), + "right_half_turn_count": int(right_canonical.flipped.sum().item()), + "selected": context.pair_rank < len(ranking.ranked_pairs), + } + if context.pair_rank >= len(ranking.ranked_pairs): + trace["reason"] = "requested_pair_rank_unavailable" + return { + "left": self._failed_arm_result(left_poses), + "right": self._failed_arm_result(right_poses), + }, trace + left_index, right_index = ranking.ranked_pairs[context.pair_rank] + left_pose = left_canonical.poses[left_index] + right_pose = right_canonical.poses[right_index] + trace.update( + { + "selected_left_index": left_index, + "selected_right_index": right_index, + "selected_pair_score": ranking.scores[context.pair_rank], + "selected_left_half_turn": bool(left_canonical.flipped[left_index]), + "selected_right_half_turn": bool(right_canonical.flipped[right_index]), + "selected_left_rotation_radians": float( + left_canonical.selected_rotation_radians[left_index] + ), + "selected_right_rotation_radians": float( + right_canonical.selected_rotation_radians[right_index] + ), + "selected_left_pose": left_pose.detach().cpu().tolist(), + "selected_right_pose": right_pose.detach().cpu().tolist(), + "selected_separation": float( + torch.linalg.vector_norm(left_pose[:3, 3] - right_pose[:3, 3]) + ), + } + ) + + def selected_arm( + arm: dict[str, Any], + poses: torch.Tensor, + index: int, + ) -> dict[str, Any]: + open_lengths = torch.as_tensor(arm["open_lengths"]) + costs = torch.as_tensor(arm["total_cost"], dtype=torch.float32) + return { + "is_success": True, + "grasp_poses": poses[index : index + 1], + "open_lengths": open_lengths[index : index + 1], + "total_cost": costs.new_zeros(1), + } + + return { + "left": selected_arm(left, left_canonical.poses, left_index), + "right": selected_arm(right, right_canonical.poses, right_index), + }, trace + + @staticmethod + def _transform_points(points: torch.Tensor, pose: torch.Tensor) -> torch.Tensor: + return points @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + @classmethod + def _filter_counts( + cls, + backend: Any, + *, + mesh_vertices: torch.Tensor, + object_pose: torch.Tensor, + arm_direction: torch.Tensor, + approach_direction: torch.Tensor, + middle_empty_ratio: float, + ) -> dict[str, int]: + pairs = backend.antipodal_pairs.to(dtype=torch.float32) + origin = cls._transform_points(pairs[:, 0], object_pose) + hit = cls._transform_points(pairs[:, 1], object_pose) + world_vertices = cls._transform_points(mesh_vertices, object_pose) + projection = torch.matmul(world_vertices, arm_direction) + extent = projection.max() - projection.min() + left_threshold = projection.min() + extent * (0.5 - middle_empty_ratio * 0.5) + right_threshold = projection.max() - extent * (0.5 - middle_empty_ratio * 0.5) + origin_projection = torch.matmul(origin, arm_direction) + hit_projection = torch.matmul(hit, arm_direction) + masks = { + "left": (origin_projection < left_threshold) + | (hit_projection < left_threshold), + "right": (origin_projection > right_threshold) + | (hit_projection > right_threshold), + } + counts: dict[str, int] = {} + for side, mask in masks.items(): + grasp_x = F.normalize(hit[mask] - origin[mask], dim=1) + cosine = torch.clamp( + torch.sum(grasp_x * approach_direction, dim=1), -1.0, 1.0 + ) + angle = torch.abs(torch.acos(cosine)) + angle_valid = torch.abs(angle - torch.pi * 0.5) <= float( + backend._max_deviation_angle + ) + counts[f"{side}_partition_pair_count"] = int(mask.sum().item()) + counts[f"{side}_angle_valid_pair_count"] = int(angle_valid.sum().item()) + return counts + + @staticmethod + def _candidate_count(result: dict[str, Any]) -> int: + poses = result.get("grasp_poses") + if not result.get("is_success", False) or not isinstance(poses, torch.Tensor): + return 0 + return int(poses.shape[0]) if poses.ndim == 3 else 0 + + def get_valid_grasp_poses( + self, + **kwargs: Any, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Apply v14's side and mid-body preference for upright pickup.""" + results = super().get_valid_grasp_poses(**kwargs) + context = self._upright_selection_context + if context is None: + return results + + vertices = torch.as_tensor(kwargs["mesh_vertices"], dtype=torch.float32) + object_poses = torch.as_tensor(kwargs["obj_poses"], dtype=torch.float32) + local_axis = context.local_axis.to( + device=vertices.device, + dtype=vertices.dtype, + ) + vertex_positions = torch.matmul(vertices, local_axis) + axis_min = vertex_positions.min() + axis_extent = vertex_positions.max() - axis_min + if float(axis_extent) <= 1.0e-6: + raise ValueError("Upright grasp axis must span non-zero object geometry.") + + ranked_results: list[tuple[torch.Tensor, torch.Tensor]] = [] + row_traces: list[dict[str, Any]] = [] + for row_index, (result, object_pose) in enumerate( + zip(results, object_poses, strict=True) + ): + grasp_poses, costs = result + grasp_poses = torch.as_tensor(grasp_poses, dtype=torch.float32) + costs = torch.as_tensor( + costs, + device=grasp_poses.device, + dtype=torch.float32, + ) + if grasp_poses.ndim == 2: + grasp_poses = grasp_poses.unsqueeze(0) + object_pose = object_pose.to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + axis = local_axis.to( + device=grasp_poses.device, + dtype=grasp_poses.dtype, + ) + world_upright = torch.matmul(object_pose[:3, :3], axis) + closing_axes = F.normalize(grasp_poses[:, :3, 0], dim=1) + axis_alignment = torch.abs( + torch.sum(closing_axes * world_upright[None], dim=1) + ) + side_compatible = axis_alignment <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_ALIGNMENT + relative_centers = grasp_poses[:, :3, 3] - object_pose[None, :3, 3] + center_axis_positions = torch.sum( + relative_centers * world_upright[None], + dim=1, + ) + center_fractions = ( + center_axis_positions - axis_min.to(grasp_poses.device) + ) / axis_extent.to(grasp_poses.device) + central_band = ( + center_fractions >= _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION + ) & (center_fractions <= _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION) + interval = ( + _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION + - _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION + ) + height_penalty = ( + torch.clamp( + _UPRIGHT_SIDE_GRASP_MIN_AXIS_FRACTION - center_fractions, + min=0.0, + ) + + torch.clamp( + center_fractions - _UPRIGHT_SIDE_GRASP_MAX_AXIS_FRACTION, + min=0.0, + ) + ) / interval + adjusted_costs = ( + torch.where( + side_compatible, + costs, + torch.full_like(costs, torch.inf), + ) + + _UPRIGHT_SIDE_GRASP_HEIGHT_COST_WEIGHT * height_penalty + ) + ranked = torch.argsort(adjusted_costs)[:_UPRIGHT_SIDE_GRASP_CANDIDATE_LIMIT] + ranked_results.append((grasp_poses[ranked], adjusted_costs[ranked])) + finite_ranked = torch.isfinite(adjusted_costs[ranked]) + best_index = int(ranked[0].item()) if bool(finite_ranked.any()) else None + row_traces.append( + { + "environment_index": row_index, + "local_axis": axis.detach().cpu().tolist(), + "candidate_count": int(grasp_poses.shape[0]), + "side_compatible_count": int(side_compatible.sum().item()), + "central_band_count": int(central_band.sum().item()), + "side_and_central_count": int( + (side_compatible & central_band).sum().item() + ), + "retained_count": int(finite_ranked.sum().item()), + "best_candidate_axis_alignment": ( + None + if best_index is None + else float(axis_alignment[best_index].item()) + ), + "best_candidate_axis_fraction": ( + None + if best_index is None + else float(center_fractions[best_index].item()) + ), + } + ) + self._last_upright_trace = ( + row_traces[0] if len(row_traces) == 1 else {"environment_rows": row_traces} + ) + return ranked_results + + def get_dual_arm_valid_grasp_poses(self, **kwargs: Any) -> list[dict | None]: + """Run the standard generator while observing its NMS/collision boundary.""" + vertices = kwargs["mesh_vertices"] + triangles = kwargs["mesh_triangles"] + backend = self._backend(vertices, triangles) + poses = self._object_poses(kwargs["obj_poses"], device=backend.device) + directions = self._approach_directions( + kwargs["approach_direction"], + batch_size=poses.shape[0], + device=backend.device, + ) + arm_direction = self._approach_directions( + kwargs["left_to_right_arm_direction"], + batch_size=1, + device=backend.device, + )[0] + ratio = float(kwargs.get("middle_empty_ratio", 0.4)) + collision_records: list[dict[str, Any]] = [] + checker = backend._collision_checker + original_query = checker.query + + def traced_query(*args: Any, **query_kwargs: Any): + colliding, distance = original_query(*args, **query_kwargs) + distance = torch.as_tensor(distance, dtype=torch.float32) + collision_records.append( + { + "nms_candidate_count": int(colliding.numel()), + "noncolliding_candidate_count": int((~colliding).sum().item()), + "minimum_signed_distance": float(distance.min().item()), + "maximum_signed_distance": float(distance.max().item()), + } + ) + return colliding, distance + + checker.query = traced_query + try: + results = super().get_dual_arm_valid_grasp_poses(**kwargs) + finally: + checker.query = original_query + + row_traces: list[dict[str, Any]] = [] + selected_results: list[dict[str, dict[str, Any]] | None] = [] + for row_index, (object_pose, approach, result) in enumerate( + zip(poses, directions, results, strict=True) + ): + filter_counts = self._filter_counts( + backend, + mesh_vertices=vertices.to(device=backend.device, dtype=torch.float32), + object_pose=object_pose, + arm_direction=arm_direction, + approach_direction=approach, + middle_empty_ratio=ratio, + ) + record_offset = 2 * row_index + records = collision_records[record_offset : record_offset + 2] + left_record = records[0] if len(records) >= 1 else {} + right_record = records[1] if len(records) >= 2 else {} + left = {} if result is None else result["left"] + right = {} if result is None else result["right"] + left_final = self._candidate_count(left) + right_final = self._candidate_count(right) + selected_result, pair_trace = self._select_pair( + result, + row_index=row_index, + ) + selected_results.append(selected_result) + row_traces.append( + { + "environment_index": row_index, + "approach_direction": approach.detach().cpu().tolist(), + "middle_empty_ratio": ratio, + "S1_grasp_pair_generation": { + "antipodal_pair_count": int(backend.antipodal_pairs.shape[0]), + }, + "S2_approach_angle_filtering": filter_counts, + "S3_nms": { + "left_candidate_count": int( + left_record.get("nms_candidate_count", 0) + ), + "right_candidate_count": int( + right_record.get("nms_candidate_count", 0) + ), + }, + "S4_collision_filtering": { + "left_candidate_count": int( + left_record.get("noncolliding_candidate_count", 0) + ), + "right_candidate_count": int( + right_record.get("noncolliding_candidate_count", 0) + ), + "left_minimum_signed_distance": left_record.get( + "minimum_signed_distance" + ), + "left_maximum_signed_distance": left_record.get( + "maximum_signed_distance" + ), + "right_minimum_signed_distance": right_record.get( + "minimum_signed_distance" + ), + "right_maximum_signed_distance": right_record.get( + "maximum_signed_distance" + ), + }, + "S5_left_right_pairing": { + "left_final_count": left_final, + "right_final_count": right_final, + "paired": left_final > 0 and right_final > 0, + }, + "pair_selection": pair_trace, + } + ) + self._last_dual_trace = ( + row_traces[0] if len(row_traces) == 1 else {"environment_rows": row_traces} + ) + return selected_results diff --git a/embodichain/gen_sim/action_engine/runtime/grounding.py b/embodichain/gen_sim/action_engine/runtime/grounding.py new file mode 100644 index 000000000..36a24518c --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/grounding.py @@ -0,0 +1,3157 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve symbolic bindings from live simulator state.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass, replace +import math +from typing import Any + +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.config import ( + RuntimePolicyCfg, + default_runtime_policy, +) +from embodichain.gen_sim.action_engine.domain import normalize_placement_relation +from embodichain.gen_sim.action_engine.orientation import ( + AlignAxisConstraint, + MatchRotationConstraint, + compile_orientation_constraint, +) +from embodichain.lab.sim.atomic_actions import ( + AntipodalAffordance, + AxisAlignAffordance, + AxisAlignGoal, + CoordinatedPickGoal, + CoordinatedPlacementGoal, + EndEffectorPoseGoal, + GraspGoal, + HeldObjectPoseGoal, + JointPositionGoal, + ObjectSemantics, + PlaceGoal, + PourGoal, + PressAffordance, + PressGoal, + SlideAffordance, + SlideGoal, + TwistAffordance, + TwistGoal, +) +from embodichain.utils.logger import log_info + +from .frames import arm_base_poses, relation_offset, robot_frame_axes +from .geometry_axes import analyze_local_geometry_axes +from .models import ExecutionProgram, GroundedAction, SemanticStep +from .motion_policy import resolve_motion_policy, with_motion_modifiers +from .robot_parts import arm_control_part +from .state import ExecutionState + +__all__ = ["ActionGrounder", "LiveArrangementPlan", "LivePlacementPlan"] + +_E2_CLEARANCE_RETREAT_DISTANCE = 0.20 +_E2_REORIENT_MINIMUM_CLEARANCE = 0.15 +DEFAULT_INTERNAL_AXIS = (0.0, 0.0, 1.0) +DEFAULT_TARGET_AXIS = (0.0, 0.0, 1.0) + + +def _batched_pose(value: Any, env: Any) -> torch.Tensor: + pose = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + if pose.shape != (int(env.num_envs), 4, 4): + raise ValueError( + "Live pose must have shape (4, 4) or " + f"({int(env.num_envs)}, 4, 4), got {tuple(pose.shape)}." + ) + return pose + + +def _object(env: Any, uid: str) -> Any: + entity = env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") + return entity + + +def _live_pose(env: Any, uid: str) -> torch.Tensor: + entity = env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") + return _batched_pose(entity.get_local_pose(to_matrix=True), env) + + +def _placement_relation(step: SemanticStep) -> str: + """Normalize a placement relation from the step's own goal contract.""" + relation = str(step.goal.get("relation", "none")) + if relation in {"none", "handover", "held_above_initial"}: + return relation + if ( + step.postcondition.get("type") == "semantic_goal" + or step.goal.get("terminal_behavior") == "place" + ): + return normalize_placement_relation(relation) + return relation + + +def _local_vertices(entity: Any, env: Any, env_id: int = 0) -> torch.Tensor: + value = entity.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (list, tuple)): + value = value[0] + vertices = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError("Rigid-object mesh vertices must have shape (N, 3).") + return vertices + + +def _world_vertices(entity: Any, env: Any, env_id: int) -> torch.Tensor: + get_vertices = getattr(entity, "get_vertices", None) + if callable(get_vertices): + vertices = _local_vertices(entity, env, env_id) + pose = _batched_pose(entity.get_local_pose(to_matrix=True), env)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + get_link_vertices = getattr(entity, "get_link_vert_face", None) + get_link_pose = getattr(entity, "get_link_pose", None) + link_names = getattr(entity, "link_names", ()) + if not callable(get_link_vertices) or not callable(get_link_pose) or not link_names: + raise ValueError("Scene entity exposes no usable collision geometry.") + world_vertices = [] + for link_name in link_names: + vertices, _ = get_link_vertices(link_name) + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=env.device, + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + continue + pose = _batched_pose( + get_link_pose(link_name, to_matrix=True), + env, + )[env_id] + world_vertices.append(vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3]) + if not world_vertices: + raise ValueError("Scene articulation exposes no usable link geometry.") + return torch.cat(world_vertices, dim=0) + + +@dataclass(frozen=True) +class _Geometry: + radius: torch.Tensor + half_height: torch.Tensor + + +class LiveArrangementPlan: + """Materialize collision-aware line slots independently in every env.""" + + def __init__( + self, + env: Any, + steps: Sequence[SemanticStep], + *, + slot_margin: float | None = None, + minimum_spacing: float | None = None, + clearance: float | None = None, + row_search_step: float | None = None, + row_search_radius: float | None = None, + ) -> None: + if not steps: + raise ValueError("An arrangement plan requires at least one step.") + self.env = env + self.steps = tuple(steps) + self.step_by_id = {step.id: step for step in steps} + self.num_envs = int(env.num_envs) + self.device = env.device + self.slot_count = len(steps) + self.axis = str(steps[0].goal.get("axis", "world_x")) + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + defaults = default_runtime_policy(profile).grounding["arrangement"] + slot_margin = defaults["slot_margin"] if slot_margin is None else slot_margin + minimum_spacing = ( + defaults["minimum_spacing"] if minimum_spacing is None else minimum_spacing + ) + self.clearance = float( + defaults["layout_clearance"] if clearance is None else clearance + ) + self.row_search_step = float( + defaults["row_search_step"] if row_search_step is None else row_search_step + ) + self.row_search_radius = float( + defaults["row_search_radius"] + if row_search_radius is None + else row_search_radius + ) + + table = _object(env, "table") + bounds = [] + for env_id in range(self.num_envs): + vertices = _world_vertices(table, env, env_id) + bounds.append( + torch.stack((vertices.min(dim=0).values, vertices.max(dim=0).values)) + ) + self.table_bounds = torch.stack(bounds) + self.table_center = self.table_bounds.mean(dim=1) + self.table_top = self.table_bounds[:, 1, 2] + if self.axis == "table_long_axis": + mean_extent = ( + self.table_bounds[:, 1, :2] - self.table_bounds[:, 0, :2] + ).mean(dim=0) + self.axis_index = int(torch.argmax(mean_extent).item()) + else: + self.axis_index = 0 if self.axis in {"x", "world_x"} else 1 + self.perpendicular_index = 1 - self.axis_index + self.geometry = {step.id: self._geometry(step) for step in self.steps} + diameters = torch.stack( + [self.geometry[step.id].radius * 2.0 for step in self.steps], + dim=1, + ) + self.spacing = torch.maximum( + diameters.max(dim=1).values + float(slot_margin), + torch.full( + (self.num_envs,), + float(minimum_spacing), + dtype=torch.float32, + device=self.device, + ), + ) + self.positions = self._make_slots() + self.reassignment_reason: list[str | None] = [None] * self.num_envs + self.reassignment_cost = torch.full( + (self.num_envs,), + float("nan"), + dtype=torch.float32, + device=self.device, + ) + self.assignments = self._initial_slot_assignments() + order_by = str(self.steps[0].goal.get("order_by", "explicit")) + direction = str(self.steps[0].goal.get("order_direction", "given")) + if order_by == "size" and not any( + step.goal.get("slot_constraint") == "free_reassignable" + for step in self.steps + ): + for env_id in range(self.num_envs): + ordered = sorted( + self.steps, + key=lambda step: float(self.geometry[step.id].radius[env_id]), + reverse=direction != "ascending", + ) + for slot_id, step in enumerate(ordered): + self.assignments[step.id][env_id] = slot_id + self.completed = { + step.id: torch.zeros( + self.num_envs, + dtype=torch.bool, + device=self.device, + ) + for step in self.steps + } + + def _initial_slot_assignments(self) -> dict[str, torch.Tensor]: + """Match free-order objects to slots in their current spatial order.""" + assignments = { + step.id: torch.full( + (self.num_envs,), + int(step.goal.get("nominal_slot_index", index)), + dtype=torch.long, + device=self.device, + ) + for index, step in enumerate(self.steps) + } + free_steps = [ + step + for step in self.steps + if step.goal.get("slot_constraint") == "free_reassignable" + ] + if not free_steps: + return assignments + required_slots = { + int(step.goal.get("nominal_slot_index", index)) + for index, step in enumerate(self.steps) + if step.goal.get("slot_constraint") != "free_reassignable" + } + available_slots = [ + slot_id + for slot_id in range(self.slot_count) + if slot_id not in required_slots + ] + if len(available_slots) != len(free_steps): + raise ValueError( + "Arrangement slot constraints do not define a one-to-one assignment." + ) + axis_positions = { + step.id: _live_pose(self.env, step.object_uid)[:, self.axis_index, 3] + for step in free_steps + } + for env_id in range(self.num_envs): + ordered_steps = sorted( + free_steps, + key=lambda step: ( + float(axis_positions[step.id][env_id]), + int(step.goal.get("nominal_slot_index", 0)), + step.id, + ), + ) + ordered_slots = sorted( + available_slots, + key=lambda slot_id: ( + float(self.positions[env_id, slot_id, self.axis_index]), + slot_id, + ), + ) + matching_cost = 0.0 + changed = False + for step, slot_id in zip(ordered_steps, ordered_slots): + nominal = int(step.goal.get("nominal_slot_index", 0)) + assignments[step.id][env_id] = slot_id + changed |= slot_id != nominal + matching_cost += abs( + float(axis_positions[step.id][env_id]) + - float(self.positions[env_id, slot_id, self.axis_index]) + ) + if changed: + self.reassignment_reason[env_id] = ( + "free arrangement initialized from live spatial order" + ) + self.reassignment_cost[env_id] = matching_cost + return assignments + + def _geometry(self, step: SemanticStep) -> _Geometry: + entity = _object(self.env, step.object_uid) + radii = [] + heights = [] + for env_id in range(self.num_envs): + vertices = _local_vertices(entity, self.env, env_id) + half_extent = ( + vertices.max(dim=0).values - vertices.min(dim=0).values + ) * 0.5 + if step.goal.get("orientation_goal", "none") in {"none", "preserve"}: + rotation = _live_pose(self.env, step.object_uid)[env_id, :3, :3] + rotated = vertices @ rotation.transpose(0, 1) + radii.append(torch.linalg.vector_norm(rotated[:, :2], dim=-1).max()) + else: + # A non-preserve target may rotate the longest local dimension + # into the table plane, so retain the conservative bound. + radii.append( + torch.linalg.vector_norm(torch.topk(half_extent, k=2).values) + ) + heights.append((vertices[:, 2].max() - vertices[:, 2].min()) * 0.5) + return _Geometry(torch.stack(radii), torch.stack(heights)) + + def _make_slots(self) -> torch.Tensor: + offsets = ( + torch.arange(self.slot_count, device=self.device, dtype=torch.float32) + - (self.slot_count - 1) / 2.0 + ) + slots = torch.empty( + self.num_envs, + self.slot_count, + 3, + dtype=torch.float32, + device=self.device, + ) + radii = torch.stack( + [self.geometry[step.id].radius for step in self.steps], + dim=1, + ) + # Free slot rematching allows any remaining object to occupy any slot. + # Size every slot for the largest member in that environment rather + # than accidentally baking the nominal object order into geometry. + slot_radii = radii.max(dim=1).values[:, None].repeat(1, self.slot_count) + obstacles = self._obstacle_bounds() + search_offsets = [0.0] + steps = int(self.row_search_radius / self.row_search_step) + for index in range(1, steps + 1): + offset = self.row_search_step * index + search_offsets.extend((offset, -offset)) + for env_id in range(self.num_envs): + chosen = None + for perpendicular in search_offsets: + candidate = self.table_center[env_id].repeat(self.slot_count, 1) + candidate[:, self.axis_index] += self.spacing[env_id] * offsets + candidate[:, self.perpendicular_index] += perpendicular + candidate[:, 2] = self.table_top[env_id] + if self._safe( + candidate, + slot_radii[env_id], + self.table_bounds[env_id], + obstacles[env_id], + ): + chosen = candidate + break + if chosen is None: + raise ValueError( + f"Environment {env_id} has no collision-free arrangement row." + ) + slots[env_id] = chosen + return slots + + def _obstacle_bounds( + self, + ) -> list[list[tuple[torch.Tensor, torch.Tensor]]]: + result: list[list[tuple[torch.Tensor, torch.Tensor]]] = [ + [] for _ in range(self.num_envs) + ] + getter = getattr(self.env.sim, "get_rigid_object_uid_list", None) + if not callable(getter): + return result + movable = {step.object_uid for step in self.steps} + for uid in getter(): + if uid == "table" or uid in movable: + continue + entity = self.env.sim.get_rigid_object(uid) + if entity is None: + continue + for env_id in range(self.num_envs): + vertices = _world_vertices(entity, self.env, env_id) + if float(vertices[:, 2].max()) < float( + self.table_top[env_id] - self.clearance + ): + continue + result[env_id].append( + ( + vertices[:, :2].min(dim=0).values, + vertices[:, :2].max(dim=0).values, + ) + ) + return result + + def _safe( + self, + slots: torch.Tensor, + radii: torch.Tensor, + table_bounds: torch.Tensor, + obstacles: Sequence[tuple[torch.Tensor, torch.Tensor]], + ) -> bool: + lower = table_bounds[0, :2] + radii[:, None] + self.clearance + upper = table_bounds[1, :2] - radii[:, None] - self.clearance + if bool(((slots[:, :2] < lower) | (slots[:, :2] > upper)).any()): + return False + for center, radius in zip(slots[:, :2], radii): + for obstacle_lower, obstacle_upper in obstacles: + closest = torch.maximum( + obstacle_lower, + torch.minimum(center, obstacle_upper), + ) + if float(torch.linalg.vector_norm(center - closest)) <= float( + radius + self.clearance + ): + return False + return True + + def target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + phase: str, + policy: Mapping[str, Any], + ) -> torch.Tensor: + """Return a live final or collision-clear staging object pose.""" + if phase not in {"staging", "final"}: + raise ValueError(f"Unsupported arrangement phase {phase!r}.") + target = object_pose.clone() + env_ids = torch.arange(self.num_envs, device=self.device) + slot_ids = self.assignments[step.id] + target[:, :2, 3] = self.positions[env_ids, slot_ids, :2] + final_z = ( + self.table_top + + self.geometry[step.id].half_height + + float(policy["surface_clearance"]) + ) + target[:, 2, 3] = final_z + if phase == "staging": + target[:, 2, 3] = final_z + float(policy["transport_clearance"]) + return target + + def mark_completed(self, step_id: str, success: torch.Tensor) -> None: + self.completed[step_id] |= success.to(self.device, dtype=torch.bool) + + def remaining(self, env_id: int) -> list[str]: + return [ + step.id for step in self.steps if not bool(self.completed[step.id][env_id]) + ] + + def available_slots(self, env_id: int) -> list[int]: + occupied = { + int(self.assignments[step.id][env_id]) + for step in self.steps + if bool(self.completed[step.id][env_id]) + } + return [index for index in range(self.slot_count) if index not in occupied] + + def assign(self, env_id: int, assignment: Mapping[str, int]) -> None: + for step_id, slot_id in assignment.items(): + self.assignments[step_id][env_id] = int(slot_id) + + def metadata(self, step: SemanticStep, env_id: int) -> dict[str, Any]: + """Describe the live slot resolution used by one environment.""" + nominal = int(step.goal.get("nominal_slot_index", 0)) + resolved = int(self.assignments[step.id][env_id]) + return { + "nominal_slot_index": nominal, + "resolved_slot_index": resolved, + "slot_constraint": str(step.goal.get("slot_constraint", "required")), + "slot_reassigned": resolved != nominal, + "reassignment_reason": self.reassignment_reason[env_id], + "matching_cost": ( + float(self.reassignment_cost[env_id]) + if torch.isfinite(self.reassignment_cost[env_id]) + else None + ), + "spacing": float(self.spacing[env_id]), + "resolved_slot_position": self.positions[env_id, resolved].tolist(), + } + + +class LivePlacementPlan: + """Allocate non-overlapping live slots for one shared container.""" + + def __init__( + self, + env: Any, + steps: Sequence[SemanticStep], + *, + clearance: float | None = None, + ) -> None: + if not steps: + raise ValueError("A placement plan requires at least one step.") + references = {step.goal.get("reference_object") for step in steps} + if len(references) != 1 or not isinstance(next(iter(references)), str): + raise ValueError("Placement-plan steps must share one reference object.") + self.env = env + self.steps = tuple(steps) + self.reference_uid = str(next(iter(references))) + self.num_envs = int(env.num_envs) + profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + default_clearance = default_runtime_policy(profile).grounding["placement"][ + "clearance" + ] + self.clearance = float(default_clearance if clearance is None else clearance) + self.positions = self._make_slots() + + def _make_slots(self) -> dict[str, torch.Tensor]: + container = _object(self.env, self.reference_uid) + positions = { + step.id: torch.empty( + self.num_envs, + 3, + dtype=torch.float32, + device=self.env.device, + ) + for step in self.steps + } + named_slots = [str(step.goal.get("slot", "auto")) for step in self.steps] + for slot in named_slots: + if slot not in {"auto", "left", "center", "right"}: + raise ValueError(f"Unsupported container slot {slot!r}.") + + for env_id in range(self.num_envs): + vertices = _world_vertices(container, self.env, env_id) + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + center = (lower + upper) * 0.5 + extent = upper[:2] - lower[:2] + axis = int(torch.argmax(extent).item()) + radii = [] + for step in self.steps: + moved_vertices = _local_vertices( + _object(self.env, step.object_uid), + self.env, + env_id, + ) + half = ( + moved_vertices.max(dim=0).values - moved_vertices.min(dim=0).values + )[:2] * 0.5 + radii.append(float(torch.linalg.vector_norm(half))) + radius = max(radii) + usable_span = float(extent[axis]) - 2.0 * (radius + self.clearance) + required_span = 2.0 * radius * max(len(self.steps) - 1, 0) + if usable_span + 1.0e-6 < required_span: + raise ValueError( + f"Environment {env_id} container {self.reference_uid!r} " + "has no non-overlapping slot plan." + ) + offsets = torch.linspace( + -required_span * 0.5, + required_span * 0.5, + len(self.steps), + device=self.env.device, + ) + named_offsets = { + "left": required_span * 0.5, + "center": 0.0, + "right": -required_span * 0.5, + } + used: list[float] = [] + for index, step in enumerate(self.steps): + slot = named_slots[index] + offset = ( + float(offsets[index]) if slot == "auto" else named_offsets[slot] + ) + if any(abs(offset - item) < 2.0 * radius for item in used): + raise ValueError( + f"Container slot {slot!r} overlaps another requested slot." + ) + used.append(offset) + target = center.clone() + target[axis] += offset + target[2] = lower[2] + positions[step.id][env_id] = target + return positions + + def target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + rotation: torch.Tensor, + *, + surface_clearance: float, + ) -> torch.Tensor: + """Return a slot pose corrected for the rotated object mesh bottom.""" + target = object_pose.clone() + target[:, :3, :3] = rotation + target[:, :2, 3] = self.positions[step.id][:, :2] + entity = _object(self.env, step.object_uid) + for env_id in range(self.num_envs): + bottom = ( + _local_vertices(entity, self.env, env_id) + @ rotation[env_id].transpose(0, 1) + )[:, 2].min() + target[env_id, 2, 3] = ( + self.positions[step.id][env_id, 2] + surface_clearance - bottom + ) + return target + + +class ActionGrounder: + """Translate one symbolic action into a public typed atomic-action target.""" + + def __init__( + self, + program: ExecutionProgram, + env: Any, + semantics_factory: Callable[[str], ObjectSemantics], + arrangement: ( + LiveArrangementPlan | Mapping[str, LiveArrangementPlan] | None + ) = None, + placements: Mapping[str, LivePlacementPlan] | None = None, + runtime_policy: RuntimePolicyCfg | None = None, + capability_registry: Any | None = None, + ) -> None: + self.program = program + self.env = env + self.semantics_factory = semantics_factory + self.capabilities = capability_registry or build_atomic_capability_registry() + self.robot_profile = str(getattr(env, "agent_robot_profile", "dual_ur10")) + self.runtime_policy = runtime_policy or default_runtime_policy( + self.robot_profile + ) + if isinstance(arrangement, Mapping): + self.arrangements = dict(arrangement) + elif arrangement is None: + self.arrangements = {} + else: + self.arrangements = { + step.id: arrangement + for step in program.semantic_steps + if step.operator in {"arrange_line", "place_in_line"} + } + self.placements = dict(placements or {}) + + def policy( + self, + action: Mapping[str, Any], + *, + extra_modifiers: tuple[tuple[str, str], ...] = (), + ) -> dict[str, Any]: + action_class = str(action.get("atomic_action_class", "")) + capability = self.capabilities.get(action_class) + motion_base = capability.motion_base or capability.name + policy_spec = action.get("motion_policy", {"modifiers": []}) + if extra_modifiers: + policy_spec = with_motion_modifiers(policy_spec, *extra_modifiers) + inline = action.get("motion_policy_config", action.get("cfg")) + return resolve_motion_policy( + self.robot_profile, + motion_base, + policy_spec, + motion_defaults=self.runtime_policy.motion_defaults, + motion_modifiers=self.runtime_policy.motion_modifiers, + inline_overrides=inline if isinstance(inline, Mapping) else None, + ) + + def _policy_value(self, policy: Mapping[str, Any], key: str) -> Any: + defaults = self.runtime_policy.grounding["semantic_defaults"] + return policy[key] if key in policy else defaults[key] + + def ground( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + _handover_workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> GroundedAction: + action_class = str(action["atomic_action_class"]) + capability = self.capabilities.require_executable(action_class) + self.capabilities.validate_binding(action) + control = str(action.get("control", "arm")) + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + raise ValueError("target_binding must be a mapping.") + kind = str(binding.get("kind", "")) + orientation = compile_orientation_constraint(step.goal) + policy = self.policy(action) + if kind == "joint_state": + joint_defaults = self.runtime_policy.grounding["joint_state"] + source = binding.get("source") + if bool(binding.get("single_release", False)): + policy["single_release"] = True + if source == "gripper_closed": + policy["sample_interval"] = int( + joint_defaults["hand_close_sample_interval"] + ) + elif source == "gripper_open": + if binding.get("coordinated_release_role") is not None: + coordinated_name = next( + name + for name in self.capabilities.executable_names() + if self.capabilities.get(name).config_materializer + == "coordinated_pickment" + ) + coordinated = self.runtime_policy.motion_defaults[coordinated_name] + policy["sample_interval"] = int( + coordinated.get( + "release_sample_interval", + joint_defaults["hand_open_sample_interval"], + ) + ) + else: + policy["sample_interval"] = int( + joint_defaults["hand_open_sample_interval"] + ) + elif source == "initial" and control == "arm": + # Returning home after release is a safety motion. If the + # collision-aware planner cannot find a route, do not silently + # replace it with collision-unaware joint interpolation. + policy["collision_safety"] = "required" + object_pose = _live_pose(self.env, step.object_uid) + if "upright_local_axis" in step.goal: + policy["upright_local_axis"] = self._upright_local_axis(step) + if capability.target_materializer == "object_grasp": + policy["obj_upright_direction"] = self._upright_local_direction(step) + reference_pose = self._reference_pose(step) + target_object_pose = None + + if capability.target_materializer_hook is not None: + grounded = capability.target_materializer_hook( + grounder=self, + action=action, + step=step, + arm=arm, + state=state, + binding=binding, + policy=policy, + object_pose=object_pose, + reference_pose=reference_pose, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + if not isinstance(grounded, GroundedAction): + raise TypeError( + f"AtomicAction {action_class!r} target materializer must " + "return GroundedAction." + ) + return replace( + grounded, + allow_yaw_search=orientation.allows_yaw_search, + ) + + if kind == "object": + semantics = self.semantics_factory( + str(binding.get("object", step.object_uid)) + ) + if capability.target_materializer == "object_grasp": + if step.operator == "pour": + semantics = self._pour_source_semantics( + step, + semantics, + object_pose, + reference_pose, + ) + target: Any = GraspGoal(semantics=semantics) + elif capability.target_materializer == "axis_align": + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise ValueError( + "AxisAlign requires an AntipodalAffordance with mesh geometry." + ) + semantics = replace( + semantics, + affordance=AxisAlignAffordance( + object_label=affordance.object_label, + custom_config=dict(affordance.custom_config), + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + internal_axis=self._upright_local_direction(step), + ), + ) + policy.setdefault("target_axis", DEFAULT_TARGET_AXIS) + policy.setdefault( + "surface_clearance", + float(self._policy_value(policy, "surface_clearance")), + ) + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = AxisAlignGoal(semantics=semantics) + elif capability.target_materializer == "coordinated_pickment": + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = CoordinatedPickGoal( + object_target_pose=target_object_pose, + semantics=semantics, + object_initial_pose=object_pose, + ) + elif capability.target_materializer == "press": + target_object_pose = object_pose.clone() + target, press_policy = self._press_goal( + step.object_uid, + object_pose, + semantics=semantics, + terminal_state=str( + step.postcondition.get("terminal_state", "activated") + ), + ) + policy.update(press_policy) + else: + raise ValueError( + f"{action_class} does not support object target bindings." + ) + elif kind == "pour_goal": + if capability.target_materializer != "pour": + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve a pour_goal." + ) + policy.setdefault("rotate_angle", math.pi / 2.0) + target = PourGoal() + elif kind == "articulation_goal": + if capability.target_materializer == "slide": + target, policy = self._slide_target( + step, + arm, + policy, + ) + elif capability.target_materializer == "twist": + target, policy = self._twist_target(step, arm, policy) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve an articulation_goal." + ) + elif kind in {"semantic_goal", "coordinated_goal"}: + phase = str(binding.get("phase", "final")) + target_object_pose = self._semantic_target( + step, + object_pose, + reference_pose, + policy, + phase=phase, + orientation_reference_pose=orientation_reference_pose, + ) + if capability.target_materializer == "coordinated_pickment": + semantics = self.semantics_factory(step.object_uid) + target = CoordinatedPickGoal( + object_target_pose=target_object_pose, + semantics=semantics, + object_initial_pose=object_pose, + ) + elif capability.target_materializer == "press": + # Press moves the TCP, not the target object. Keep the object's + # live pose as the postcondition reference while grounding a + # downward contact point from its current surface geometry. + target_object_pose = object_pose.clone() + target, press_policy = self._press_goal( + step.object_uid, + object_pose, + terminal_state=str( + step.postcondition.get("terminal_state", "activated") + ), + ) + policy.update(press_policy) + elif capability.target_materializer == "semantic_held_object": + target = HeldObjectPoseGoal(object_target_pose=target_object_pose) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} cannot " + f"resolve {kind!r}." + ) + elif kind == "coordinated_placement_goal": + support_uid = binding.get( + "support_object", + step.goal.get("support_object"), + ) + placing_uid = binding.get("placing_object", step.object_uid) + if not isinstance(placing_uid, str) or not placing_uid: + raise ValueError("coordinated_placement_goal requires placing_object.") + if not isinstance(support_uid, str) or not support_uid: + raise ValueError("coordinated_placement_goal requires support_object.") + support_pose = _live_pose(self.env, support_uid) + target_object_pose = self._semantic_target( + step, + object_pose, + support_pose, + policy, + phase="final", + orientation_reference_pose=orientation_reference_pose, + ) + target = CoordinatedPlacementGoal( + placing_object_target_pose=target_object_pose, + support_object_target_pose=support_pose, + release=bool(step.goal.get("release", True)), + ) + elif kind == "current_held_pose": + if state.get_held_object(arm_control_part(self.env, arm)) is None: + raise ValueError("Place requires a held object from a prior PickUp.") + target = PlaceGoal( + xpos=( + reference_eef_pose + if reference_eef_pose is not None + else self._current_eef_pose(arm) + ) + ) + elif kind == "policy_pose": + source = binding.get("source") + operation = binding.get("operation") + retreat_reference = self._retreat_reference_pose( + arm, + reference_eef_pose, + ) + if operation in {"retreat", "lift_clear", "retreat_after_lift"}: + policy["retreat_reachability_search"] = True + policy["retreat_reference_pose"] = retreat_reference.clone() + if operation == "lift_clear": + policy["retreat_search_mode"] = "vertical_only" + if bool(binding.get("verify_lift_clear", False)): + policy["verify_lift_clear"] = True + policy["minimum_clearance"] = max( + float(policy.get("minimum_clearance", 0.0)), + _E2_REORIENT_MINIMUM_CLEARANCE, + ) + if operation == "reorient_tool_down": + policy["reorient_tool_down"] = True + policy["reorient_reference_pose"] = retreat_reference.clone() + policy["reorient_waypoint_count"] = 5 + policy["reorient_yaw_degrees"] = [ + 0.0, + 45.0, + -45.0, + 90.0, + -90.0, + 180.0, + ] + policy["minimum_clearance"] = _E2_REORIENT_MINIMUM_CLEARANCE + if operation == "retreat_after_lift": + policy["retreat_distance"] = _E2_CLEARANCE_RETREAT_DISTANCE + policy["minimum_retreat_distance"] = 0.05 + policy["retreat_search_mode"] = "horizontal_only" + policy["retreat_search_samples"] = 4 + if source in {"release", "handover"}: + policy["clearance_object_uid"] = step.object_uid + policy["collision_safety"] = "required" + contact_uids = [step.object_uid] + reference_uid = step.goal.get("reference_object") + if isinstance(reference_uid, str) and reference_uid: + contact_uids.append(reference_uid) + policy["collision_exclusion_uids"] = list(dict.fromkeys(contact_uids)) + if source == "handover": + policy.update(self.runtime_policy.grounding["handover"]) + policy["transfer_arm"] = arm + policy["transfer_role_axis"] = self._handover_role_axis( + arm, + dtype=object_pose.dtype, + device=object_pose.device, + ) + target = EndEffectorPoseGoal( + xpos=( + retreat_reference.clone() + if operation == "reorient_tool_down" + else self._retreat_pose( + arm, + policy, + retreat_reference, + clear_exchange=source == "handover", + retreat_after_lift=operation == "retreat_after_lift", + ) + ) + ) + elif kind == "visual_constraint": + visual_pose = self._visual_target(binding, arm) + if capability.target_materializer == "semantic_held_object": + target_object_pose = object_pose.clone() + target_object_pose[:, :3, 3] = visual_pose[:, :3, 3] + target = HeldObjectPoseGoal(object_target_pose=target_object_pose) + elif capability.target_materializer == "eef_pose": + target = EndEffectorPoseGoal(xpos=visual_pose) + else: + raise ValueError( + f"Target materializer {capability.target_materializer!r} " + "cannot resolve a visual_constraint." + ) + elif kind == "joint_state": + if bool(binding.get("required_home", False)): + policy["verify_required_home"] = True + target = JointPositionGoal( + target=self._joint_target( + arm, + control, + str(binding.get("source", "initial")), + binding, + ) + ) + elif kind in {"eef_pose", "pose"}: + target = EndEffectorPoseGoal(xpos=self._explicit_pose(binding, object_pose)) + elif kind == "handover_goal": + target, target_object_pose, policy = self._handover_target( + step, + binding, + object_pose, + reference_pose, + policy, + state, + orientation_reference_pose=orientation_reference_pose, + workspace=_handover_workspace, + ) + elif kind == "handover_staging": + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + receive_arm = str(binding.get("receive_arm", "right_arm")) + middle, _ = self._handover_workspace_poses( + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + step=step, + orientation_reference_pose=orientation_reference_pose, + ) + middle[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + target_object_pose = middle + target = HeldObjectPoseGoal(object_target_pose=middle) + else: + raise ValueError(f"Unsupported target binding kind {kind!r}.") + return GroundedAction( + action_class=action_class, + arm=arm, + control=control, + target=target, + cfg=policy, + object_pose=object_pose, + reference_pose=reference_pose, + target_object_pose=target_object_pose, + motion_policy=policy, + object_uid=step.object_uid, + allow_yaw_search=orientation.allows_yaw_search, + ) + + def _handover_role_axis( + self, + transfer_arm: str, + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Return the world-space axis from the receiver base to transfer base.""" + if transfer_arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Unknown handover arm {transfer_arm!r}.") + _, lateral = robot_frame_axes(self.env) + horizontal = lateral if transfer_arm == "left_arm" else -lateral + return torch.cat( + ( + horizontal.to(dtype=dtype, device=device), + torch.zeros( + (int(self.env.num_envs), 1), + dtype=dtype, + device=device, + ), + ), + dim=1, + ) + + def ground_candidates( + self, + action: Mapping[str, Any], + step: SemanticStep, + *, + arm: str, + state: ExecutionState, + reference_eef_pose: torch.Tensor | None = None, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[GroundedAction, ...]: + """Return deterministic grounding candidates for an opt-in capability.""" + binding = action.get("target_binding", {}) + if not isinstance(binding, Mapping): + return ( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ), + ) + placement_support_uid = self._placement_support_uid(step) + placement_relation = _placement_relation(step) + is_on_placement = ( + binding.get("kind") == "semantic_goal" + and binding.get("phase", "final") != "staging" + and placement_relation in {"on", "on_top", "on_top_of"} + and placement_support_uid is not None + ) + if is_on_placement: + base = self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ) + return self._placement_grounding_candidates( + base, + step, + support_uid=placement_support_uid, + ) + if binding.get("kind") != "handover_goal": + return ( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + ), + ) + policy = self.policy(action) + object_pose = _live_pose(self.env, step.object_uid) + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + workspaces = self._handover_workspace_candidates( + step, + object_pose, + transfer_arm=str(binding.get("transfer_arm", "left_arm")), + receive_arm=str(binding.get("receive_arm", "right_arm")), + policy=policy, + rotation=rotation, + ) + return tuple( + self.ground( + action, + step, + arm=arm, + state=state, + reference_eef_pose=reference_eef_pose, + orientation_reference_pose=orientation_reference_pose, + _handover_workspace=workspace, + ) + for workspace in workspaces + ) + + def _placement_grounding_candidates( + self, + base: GroundedAction, + step: SemanticStep, + *, + support_uid: str, + ) -> tuple[GroundedAction, ...]: + """Sample bounded support-relative poses from live object geometry.""" + if base.target_object_pose is None or not isinstance( + base.target, HeldObjectPoseGoal + ): + return (base,) + support = _object(self.env, support_uid) + moved = _object(self.env, step.object_uid) + placement = self.runtime_policy.grounding["placement"] + count = int(placement["candidate_count"]) + fraction = float(placement["candidate_offset_fraction"]) + margin = float(placement["support_margin"]) + patterns = ( + (0.0, 0.0), + (1.0, 0.0), + (-1.0, 0.0), + (0.0, 1.0), + (0.0, -1.0), + (1.0, 1.0), + (1.0, -1.0), + (-1.0, 1.0), + (-1.0, -1.0), + )[:count] + candidates: list[GroundedAction] = [] + seen_offsets: list[torch.Tensor] = [] + for candidate_index, pattern in enumerate(patterns): + target_pose = base.target_object_pose.clone() + offsets = target_pose.new_zeros((int(self.env.num_envs), 2)) + for env_id in range(int(self.env.num_envs)): + support_vertices = _world_vertices(support, self.env, env_id) + moved_local = _local_vertices(moved, self.env, env_id) + rotated = moved_local @ target_pose[env_id, :3, :3].transpose(0, 1) + support_lower = support_vertices[:, :2].min(dim=0).values + support_upper = support_vertices[:, :2].max(dim=0).values + moved_lower = rotated[:, :2].min(dim=0).values + moved_upper = rotated[:, :2].max(dim=0).values + allowed_lower = support_lower + margin - moved_lower + allowed_upper = support_upper - margin - moved_upper + if bool(torch.all(allowed_lower <= allowed_upper)): + base_xy = target_pose[env_id, :2, 3].clone() + center = torch.minimum( + torch.maximum(base_xy, allowed_lower), + allowed_upper, + ) + direction = target_pose.new_tensor(pattern) + room = torch.where( + direction >= 0.0, + allowed_upper - center, + center - allowed_lower, + ) + candidate_xy = center + direction * room * fraction + offsets[env_id] = candidate_xy - base_xy + target_pose[env_id, :2, 3] = candidate_xy + + footprint_lower = target_pose[env_id, :2, 3] + moved_lower + footprint_upper = target_pose[env_id, :2, 3] + moved_upper + local_mask = torch.all( + (support_vertices[:, :2] >= footprint_lower - margin) + & (support_vertices[:, :2] <= footprint_upper + margin), + dim=1, + ) + if bool(local_mask.any()): + support_height = support_vertices[local_mask, 2].max() + else: + distances = torch.linalg.vector_norm( + support_vertices[:, :2] - target_pose[env_id, :2, 3], + dim=1, + ) + nearest_count = min(8, int(support_vertices.shape[0])) + nearest = torch.topk( + distances, + nearest_count, + largest=False, + ).indices + support_height = support_vertices[nearest, 2].max() + target_pose[env_id, 2, 3] = ( + support_height + + float(self._policy_value(base.motion_policy, "surface_clearance")) + - rotated[:, 2].min() + ) + if any(torch.allclose(offsets, prior) for prior in seen_offsets): + continue + seen_offsets.append(offsets) + candidates.append( + replace( + base, + target=replace(base.target, object_target_pose=target_pose), + target_object_pose=target_pose, + motion_policy={ + **base.motion_policy, + "placement_candidate_index": candidate_index, + "placement_xy_offset": offsets, + }, + ) + ) + return tuple(candidates) or (base,) + + @staticmethod + def _placement_support_uid(step: SemanticStep) -> str | None: + value = step.goal.get("reference_object", step.goal.get("support_object")) + if isinstance(value, str) and value: + return value + if ( + step.postcondition.get("type") == "stack_layer_supported" + and int(step.goal.get("layer_index", -1)) == 0 + ): + return "table" + return None + + def _visual_target( + self, + binding: Mapping[str, Any], + arm: str, + ) -> torch.Tensor: + """Unproject one normalized image keypoint using live camera depth.""" + camera_uid = str(binding.get("camera_uid", "")) + sensor = self.env.sim.get_sensor(camera_uid) + if sensor is None: + raise ValueError(f"Unknown visual-constraint camera {camera_uid!r}.") + keypoint_value = binding.get("normalized_keypoint") + if keypoint_value is None: + bbox = binding.get("normalized_bbox") + if isinstance(bbox, Sequence) and len(bbox) == 4: + keypoint_value = [ + (float(bbox[0]) + float(bbox[2])) * 0.5, + (float(bbox[1]) + float(bbox[3])) * 0.5, + ] + if keypoint_value is None: + raise ValueError( + "visual_constraint requires a normalized keypoint or bbox in [0, 1]." + ) + keypoint = torch.as_tensor( + keypoint_value, + dtype=torch.float32, + device=self.env.device, + ).flatten() + if keypoint.numel() != 2 or bool( + ((~torch.isfinite(keypoint)) | (keypoint < 0.0) | (keypoint > 1.0)).any() + ): + raise ValueError( + "visual_constraint requires a normalized keypoint or bbox in [0, 1]." + ) + data = sensor.get_data() + if "depth" not in data: + raise ValueError( + f"Camera {camera_uid!r} must provide depth for visual Grounding." + ) + depth = torch.as_tensor(data["depth"], device=self.env.device).squeeze(-1) + if depth.ndim == 2: + depth = depth.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + if depth.ndim != 3 or depth.shape[0] != int(self.env.num_envs): + raise ValueError("Camera depth must have shape (N, H, W) or (N, H, W, 1).") + height, width = depth.shape[-2:] + pixel_x = min(max(int(round(float(keypoint[0]) * (width - 1))), 0), width - 1) + pixel_y = min(max(int(round(float(keypoint[1]) * (height - 1))), 0), height - 1) + distance = depth[:, pixel_y, pixel_x].to(torch.float32) + if bool((~torch.isfinite(distance) | (distance <= 0.0)).any()): + raise ValueError("visual_constraint keypoint has no valid live depth.") + intrinsics = torch.as_tensor( + sensor.get_intrinsics(), + dtype=torch.float32, + device=self.env.device, + ) + if intrinsics.ndim == 2: + intrinsics = intrinsics.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + camera_pose = torch.as_tensor( + sensor.get_arena_pose(to_matrix=True), + dtype=torch.float32, + device=self.env.device, + ) + if camera_pose.ndim == 2: + camera_pose = camera_pose.unsqueeze(0).repeat(int(self.env.num_envs), 1, 1) + fx = intrinsics[:, 0, 0] + fy = intrinsics[:, 1, 1] + cx = intrinsics[:, 0, 2] + cy = intrinsics[:, 1, 2] + point = torch.stack( + ( + (float(pixel_x) - cx) * distance / fx, + (float(pixel_y) - cy) * distance / fy, + distance, + torch.ones_like(distance), + ), + dim=1, + ) + world = torch.bmm(camera_pose, point.unsqueeze(-1)).squeeze(-1) + target = self._current_eef_pose(arm).clone() + target[:, :3, 3] = world[:, :3] + return target + + def _handover_target( + self, + step: SemanticStep, + binding: Mapping[str, Any], + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + policy: Mapping[str, Any], + state: ExecutionState, + *, + orientation_reference_pose: torch.Tensor | None, + workspace: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> tuple[GraspGoal, torch.Tensor, dict[str, Any]]: + transfer_arm = str(binding.get("transfer_arm", "left_arm")) + receive_arm = str( + binding.get( + "receive_arm", + "right_arm" if transfer_arm == "left_arm" else "left_arm", + ) + ) + if transfer_arm == receive_arm or {transfer_arm, receive_arm} != { + "left_arm", + "right_arm", + }: + raise ValueError("HandOver requires distinct left_arm/right_arm roles.") + transfer_part = arm_control_part(self.env, transfer_arm) + held = state.get_held_object(transfer_part) + if held is None: + raise ValueError( + f"HandOver requires {transfer_arm} to hold {step.object_uid!r}." + ) + + del reference_pose + if workspace is None: + middle, final = self._handover_workspace_poses( + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + step=step, + orientation_reference_pose=orientation_reference_pose, + ) + else: + middle, final = (item.clone() for item in workspace) + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + middle[:, :3, :3] = rotation + final[:, :3, :3] = rotation + semantics = self.semantics_factory(step.object_uid) + grounded_policy = dict(policy) + grounded_policy.update( + { + "transfer_arm": transfer_arm, + "receive_arm": receive_arm, + "middle_object_pose": middle, + "final_object_pose": final, + } + ) + return ( + GraspGoal(semantics=semantics), + middle, + grounded_policy, + ) + + def _handover_workspace_poses( + self, + object_pose: torch.Tensor, + *, + transfer_arm: str, + receive_arm: str, + policy: Mapping[str, Any], + step: SemanticStep, + orientation_reference_pose: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Choose the highest-ranked collision-aware handover workspace.""" + rotation = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + candidates = self._handover_workspace_candidates( + step, + object_pose, + transfer_arm=transfer_arm, + receive_arm=receive_arm, + policy=policy, + rotation=rotation, + ) + return candidates[0] + + def _handover_workspace_candidates( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + transfer_arm: str, + receive_arm: str, + policy: Mapping[str, Any], + rotation: torch.Tensor, + ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + """Rank exchange poses inside the two arm workspaces and above obstacles.""" + if transfer_arm == receive_arm or {transfer_arm, receive_arm} != { + "left_arm", + "right_arm", + }: + raise ValueError("Handover workspace requires distinct arm roles.") + table = self.env.sim.get_rigid_object("table") + if table is not None and hasattr(table, "get_vertices"): + centers = [] + tops = [] + bounds = [] + for env_id in range(int(self.env.num_envs)): + vertices = _world_vertices(table, self.env, env_id) + lower = vertices[:, :2].min(dim=0).values + upper = vertices[:, :2].max(dim=0).values + centers.append((lower + upper) * 0.5) + tops.append(vertices[:, 2].max()) + bounds.append(torch.stack((lower, upper))) + center = torch.stack(centers) + table_top = torch.stack(tops) + table_bounds = torch.stack(bounds) + else: + left = self._current_eef_pose("left_arm") + right = self._current_eef_pose("right_arm") + center = (left[:, :2, 3] + right[:, :2, 3]) * 0.5 + table_top = object_pose[:, 2, 3] + extent = float(policy.get("exchange_candidate_offset", 0.16)) * 2.0 + table_bounds = torch.stack((center - extent, center + extent), dim=1) + + forward, lateral = robot_frame_axes(self.env) + left_base, right_base = arm_base_poses(self.env) + transfer_base = left_base if transfer_arm == "left_arm" else right_base + receive_base = right_base if receive_arm == "right_arm" else left_base + base_midpoint = (transfer_base[:, :2, 3] + receive_base[:, :2, 3]) * 0.5 + table_forward = torch.sum((center - base_midpoint) * forward, dim=1) + shared_center = base_midpoint + forward * table_forward[:, None] + offset = float(policy.get("exchange_candidate_offset", 0.16)) + obstacle_clearance = float(policy.get("exchange_obstacle_clearance", 0.04)) + tool_horizontal_envelope = float( + policy.get("exchange_gripper_horizontal_envelope", 0.035) + ) + float(policy.get("exchange_wrist_horizontal_envelope", 0.055)) + tool_vertical_envelope = float( + policy.get("exchange_gripper_vertical_envelope", 0.025) + ) + float(policy.get("exchange_wrist_vertical_envelope", 0.04)) + minimum_reach = float(policy.get("exchange_minimum_reach", 0.10)) + maximum_reach = float(policy.get("exchange_maximum_reach", 1.00)) + if not 0.0 <= minimum_reach < maximum_reach: + raise ValueError("Handover reach bounds require 0 <= minimum < maximum.") + requested_count = max(1, int(policy.get("exchange_candidate_count", 4))) + object_clearance = float(policy.get("exchange_clearance", 0.06)) + if ( + min( + obstacle_clearance, + tool_horizontal_envelope, + tool_vertical_envelope, + object_clearance, + ) + < 0.0 + ): + raise ValueError("Handover geometry clearances must be non-negative.") + xy_coefficients = ( + (0.0, 0.0), + (1.0, 0.0), + (-1.0, 0.0), + (2.0, 0.0), + (-2.0, 0.0), + (0.0, 0.5), + (0.0, -0.5), + ) + ranked_by_env: list[list[tuple[float, torch.Tensor]]] = [] + moved = _object(self.env, step.object_uid) + obstacle_uids = ( + self.env.sim.get_rigid_object_uid_list() + if hasattr(self.env.sim, "get_rigid_object_uid_list") + else [] + ) + for env_id in range(int(self.env.num_envs)): + local_vertices = _local_vertices(moved, self.env, env_id) + rotated = local_vertices @ rotation[env_id].transpose(0, 1) + half_xy = ( + rotated[:, :2].max(dim=0).values - rotated[:, :2].min(dim=0).values + ) * 0.5 + bottom = rotated[:, 2].min() + margin = half_xy + obstacle_clearance + tool_horizontal_envelope + lower_limit = table_bounds[env_id, 0] + margin + upper_limit = table_bounds[env_id, 1] - margin + options: list[tuple[float, torch.Tensor]] = [] + for forward_scale, lateral_scale in xy_coefficients: + xy = ( + shared_center[env_id] + + forward[env_id] * (offset * forward_scale) + + lateral[env_id] * (offset * lateral_scale) + ) + if bool(((xy < lower_limit) | (xy > upper_limit)).any()): + continue + transfer_distance = torch.linalg.vector_norm( + xy - transfer_base[env_id, :2, 3] + ) + receive_distance = torch.linalg.vector_norm( + xy - receive_base[env_id, :2, 3] + ) + if not ( + minimum_reach <= float(transfer_distance) <= maximum_reach + and minimum_reach <= float(receive_distance) <= maximum_reach + ): + continue + obstacle_score, nearby_obstacle_top = self._handover_obstacle_metrics( + xy, + env_id=env_id, + object_uid=step.object_uid, + obstacle_uids=obstacle_uids, + half_xy=half_xy, + clearance=obstacle_clearance + tool_horizontal_envelope, + ) + center_cost = float( + torch.linalg.vector_norm(xy - shared_center[env_id]) + ) + pose = object_pose[env_id].clone() + pose[:3, :3] = rotation[env_id] + pose[:2, 3] = xy + safety_floor = torch.maximum(table_top[env_id], nearby_obstacle_top) + safe_z = ( + safety_floor + object_clearance + tool_vertical_envelope - bottom + ) + pose[2, 3] = torch.maximum(object_pose[env_id, 2, 3], safe_z) + lift_cost = max( + 0.0, + float(pose[2, 3] - object_pose[env_id, 2, 3]), + ) + options.append( + (obstacle_score + center_cost * 0.25 + lift_cost * 0.1, pose) + ) + if not options: + raise ValueError( + "No handover exchange pose lies inside the table bounds and " + "the reachable intersection of both arm bases." + ) + options.sort(key=lambda item: item[0]) + ranked_by_env.append(options[:requested_count]) + + candidate_count = min( + requested_count, + max(len(options) for options in ranked_by_env), + ) + candidates = [] + for candidate_index in range(candidate_count): + middle = object_pose.clone() + for env_id, options in enumerate(ranked_by_env): + middle[env_id] = options[min(candidate_index, len(options) - 1)][1] + # The built-in HandOver primitive plans its final transfer/receiver + # phase concurrently. An exchange-to-exchange target makes that + # receiver path stationary; graph-level retreat/home nodes then + # clear the transfer arm before any receiver-side continuation. + final = middle.clone() + candidates.append((middle, final)) + return tuple(candidates) + + def _handover_obstacle_metrics( + self, + xy: torch.Tensor, + *, + env_id: int, + object_uid: str, + obstacle_uids: Sequence[str], + half_xy: torch.Tensor, + clearance: float, + ) -> tuple[float, torch.Tensor]: + score = 0.0 + highest_top = torch.tensor( + -torch.inf, + dtype=xy.dtype, + device=xy.device, + ) + for uid in obstacle_uids: + if uid in {"table", object_uid}: + continue + obstacle = self.env.sim.get_rigid_object(uid) + if obstacle is None or not hasattr(obstacle, "get_vertices"): + continue + vertices = _world_vertices(obstacle, self.env, env_id) + lower = vertices[:, :2].min(dim=0).values - half_xy - clearance + upper = vertices[:, :2].max(dim=0).values + half_xy + clearance + outside = torch.maximum( + torch.maximum(lower - xy, xy - upper), + torch.zeros_like(xy), + ) + if bool((outside > 0.0).any()): + distance = float(torch.linalg.vector_norm(outside)) + score += 1.0 / max(distance, 1.0e-3) + else: + score += 1.0e3 + highest_top = torch.maximum(highest_top, vertices[:, 2].max()) + return score, highest_top + + def _handover_receiver_exit( + self, + middle: torch.Tensor, + receive_arm: str, + policy: Mapping[str, Any], + ) -> torch.Tensor: + final = middle.clone() + receive_pose = self._current_eef_pose(receive_arm) + direction = receive_pose[:, :2, 3] - middle[:, :2, 3] + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + fallback = direction.new_zeros(direction.shape) + fallback[:, 1] = -1.0 if receive_arm == "right_arm" else 1.0 + direction = torch.where( + norm > 1.0e-6, direction / norm.clamp_min(1.0e-6), fallback + ) + final[:, :2, 3] += direction * min( + 0.12, + float(self._policy_value(policy, "relation_distance")) * 0.5, + ) + return final + + def _reference_pose(self, step: SemanticStep) -> torch.Tensor | None: + uid = step.goal.get("reference_object", step.goal.get("support_object")) + if not isinstance(uid, str) or not uid: + return None + if step.goal.get("reference_state") == "initial": + initial = getattr(self.env, "agent_initial_object_poses", {}).get(uid) + if initial is None: + raise ValueError(f"Initial pose for {uid!r} is unavailable.") + return _batched_pose(initial, self.env) + return _live_pose(self.env, uid) + + def _semantic_target( + self, + step: SemanticStep, + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + policy: Mapping[str, Any], + *, + phase: str, + orientation_reference_pose: torch.Tensor | None = None, + ) -> torch.Tensor: + if phase == "handover_exit": + receive_arm = str(step.goal.get("receive_arm", "")) + if receive_arm not in {"left_arm", "right_arm"}: + raise ValueError("handover_exit requires a concrete receive_arm.") + target = self._handover_receiver_exit(object_pose, receive_arm, policy) + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + return target + if step.operator in {"arrange_line", "place_in_line"}: + arrangement = self.arrangements.get(step.id) + if arrangement is None: + raise ValueError("arrange_line requires a live arrangement plan.") + target = arrangement.target( + step, + object_pose, + phase=phase, + policy=policy, + ) + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + arrangement.table_top[env_id] + + float(policy["surface_clearance"]) + - bottom + ) + if phase == "staging": + target[:, 2, 3] += float(policy["transport_clearance"]) + return target + placement = self.placements.get(step.id) + if placement is not None: + target = placement.target( + step, + object_pose, + self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ), + surface_clearance=float(policy["surface_clearance"]), + ) + if phase == "staging": + target[:, 2, 3] += float(policy["transport_clearance"]) + return target + if step.goal.get("position_anchor") in {"initial_xy", "live_xy"} and isinstance( + step.goal.get("support_object"), str + ): + initial = None + if step.goal.get("position_anchor", "initial_xy") == "initial_xy": + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + target = ( + _batched_pose(initial, self.env).clone() + if initial is not None + else object_pose.clone() + ) + target[:, :3, :3] = self._target_rotation( + step, + target, + orientation_reference_pose=orientation_reference_pose, + ) + support_uid = str(step.goal.get("support_object", "table")) + support = _object(self.env, support_uid) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + float(policy["surface_clearance"]) - bottom + ) + if phase == "staging": + target[:, 2, 3] += float(policy["staging_lift_height"]) + return target + if step.operator == "pour": + if phase == "return": + initial = getattr(self.env, "agent_initial_object_poses", {}).get( + step.object_uid + ) + if initial is None: + raise ValueError( + f"Pour return requires an initial pose for {step.object_uid!r}." + ) + return _batched_pose(initial, self.env).clone() + if reference_pose is None: + raise ValueError("Pour requires a live target-container pose.") + target = object_pose.clone() + source = _object(self.env, step.object_uid) + receiver_uid = str(step.goal.get("reference_object", "")) + receiver = _object(self.env, receiver_uid) + clearance = float(self._policy_value(policy, "transport_clearance")) + for env_id in range(int(self.env.num_envs)): + receiver_vertices = _world_vertices(receiver, self.env, env_id) + target[env_id, :2, 3] = ( + receiver_vertices[:, :2].min(dim=0).values + + receiver_vertices[:, :2].max(dim=0).values + ) * 0.5 + source_vertices = _local_vertices(source, self.env, env_id) + rotation_radius = torch.linalg.vector_norm(source_vertices, dim=1).max() + target[env_id, 2, 3] = ( + receiver_vertices[:, 2].max() + clearance + rotation_radius + ) + return target + target = object_pose.clone() + if reference_pose is not None: + target[:, :3, 3] = reference_pose[:, :3, 3] + # Operators without a relational goal (for example press or a + # direction-only coordinated transport) must preserve the live origin + # instead of being silently projected onto a synthetic table support. + relation = _placement_relation(step) + distance = float(self._policy_value(policy, "relation_distance")) + relation_frame = str(step.goal.get("relation_frame", "world")) + forward_distance = distance + lateral_distance = distance + if relation_frame == "robot" and reference_pose is not None: + nominal = float(policy.get("robot_relative_distance", 0.10)) + clearance = float(policy.get("relation_clearance", 0.02)) + reference_uid = str(step.goal.get("reference_object", "")) + if reference_uid: + forward_axis, lateral_axis = robot_frame_axes(self.env) + forward_distance = self._relative_object_spacing( + step.object_uid, + reference_uid, + axis=forward_axis, + nominal=nominal, + clearance=clearance, + ) + lateral_distance = self._relative_object_spacing( + step.object_uid, + reference_uid, + axis=lateral_axis, + nominal=nominal, + clearance=clearance, + ) + directional_offset = relation_offset( + self.env, + relation, + frame=relation_frame, + forward_distance=forward_distance, + lateral_distance=lateral_distance, + dtype=target.dtype, + device=target.device, + ) + offsets = { + "above": (0.0, 0.0, float(self._policy_value(policy, "hover_height"))), + "held_above_initial": ( + 0.0, + 0.0, + float(self._policy_value(policy, "hover_height")), + ), + } + if directional_offset is not None: + target[:, :3, 3] += directional_offset + elif (offset := offsets.get(relation)) is not None: + target[:, :3, 3] += torch.tensor( + offset, + dtype=target.dtype, + device=target.device, + ) + slot = str(step.goal.get("slot", "auto")) + if relation in {"on", "on_top", "on_top_of", "inside"} and slot in { + "left", + "right", + }: + slot_offset = relation_offset( + self.env, + slot, + frame=relation_frame, + forward_distance=forward_distance, + lateral_distance=lateral_distance, + dtype=target.dtype, + device=target.device, + ) + if slot_offset is not None: + target[:, :3, 3] += slot_offset + direction = str(step.goal.get("direction", "none")) + direction_offsets = { + "world_x": (distance, 0.0, 0.0), + "world_y": (0.0, distance, 0.0), + "up": (0.0, 0.0, distance), + "down": (0.0, 0.0, -distance), + } + planar_direction_offset = relation_offset( + self.env, + direction, + frame=relation_frame, + forward_distance=distance, + lateral_distance=distance, + dtype=target.dtype, + device=target.device, + ) + if planar_direction_offset is not None: + target[:, :3, 3] += planar_direction_offset + elif direction in direction_offsets: + target[:, :3, 3] += torch.tensor( + direction_offsets[direction], + dtype=target.dtype, + device=target.device, + ) + + root_stack_layer = ( + step.operator == "build_stack" + and int(step.goal.get("layer_index", 0)) == 0 + and reference_pose is None + ) + if root_stack_layer: + table = _object(self.env, "table") + for env_id in range(int(self.env.num_envs)): + vertices = _world_vertices(table, self.env, env_id) + target[env_id, :2, 3] = ( + vertices[:, :2].min(dim=0).values + + vertices[:, :2].max(dim=0).values + ) * 0.5 + + target[:, :3, :3] = self._target_rotation( + step, + object_pose, + orientation_reference_pose=orientation_reference_pose, + ) + if ( + step.actor.get("mode") == "coordinated" + and "terminal_behavior" in step.goal + and relation not in {"on", "on_top", "on_top_of", "inside"} + and direction not in {"up", "down"} + ): + release = str(step.goal.get("terminal_behavior", "hold")) == "place" + if not release: + target[:, 2, 3] = object_pose[:, 2, 3] + float( + self._policy_value(policy, "transport_clearance") + ) + else: + table = _object(self.env, "table") + moved = _object(self.env, step.object_uid) + clearance = float(self._policy_value(policy, "surface_clearance")) + for env_id in range(int(self.env.num_envs)): + table_top = _world_vertices(table, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = table_top + clearance - bottom + if relation in {"on", "on_top", "on_top_of"} or root_stack_layer: + support_uid = ( + step.goal.get("reference_object") + or step.goal.get("support_object") + or "table" + ) + support = _object(self.env, str(support_uid)) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + + float(self._policy_value(policy, "surface_clearance")) + - bottom + ) + elif relation == "inside" and reference_pose is not None: + # Grounding the final move happens after the staging lift. Preserve + # the pre-pick supported height rather than the lifted live height. + supported_pose = orientation_reference_pose + if supported_pose is None: + supported_pose = object_pose + supported_pose = _batched_pose(supported_pose, self.env) + target[:, 2, 3] = supported_pose[:, 2, 3] + elif isinstance(step.goal.get("support_object"), str) and relation not in { + "none", + "handover", + "above", + "held_above_initial", + }: + support = _object(self.env, str(step.goal["support_object"])) + moved = _object(self.env, step.object_uid) + for env_id in range(int(self.env.num_envs)): + support_top = _world_vertices(support, self.env, env_id)[:, 2].max() + bottom = self._rotated_local_z_min( + moved, + target[env_id, :3, :3], + env_id, + ) + target[env_id, 2, 3] = ( + support_top + + float(self._policy_value(policy, "surface_clearance")) + - bottom + ) + if phase == "staging": + # Staging is a runtime waypoint, not a persisted coordinate. This + # keeps in-place orientation robust to the object's live height. + target[:, 2, 3] += float(self._policy_value(policy, "transport_clearance")) + return target + + def _pour_source_semantics( + self, + step: SemanticStep, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + reference_pose: torch.Tensor | None, + ) -> ObjectSemantics: + """Attach the target-directed local rotation axis used by Pour.""" + affordance = semantics.affordance + if not isinstance(affordance, AntipodalAffordance): + raise ValueError("Pour requires an AntipodalAffordance for pickup.") + if reference_pose is None: + raise ValueError("Pour requires a live target-container pose.") + direction = reference_pose[:, :3, 3] - object_pose[:, :3, 3] + direction[:, 2] = 0.0 + norm = torch.linalg.vector_norm(direction, dim=1) + if torch.any(norm <= 1.0e-6): + raise ValueError( + "Pour source and target must define a non-zero horizontal direction." + ) + direction = direction / norm.unsqueeze(1) + world_up = direction.new_tensor([0.0, 0.0, 1.0]).expand_as(direction) + world_axis = torch.linalg.cross(world_up, direction, dim=1) + local_axes = torch.bmm( + object_pose[:, :3, :3].transpose(1, 2), + world_axis.unsqueeze(2), + ).squeeze(2) + local_axes = torch.nn.functional.normalize(local_axes, dim=1) + if not torch.allclose( + local_axes, + local_axes[:1].expand_as(local_axes), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched Pour environments require one shared object-local " + "target-directed rotation axis." + ) + return replace( + semantics, + affordance=AxisAlignAffordance( + object_label=affordance.object_label, + custom_config=dict(affordance.custom_config), + mesh_vertices=affordance.mesh_vertices, + mesh_triangles=affordance.mesh_triangles, + internal_axis=local_axes[0], + ), + ) + + def _slide_target( + self, + step: SemanticStep, + arm: str, + policy: Mapping[str, Any], + ) -> tuple[SlideGoal, dict[str, Any]]: + """Build one Slide goal from live prismatic-joint metadata.""" + articulation = getattr(self.env.sim, "get_articulation", lambda _uid: None)( + step.object_uid + ) + if articulation is None: + raise ValueError( + f"Articulation action requires live articulation {step.object_uid!r}." + ) + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + backend = backend_entities[0] + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == "prismatic": + candidates.append((int(joint_id), joint_name, info)) + if len(candidates) != 1: + raise ValueError( + "Slide grounding requires exactly one active prismatic joint; " + f"found {[name for _, name, _ in candidates]}." + ) + joint_id, joint_name, joint_info = candidates[0] + child_link = str(getattr(joint_info, "child_link_name", "")) + parent_link = str(getattr(joint_info, "parent_link_name", "")) + if ( + child_link not in articulation.link_names + or parent_link not in articulation.link_names + ): + raise ValueError( + "Prismatic joint metadata must identify live parent and child links." + ) + contact_links = [] + for candidate_name in getattr(articulation, "all_joint_names", ()): + candidate_info = backend.get_joint_info(str(candidate_name)) + candidate_type = ( + str( + getattr( + getattr(candidate_info, "joint_type", None), + "name", + candidate_info.joint_type, + ) + ) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + candidate_child = str(getattr(candidate_info, "child_link_name", "")) + if ( + candidate_type == "fixed" + and str(getattr(candidate_info, "parent_link_name", "")) == child_link + and candidate_child in articulation.link_names + ): + contact_links.append(candidate_child) + if len(contact_links) != 1: + raise ValueError( + "Slide grounding requires exactly one fixed contact endpoint " + f"on prismatic child link {child_link!r}; found {contact_links}." + ) + contact_link = contact_links[0] + vertices, triangles = articulation.get_link_vert_face(contact_link) + vertices = torch.as_tensor( + vertices, dtype=torch.float32, device=self.env.device + ) + triangles = torch.as_tensor( + triangles, dtype=torch.int64, device=self.env.device + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): + raise ValueError("Slide contact link has no valid grasp geometry.") + if triangles.ndim != 2 or triangles.shape[-1] != 3 or not triangles.numel(): + raise ValueError("Slide contact link has no valid triangle geometry.") + + contact_pose = _batched_pose( + articulation.get_link_pose(contact_link, to_matrix=True), self.env + ) + parent_pose = _batched_pose( + articulation.get_link_pose(parent_link, to_matrix=True), self.env + ) + raw_axis = joint_info.axis + raw_origin = joint_info.origin_pose + axis = torch.tensor( + raw_axis.tolist() if hasattr(raw_axis, "tolist") else raw_axis, + dtype=torch.float32, + device=self.env.device, + ).reshape(3) + origin = torch.tensor( + raw_origin.tolist() if hasattr(raw_origin, "tolist") else raw_origin, + dtype=torch.float32, + device=self.env.device, + ).reshape(4, 4) + opening_world = torch.matmul( + torch.matmul(parent_pose[:, :3, :3], origin[:3, :3]), axis + ) + push_world = -torch.nn.functional.normalize(opening_world, dim=1) + push_local = torch.bmm( + contact_pose[:, :3, :3].transpose(1, 2), + push_world.unsqueeze(2), + ).squeeze(2) + push_local = torch.nn.functional.normalize(push_local, dim=1) + if not torch.allclose( + push_local, + push_local[:1].expand_as(push_local), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched Slide environments require one shared child-link axis." + ) + + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + qpos = articulation.get_qpos()[:, joint_id] + if limits.shape != (int(self.env.num_envs), 2): + raise ValueError("Prismatic joint limits have an invalid batch shape.") + if not torch.isfinite(limits).all() or torch.any(limits[:, 0] >= limits[:, 1]): + raise ValueError("Slide requires finite ordered prismatic joint limits.") + if step.operator not in {"pull_articulated_part", "push_articulated_part"}: + raise ValueError("Slide grounding requires a pull or push operator.") + direction = "pull" if step.operator == "pull_articulated_part" else "push" + target_qpos = limits[:, 1] if direction == "pull" else limits[:, 0] + distances = torch.abs(target_qpos - qpos) + if torch.any(distances <= 1.0e-5): + raise ValueError("Articulation joint is already at the requested target.") + if not torch.allclose( + distances, + distances[:1].expand_as(distances), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched Slide environments require one shared translation distance." + ) + scoped_policy = dict(policy) + scoped_policy.update( + { + "direction": direction, + "translation_distance": float(distances[0]), + "articulation_joint_name": joint_name, + "articulation_joint_id": joint_id, + "articulation_initial_qpos": qpos, + "articulation_target_qpos": target_qpos, + "articulation_push_axis_world": push_world, + } + ) + left_base, right_base = arm_base_poses(self.env) + arm_base = left_base if arm == "left_arm" else right_base + current_eef = self._current_eef_pose(arm) + log_info( + f"Slide grounding {step.id}/{arm}: contact_link={contact_link!r}, " + f"contact_position={contact_pose[0, :3, 3].detach().cpu().tolist()}, " + f"push_axis={push_world[0].detach().cpu().tolist()}, " + f"eef_position={current_eef[0, :3, 3].detach().cpu().tolist()}, " + f"arm_base_position={arm_base[0, :3, 3].detach().cpu().tolist()}." + ) + contact_entity_id = f"{step.object_uid}:{contact_link}" + affordance = SlideAffordance( + object_label=contact_entity_id, + mesh_vertices=vertices, + mesh_triangles=triangles, + translation_axis=push_local[0], + joint_name=joint_name, + joint_limits=(float(limits[0, 0]), float(limits[0, 1])), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={"mesh_vertices": vertices, "mesh_triangles": triangles}, + entity_id=contact_entity_id, + label=contact_entity_id, + ) + return ( + SlideGoal( + semantics=semantics, + target_pose=contact_pose, + ), + scoped_policy, + ) + + def _twist_target( + self, + step: SemanticStep, + arm: str, + policy: Mapping[str, Any], + ) -> tuple[TwistGoal, dict[str, Any]]: + """Build one Twist goal from a live revolute joint and setting map.""" + articulation = getattr(self.env.sim, "get_articulation", lambda _uid: None)( + step.object_uid + ) + if articulation is None: + raise ValueError( + f"TurnKnob requires live articulation {step.object_uid!r}." + ) + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + backend = backend_entities[0] + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == "revolute": + candidates.append((int(joint_id), joint_name, info)) + if len(candidates) != 1: + raise ValueError( + "TurnKnob grounding requires exactly one active revolute joint; " + f"found {[name for _, name, _ in candidates]}." + ) + joint_id, joint_name, joint_info = candidates[0] + child_link = str(getattr(joint_info, "child_link_name", "")) + parent_link = str(getattr(joint_info, "parent_link_name", "")) + if ( + child_link not in articulation.link_names + or parent_link not in articulation.link_names + ): + raise ValueError( + "Revolute joint metadata must identify live parent and child links." + ) + vertices, _ = articulation.get_link_vert_face(child_link) + vertices = torch.as_tensor( + vertices, dtype=torch.float32, device=self.env.device + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): + raise ValueError("TurnKnob child link has no valid grasp geometry.") + + child_pose = _batched_pose( + articulation.get_link_pose(child_link, to_matrix=True), self.env + ) + parent_pose = _batched_pose( + articulation.get_link_pose(parent_link, to_matrix=True), self.env + ) + raw_axis = joint_info.axis + raw_origin = joint_info.origin_pose + axis = torch.tensor( + raw_axis.tolist() if hasattr(raw_axis, "tolist") else raw_axis, + dtype=torch.float32, + device=self.env.device, + ).reshape(3) + origin = torch.tensor( + raw_origin.tolist() if hasattr(raw_origin, "tolist") else raw_origin, + dtype=torch.float32, + device=self.env.device, + ).reshape(4, 4) + joint_pose = torch.matmul(parent_pose, origin) + world_axis = torch.matmul(joint_pose[:, :3, :3], axis) + local_axis = torch.bmm( + child_pose[:, :3, :3].transpose(1, 2), + world_axis.unsqueeze(2), + ).squeeze(2) + local_axis = torch.nn.functional.normalize(local_axis, dim=1) + joint_origin_local = torch.bmm(torch.linalg.inv(child_pose), joint_pose)[ + :, :3, 3 + ] + if not torch.allclose( + local_axis, + local_axis[:1].expand_as(local_axis), + atol=1.0e-4, + rtol=1.0e-4, + ) or not torch.allclose( + joint_origin_local, + joint_origin_local[:1].expand_as(joint_origin_local), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched TurnKnob environments require shared local joint geometry." + ) + + agent_config = getattr(self.env, "agent_config", {}) + articulation_settings = ( + agent_config.get("articulation_settings", {}) + if isinstance(agent_config, Mapping) + else {} + ) + per_articulation = ( + articulation_settings.get(step.object_uid, {}) + if isinstance(articulation_settings, Mapping) + else {} + ) + setting_values = ( + per_articulation.get(joint_name, ()) + if isinstance(per_articulation, Mapping) + else () + ) + if ( + not isinstance(setting_values, Sequence) + or isinstance(setting_values, (str, bytes, bytearray)) + or not setting_values + ): + raise ValueError( + "TurnKnob requires explicit setting_values; ordinal settings " + "cannot be guessed from joint limits." + ) + values = torch.as_tensor( + setting_values, dtype=torch.float32, device=self.env.device + ) + if values.ndim != 1 or not torch.isfinite(values).all(): + raise ValueError("TurnKnob setting_values must be a finite list.") + setting = int(step.goal.get("target_setting", -1)) + if setting < 0 or setting >= values.numel(): + raise ValueError("TurnKnob target_setting is outside setting_values.") + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + target_qpos = values[setting] + if torch.any(target_qpos < limits[:, 0]) or torch.any( + target_qpos > limits[:, 1] + ): + raise ValueError("TurnKnob target setting violates revolute joint limits.") + qpos = articulation.get_qpos()[:, joint_id] + twist_angles = target_qpos - qpos + if not torch.allclose( + twist_angles, + twist_angles[:1].expand_as(twist_angles), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError( + "Batched TurnKnob environments require one shared twist angle." + ) + + axis_local = local_axis[0] + origin_local = joint_origin_local[0] + geometry_center = ( + vertices.min(dim=0).values + vertices.max(dim=0).values + ) * 0.5 + projections = torch.matmul(vertices - origin_local, axis_local) + axial_value = ( + projections.max() + if torch.abs(projections.max()) >= torch.abs(projections.min()) + else projections.min() + ) + center_projection = torch.dot(geometry_center - origin_local, axis_local) + grasp_position = geometry_center + ( + 0.8 * (axial_value - center_projection) * axis_local + ) + affordance = TwistAffordance( + object_label=f"{step.object_uid}:{child_link}", + grasp_position=tuple(float(value) for value in grasp_position), + axis_origin=tuple(float(value) for value in origin_local), + twist_axis=axis_local, + joint_name=joint_name, + joint_limits=(float(limits[0, 0]), float(limits[0, 1])), + ) + semantics = ObjectSemantics( + affordance=affordance, + geometry={"mesh_vertices": vertices}, + entity_id=f"{step.object_uid}:{child_link}", + label=f"{step.object_uid}:{child_link}", + ) + scoped_policy = dict(policy) + scoped_policy.update( + { + "twist_angle": float(twist_angles[0]), + "articulation_joint_name": joint_name, + "articulation_joint_id": joint_id, + "articulation_initial_qpos": qpos, + "articulation_target_qpos": target_qpos, + } + ) + log_info( + f"Twist grounding {step.id}/{arm}: joint={joint_name!r}, " + f"current={float(qpos[0]):.4f}, target={float(target_qpos):.4f}." + ) + return TwistGoal(semantics=semantics, target_pose=child_pose), scoped_policy + + def _relative_object_spacing( + self, + moved_uid: str, + reference_uid: str, + *, + axis: int | torch.Tensor, + nominal: float, + clearance: float, + ) -> float: + """Return deterministic center spacing from live object extents.""" + moved = _object(self.env, moved_uid) + reference = _object(self.env, reference_uid) + required = float(nominal) + for env_id in range(int(self.env.num_envs)): + moved_vertices = _world_vertices(moved, self.env, env_id) + reference_vertices = _world_vertices(reference, self.env, env_id) + if isinstance(axis, torch.Tensor): + direction = axis[env_id].to( + dtype=moved_vertices.dtype, + device=moved_vertices.device, + ) + moved_axis = moved_vertices[:, :2] @ direction + reference_axis = reference_vertices[:, :2] @ direction + else: + moved_axis = moved_vertices[:, axis] + reference_axis = reference_vertices[:, axis] + moved_half = (moved_axis.max() - moved_axis.min()) * 0.5 + reference_half = (reference_axis.max() - reference_axis.min()) * 0.5 + required = max( + required, + float(moved_half + reference_half) + float(clearance), + ) + return required + + def _upright_local_direction(self, step: SemanticStep) -> torch.Tensor: + axis = self._upright_local_axis(step) + if axis == "z": + return torch.tensor( + DEFAULT_INTERNAL_AXIS, + dtype=torch.float32, + device=self.env.device, + ) + entity = _object(self.env, step.object_uid) + vertices = _local_vertices(entity, self.env, 0) + if axis in {"long_axis", "short_axis"}: + axes = analyze_local_geometry_axes(vertices) + axis_index = ( + axes.long_axis_index if axis == "long_axis" else axes.short_axis_index + ) + else: + axis_index = {"x": 0, "y": 1, "z": 2}[axis] + direction = torch.zeros(3, dtype=torch.float32, device=self.env.device) + direction[axis_index] = 1.0 + return direction + + @staticmethod + def _upright_local_axis(step: SemanticStep) -> str: + align_terms = tuple( + term + for term in compile_orientation_constraint(step.goal).terms + if isinstance(term, AlignAxisConstraint) + ) + if align_terms: + return align_terms[0].local_axis + axis = str(step.goal.get("upright_local_axis", "auto")) + return "long_axis" if axis == "auto" else axis + + def _target_rotation( + self, + step: SemanticStep, + object_pose: torch.Tensor, + *, + orientation_reference_pose: torch.Tensor | None = None, + ) -> torch.Tensor: + constraint = compile_orientation_constraint(step.goal) + if not constraint.terms: + return object_pose[:, :3, :3].clone() + if ( + len(constraint.terms) == 1 + and isinstance(constraint.terms[0], MatchRotationConstraint) + and constraint.terms[0].reference == "step_start" + ): + if orientation_reference_pose is not None: + reference = _batched_pose(orientation_reference_pose, self.env) + return reference[:, :3, :3].clone() + return object_pose[:, :3, :3].clone() + goal = str(step.goal.get("orientation_goal", "none")) + align_term = next( + ( + term + for term in constraint.terms + if isinstance(term, AlignAxisConstraint) + ), + None, + ) + if align_term is not None: + goal = "upright" + if goal not in {"upright", "lay_flat", "axis_align"}: + raise ValueError(f"Unsupported orientation_goal {goal!r}.") + + entity = _object(self.env, step.object_uid) + rotations = [] + for env_id in range(int(self.env.num_envs)): + vertices = _local_vertices(entity, self.env, env_id) + longest_to_shortest = list( + analyze_local_geometry_axes(vertices).ordered_axis_indices + ) + if goal == "upright": + upright_axis = ( + align_term.local_axis + if align_term is not None + else self._upright_local_axis(step) + ) + vertical_axis = ( + int(longest_to_shortest[0]) + if upright_axis == "long_axis" + else ( + int(longest_to_shortest[-1]) + if upright_axis == "short_axis" + else {"x": 0, "y": 1, "z": 2}[upright_axis] + ) + ) + horizontal_axis = next( + int(axis) + for axis in longest_to_shortest + if int(axis) != vertical_axis + ) + elif goal == "lay_flat": + vertical_axis = int(longest_to_shortest[-1]) + horizontal_axis = int(longest_to_shortest[0]) + else: + horizontal_axis = self._aligned_local_axis( + step, + longest_to_shortest, + ) + vertical_axis = next( + int(axis) + for axis in reversed(longest_to_shortest) + if int(axis) != horizontal_axis + ) + direction = self._horizontal_orientation( + step, + object_pose, + env_id, + horizontal_axis, + ) + rotations.append( + self._world_aligned_rotation( + direction, + horizontal_axis=horizontal_axis, + vertical_axis=vertical_axis, + ) + ) + return torch.stack(rotations) + + @staticmethod + def _aligned_local_axis( + step: SemanticStep, + longest_to_shortest: Sequence[int], + ) -> int: + axis = str(step.goal.get("orientation_axis", "long_axis")) + if axis == "x": + return 0 + if axis == "y": + return 1 + if axis == "long_axis": + return int(longest_to_shortest[0]) + if axis == "short_axis": + return int(longest_to_shortest[-1]) + raise ValueError(f"Unsupported axis_align orientation_axis {axis!r}.") + + def _horizontal_orientation( + self, + step: SemanticStep, + object_pose: torch.Tensor, + env_id: int, + local_axis: int, + ) -> torch.Tensor: + align_to = step.goal.get("orientation_reference_object") + if isinstance(align_to, str) and align_to: + reference = _object(self.env, align_to) + vertices = _local_vertices(reference, self.env, env_id) + ordered = analyze_local_geometry_axes(vertices).ordered_axis_indices + requested = str(step.goal.get("orientation_axis", "long_axis")) + reference_axis = int( + ordered[-1] if requested == "short_axis" else ordered[0] + ) + reference_pose = _live_pose(self.env, align_to) + direction = reference_pose[env_id, :3, reference_axis].clone() + elif step.operator in {"arrange_line", "place_in_line"}: + arrangement = self.arrangements.get(step.id) + axis_index = 0 if arrangement is None else arrangement.axis_index + direction = torch.zeros( + 3, + dtype=object_pose.dtype, + device=object_pose.device, + ) + direction[axis_index] = 1.0 + elif str(step.goal.get("orientation_axis", "")) in {"y", "world_y"}: + direction = object_pose.new_tensor([0.0, 1.0, 0.0]) + elif str(step.goal.get("orientation_axis", "")) in {"x", "world_x"}: + direction = object_pose.new_tensor([1.0, 0.0, 0.0]) + else: + direction = object_pose[env_id, :3, local_axis].clone() + direction[2] = 0.0 + norm = torch.linalg.vector_norm(direction) + if float(norm) < 1.0e-6: + return object_pose.new_tensor([1.0, 0.0, 0.0]) + return direction / norm + + @staticmethod + def _world_aligned_rotation( + horizontal_direction: torch.Tensor, + *, + horizontal_axis: int, + vertical_axis: int, + ) -> torch.Tensor: + world_up = horizontal_direction.new_tensor([0.0, 0.0, 1.0]) + remaining_axis = ({0, 1, 2} - {horizontal_axis, vertical_axis}).pop() + columns = [torch.zeros_like(world_up) for _ in range(3)] + columns[horizontal_axis] = horizontal_direction + columns[vertical_axis] = world_up + columns[remaining_axis] = torch.linalg.cross( + world_up, + horizontal_direction, + ) + rotation = torch.stack(columns, dim=1) + if float(torch.linalg.det(rotation)) < 0.0: + rotation[:, remaining_axis] *= -1.0 + return rotation + + def _rotated_local_z_min( + self, + entity: Any, + rotation: torch.Tensor, + env_id: int, + ) -> torch.Tensor: + vertices = _local_vertices(entity, self.env, env_id) + return (vertices @ rotation.transpose(0, 1))[:, 2].min() + + def _current_eef_pose(self, arm: str) -> torch.Tensor: + """Return the live TCP pose for one logical Action Engine arm.""" + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a physical arm, got {arm!r}.") + if hasattr(self.env, "get_current_xpos_agent"): + left, right = self.env.get_current_xpos_agent() + value = left if arm == "left_arm" else right + if value is not None: + return _batched_pose(value, self.env) + + is_left = arm == "left_arm" + if not hasattr(self.env, "get_agent_arm_control_part"): + raise ValueError("Coordinated placement requires live TCP poses.") + part = self.env.get_agent_arm_control_part(is_left) + qpos = self._arm_qpos(arm) + return _batched_pose( + self.env.robot.compute_fk(qpos=qpos, name=part, to_matrix=True), + self.env, + ) + + def _press_goal( + self, + uid: str, + object_pose: torch.Tensor, + *, + semantics: ObjectSemantics | None = None, + terminal_state: str = "activated", + ) -> tuple[PressGoal, dict[str, Any]]: + """Ground the live top surface into the typed press-affordance contract.""" + if semantics is None: + semantics = self.semantics_factory(uid) + articulation = getattr( + self.env.sim, + "get_articulation", + lambda _uid: None, + )(uid) + if articulation is not None: + return self._articulation_press_goal( + uid, + articulation, + semantics, + terminal_state=terminal_state, + ) + entity = _object(self.env, uid) + reference_pose = object_pose[0] + world_position = reference_pose[:3, 3].clone() + world_position[2] = _world_vertices(entity, self.env, 0)[:, 2].max() + rotation = reference_pose[:3, :3] + local_position = rotation.transpose(0, 1) @ ( + world_position - reference_pose[:3, 3] + ) + local_axis = rotation.transpose(0, 1) @ torch.tensor( + [0.0, 0.0, -1.0], + dtype=rotation.dtype, + device=rotation.device, + ) + press_semantics = replace( + semantics, + affordance=PressAffordance( + press_axis=local_axis, + press_position=tuple(float(value) for value in local_position), + ), + ) + return ( + PressGoal( + semantics=press_semantics, + target_pose=object_pose.clone(), + ), + {}, + ) + + def _articulation_press_goal( + self, + uid: str, + articulation: Any, + semantics: ObjectSemantics, + *, + terminal_state: str, + ) -> tuple[PressGoal, dict[str, Any]]: + """Ground a calibrated prismatic button from live joint metadata.""" + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Button articulation exposes no joint metadata.") + backend = backend_entities[0] + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == "prismatic": + candidates.append((int(joint_id), joint_name, info)) + if len(candidates) != 1: + raise ValueError("Press requires exactly one active prismatic joint.") + joint_id, joint_name, joint_info = candidates[0] + child_link = str(getattr(joint_info, "child_link_name", "")) + parent_link = str(getattr(joint_info, "parent_link_name", "")) + if ( + child_link not in articulation.link_names + or parent_link not in articulation.link_names + ): + raise ValueError("Button joint must identify live parent and child links.") + + settings = getattr(self.env, "agent_config", {}).get( + "articulation_settings", {} + ) + per_articulation = ( + settings.get(uid, {}) if isinstance(settings, Mapping) else {} + ) + values = ( + per_articulation.get(joint_name, ()) + if isinstance(per_articulation, Mapping) + else () + ) + if ( + not isinstance(values, Sequence) + or isinstance(values, (str, bytes, bytearray)) + or len(values) < 2 + ): + raise ValueError( + "Press requires explicit inactive/activated joint settings." + ) + values_tensor = torch.as_tensor( + values, dtype=torch.float32, device=self.env.device + ) + if not torch.isfinite(values_tensor).all(): + raise ValueError("Press joint settings must be finite.") + if terminal_state == "activated": + target_qpos = values_tensor[-1] + elif terminal_state == "inactive": + target_qpos = values_tensor[0] + else: + raise ValueError(f"Unsupported button terminal state {terminal_state!r}.") + + qpos = articulation.get_qpos()[:, joint_id] + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + if torch.any(target_qpos < limits[:, 0]) or torch.any( + target_qpos > limits[:, 1] + ): + raise ValueError("Button target setting violates prismatic limits.") + distances = torch.abs(target_qpos - qpos) + if torch.any(distances <= 1.0e-5): + raise ValueError("Button is already at the requested terminal state.") + if not torch.allclose( + distances, + distances[:1].expand_as(distances), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError("Batched Press requires one shared press distance.") + + child_pose = _batched_pose( + articulation.get_link_pose(child_link, to_matrix=True), self.env + ) + parent_pose = _batched_pose( + articulation.get_link_pose(parent_link, to_matrix=True), self.env + ) + raw_axis = joint_info.axis + raw_origin = joint_info.origin_pose + axis = torch.tensor( + raw_axis.tolist() if hasattr(raw_axis, "tolist") else raw_axis, + dtype=torch.float32, + device=self.env.device, + ).reshape(3) + origin = torch.tensor( + raw_origin.tolist() if hasattr(raw_origin, "tolist") else raw_origin, + dtype=torch.float32, + device=self.env.device, + ).reshape(4, 4) + joint_pose = torch.matmul(parent_pose, origin) + world_axis = torch.matmul(joint_pose[:, :3, :3], axis) + direction_sign = torch.sign(target_qpos - qpos) + movement_world = torch.nn.functional.normalize( + world_axis * direction_sign[:, None], dim=1 + ) + movement_local = torch.bmm( + child_pose[:, :3, :3].transpose(1, 2), + movement_world.unsqueeze(2), + ).squeeze(2) + if not torch.allclose( + movement_local, + movement_local[:1].expand_as(movement_local), + atol=1.0e-4, + rtol=1.0e-4, + ): + raise ValueError("Batched Press requires one shared movement axis.") + + vertices, _ = articulation.get_link_vert_face(child_link) + vertices = torch.as_tensor( + vertices, dtype=torch.float32, device=self.env.device + ) + if vertices.ndim != 2 or vertices.shape[-1] != 3 or not vertices.numel(): + raise ValueError("Button child link has no contact geometry.") + press_axis = movement_local[0] + geometry_center = ( + vertices.min(dim=0).values + vertices.max(dim=0).values + ) * 0.5 + projections = torch.matmul(vertices, press_axis) + contact_projection = projections.min() + center_projection = torch.dot(geometry_center, press_axis) + press_position = ( + geometry_center + (contact_projection - center_projection) * press_axis + ) + press_semantics = replace( + semantics, + affordance=PressAffordance( + press_axis=press_axis, + press_position=tuple(float(value) for value in press_position), + ), + ) + return ( + PressGoal(semantics=press_semantics, target_pose=child_pose), + { + "press_distance": float(distances[0]), + "articulation_joint_name": joint_name, + "articulation_joint_id": joint_id, + "articulation_initial_qpos": qpos, + "articulation_target_qpos": target_qpos, + }, + ) + + def _retreat_pose( + self, + arm: str, + policy: Mapping[str, Any], + reference: torch.Tensor | None, + *, + clear_exchange: bool = False, + retreat_after_lift: bool = False, + ) -> torch.Tensor: + target = self._retreat_reference_pose(arm, reference).clone() + if retreat_after_lift: + direction: torch.Tensor | None = None + clearance_uid = policy.get("clearance_object_uid") + if isinstance(clearance_uid, str) and clearance_uid: + entity = _object(self.env, clearance_uid) + object_pose = _batched_pose( + entity.get_local_pose(to_matrix=True), + self.env, + ) + direction = target[:, :2, 3] - object_pose[:, :2, 3] + if direction is None: + direction = torch.zeros_like(target[:, :2, 3]) + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + unresolved = norm <= 1.0e-6 + if unresolved.any(): + left_base, right_base = arm_base_poses(self.env) + base = left_base if arm == "left_arm" else right_base + baseward = base[:, :2, 3] - target[:, :2, 3] + baseward_norm = torch.linalg.vector_norm(baseward, dim=1, keepdim=True) + baseward = baseward / torch.clamp(baseward_norm, min=1.0e-6) + direction = torch.where(unresolved, baseward, direction) + norm = torch.linalg.vector_norm(direction, dim=1, keepdim=True) + direction = direction / torch.clamp(norm, min=1.0e-6) + target[:, :2, 3] += direction * float(policy.get("retreat_distance", 0.10)) + return target + desired = float(self._policy_value(policy, "retreat_height")) + if clear_exchange: + _, lateral = robot_frame_axes(self.env) + direction = lateral if arm == "left_arm" else -lateral + target[:, :2, 3] += direction.to( + dtype=target.dtype, + device=target.device, + ) * float(policy.get("retreat_distance", 0.10)) + desired = max( + desired, + float(self._policy_value(policy, "minimum_retreat_height")), + ) + ceiling = float(self._policy_value(policy, "maximum_eef_height")) + height = torch.clamp(ceiling - target[:, 2, 3], min=0.0, max=desired) + target[:, 2, 3] += height + return target + + def _retreat_reference_pose( + self, + arm: str, + reference: torch.Tensor | None, + ) -> torch.Tensor: + """Resolve the live or speculative TCP pose from which retreat starts.""" + pose = reference + if pose is None and hasattr(self.env, "get_current_xpos_agent"): + left, right = self.env.get_current_xpos_agent() + pose = left if arm == "left_arm" else right + if pose is None: + raise ValueError("Retreat grounding requires a live end-effector pose.") + return _batched_pose(pose, self.env) + + def _joint_target( + self, + arm: str, + control: str, + source: str, + binding: Mapping[str, Any], + ) -> torch.Tensor: + if source in {"gripper_closed", "gripper_open"}: + value = ( + getattr(self.env, "close_state") + if source == "gripper_closed" + else getattr(self.env, "open_state") + ) + return torch.as_tensor( + value, + dtype=torch.float32, + device=self.env.device, + ) + if source == "joint_delta": + current = self._arm_qpos(arm).clone() + index = int(binding["joint_index"]) + current[:, index] += torch.deg2rad( + torch.tensor( + float(binding.get("delta_degrees", 0.0)), + device=current.device, + ) + ) + return current + initial = getattr(self.env, "init_qpos", self.env.robot.get_qpos()) + joint_ids = self._joint_ids(arm, control) + return torch.as_tensor(initial, device=self.env.device)[:, joint_ids] + + def _arm_qpos(self, arm: str) -> torch.Tensor: + if hasattr(self.env, "get_current_qpos_agent"): + left, right = self.env.get_current_qpos_agent() + return torch.as_tensor( + left if arm == "left_arm" else right, + dtype=torch.float32, + device=self.env.device, + ) + return self.env.robot.get_qpos()[:, self._joint_ids(arm, "arm")] + + def _joint_ids(self, arm: str, control: str) -> list[int]: + side = "left" if arm == "left_arm" else "right" + key = f"{side}_{'eef' if control == 'hand' else 'arm'}_joints" + return list(getattr(self.env, key, ())) + + def _explicit_pose( + self, + binding: Mapping[str, Any], + object_pose: torch.Tensor, + ) -> torch.Tensor: + reference = str(binding.get("reference", "absolute")) + target = object_pose.clone() + if reference == "absolute": + values = binding.get("position_by_env", binding.get("position")) + position = torch.as_tensor( + values, + dtype=target.dtype, + device=target.device, + ) + if position.ndim == 1: + position = position.unsqueeze(0).repeat(int(self.env.num_envs), 1) + target[:, :3, 3] = position + return target + offset = torch.as_tensor( + binding.get("offset", (0.0, 0.0, 0.0)), + dtype=target.dtype, + device=target.device, + ) + target[:, :3, 3] += offset + return target + + def _coordinated_grasps( + self, + semantics: ObjectSemantics, + object_pose: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build a deterministic opposing pair along the object's longest XY axis.""" + vertices = semantics.geometry.get("mesh_vertices") + vertices = torch.as_tensor( + vertices, + dtype=torch.float32, + device=self.env.device, + ) + lower = vertices.min(dim=0).values + upper = vertices.max(dim=0).values + axis = int(torch.argmax(upper[:2] - lower[:2]).item()) + center = (lower + upper) * 0.5 + grasp_policy = self.runtime_policy.grounding["coordinated_grasp"] + inset = max( + float(grasp_policy["minimum_inset"]), + float((upper[axis] - lower[axis]) * grasp_policy["inset_fraction"]), + ) + left = torch.eye(4, dtype=torch.float32, device=self.env.device) + right = left.clone() + left[:3, 3] = center + right[:3, 3] = center + left[axis, 3] = lower[axis] + inset + right[axis, 3] = upper[axis] - inset + # Keep TCP z horizontal and facing the object from opposite sides. + if axis == 0: + left[:3, :3] = torch.tensor( + [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=self.env.device, + ) + right[:3, :3] = torch.tensor( + [[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=self.env.device, + ) + else: + left[:3, :3] = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, -1.0, 0.0]], + device=self.env.device, + ) + right[:3, :3] = torch.tensor( + [[-1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, -1.0, 0.0]], + device=self.env.device, + ) + batch = int(self.env.num_envs) + return left.unsqueeze(0).repeat(batch, 1, 1), right.unsqueeze(0).repeat( + batch, 1, 1 + ) diff --git a/embodichain/gen_sim/action_engine/runtime/predicates.py b/embodichain/gen_sim/action_engine/runtime/predicates.py new file mode 100644 index 000000000..84006d57d --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/predicates.py @@ -0,0 +1,891 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Evaluate canonical closed-loop predicates against live environment state.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +import torch + +from .geometry_axes import analyze_local_geometry_axes + +from embodichain.gen_sim.action_engine.config import default_runtime_policy + +from .frames import relation_axes +from .robot_parts import arm_control_part + +__all__ = ["PREDICATE_TYPES", "evaluate_predicate"] + +PREDICATE_TYPES = frozenset( + { + "both_arms_at_initial_qpos", + "both_grippers_open", + "coordinated_placed", + "grippers_clear_of_object", + "held_by_both_grippers", + "object_axis_near", + "object_axis_offset_near", + "object_held", + "object_held_by_both_grippers", + "object_held_by_gripper", + "object_in_container", + "object_lifted", + "object_not_fallen", + "object_on_object", + "object_supported_by", + "object_position_near", + "object_relative_position", + "object_upright", + "object_xy_near", + "objects_collinear", + "objects_ordered", + "pressed", + "poured", + } +) +_DEFAULT_PREDICATE_FALLBACKS = default_runtime_policy("dual_ur10").predicate_fallbacks + + +def _predicate_fallbacks(env: Any) -> Mapping[str, Any]: + policy = getattr(env, "runtime_policy", None) + value = getattr(policy, "predicate_fallbacks", None) + return value if isinstance(value, Mapping) else _DEFAULT_PREDICATE_FALLBACKS + + +def _constant(env: Any, value: bool) -> torch.Tensor: + return torch.full( + (int(env.num_envs),), + value, + dtype=torch.bool, + device=env.device, + ) + + +def _entity(env: Any, uid: str) -> Any: + entity = env.sim.get_rigid_object(uid) + if entity is None: + entity = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if entity is None: + raise ValueError(f"Unknown scene entity {uid!r}.") + return entity + + +def _pose(env: Any, uid: str) -> torch.Tensor: + entity = _entity(env, uid) + pose = torch.as_tensor( + entity.get_local_pose(to_matrix=True), + dtype=torch.float32, + device=env.device, + ) + if pose.ndim == 2: + pose = pose.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + return pose + + +def _position(env: Any, uid: str) -> torch.Tensor: + return _pose(env, uid)[:, :3, 3] + + +def _world_vertices(env: Any, uid: str, env_id: int) -> torch.Tensor: + entity = _entity(env, uid) + value = entity.get_vertices(env_ids=[env_id], scale=True) + if isinstance(value, (tuple, list)): + value = value[0] + vertices = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if vertices.ndim == 3 and vertices.shape[0] == 1: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") + pose = _pose(env, uid)[env_id] + return vertices @ pose[:3, :3].transpose(0, 1) + pose[:3, 3] + + +def _projected_center_of_mass( + env: Any, + uid: str, + env_id: int, + world_vertices: torch.Tensor, +) -> torch.Tensor: + """Return the live COM projection, with a geometry-center fallback.""" + entity = _entity(env, uid) + body_data = None if entity is None else getattr(entity, "body_data", None) + com_pose = None if body_data is None else getattr(body_data, "com_pose", None) + if callable(com_pose): + com_pose = com_pose() + if com_pose is not None: + local_com = torch.as_tensor( + com_pose, + dtype=torch.float32, + device=env.device, + ) + if local_com.ndim == 1: + local_com = local_com.unsqueeze(0).repeat(int(env.num_envs), 1) + if local_com.ndim == 2 and local_com.shape[0] == int(env.num_envs): + pose = _pose(env, uid)[env_id] + return (pose[:3, :3] @ local_com[env_id, :3] + pose[:3, 3])[:2] + return ( + world_vertices[:, :2].min(dim=0).values + + world_vertices[:, :2].max(dim=0).values + ) * 0.5 + + +def _object_supported_by( + env: Any, + spec: Mapping[str, Any], + defaults: Mapping[str, Any], +) -> torch.Tensor: + """Evaluate one-frame geometric support without advancing simulation.""" + object_uid = _object(spec) + support_uid = str( + spec.get( + "support", + spec.get("reference_object", spec.get("reference", "")), + ) + ) + if not support_uid: + raise ValueError("Support predicate requires a support object uid.") + margin = float(spec.get("com_margin", defaults["support_com_margin"])) + max_gap = float(spec.get("max_vertical_gap", defaults["support_max_vertical_gap"])) + max_penetration = float( + spec.get("max_penetration", defaults["support_max_penetration"]) + ) + min_overlap = float( + spec.get("min_overlap_ratio", defaults["support_min_overlap_ratio"]) + ) + result = _constant(env, False) + for env_id in range(int(env.num_envs)): + moved = _world_vertices(env, object_uid, env_id) + support = _world_vertices(env, support_uid, env_id) + moved_lower = moved[:, :2].min(dim=0).values + moved_upper = moved[:, :2].max(dim=0).values + support_lower = support[:, :2].min(dim=0).values + support_upper = support[:, :2].max(dim=0).values + overlap_extent = torch.clamp( + torch.minimum(moved_upper, support_upper) + - torch.maximum(moved_lower, support_lower), + min=0.0, + ) + moved_extent = torch.clamp(moved_upper - moved_lower, min=1e-6) + overlap_ratio = torch.prod(overlap_extent) / torch.prod(moved_extent) + projected_center = _projected_center_of_mass( + env, + object_uid, + env_id, + moved, + ) + center_supported = torch.all( + projected_center >= support_lower + margin + ) & torch.all(projected_center <= support_upper - margin) + local_mask = torch.all( + (support[:, :2] >= moved_lower - margin) + & (support[:, :2] <= moved_upper + margin), + dim=1, + ) + if bool(local_mask.any()): + local_support_height = support[local_mask, 2].max() + else: + # Sparse meshes may have no vertex exactly under a small payload. + # Nearest vertices are a local fallback; using the mesh-wide peak + # would confuse a remote protrusion with the candidate support pose. + distances = torch.linalg.vector_norm( + support[:, :2] - projected_center, + dim=1, + ) + count = min(8, int(support.shape[0])) + local_support_height = support[ + torch.topk(distances, count, largest=False).indices, 2 + ].max() + vertical_gap = moved[:, 2].min() - local_support_height + result[env_id] = bool( + center_supported + and overlap_ratio >= min_overlap + and vertical_gap >= -max_penetration + and vertical_gap <= max_gap + ) + return result + + +def _objects(spec: Mapping[str, Any]) -> list[str]: + values = spec.get("objects", spec.get("object_uids")) + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + raise ValueError("Predicate requires a non-empty objects list.") + return [str(value) for value in values] + + +def _object(spec: Mapping[str, Any]) -> str: + value = spec.get("object", spec.get("object_uid")) + if not isinstance(value, str) or not value: + raise ValueError("Predicate requires a non-empty object uid.") + return value + + +def _local_axis_index(env: Any, uid: str, axis: Any) -> int: + name = str(axis).lower() + if name in {"x", "y", "z"}: + return {"x": 0, "y": 1, "z": 2}[name] + geometry_axis = { + "long": "long", + "long_axis": "long", + "longest": "long", + "short": "short", + "short_axis": "short", + "shortest": "short", + }.get(name) + if geometry_axis is None: + raise ValueError(f"Unsupported upright local axis {axis!r}.") + entity = _entity(env, uid) + vertices = entity.get_vertices(env_ids=[0], scale=True) + if isinstance(vertices, (tuple, list)): + vertices = vertices[0] + vertices = torch.as_tensor(vertices, dtype=torch.float32, device=env.device) + if vertices.ndim == 3: + vertices = vertices[0] + if vertices.ndim != 2 or vertices.shape[-1] != 3 or vertices.numel() == 0: + raise ValueError(f"Rigid object {uid!r} has invalid mesh vertices.") + axes = analyze_local_geometry_axes(vertices) + return axes.long_axis_index if geometry_axis == "long" else axes.short_axis_index + + +def _arm_values( + env: Any, kind: str +) -> tuple[torch.Tensor | None, torch.Tensor | None] | None: + getter = getattr(env, f"get_current_{kind}_agent", None) + if callable(getter): + left, right = getter() + values = [] + for side, value in zip(("left", "right"), (left, right)): + if value is None: + values.append(None) + continue + item = torch.as_tensor(value, device=env.device) + if kind == "xpos" and item.ndim == 2: + item = item.unsqueeze(0).repeat(int(env.num_envs), 1, 1) + elif kind == "gripper_state" and item.ndim == 1: + item = item.unsqueeze(0) + if kind == "gripper_state": + configured = getattr( + env, + "agent_gripper_state_joint_indices", + {}, + ) + indices = ( + configured.get(side) if isinstance(configured, Mapping) else None + ) + if indices is not None: + item = item[:, list(indices)] + values.append(item) + return values[0], values[1] + if kind != "gripper_state": + return None + qpos = env.robot.get_qpos() + values = [] + for side in ("left", "right"): + ids = list(getattr(env, f"{side}_eef_joints", ())) + if not ids: + return None + item = qpos[:, ids] + configured = getattr(env, "agent_gripper_state_joint_indices", {}) + indices = configured.get(side) if isinstance(configured, Mapping) else None + if indices is not None: + item = item[:, list(indices)] + values.append(item) + return values[0], values[1] + + +def _gripper_has_closed( + env: Any, + gripper: torch.Tensor, + *, + tolerance: float, +) -> torch.Tensor: + """Check closure intent without requiring an impossible empty-gripper pose.""" + gripper = gripper.to(device=env.device, dtype=torch.float32) + open_state = getattr(env, "open_state", None) + close_state = getattr(env, "close_state", None) + reference = open_state if open_state is not None else close_state + if reference is None: + return _constant(env, False) + expected = torch.as_tensor( + reference, + dtype=torch.float32, + device=env.device, + ).flatten() + repeats = (gripper.shape[-1] + expected.numel() - 1) // expected.numel() + expected = expected.repeat(repeats)[: gripper.shape[-1]] + distance = torch.linalg.vector_norm(gripper - expected, dim=-1) + if open_state is not None: + return distance > tolerance + return distance <= tolerance + + +def _object_held( + env: Any, + uid: str, + *, + owners: Mapping[str, Sequence[str | None]] | None, + states: Mapping[tuple[str, str], Any] | None, + position_tolerance: float, + gripper_tolerance: float, + required_arm: str | None = None, +) -> torch.Tensor: + """Verify registry ownership against live object, TCP, and gripper state.""" + result = _constant(env, False) + if owners is None or states is None or uid not in owners: + return result + eef_values = _arm_values(env, "xpos") + gripper_values = _arm_values(env, "gripper_state") + if eef_values is None or gripper_values is None: + return result + + object_pose = _pose(env, uid) + for arm_index, arm in enumerate(("left_arm", "right_arm")): + if required_arm is not None and arm != required_arm: + continue + state = states.get((uid, arm)) + held = ( + None if state is None else state.get_held_object(arm_control_part(env, arm)) + ) + actual_eef = eef_values[arm_index] + gripper = gripper_values[arm_index] + if held is None or actual_eef is None or gripper is None: + continue + entity_id = getattr(held.semantics, "entity_id", None) + if entity_id != uid: + continue + actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) + expected_eef = torch.bmm( + object_pose, + held.object_to_eef.to(device=env.device, dtype=object_pose.dtype), + ) + position_ok = ( + torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], dim=-1 + ) + <= position_tolerance + ) + closed = _gripper_has_closed( + env, + gripper, + tolerance=gripper_tolerance, + ) + owned = torch.tensor( + [item == arm for item in owners[uid]], + dtype=torch.bool, + device=env.device, + ) + result |= owned & position_ok & closed + return result + + +def _coordinated_held( + env: Any, + uid: str, + state: Any, + *, + position_tolerance: float, + gripper_tolerance: float, +) -> torch.Tensor: + result = _constant(env, False) + if state is None: + return result + held_relations = tuple( + state.get_held_object(arm_control_part(env, arm)) + for arm in ("left_arm", "right_arm") + ) + if any(held is None for held in held_relations): + return result + for held in held_relations: + assert held is not None + entity_id = getattr(held.semantics, "entity_id", None) + if entity_id != uid: + return result + eef_values = _arm_values(env, "xpos") + gripper_values = _arm_values(env, "gripper_state") + if eef_values is None or gripper_values is None: + return result + + object_pose = _pose(env, uid) + result = _constant(env, True) + for arm_index, held in enumerate(held_relations): + assert held is not None + if held.env_mask is not None: + result &= held.env_mask.to(device=env.device) + actual_eef = eef_values[arm_index] + gripper = gripper_values[arm_index] + if actual_eef is None or gripper is None: + return _constant(env, False) + transform = held.object_to_eef.to( + device=env.device, + dtype=object_pose.dtype, + ) + expected_eef = torch.bmm(object_pose, transform) + actual_eef = actual_eef.to(device=env.device, dtype=object_pose.dtype) + position_ok = ( + torch.linalg.vector_norm( + actual_eef[:, :3, 3] - expected_eef[:, :3, 3], + dim=-1, + ) + <= position_tolerance + ) + closed = _gripper_has_closed( + env, + gripper, + tolerance=gripper_tolerance, + ) + result &= position_ok & closed + return result + + +def evaluate_predicate( + env: Any, + spec: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None, + *, + held_owners: Mapping[str, Sequence[str | None]] | None = None, + held_states: Mapping[tuple[str, str], Any] | None = None, + coordinated_state: Any | None = None, +) -> torch.Tensor: + """Evaluate one typed predicate or a boolean predicate tree.""" + runtime = { + "held_owners": held_owners, + "held_states": held_states, + "coordinated_state": coordinated_state, + } + defaults = _predicate_fallbacks(env) + if spec is None: + return _constant(env, True) + if isinstance(spec, Sequence) and not isinstance(spec, (str, bytes, Mapping)): + result = _constant(env, True) + for term in spec: + result &= evaluate_predicate(env, term, **runtime) + return result + if not isinstance(spec, Mapping): + raise TypeError("Predicate must be a mapping or a sequence of mappings.") + op = str(spec.get("op", "")).lower() + if not op and "terms" in spec: + op = "all" + if op in {"all", "and"}: + return evaluate_predicate(env, list(spec.get("terms", ())), **runtime) + if op in {"any", "or"}: + result = _constant(env, False) + for term in spec.get("terms", ()): + result |= evaluate_predicate(env, term, **runtime) + return result + if op == "not": + return ~evaluate_predicate(env, spec.get("term"), **runtime) + + kind = str(spec.get("type", spec.get("kind", ""))).lower() + if kind in {"semantic_goal", "line_member_placed", "stack_layer_supported"}: + raise ValueError( + f"Predicate {kind!r} is a compiler marker and requires the " + "executor's grounded target." + ) + if kind in {"object_held", "object_held_by_gripper"}: + required_arm = spec.get("arm") + if required_arm in {"left", "right"}: + required_arm = f"{required_arm}_arm" + return _object_held( + env, + _object(spec), + owners=held_owners, + states=held_states, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + required_arm=str(required_arm) if required_arm else None, + ) + if kind == "handover_complete": + required_arm = spec.get("arm", "right_arm") + return _object_held( + env, + _object(spec), + owners=held_owners, + states=held_states, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + required_arm=str(required_arm), + ) + if kind in {"held_by_both_grippers", "object_held_by_both_grippers"}: + return _coordinated_held( + env, + _object(spec), + coordinated_state, + position_tolerance=float( + spec.get("position_tolerance", defaults["held_position_tolerance"]) + ), + gripper_tolerance=float( + spec.get("gripper_tolerance", defaults["held_gripper_tolerance"]) + ), + ) + if kind in {"object_position_near", "position_near"}: + position = _position(env, _object(spec)) + target = torch.as_tensor( + spec.get("target_position", spec.get("target")), + dtype=position.dtype, + device=position.device, + ) + if target.ndim == 1: + target = target.unsqueeze(0) + return torch.linalg.vector_norm(position - target, dim=-1) <= float( + spec.get("tolerance", defaults["position_tolerance"]) + ) + if kind in {"object_xy_near", "xy_near"}: + position = _position(env, _object(spec))[:, :2] + target = torch.as_tensor( + spec.get("target_xy", spec.get("target")), + dtype=position.dtype, + device=position.device, + ).reshape(-1, 2) + return torch.linalg.vector_norm(position - target, dim=-1) <= float( + spec.get("tolerance", defaults["xy_tolerance"]) + ) + if kind in {"object_relative_position", "relative_position"}: + reference_uid = spec.get("reference_object", spec.get("reference")) + if not isinstance(reference_uid, str) or not reference_uid: + raise ValueError("Relative-position predicate requires a reference object.") + relation = str(spec.get("relation", "")) + axes = relation_axes( + env, + relation, + frame=str(spec.get("relation_frame", "world")), + ) + if not axes: + raise ValueError(f"Unsupported directional relation {relation!r}.") + delta = ( + _position(env, _object(spec))[:, :2] - _position(env, reference_uid)[:, :2] + ) + minimum_distance = float(spec.get("minimum_distance", 0.0)) + result = _constant(env, True) + for axis in axes: + projection = torch.sum( + delta * axis.to(dtype=delta.dtype, device=delta.device), dim=1 + ) + result &= projection >= minimum_distance + return result + if kind in {"object_in_container", "inside"}: + position = _position(env, _object(spec)) + container = _position( + env, str(spec.get("container", spec.get("reference_object"))) + ) + xy = torch.linalg.vector_norm(position[:, :2] - container[:, :2], dim=-1) + z = position[:, 2] - container[:, 2] + return ( + (xy <= float(spec.get("xy_radius", defaults["container_xy_radius"]))) + & (z >= float(spec.get("min_z_offset", defaults["container_min_z_offset"]))) + & (z <= float(spec.get("max_z_offset", defaults["container_max_z_offset"]))) + ) + if kind in {"object_supported_by", "object_on_object", "on"}: + return _object_supported_by(env, spec, defaults) + if kind == "object_not_fallen": + axis = _pose(env, _object(spec))[:, :3, 2] + cosine = axis[:, 2].clamp(-1.0, 1.0) + return torch.arccos(cosine) <= float( + spec.get("max_tilt", defaults["not_fallen_max_tilt"]) + ) + if kind == "object_upright": + uid = _object(spec) + local_axis = spec.get("local_axis", "long_axis") + axis_index = _local_axis_index( + env, + uid, + local_axis, + ) + axis = _pose(env, uid)[:, :3, axis_index] + cosine = axis[:, 2].clamp(-1.0, 1.0) + directed = spec.get( + "directed", + str(local_axis).lower() + not in { + "long", + "long_axis", + "longest", + "short", + "short_axis", + "shortest", + }, + ) + if not isinstance(directed, bool): + raise ValueError("object_upright directed must be a boolean.") + if not directed: + cosine = cosine.abs() + return torch.arccos(cosine) <= float( + spec.get("max_tilt", defaults["upright_max_tilt"]) + ) + if kind in {"object_axis_offset_near", "object_axis_near"}: + object_position = _position(env, _object(spec)) + axis = _axis_index(spec.get("axis", "x")) + reference_uid = spec.get( + "reference_object", + spec.get("reference", spec.get("support")), + ) + if isinstance(reference_uid, str) and reference_uid: + values = object_position[:, axis] - _position(env, reference_uid)[:, axis] + else: + values = object_position[:, axis] + target = spec.get( + "target_offset", + spec.get("offset", spec.get("target", 0.0)), + ) + target_value = torch.as_tensor( + target, + dtype=values.dtype, + device=values.device, + ) + return torch.abs(values - target_value) <= float( + spec.get("tolerance", defaults["axis_tolerance"]) + ) + if kind in {"objects_collinear", "collinear"}: + positions = torch.stack( + [_position(env, uid) for uid in _objects(spec)], + dim=1, + ) + axis = 0 if str(spec.get("axis", "x")) in {"x", "world_x"} else 1 + values = positions[:, :, 1 - axis] + return values.max(dim=1).values - values.min(dim=1).values <= float( + spec.get("tolerance", defaults["collinearity_tolerance"]) + ) + if kind in {"objects_ordered", "ordered"}: + positions = torch.stack( + [_position(env, uid) for uid in _objects(spec)], + dim=1, + ) + axis = 0 if str(spec.get("axis", "x")) in {"x", "world_x"} else 1 + differences = torch.diff(positions[:, :, axis], dim=1) + tolerance = float(spec.get("tolerance", defaults["ordering_tolerance"])) + if str(spec.get("direction", "ascending")) == "descending": + return torch.all(differences <= tolerance, dim=1) + return torch.all(differences >= -tolerance, dim=1) + if kind == "object_lifted": + position = _position(env, _object(spec))[:, 2] + initial = spec.get("initial_height") + if initial is None: + initial_pose = getattr(env, "agent_initial_object_poses", {}).get( + _object(spec) + ) + if initial_pose is None: + raise ValueError("object_lifted requires an initial object pose.") + initial = initial_pose[:, 2, 3] + initial = torch.as_tensor(initial, device=position.device) + return position >= initial + float( + spec.get("min_height", defaults["minimum_lift_height"]) + ) + if kind in {"both_arms_at_initial_qpos", "arms_home"}: + current = env.robot.get_qpos() + initial = getattr(env, "init_qpos", current) + return torch.all( + torch.abs(current - initial) + <= float(spec.get("tolerance", defaults["arm_initial_qpos_tolerance"])), + dim=-1, + ) + if kind in {"both_grippers_open", "grippers_open"}: + gripper_values = _arm_values(env, "gripper_state") + if gripper_values is None: + return _constant(env, False) + results = [] + for side, value in zip(("left", "right"), gripper_values): + if value is None: + return _constant(env, False) + value = torch.as_tensor(value, dtype=torch.float32, device=env.device) + if value.ndim == 1: + value = value.unsqueeze(0).repeat(int(env.num_envs), 1) + expected = getattr(env, f"{side}_arm_init_gripper_state", env.open_state) + expected = torch.as_tensor( + expected, + dtype=torch.float32, + device=env.device, + ) + if expected.ndim == 1: + expected = expected.unsqueeze(0).repeat(int(env.num_envs), 1) + configured = getattr(env, "agent_gripper_state_joint_indices", {}) + indices = configured.get(side) if isinstance(configured, Mapping) else None + if indices is not None: + expected = expected[:, list(indices)] + results.append( + torch.linalg.vector_norm(value - expected, dim=-1) + <= float(spec.get("tolerance", defaults["gripper_state_tolerance"])) + ) + return results[0] & results[1] + if kind == "grippers_clear_of_object": + eef_values = _arm_values(env, "xpos") + if eef_values is None: + return _constant(env, False) + object_position = _position(env, _object(spec)) + clearance = float( + spec.get( + "min_distance", + spec.get("clearance", defaults["gripper_clear_min_distance"]), + ) + ) + result = _constant(env, True) + for eef in eef_values: + if eef is None: + return _constant(env, False) + result &= ( + torch.linalg.vector_norm( + eef[:, :3, 3] - object_position, + dim=-1, + ) + >= clearance + ) + return result + if kind == "pressed": + checker = getattr(env, "is_object_pressed", None) + if callable(checker): + value = checker(_object(spec), spec.get("terminal_state", "activated")) + result = torch.as_tensor(value, dtype=torch.bool, device=env.device) + return ( + result.repeat(int(env.num_envs)) + if result.ndim == 0 + else result.reshape(-1) + ) + return _constant(env, False) + if kind == "poured": + if spec.get("verification") == "action_completion": + # Reaching semantic-step verification means every required pour edge + # already completed without a fatal planning or execution failure. + return _constant(env, True) + + raw_contents = spec.get("contents", ()) + if not isinstance(raw_contents, Sequence) or isinstance( + raw_contents, (str, bytes, bytearray) + ): + raise ValueError("poured contents must be a list of observable objects.") + contents = [ + item.get("object") if isinstance(item, Mapping) else item + for item in raw_contents + ] + if not contents or any(not isinstance(uid, str) or not uid for uid in contents): + raise ValueError( + "poured requires at least one independently observable content object." + ) + if len(contents) != len(set(contents)): + raise ValueError("poured content objects must be unique.") + target = spec.get("reference_object", spec.get("container")) + if not isinstance(target, str) or not target: + raise ValueError("poured requires a target reference_object.") + transferred = _constant(env, True) + for uid in contents: + transferred &= evaluate_predicate( + env, + { + "type": "object_in_container", + "object": uid, + "container": target, + }, + **runtime, + ) + return transferred + if kind == "articulation_joint_near": + uid = _object(spec) + articulation = getattr(env.sim, "get_articulation", lambda _uid: None)(uid) + if articulation is None: + raise ValueError(f"Unknown articulation {uid!r}.") + backend_entities = getattr( + articulation, + "_entities", + getattr(articulation, "entities", ()), + ) + if not backend_entities: + raise ValueError("Articulation backend does not expose joint metadata.") + backend = backend_entities[0] + expected_joint_type = "revolute" if "target_setting" in spec else "prismatic" + candidates = [] + for joint_id in getattr( + articulation, + "active_joint_ids", + range(len(articulation.joint_names)), + ): + joint_name = str(articulation.joint_names[int(joint_id)]) + info = backend.get_joint_info(joint_name) + joint_type = ( + str(getattr(getattr(info, "joint_type", None), "name", info.joint_type)) + .rsplit(".", maxsplit=1)[-1] + .lower() + ) + if joint_type == expected_joint_type: + candidates.append((int(joint_id), joint_name)) + requested = spec.get("joint_name") + if requested is not None: + candidates = [item for item in candidates if item[1] == str(requested)] + if len(candidates) != 1: + raise ValueError( + "articulation_joint_near requires exactly one matching " + f"{expected_joint_type} joint." + ) + joint_id, _ = candidates[0] + limits = articulation.get_qpos_limits(joint_ids=[joint_id])[:, 0] + qpos = articulation.get_qpos()[:, joint_id] + target_state = spec.get("target_state") + if target_state == "open": + target = limits[:, 1] + elif target_state == "closed": + target = limits[:, 0] + elif "target_qpos" in spec: + target = torch.as_tensor( + spec["target_qpos"], dtype=torch.float32, device=env.device + ).expand_as(qpos) + elif "target_setting" in spec: + raw_values = spec.get("setting_values", ()) + if ( + not isinstance(raw_values, Sequence) + or isinstance(raw_values, (str, bytes, bytearray)) + or not raw_values + ): + raise ValueError( + "articulation_joint_near target_setting requires setting_values." + ) + values = torch.as_tensor(raw_values, dtype=torch.float32, device=env.device) + setting = int(spec["target_setting"]) + if setting < 0 or setting >= values.numel(): + raise ValueError( + "articulation_joint_near target_setting is outside setting_values." + ) + target = values[setting].expand_as(qpos) + else: + raise ValueError( + "articulation_joint_near requires open/closed target_state or " + "target_setting with setting_values." + ) + tolerance = float(spec.get("tolerance", defaults["axis_tolerance"])) + return torch.isfinite(qpos) & (torch.abs(qpos - target) <= tolerance) + if kind == "coordinated_placed": + relation = str(spec.get("relation", "on")) + reference = spec.get("support_object", spec.get("reference_object")) + translated = { + "type": ( + "object_in_container" if relation == "inside" else "object_supported_by" + ), + "object": _object(spec), + ("container" if relation == "inside" else "support"): reference, + } + return evaluate_predicate(env, translated, **runtime) + raise ValueError(f"Unsupported execution predicate {kind!r}.") + + +def _axis_index(value: Any) -> int: + axis = str(value).lower().replace("world_", "") + if axis not in {"x", "y", "z"}: + raise ValueError(f"Unsupported predicate axis {value!r}.") + return {"x": 0, "y": 1, "z": 2}[axis] diff --git a/embodichain/gen_sim/action_engine/runtime/robot_parts.py b/embodichain/gen_sim/action_engine/runtime/robot_parts.py new file mode 100644 index 000000000..fe6f71c58 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/robot_parts.py @@ -0,0 +1,34 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve semantic Action Engine arms to physical robot control parts.""" + +from __future__ import annotations + +from typing import Any + +__all__ = ["arm_control_part"] + + +def arm_control_part(env: Any, arm: str) -> str: + """Return the physical arm control part for a semantic arm name.""" + if arm not in {"left_arm", "right_arm"}: + raise ValueError(f"Expected a semantic arm, got {arm!r}.") + if hasattr(env, "get_agent_arm_control_part"): + part = env.get_agent_arm_control_part(arm == "left_arm") + if part: + return str(part) + return arm diff --git a/embodichain/gen_sim/action_engine/runtime/solver_compat.py b/embodichain/gen_sim/action_engine/runtime/solver_compat.py new file mode 100644 index 000000000..810bc8cf7 --- /dev/null +++ b/embodichain/gen_sim/action_engine/runtime/solver_compat.py @@ -0,0 +1,234 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Install solver compatibility corrections scoped to Action Engine.""" + +from __future__ import annotations + +from collections.abc import Mapping +import functools +import threading +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.sim.solvers import PytorchSolver, URSolver, URSolverCfg + +__all__ = [ + "install_action_engine_solver_compat", + "install_pytorch_solver_tcp_compat", + "install_ur5_solver_frame_compat", + "repair_action_engine_ur5_solver_cfg", +] + +_PYTORCH_INSTALL_MARKER = "_action_engine_tcp_inverse_compat_installed" +_UR5_INSTALL_MARKER = "_action_engine_ur5_frame_compat_installed" +_UR5_ANALYTIC_TO_URDF_EE = np.eye(4, dtype=np.float32) +_UR5_ANALYTIC_TO_URDF_EE[0, 3] = -0.01 +_UR_DH_FIELDS = ("d1", "a2", "a3", "d4", "d5", "d6") + + +def repair_action_engine_ur5_solver_cfg(robot_cfg: Any) -> int: + """Repair stale UR10 DH defaults before Action Engine creates a UR5 robot. + + ``SolverCfg.from_dict`` constructs a UR10 config before assigning a + non-default ``ur_type``. Generated UR5 Action Engine configs therefore + reach the environment with UR10 DH values. Repair only that exact stale + signature so explicitly calibrated parameters remain untouched. + + Args: + robot_cfg: Robot configuration whose solver configs will be inspected. + + Returns: + Number of unique solver configs repaired by this call. + """ + configured = getattr(robot_cfg, "solver_cfg", None) + candidates = ( + configured.values() if isinstance(configured, Mapping) else (configured,) + ) + stale_defaults = URSolverCfg() + stale_dh = tuple(float(getattr(stale_defaults, name)) for name in _UR_DH_FIELDS) + + repaired = 0 + visited: set[int] = set() + for solver_cfg in candidates: + cfg_id = id(solver_cfg) + if cfg_id in visited: + continue + visited.add(cfg_id) + if not isinstance(solver_cfg, URSolverCfg): + continue + ur_type = str(getattr(solver_cfg, "ur_type", "")) + if ur_type != "ur5": + continue + current_dh = tuple(float(getattr(solver_cfg, name)) for name in _UR_DH_FIELDS) + if not np.allclose(current_dh, stale_dh, rtol=0.0, atol=1.0e-12): + continue + canonical = URSolverCfg(ur_type=ur_type) + for name in _UR_DH_FIELDS: + setattr(solver_cfg, name, getattr(canonical, name)) + repaired += 1 + return repaired + + +def install_action_engine_solver_compat(robot: Any) -> int: + """Install all solver corrections required by the Action Engine runtime.""" + return install_pytorch_solver_tcp_compat(robot) + install_ur5_solver_frame_compat( + robot + ) + + +def install_pytorch_solver_tcp_compat(robot: Any) -> int: + """Correct TCP inversion on every PytorchSolver owned by ``robot``. + + The shared solver currently transposes a rotation into an overlapping + tensor view. This wrapper transforms the requested TCP pose with a proper + matrix inverse, temporarily presents an identity TCP to the original + implementation, and otherwise preserves its sampling and ranking behavior. + + Args: + robot: Initialized robot containing its private solver registry. + + Returns: + Number of solver instances wrapped by this call. + """ + solvers = getattr(robot, "_solvers", None) + if not isinstance(solvers, Mapping): + return 0 + + installed = 0 + visited: set[int] = set() + for solver in solvers.values(): + solver_id = id(solver) + if solver_id in visited: + continue + visited.add(solver_id) + if not isinstance(solver, PytorchSolver) or bool( + getattr(solver, _PYTORCH_INSTALL_MARKER, False) + ): + continue + _wrap_solver(solver) + installed += 1 + return installed + + +def _wrap_solver(solver: PytorchSolver) -> None: + original_get_ik = solver.get_ik + call_lock = threading.RLock() + + @functools.wraps(original_get_ik) + def corrected_get_ik( + target_xpos: torch.Tensor | np.ndarray, + *args: Any, + **kwargs: Any, + ) -> Any: + target = torch.as_tensor( + target_xpos, + dtype=torch.float32, + device=solver.device, + ) + tcp = torch.as_tensor( + solver.tcp_xpos, + dtype=torch.float32, + device=solver.device, + ) + link_target = target @ torch.linalg.inv(tcp) + + # The solver instance is shared by vectorized environments. Protect the + # temporary TCP substitution in case a caller plans from another thread. + with call_lock: + active_tcp = solver.tcp_xpos + solver.tcp_xpos = np.eye(4, dtype=np.float32) + try: + return original_get_ik( + target_xpos=link_target, + *args, + **kwargs, + ) + finally: + solver.tcp_xpos = active_tcp + + solver.get_ik = corrected_get_ik + setattr(solver, _PYTORCH_INSTALL_MARKER, True) + + +def install_ur5_solver_frame_compat(robot: Any) -> int: + """Align UR5 analytic IK targets with the URDF ``ee_link`` frame. + + The UR5 asset carries a fixed ``-0.01 m`` local-x offset on ``ee_link`` + that is absent from the analytic DH model. The correction is installed + only for UR5 solvers owned by an Action Engine environment. + + Args: + robot: Initialized robot containing its private solver registry. + + Returns: + Number of solver instances wrapped by this call. + """ + solvers = getattr(robot, "_solvers", None) + if not isinstance(solvers, Mapping): + return 0 + + installed = 0 + visited: set[int] = set() + for solver in solvers.values(): + solver_id = id(solver) + if solver_id in visited: + continue + visited.add(solver_id) + if ( + not isinstance(solver, URSolver) + or str(getattr(getattr(solver, "cfg", None), "ur_type", "")) != "ur5" + or bool(getattr(solver, _UR5_INSTALL_MARKER, False)) + ): + continue + _wrap_ur5_solver(solver) + installed += 1 + return installed + + +def _wrap_ur5_solver(solver: URSolver) -> None: + original_get_ik = solver.get_ik + + @functools.wraps(original_get_ik) + def corrected_get_ik( + target_xpos: torch.Tensor | np.ndarray, + *args: Any, + **kwargs: Any, + ) -> Any: + target = torch.as_tensor( + target_xpos, + dtype=torch.float32, + device=solver.device, + ) + tcp = torch.as_tensor( + solver.tcp_xpos, + dtype=torch.float32, + device=solver.device, + ) + analytic_to_urdf = torch.as_tensor( + _UR5_ANALYTIC_TO_URDF_EE, + dtype=torch.float32, + device=solver.device, + ) + corrected_target = ( + target @ torch.linalg.inv(tcp) @ torch.linalg.inv(analytic_to_urdf) @ tcp + ) + return original_get_ik(corrected_target, *args, **kwargs) + + solver.get_ik = corrected_get_ik + setattr(solver, _UR5_INSTALL_MARKER, True) diff --git a/embodichain/gen_sim/action_engine/runtime/state.py b/embodichain/gen_sim/action_engine/runtime/state.py index dbafc4748..7fdafebd7 100644 --- a/embodichain/gen_sim/action_engine/runtime/state.py +++ b/embodichain/gen_sim/action_engine/runtime/state.py @@ -19,18 +19,81 @@ from __future__ import annotations from dataclasses import dataclass, field +from types import MappingProxyType from typing import Mapping import torch from embodichain.lab.sim.atomic_actions import ( HeldObjectState, + SceneSnapshot, TaskState, ) __all__ = ["ExecutionState"] +@dataclass(frozen=True, slots=True, eq=False) +class _CollisionOverrideSceneSnapshot(SceneSnapshot): + """Keep semantic entity poses live while overriding collision poses.""" + + collision_pose_overrides: Mapping[str, torch.Tensor] = field(default_factory=dict) + + def __post_init__(self) -> None: + SceneSnapshot.__post_init__(self) + normalized: dict[str, torch.Tensor] = {} + for entity_id, pose in self.collision_pose_overrides.items(): + if entity_id not in self.collision_entity_ids: + raise ValueError( + "Collision pose overrides must reference collision entities." + ) + if ( + not isinstance(pose, torch.Tensor) + or not pose.is_floating_point() + or pose.dim() not in (2, 3) + or pose.shape[-2:] != (4, 4) + or not bool(torch.isfinite(pose).all().item()) + ): + raise ValueError( + "Collision pose overrides must be finite floating tensors " + "with shape (4, 4) or (B, 4, 4)." + ) + normalized[entity_id] = pose.detach().clone() + object.__setattr__( + self, + "collision_pose_overrides", + MappingProxyType(normalized), + ) + + def collision_obstacle_poses( + self, + *, + batch_size: int, + device: torch.device, + dtype: torch.dtype, + ) -> Mapping[str, torch.Tensor]: + """Return planner poses with intentional-contact rows parked.""" + poses = dict( + SceneSnapshot.collision_obstacle_poses( + self, + batch_size=batch_size, + device=device, + dtype=dtype, + ) + ) + for entity_id, override in self.collision_pose_overrides.items(): + pose = override.to(device=device, dtype=dtype) + if pose.shape == (4, 4): + pose = pose.unsqueeze(0).expand(batch_size, -1, -1) + elif pose.shape != (batch_size, 4, 4): + raise ValueError( + f"Collision override {entity_id!r} must match planning " + f"batch size {batch_size}." + ) + poses[entity_id] = pose.clone() + return MappingProxyType(poses) + + @dataclass(slots=True, eq=False) class ExecutionState: """Projected task state paired with the next full-robot planning seed. diff --git a/tests/gen_sim/action_engine/capabilities/__init__.py b/tests/gen_sim/action_engine/capabilities/__init__.py new file mode 100644 index 000000000..046cb429b --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Action Engine capability tests.""" diff --git a/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py new file mode 100644 index 000000000..7c27c824f --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/test_atomic_v2.py @@ -0,0 +1,469 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from copy import deepcopy +from dataclasses import dataclass +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + AtomicCapability, + StateAtom, + build_atomic_capability_registry, +) +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.grounding import ActionGrounder +from embodichain.gen_sim.action_engine.runtime.loader import load_execution_program +from embodichain.gen_sim.action_engine.runtime.models import GroundedAction +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + +from ..task_fixtures import make_task_spec +from embodichain.gen_sim.action_engine.planning.linker import link_seed_graph +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionOptions, + ActionPlan, + EndEffectorPoseGoal, + PlannerDiagnostics, + RuntimeCommandFrame, + TimedTrajectory, + TimedCommandSequence, + TrackingPolicy, +) + + +@dataclass(frozen=True, slots=True) +class _TestOptions(ActionOptions): + marker: str = "test" + + +class _TestAction: + skill_id = "test_retreat" + end_effector_roles: tuple[str, ...] = () + + +class _TestEngine: + binding_owner_id = "test-engine" + + def bind_control_parts(self, _skill_id, _endpoints): + return ActionBinding(owner_id=self.binding_owner_id) + + def plan(self, invocation, context): + assert isinstance(invocation.skill_options, _TestOptions) + positions = context.robot.qpos[:, None, :] + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=context.env_ids, + step_dt=context.require_control_dt(), + ) + return ActionPlan( + skill_id=invocation.skill_id, + plan_success=torch.ones(context.batch_size, dtype=torch.bool), + commands=TimedCommandSequence( + frames=( + RuntimeCommandFrame( + commands=(), + active_mask=torch.ones( + context.batch_size, + dtype=torch.bool, + ), + env_ids=context.env_ids, + hold_duration=trajectory.dt[:, 0], + ), + ), + env_ids=context.env_ids, + ), + joint_trajectory=trajectory, + recovery_policy=invocation.recovery_policy, + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=context.scene.version, + planned_collision_world_revision=(0,) * context.batch_size, + diagnostics=PlannerDiagnostics(backend="test"), + ) + + +class _Robot: + dof = 2 + uid = "test_robot" + control_parts = {"left_arm": [0], "right_arm": [1]} + + def get_qpos(self): + return torch.zeros((1, 2)) + + def get_joint_ids(self, *, name: str): + return self.control_parts.get(name, []) + + +class _Entity: + def get_local_pose(self, *, to_matrix: bool): + assert to_matrix + return torch.eye(4).unsqueeze(0) + + +class _Sim: + def get_rigid_object(self, _uid: str): + return _Entity() + + +def test_axis_align_uses_its_tutorial_motion_policy_base() -> None: + capability = build_atomic_capability_registry().get("AxisAlign") + + assert capability.motion_base == "AxisAlign" + + +def test_axis_align_retains_ownership_until_explicit_release() -> None: + capability = build_atomic_capability_registry().get("AxisAlign") + contract = capability.resolve_contract( + { + "object_uid": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + } + ) + + assert capability.state_effect == "hold" + assert capability.verifier_hook is None + assert contract.requires == ( + StateAtom("arm_free", arm="left_arm"), + StateAtom("object_free", object_uid="can"), + ) + assert contract.effects[-1].atom == StateAtom( + "object_held", + object_uid="can", + arm="left_arm", + ) + + +def test_single_arm_release_contract_frees_the_object_and_arm() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + contract = capability.resolve_contract( + { + "object_uid": "can", + "actor": {"mode": "required", "arm": "left_arm"}, + "control": "hand", + "target_binding": { + "kind": "joint_state", + "source": "gripper_open", + "single_release": True, + }, + } + ) + + assert contract.requires == ( + StateAtom("object_held", object_uid="can", arm="left_arm"), + ) + assert [(effect.op, effect.atom.predicate) for effect in contract.effects] == [ + ("delete", "object_held"), + ("add", "arm_free"), + ("add", "object_free"), + ] + assert contract.failure_policy == "task_required" + + +def test_single_arm_release_verifies_the_selected_gripper_is_open() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + executor = SimpleNamespace( + env=SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + open_state=(0.0, 0.0), + close_state=(1.0, 1.0), + get_current_gripper_state_agent=lambda: ( + torch.tensor([[0.0, 0.0], [0.1, 0.1]]), + torch.ones(2, 2), + ), + ), + runtime_policy=SimpleNamespace( + predicate_fallbacks={"gripper_state_tolerance": 1.0e-3} + ), + ) + outcome = SimpleNamespace( + grounded=SimpleNamespace(motion_policy={"single_release": True}) + ) + + verified = capability.verifier_hook( + executor=executor, + step=SimpleNamespace(), + arm="left_arm", + outcome=outcome, + attempted=torch.tensor([True, True]), + ) + + assert verified.tolist() == [True, False] + + +def test_single_arm_release_ignores_passive_mimic_joint_residuals() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + executor = SimpleNamespace( + env=SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + open_state=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + close_state=(0.7, -0.7, 0.7, -0.7, -0.7, 0.7), + agent_gripper_model="robotiq", + agent_gripper_state_joint_indices={"left": (0,), "right": (0,)}, + get_current_gripper_state_agent=lambda: ( + torch.tensor( + [ + [0.0, 0.01, -0.02, 0.03, -0.01, 0.02], + [0.1, 0.0, 0.0, 0.0, 0.0, 0.0], + ] + ), + torch.zeros(2, 6), + ), + ), + runtime_policy=SimpleNamespace( + predicate_fallbacks={"gripper_state_tolerance": 1.0e-3} + ), + ) + outcome = SimpleNamespace( + grounded=SimpleNamespace(motion_policy={"single_release": True}) + ) + + verified = capability.verifier_hook( + executor=executor, + step=SimpleNamespace(), + arm="left_arm", + outcome=outcome, + attempted=torch.tensor([True, True]), + ) + + assert verified.tolist() == [True, False] + + +def test_single_arm_release_uses_normalized_opening_and_physical_support() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + executor = SimpleNamespace( + env=SimpleNamespace( + device=torch.device("cpu"), + num_envs=2, + open_state=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + close_state=(0.7, -0.7, 0.7, -0.7, -0.7, 0.7), + agent_gripper_model="robotiq", + agent_gripper_state_joint_indices={"left": (0,), "right": (0,)}, + get_current_gripper_state_agent=lambda: ( + torch.tensor( + [ + [0.015, 0.2, -0.2, 0.2, -0.2, 0.2], + [0.015, 0.2, -0.2, 0.2, -0.2, 0.2], + ] + ), + torch.zeros(2, 6), + ), + ), + _support_reference_uid=lambda _step: "table", + _support_stable_for=lambda _step, _support, _active: torch.tensor( + [True, False] + ), + _entity_pose=lambda _uid: torch.eye(4).repeat(2, 1, 1), + _placement_orientation_satisfied=lambda _step, _pose: torch.tensor( + [False, False] + ), + runtime_policy=SimpleNamespace( + predicate_fallbacks={"gripper_state_tolerance": 1.0e-3} + ), + ) + planner_trace: dict[str, object] = {} + outcome = SimpleNamespace( + grounded=SimpleNamespace(motion_policy={"single_release": True}), + planner_trace=planner_trace, + ) + + verified = capability.verifier_hook( + executor=executor, + step=SimpleNamespace(object_uid="can"), + arm="left_arm", + outcome=outcome, + attempted=torch.tensor([True, True]), + ) + + assert verified.tolist() == [True, False] + release = planner_trace["release_verification"] + assert release["open_error_fraction"] == pytest.approx([0.015 / 0.7] * 2) + assert release["gripper_open"] == [True, True] + assert release["support_stable"] == [True, False] + assert release["upright"] == [False, False] + assert release["accepted"] == [True, False] + + +def test_explicit_required_home_is_safety_required_for_any_task_type() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + base = { + "atomic_action": "MoveJoints", + "object_uid": "can", + "actor": {"mode": "required", "arm": "right_arm"}, + "control": "arm", + "role": "cleanup", + "target_binding": {"kind": "joint_state", "source": "initial"}, + } + + generic = capability.resolve_contract(base) + required_home = capability.resolve_contract( + { + **base, + "task_type": "test_carrier_consumer", + "target_binding": { + **base["target_binding"], + "operation": "custom_home", + "required_home": True, + }, + } + ) + + assert generic.failure_policy == "best_effort" + assert required_home.failure_policy == "safety_required" + + +def test_coordinated_release_contract_uses_binding_not_task_number() -> None: + capability = build_atomic_capability_registry().get("MoveJoints") + node = { + "atomic_action": "MoveJoints", + "task_type": "test_carrier_consumer", + "object_uid": "tray", + "actor": {"mode": "required", "arm": "left_arm"}, + "control": "hand", + "role": "primary", + "sync_group": "release_pair", + "target_binding": { + "kind": "joint_state", + "source": "gripper_open", + "coordinated_release_role": "participant", + }, + } + + contract = capability.resolve_contract(node) + + assert contract.requires == ( + StateAtom("object_coordinated_held", object_uid="tray"), + ) + + +def test_new_descriptor_reuses_loader_and_adapter_without_dispatch_changes() -> None: + registry = build_atomic_capability_registry() + calls = [] + + def target_hook(**kwargs): + calls.append("target") + pose = kwargs["object_pose"].clone() + return GroundedAction( + action_class="TestRetreat", + arm=kwargs["arm"], + control="arm", + target=EndEffectorPoseGoal(xpos=pose), + cfg=kwargs["policy"], + object_pose=pose, + target_object_pose=pose, + motion_policy=kwargs["policy"], + ) + + def config_hook(**_kwargs): + calls.append("config") + return _TestOptions() + + registry.register( + AtomicCapability( + "TestRetreat", + _TestAction, + _TestOptions, + frozenset({"policy_pose"}), + frozenset({"arm"}), + "single_arm", + "preserve", + "eef_pose", + motion_base="MoveEndEffector", + target_materializer_hook=target_hook, + config_materializer_hook=config_hook, + contract_resolver_hook=registry.get( + "MoveEndEffector" + ).contract_resolver_hook, + ) + ) + task, requirements = make_task_spec("E1") + bindings = { + item["role_id"]: f"uid_{item['role_id']}" for item in requirements["objects"] + } + graph = instantiate_seed_graph(task, bindings) + graph = deepcopy(graph) + graph["capability_catalog_hash"] = registry.catalog_hash() + cleanup = next( + node for node in graph["nodes"] if node["atomic_action"] == "MoveEndEffector" + ) + cleanup["atomic_action"] = "TestRetreat" + cleanup.pop("contract") + for group in graph["task_groups"]: + group.pop("contract") + graph["metadata"].pop("action_contract_linker") + graph = link_seed_graph(graph, registry=registry) + + program = load_execution_program(graph, registry=registry) + assert any( + action["atomic_action_class"] == "TestRetreat" + for edge in program.edges + for action in edge.actions + ) + + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + robot=_Robot(), + sim=_Sim(), + agent_robot_profile="dual_ur10", + get_agent_arm_control_part=lambda is_left: ( + "left_arm" if is_left else "right_arm" + ), + get_agent_eef_control_part=lambda _is_left: None, + ) + adapter = AtomicActionAdapter( + env, + grasp_policy={}, + capability_registry=registry, + ) + adapter._atomic_engine = _TestEngine() + step = next( + step + for step in program.semantic_steps + if any( + action["atomic_action_class"] == "TestRetreat" + for edge_id in step.edge_ids + for action in next( + edge for edge in program.edges if edge.id == edge_id + ).actions + ) + ) + action = next( + action + for edge in program.edges + for action in edge.actions + if action["atomic_action_class"] == "TestRetreat" + ) + grounder = ActionGrounder( + program, + env, + lambda _uid: None, + capability_registry=registry, + ) + state = ExecutionState(last_qpos=torch.zeros((1, 2))) + grounded = grounder.ground(action, step, arm="left_arm", state=state) + outcome = adapter.plan( + grounded, + state, + ) + assert outcome.success.tolist() == [True] + assert calls == ["target", "config"] diff --git a/tests/gen_sim/action_engine/capabilities/test_held_hand_over.py b/tests/gen_sim/action_engine/capabilities/test_held_hand_over.py new file mode 100644 index 000000000..113bcb631 --- /dev/null +++ b/tests/gen_sim/action_engine/capabilities/test_held_hand_over.py @@ -0,0 +1,311 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Focused contracts for GenSim's receiver-hold handover action.""" + +from __future__ import annotations + +from unittest.mock import Mock + +import pytest +import torch + +from embodichain.gen_sim.action_engine.capabilities import ( + HeldObjectHandOver, + HeldObjectHandOverOptions, + build_atomic_capability_registry, +) +from embodichain.lab.sim.atomic_actions import ( + ActionInvocation, + AntipodalAffordance, + AtomicActionEngine, + ControlPartCommandProfile, + GraspGoal, + HeldObjectState, + MotionPolicy, + ObjectSemantics, + PlanningContext, + RobotObservation, + SceneSnapshot, + TaskState, +) +from embodichain.lab.sim.planners import MotionGenerator +from embodichain.toolkits.graspkit import ( + ParallelJawGraspPoseGenerator, + ParallelJawGripperModelCfg, +) + +_HAND_DOF = 1 +_ROBOT_DOF = 6 +_CONTROL_DT = 1.0 / 60.0 + + +class _GraspGenerator(ParallelJawGraspPoseGenerator): + def __init__(self) -> None: + super().__init__(ParallelJawGripperModelCfg(model_id="test_gripper")) + + def get_valid_grasp_poses( + self, + *, + mesh_vertices: torch.Tensor, + mesh_triangles: torch.Tensor, + obj_poses: torch.Tensor, + approach_direction: torch.Tensor, + obj_longest_axis: torch.Tensor | None = None, + is_positive_part: bool | torch.Tensor = True, + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + del ( + mesh_vertices, + mesh_triangles, + approach_direction, + obj_longest_axis, + is_positive_part, + ) + return [ + (torch.eye(4).unsqueeze(0), torch.zeros(1)) + for _ in range(obj_poses.shape[0]) + ] + + def get_best_grasp_poses(self, **kwargs: object): + poses = kwargs["obj_poses"] + assert isinstance(poses, torch.Tensor) + return ( + torch.ones(poses.shape[0], dtype=torch.bool), + poses, + torch.zeros(poses.shape[0]), + ) + + def get_dual_arm_valid_grasp_poses(self, **kwargs: object): + del kwargs + raise AssertionError("Receiver-hold handover uses one destination grasp.") + + +def _motion_generator() -> MotionGenerator: + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = _ROBOT_DOF + robot.control_parts = { + "left_arm": object(), + "left_hand": object(), + "right_arm": object(), + "right_hand": object(), + } + joint_ids = { + "left_arm": [0, 1], + "left_hand": [2], + "right_arm": [3, 4], + "right_hand": [5], + } + robot.get_joint_ids.side_effect = lambda name: list(joint_ids[name]) + robot.get_qpos.return_value = torch.zeros(1, _ROBOT_DOF) + robot.compute_fk.side_effect = lambda qpos, **_kwargs: torch.eye(4).repeat( + qpos.shape[0], 1, 1 + ) + + generator = object.__new__(MotionGenerator) + generator.robot = robot + generator.device = torch.device("cpu") + generator.planner = Mock() + generator.planner.cfg.planner_type = "stub" + generator.planner.collision_world_info = None + generator.planner.preserve_plan_samples = False + return generator + + +def _engine() -> AtomicActionEngine: + generator = _motion_generator() + profiles = { + hand: ControlPartCommandProfile.joint_positions( + open=torch.zeros(_HAND_DOF), + grasp=torch.ones(_HAND_DOF), + ) + for hand in ("left_hand", "right_hand") + } + grasp = _GraspGenerator() + engine = AtomicActionEngine( + generator, + control_profiles=profiles, + grasp_pose_generators={"left_hand": grasp, "right_hand": grasp}, + load_builtins=False, + ) + engine.register(HeldObjectHandOver()) + return engine + + +def _semantics(label: str = "can") -> ObjectSemantics: + return ObjectSemantics( + label=label, + entity_id=label, + geometry={}, + affordance=AntipodalAffordance( + object_label=label, + mesh_vertices=torch.tensor( + [[0.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.0, 0.1, 0.0]] + ), + mesh_triangles=torch.tensor([[0, 1, 2]]), + ), + ) + + +def _context(semantics: ObjectSemantics) -> PlanningContext: + relation = torch.eye(4).unsqueeze(0) + task = TaskState( + batch_size=1, + device="cpu", + held_objects={ + "left_arm": HeldObjectState( + semantics=semantics, + object_to_eef=relation, + grasp_xpos=relation, + ) + }, + ) + qpos = torch.zeros(1, _ROBOT_DOF) + return PlanningContext( + robot=RobotObservation(0.0, qpos, torch.zeros_like(qpos)), + task=task, + scene=SceneSnapshot.empty(), + env_ids=torch.tensor([0]), + control_dt=_CONTROL_DT, + ) + + +def _invocation( + engine: AtomicActionEngine, + semantics: ObjectSemantics, + *, + final_x: float = 0.0, +) -> ActionInvocation: + middle = torch.eye(4) + final = middle.clone() + final[0, 3] = final_x + binding = engine.bind_control_parts( + "hand_over", + { + "source": {"motion": "left_arm", "grasp": "left_hand"}, + "destination": {"motion": "right_arm", "grasp": "right_hand"}, + }, + ) + return ActionInvocation( + skill_id="hand_over", + goal=GraspGoal(semantics=semantics), + binding=binding, + motion_policy=MotionPolicy(sample_count=24), + skill_options=HeldObjectHandOverOptions( + middle_object_pose=middle, + final_object_pose=final, + hand_interp_steps=2, + hold_steps=2, + retreat_steps=4, + ), + ) + + +def _install_planner(monkeypatch: pytest.MonkeyPatch) -> None: + def plan( + _generator: MotionGenerator, + _control_part: str, + start_qpos: torch.Tensor, + _target_poses: torch.Tensor, + n_waypoints: int, + _motion_policy: MotionPolicy, + _control_dt: float | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.ones(start_qpos.shape[0], dtype=torch.bool), + start_qpos.unsqueeze(1).repeat(1, n_waypoints, 1), + ) + + monkeypatch.setattr( + "embodichain.gen_sim.action_engine.capabilities.held_hand_over." + "plan_named_arm_trajectory", + plan, + ) + + +def test_handover_capability_matches_installed_receiver_hold_action() -> None: + capability = build_atomic_capability_registry().get("HandOver") + + assert capability.action_type is HeldObjectHandOver + assert capability.config_type is HeldObjectHandOverOptions + assert capability.action_type.GoalType is GraspGoal + + +def test_receiver_hold_plan_transfers_ownership_and_preserves_phase_order( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_planner(monkeypatch) + engine = _engine() + semantics = _semantics() + context = _context(semantics) + + plan = engine.plan(_invocation(engine, semantics), context) + projected = plan.expected_effects.apply(context.task, plan.plan_success) + + assert plan.plan_success.tolist() == [True] + assert projected.get_held_object("left_arm") is None + received = projected.get_held_object("right_arm") + assert received is not None + assert received.semantics.entity_id == "can" + assert [segment.name for segment in plan.segments] == [ + "transfer", + "approach", + "close", + "hold", + "release", + "retreat", + ] + assert plan.joint_trajectory is not None + positions = plan.joint_trajectory.positions + close = plan.segment("close") + release = plan.segment("release") + retreat = plan.segment("retreat") + torch.testing.assert_close(positions[:, close.stop - 1, 5], torch.ones(1)) + torch.testing.assert_close(positions[:, release.stop - 1, 2], torch.zeros(1)) + torch.testing.assert_close( + positions[:, release.start :, 5], + torch.ones_like(positions[:, release.start :, 5]), + ) + torch.testing.assert_close( + positions[:, retreat.start : retreat.stop, 3:5], + positions[:, retreat.start : retreat.start + 1, 3:5].expand_as( + positions[:, retreat.start : retreat.stop, 3:5] + ), + ) + + +def test_receiver_hold_rejects_delivery_away_from_exchange( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_planner(monkeypatch) + engine = _engine() + semantics = _semantics() + + with pytest.raises(ValueError, match="receiver remains stationary"): + engine.plan(_invocation(engine, semantics, final_x=0.1), _context(semantics)) + + +def test_receiver_hold_rejects_a_different_requested_object( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _install_planner(monkeypatch) + engine = _engine() + + with pytest.raises(ValueError, match="object held by the source"): + engine.plan( + _invocation(engine, _semantics("other")), + _context(_semantics("held")), + ) diff --git a/tests/gen_sim/action_engine/runtime/__init__.py b/tests/gen_sim/action_engine/runtime/__init__.py new file mode 100644 index 000000000..66ef3ba1a --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +"""Runtime contract tests for Action Engine.""" diff --git a/tests/gen_sim/action_engine/runtime/test_actions.py b/tests/gen_sim/action_engine/runtime/test_actions.py new file mode 100644 index 000000000..6799b6298 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_actions.py @@ -0,0 +1,1834 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Focused contracts for the public atomic-action adapter.""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import replace +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +from embodichain.gen_sim.action_engine.capabilities import HeldObjectHandOver +from embodichain.gen_sim.action_engine.runtime import actions +from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ActionEngineMoveJoints, + ExactTargetMoveHeldObject, +) +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.models import ( + ActionOutcome, + GroundedAction, +) +from embodichain.gen_sim.action_engine.runtime.state import ExecutionState +from embodichain.gen_sim.action_engine.gripper_profiles import get_gripper_profile +from embodichain.lab.sim.atomic_actions import ( + Affordance, + ActionBinding, + ActionPlan, + AntipodalAffordance, + AxisAlignAffordance, + AxisAlignGoal, + CoordinatedPickGoal, + EndEffectorPoseGoal, + EntityState, + GraspGoal, + HeldObjectPoseGoal, + HeldObjectState, + JointPositionGoal, + ObjectSemantics, + PlannerDiagnostics, + PlanningFailure, + RecoveryPolicy, + RuntimeCommandFrame, + SceneSnapshot, + StateDelta, + TimedCommandSequence, + TimedTrajectory, + TrackingPolicy, +) +from embodichain.lab.sim.planners import CuroboPlannerCfg, ToppraPlannerCfg +from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator +from embodichain.utils.logger import log_warning + + +class _MeshEntity: + def get_vertices(self, *, env_ids: list[int], scale: bool) -> torch.Tensor: + assert env_ids == [0] + assert scale + return torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + ], + dtype=torch.float32, + ) + + def get_triangles(self, *, env_ids: list[int]) -> torch.Tensor: + assert env_ids == [0] + return torch.tensor([[0, 1, 2]], dtype=torch.int64) + + +def _cuboid_vertices(x: float, y: float, z: float) -> torch.Tensor: + return torch.tensor( + [ + [sx * x, sy * y, sz * z] + for sx in (-1.0, 1.0) + for sy in (-1.0, 1.0) + for sz in (-1.0, 1.0) + ], + dtype=torch.float32, + ) + + +def _rotation_x(degrees: float) -> torch.Tensor: + angle = torch.deg2rad(torch.tensor(degrees, dtype=torch.float32)) + rotation = torch.eye(3) + rotation[1, 1] = torch.cos(angle) + rotation[1, 2] = -torch.sin(angle) + rotation[2, 1] = torch.sin(angle) + rotation[2, 2] = torch.cos(angle) + return rotation + + +def _rotation_z(degrees: float) -> torch.Tensor: + angle = torch.deg2rad(torch.tensor(degrees, dtype=torch.float32)) + rotation = torch.eye(3) + rotation[0, 0] = torch.cos(angle) + rotation[0, 1] = -torch.sin(angle) + rotation[1, 0] = torch.sin(angle) + rotation[1, 1] = torch.cos(angle) + return rotation + + +class _PoseEntity: + def __init__(self, pose: torch.Tensor) -> None: + self.pose = pose + + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return self.pose.clone() + + +class _PlannerRobot: + uid = "test_robot" + dof = 8 + + _ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": [2, 3], + "physical_right_arm": [4, 5], + "physical_right_eef": [6, 7], + } + control_parts = _ids + + def get_joint_ids(self, *, name: str) -> list[int]: + return list(self._ids[name]) + + def get_control_part_base_pose(self, *, name: str, to_matrix: bool) -> torch.Tensor: + assert to_matrix + pose = torch.eye(4).repeat(2, 1, 1) + pose[:, 1, 3] = 0.3 if name == "physical_left_arm" else -0.3 + return pose + + def compute_batch_ik( + self, + *, + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + del pose, name + return torch.ones(joint_seed.shape[:2], dtype=torch.bool), joint_seed.clone() + + +def _commands_for(trajectory: TimedTrajectory) -> TimedCommandSequence: + """Build timing-only frames for retained test trajectories.""" + active = torch.ones( + trajectory.batch_size, + dtype=torch.bool, + device=trajectory.positions.device, + ) + frames = tuple( + RuntimeCommandFrame( + commands=(), + active_mask=active, + env_ids=trajectory.env_ids, + hold_duration=trajectory.dt[:, index], + ) + for index in range(trajectory.waypoint_count) + ) + return TimedCommandSequence(frames=frames, env_ids=trajectory.env_ids) + + +def _planner_diagnostics(success: torch.Tensor) -> PlannerDiagnostics: + failure = None if bool(success.all()) else PlanningFailure("planning_failed") + return PlannerDiagnostics(backend="fake", failure=failure) + + +class _FakeEngine: + """Minimal endpoint-binding and planning surface for adapter unit tests.""" + + binding_owner_id = "action-engine-test" + + def __init__(self, plan=None) -> None: + self._plan = plan + + def bind_control_parts(self, _skill_id, _endpoints) -> ActionBinding: + return ActionBinding(owner_id=self.binding_owner_id) + + def plan(self, invocation, context) -> ActionPlan: + if self._plan is None: + raise AssertionError("This fake engine has no planning callback.") + return self._plan(invocation, context) + + +def test_adapter_registers_gen_sim_compat_actions( + monkeypatch: Any, +) -> None: + registered: list[tuple[type, bool]] = [] + + class Engine: + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + pass + + def register(self, action: Any, *, replace: bool = False) -> None: + registered.append((type(action), replace)) + + adapter = AtomicActionAdapter(_planner_env()) + monkeypatch.setattr(actions, "AtomicActionEngine", Engine) + monkeypatch.setattr(adapter, "_generator", lambda: object()) + monkeypatch.setattr(adapter, "_control_profiles", lambda: {}) + monkeypatch.setattr(adapter, "_grasp_pose_generators", lambda **_kwargs: {}) + + engine = adapter._engine() + + assert isinstance(engine, Engine) + assert registered == [ + (ExactTargetMoveHeldObject, True), + (ActionEngineMoveJoints, True), + (HeldObjectHandOver, True), + ] + + +def test_free_yaw_search_uses_an_internal_reachability_sample_set( + monkeypatch: pytest.MonkeyPatch, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + target = torch.eye(4).repeat(2, 1, 1) + target[:, :3, 3] = torch.tensor([0.1, -0.2, 0.9]) + semantics = ObjectSemantics( + label="can", + entity_id="can", + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + state = ExecutionState( + last_qpos=torch.zeros(2, 8), + held_objects={"physical_left_arm": held}, + ) + grounded = GroundedAction( + "MoveHeldObject", + "left_arm", + "arm", + HeldObjectPoseGoal(object_target_pose=target), + {}, + target_object_pose=target, + allow_yaw_search=True, + ) + + def compute_batch_ik( + *, + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + assert name == "physical_left_arm" + assert pose.shape[:2] == (2, 8) + success = torch.zeros(2, 8, dtype=torch.bool) + success[:, 3] = True + return success, joint_seed.clone() + + monkeypatch.setattr(env.robot, "compute_batch_ik", compute_batch_ik) + + selected = adapter._select_transport_yaw(grounded, state) + + assert selected.target_object_pose is not None + torch.testing.assert_close(selected.target_object_pose[:, :3, 3], target[:, :3, 3]) + torch.testing.assert_close( + selected.target_object_pose[:, :3, :3], + torch.tensor([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]).repeat( + 2, 1, 1 + ), + atol=1.0e-6, + rtol=1.0e-6, + ) + + def all_yaws_reachable( + *, + pose: torch.Tensor, + name: str, + joint_seed: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + del pose, name + qpos = joint_seed.clone() + qpos[:, 0] += 1.0 + return torch.ones(2, 8, dtype=torch.bool), qpos + + monkeypatch.setattr(env.robot, "compute_batch_ik", all_yaws_reachable) + + minimum_rotation = adapter._select_transport_yaw(grounded, state) + + assert minimum_rotation.target_object_pose is not None + torch.testing.assert_close(minimum_rotation.target_object_pose, target) + + +def test_adapter_lowers_axis_align_from_live_pose_with_a_stable_seed() -> None: + vertices = torch.tensor( + [[x, y, z] for x in (-0.03, 0.03) for y in (-0.06, 0.06) for z in (-0.03, 0.03)] + ) + semantics = ObjectSemantics( + label="can", + entity_id="can", + geometry={}, + affordance=AxisAlignAffordance( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2]]), + internal_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + ) + + live_pose = torch.eye(4).repeat(2, 1, 1) + live_pose[:, 2, 3] = 1.10 + parked_pose = live_pose.clone() + parked_pose[:, 2, 3] -= 100.0 + sampled_poses: list[torch.Tensor] = [] + sampled_random_values: list[float] = [] + + class Generator: + def get_valid_grasp_poses(self, **kwargs: Any): + poses = kwargs["obj_poses"].clone() + sampled_poses.append(poses) + sampled_random_values.append(float(torch.rand(()))) + return [ + (pose.unsqueeze(0), torch.tensor([0.1])) for pose in poses.unbind(dim=0) + ] + + adapter = AtomicActionAdapter(_planner_env()) + adapter._atomic_engine = SimpleNamespace( + grasp_pose_generators={"physical_left_eef": Generator()} + ) + grounded = GroundedAction( + "AxisAlign", + "left_arm", + "arm", + AxisAlignGoal(semantics=semantics), + {}, + object_pose=live_pose, + object_uid="can", + ) + contexts = [ + SimpleNamespace( + robot=SimpleNamespace(qpos=torch.zeros(2, 8)), + scene=SceneSnapshot( + timestamp=0.0, + version=version, + entities={"can": EntityState(parked_pose)}, + ), + ) + for version in (3, 97) + ] + + adaptations = [ + adapter._adapt_axis_align_body_grasps( + grounded, + context, + adapter.capabilities.get("AxisAlign"), + ) + for context in contexts + ] + adapted_items = adaptations[0] + adapted = adapted_items[0] + + assert len(adapted_items) == 1 + assert isinstance(adapted.target, AxisAlignGoal) + assert adapted.target.semantics.entity_id == "can" + assert adapted.target.grasp_xpos is not None + assert adapted.motion_policy["body_grasp"]["long_axis_index"] == 1 + assert adapted.motion_policy["body_grasp"]["candidate_counts"] == [1, 1] + assert len(sampled_poses) == 2 + torch.testing.assert_close(sampled_poses[0], live_pose) + torch.testing.assert_close(sampled_poses[1], live_pose) + assert not torch.equal(sampled_poses[0], parked_pose) + assert sampled_random_values[0] == sampled_random_values[1] + torch.testing.assert_close( + adaptations[0][0].target.grasp_xpos, + adaptations[1][0].target.grasp_xpos, + ) + + +def test_axis_align_body_grasp_does_not_fall_back_to_scene_snapshot_pose() -> None: + vertices = torch.tensor( + [[x, y, z] for x in (-0.03, 0.03) for y in (-0.06, 0.06) for z in (-0.03, 0.03)] + ) + grounded = GroundedAction( + "AxisAlign", + "left_arm", + "arm", + AxisAlignGoal( + semantics=ObjectSemantics( + label="can", + entity_id="can", + geometry={}, + affordance=AxisAlignAffordance( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2]]), + internal_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + ) + ), + {}, + object_uid="can", + ) + parked_pose = torch.eye(4).repeat(2, 1, 1) + parked_pose[:, 2, 3] = -98.9 + context = SimpleNamespace( + scene=SceneSnapshot( + timestamp=0.0, + version=3, + entities={"can": EntityState(parked_pose)}, + ) + ) + adapter = AtomicActionAdapter(_planner_env()) + + with pytest.raises(ValueError, match="grounded live object pose"): + adapter._adapt_axis_align_body_grasps( + grounded, + context, + adapter.capabilities.get("AxisAlign"), + ) + + +def _planner_env( + *, + table: Any | None = None, + rigid_objects: dict[str, Any] | None = None, + gripper_model: str = "pgi", +) -> SimpleNamespace: + entities = dict(rigid_objects or {}) + if table is not None: + entities["table"] = table + return SimpleNamespace( + num_envs=2, + device=torch.device("cpu"), + robot=_PlannerRobot(), + sim=SimpleNamespace(get_rigid_object=entities.get), + left_arm_joints=[0, 1], + left_eef_joints=[2, 3], + right_arm_joints=[4, 5], + right_eef_joints=[6, 7], + open_state=torch.zeros(2), + close_state=torch.ones(2), + agent_gripper_model=gripper_model, + get_agent_arm_control_part=lambda is_left: ( + "physical_left_arm" if is_left else "physical_right_arm" + ), + get_agent_eef_control_part=lambda is_left: ( + "physical_left_eef" if is_left else "physical_right_eef" + ), + ) + + +def test_semantics_builds_one_mesh_only_affordance() -> None: + entity = _MeshEntity() + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + sim=SimpleNamespace( + get_rigid_object=lambda uid: entity if uid == "cube" else None + ), + ) + + adapter = AtomicActionAdapter(env) + first = adapter.semantics("cube") + second = adapter.semantics("cube") + + assert first is second + assert first.entity_id == "cube" + assert isinstance(first.affordance, AntipodalAffordance) + assert first.affordance.mesh_vertices is not None + assert first.affordance.mesh_triangles is not None + assert first.affordance.mesh_vertices.dtype == torch.float32 + assert first.affordance.mesh_triangles.dtype == torch.int64 + + +def test_planner_policy_uses_curobo_for_single_arm_and_ik_for_dual_arm() -> None: + adapter = AtomicActionAdapter(_planner_env()) + adapter._atomic_engine = _FakeEngine() + goal = JointPositionGoal(target=torch.zeros(2, 2)) + + single = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "arm", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + coordinated_goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="tray", + entity_id="tray", + geometry={}, + affordance=AntipodalAffordance(), + ), + object_target_pose=torch.eye(4), + object_initial_pose=torch.eye(4), + ) + coordinated = adapter._invocation( + GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + coordinated_goal, + {}, + ), + adapter.capabilities.get("CoordinatedPickment"), + ) + hand = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "hand", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + + assert adapter.planner_policy["backend"] == "curobo" + assert single.motion_policy.strategy == "motion_gen" + assert coordinated.motion_policy.strategy == "ik_interp" + assert torch.allclose( + coordinated.skill_options.left_to_right_arm_direction, + torch.tensor([0.0, -1.0, 0.0]), + ) + assert hand.motion_policy.strategy == "ik_interp" + + +@pytest.mark.parametrize( + ("planner_policy", "single_strategy"), + [ + ( + { + "backend": "toppra", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "dynamic_collision": False, + }, + "motion_gen", + ), + ( + { + "backend": "toppra", + "single_arm_strategy": "ik_interp", + "coordinated_strategy": "ik_interp", + "dynamic_collision": False, + }, + "ik_interp", + ), + ], +) +def test_toppra_and_ik_interp_policies_reach_runtime_factory_and_strategy( + planner_policy: dict[str, object], + single_strategy: str, +) -> None: + adapter = AtomicActionAdapter(_planner_env(), planner_policy=planner_policy) + adapter._atomic_engine = _FakeEngine() + goal = JointPositionGoal(target=torch.zeros(2, 2)) + + invocation = adapter._invocation( + GroundedAction("MoveJoints", "left_arm", "arm", goal, {}), + adapter.capabilities.get("MoveJoints"), + ) + + assert isinstance(adapter._motion_generator_cfg().planner_cfg, ToppraPlannerCfg) + assert invocation.motion_policy.strategy == single_strategy + assert invocation.motion_policy.dynamic_collision_mode.value == "off" + + +def test_runtime_rejects_requested_and_effective_backend_mismatch( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter( + env, + planner_policy={ + "backend": "toppra", + "single_arm_strategy": "motion_gen", + "coordinated_strategy": "ik_interp", + "dynamic_collision": False, + }, + ) + trajectory = TimedTrajectory.from_uniform_step( + torch.zeros(2, 2, 8), + env_ids=torch.arange(2), + step_dt=0.01, + ) + wrong_backend_plan = ActionPlan( + skill_id="move_joints", + plan_success=torch.ones(2, dtype=torch.bool), + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="curobo"), + expected_effects=StateDelta(), + ) + monkeypatch.setattr( + adapter, + "_engine", + lambda: _FakeEngine(lambda *_args: wrong_backend_plan), + ) + + with pytest.raises(RuntimeError, match="requested 'toppra'.*reported 'curobo'"): + adapter.plan( + GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ), + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + +def test_coordinated_pickment_uses_engine_scoped_grasp_generator() -> None: + adapter = AtomicActionAdapter(_planner_env()) + adapter._atomic_engine = _FakeEngine() + affordance = AntipodalAffordance() + goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="tray", + entity_id="tray", + geometry={}, + affordance=affordance, + ), + object_target_pose=torch.eye(4), + object_initial_pose=torch.eye(4), + ) + grounded = GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + goal, + { + "middle_empty_ratio": 0.7, + "is_filter_ground_collision": False, + }, + ) + + invocation = adapter._invocation( + grounded, + adapter.capabilities.get("CoordinatedPickment"), + ) + + scoped_affordance = invocation.goal.semantics.affordance + assert isinstance(scoped_affordance, AntipodalAffordance) + assert scoped_affordance is affordance + assert invocation.skill_options.middle_empty_ratio == pytest.approx(0.7) + + +def _coordinated_grounded( + rotation: torch.Tensor, + *, + vertices: torch.Tensor | None = None, +) -> GroundedAction: + object_pose = torch.eye(4).repeat(2, 1, 1) + object_pose[:, :3, :3] = rotation + object_pose[:, :3, 3] = torch.tensor([0.05, 0.0, 0.75]) + mesh_vertices = _cuboid_vertices(0.03, 0.04, 0.20) if vertices is None else vertices + goal = CoordinatedPickGoal( + semantics=ObjectSemantics( + label="test_object", + entity_id="test_object", + geometry={}, + affordance=AntipodalAffordance( + object_label="test_object", + mesh_vertices=mesh_vertices, + mesh_triangles=torch.tensor([[0, 1, 2]], dtype=torch.int64), + ), + ), + object_target_pose=object_pose.clone(), + object_initial_pose=object_pose.clone(), + ) + return GroundedAction( + "CoordinatedPickment", + "coordinated", + "arm", + goal, + {"middle_empty_ratio": 0.4}, + object_pose=object_pose, + object_uid="test_object", + ) + + +def test_coordinated_pickment_geometry_candidates_are_live_and_continuous() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + + vertical = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(torch.eye(3)), capability + ) + tilted = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(_rotation_x(45.0)), capability + ) + horizontal = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(_rotation_x(90.0)), capability + ) + yawed = adapter._adapt_coordinated_pickment_grasps( + _coordinated_grounded(_rotation_z(90.0) @ _rotation_x(90.0)), + capability, + ) + + preferred = [ + next( + candidate.cfg["middle_empty_ratio"] + for candidate in candidates + if candidate.motion_policy["coordinated_grasp"]["approach_candidate_label"] + == "robot_forward_down" + ) + for candidates in (vertical, tilted, horizontal) + ] + assert preferred[0] < preferred[1] < preferred[2] + yawed_forward_down = next( + candidate + for candidate in yawed + if candidate.motion_policy["coordinated_grasp"]["approach_candidate_label"] + == "robot_forward_down" + ) + assert yawed_forward_down.cfg["middle_empty_ratio"] == pytest.approx(preferred[0]) + for candidates in (vertical, tilted, horizontal, yawed): + assert candidates + assert torch.allclose( + candidates[0].cfg["left_to_right_arm_direction"], + torch.tensor([0.0, -1.0, 0.0]), + ) + assert torch.allclose( + candidates[0].cfg["approach_direction"], + torch.tensor([2**-0.5, 0.0, -(2**-0.5)]), + ) + assert any( + torch.allclose( + candidate.cfg["approach_direction"], + torch.tensor([0.0, 0.0, -1.0]), + ) + for candidate in candidates + ) + + +def test_coordinated_pickment_approach_family_follows_live_shared_reach() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + grounded = _coordinated_grounded( + _rotation_z(90.0), + vertices=_cuboid_vertices(0.20, 0.08, 0.02), + ) + left = torch.eye(4).repeat(2, 1, 1) + right = torch.eye(4).repeat(2, 1, 1) + left[:, 0, 3] = -0.3 + right[:, 0, 3] = 0.3 + left[:, 1, 3] = right[:, 1, 3] = -0.5 + adapter._coordinated_arm_bases = lambda: (left, right) + + candidates = adapter._adapt_coordinated_pickment_grasps(grounded, capability) + + trace = candidates[0].motion_policy["coordinated_grasp"] + assert trace["approach_candidate_label"] == "robot_forward_down" + torch.testing.assert_close( + candidates[0].cfg["approach_direction"], + torch.tensor([0.0, 2**-0.5, -(2**-0.5)]), + ) + assert trace["grasp_seed"] == 17_393 + assert [item["label"] for item in trace["approach_candidates"][:3]] == [ + "robot_forward_down", + "current", + "robot_forward", + ] + assert [candidate.cfg["grasp_pair_rank"] for candidate in candidates[:3]] == [ + 0, + 1, + 2, + ] + + +def test_coordinated_pickment_rejects_boolean_pair_candidate_count() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + grounded = _coordinated_grounded(torch.eye(3)) + grounded = replace( + grounded, + cfg={**grounded.cfg, "grasp_pair_candidate_count": True}, + ) + + with pytest.raises(ValueError, match="must be an integer"): + adapter._adapt_coordinated_pickment_grasps(grounded, capability) + + +def test_coordinated_pickment_continues_after_pair_trajectory_audit_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + positions = torch.zeros(2, 2, env.robot.dof) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + successful_plan = ActionPlan( + skill_id="coordinated_pickment", + plan_success=torch.ones(2, dtype=torch.bool), + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics(backend="fake"), + ) + engine = _FakeEngine(lambda *_args: successful_plan) + engine.grasp_pose_generators = {} + monkeypatch.setattr(adapter, "_engine_for", lambda *_args: engine) + monkeypatch.setattr( + adapter, + "_coordinated_pair_selection_context", + lambda *_args: nullcontext(), + ) + monkeypatch.setattr( + adapter, + "_latest_coordinated_grasp_trace", + lambda *_args: { + "pair_selection": { + "selected": True, + "selected_left_pose": torch.eye(4).tolist(), + "selected_right_pose": torch.eye(4).tolist(), + } + }, + ) + audited_ranks: list[int] = [] + + def audit(candidate, _invocation, plan, _context, _stages): + rank = int(candidate.cfg["grasp_pair_rank"]) + audited_ranks.append(rank) + success = torch.full((2,), rank == 1, dtype=torch.bool) + diagnostics = ( + plan.diagnostics + if bool(success.all()) + else PlannerDiagnostics( + backend="fake", + failure=PlanningFailure("trajectory_safety_failed"), + ) + ) + return ( + replace(plan, plan_success=success, diagnostics=diagnostics), + {"success": success.tolist()}, + ) + + monkeypatch.setattr(adapter, "_audit_coordinated_trajectory", audit) + grounded = _coordinated_grounded(torch.eye(3)) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, env.robot.dof)), + ) + + assert audited_ranks == [0, 1] + assert outcome.success.tolist() == [True, True] + trace = outcome.planner_trace["coordinated_grasp"] + assert trace["grasp_pair_rank"] == 1 + assert len(trace["search_attempts"]) == 2 + + +def test_coordinated_pickment_geometry_candidates_are_deterministic_for_tray() -> None: + adapter = AtomicActionAdapter(_planner_env()) + capability = adapter.capabilities.get("CoordinatedPickment") + grounded = _coordinated_grounded( + _rotation_z(31.0), + vertices=_cuboid_vertices(0.20, 0.14, 0.02), + ) + + first = adapter._adapt_coordinated_pickment_grasps(grounded, capability) + second = adapter._adapt_coordinated_pickment_grasps(grounded, capability) + + assert [item.cfg["middle_empty_ratio"] for item in first] == pytest.approx( + [item.cfg["middle_empty_ratio"] for item in second] + ) + assert ( + first[0].motion_policy["coordinated_grasp"] + == second[0].motion_policy["coordinated_grasp"] + ) + + +def test_grasp_generators_follow_mainline_service_contract() -> None: + adapter = AtomicActionAdapter(_planner_env(gripper_model="pgi")) + + generators = adapter._grasp_pose_generators() + + assert set(generators) == {"physical_left_eef", "physical_right_eef"} + generator = generators["physical_left_eef"] + assert generators["physical_right_eef"] is generator + assert isinstance(generator, AntipodalGraspPoseGenerator) + assert generator.gripper_model.model_id == "dh_pgi_140_80" + assert generator.gripper_model.max_opening_width == pytest.approx(0.100) + assert generator.gripper_model.finger_length == pytest.approx(0.10) + assert generator.collision_cfg.opening_margin == pytest.approx(0.03) + assert generator.algorithm_cfg.sample_count == 10000 + assert generator.algorithm_cfg.approach_direction_samples == 4 + assert generator.algorithm_cfg.max_candidates == 500 + assert generator.collision_cfg.max_decomposition_hulls == 16 + assert generator.collision_cfg.filter_ground_collision is True + + +def test_robotiq_grasp_generator_preserves_existing_geometry() -> None: + adapter = AtomicActionAdapter(_planner_env(gripper_model="robotiq")) + + generators = adapter._grasp_pose_generators() + generator = generators["physical_left_eef"] + + assert generators["physical_right_eef"] is generator + assert generator.gripper_model.model_id == "robotiq_arg2f_140" + assert generator.gripper_model.max_opening_width == pytest.approx(0.15) + assert generator.gripper_model.finger_length == pytest.approx(0.13) + assert generator.collision_cfg.opening_margin == pytest.approx(0.01) + + +def test_coordinated_robotiq_generator_uses_e5_opening_margin_only() -> None: + adapter = AtomicActionAdapter(_planner_env(gripper_model="robotiq")) + + coordinated = adapter._grasp_pose_generators( + filter_ground_collision=False, + opening_margin=0.02, + )["physical_left_eef"] + ordinary = adapter._grasp_pose_generators()["physical_left_eef"] + + assert coordinated.collision_cfg.opening_margin == pytest.approx(0.02) + assert ordinary.collision_cfg.opening_margin == pytest.approx(0.01) + + +def test_adapter_validates_declared_runtime_ik_solver_classes() -> None: + pytorch_solver_type = type("PytorchSolver", (), {}) + env = _planner_env() + env.agent_ik_solver = "pytorch" + env.robot.get_solver = lambda **_kwargs: pytorch_solver_type() + + adapter = AtomicActionAdapter(env) + + assert adapter.ik_solver == "pytorch" + assert adapter.ik_solver_classes == { + "left_arm": "PytorchSolver", + "right_arm": "PytorchSolver", + } + + env.agent_ik_solver = "ur" + with pytest.raises(ValueError, match="left_arm.*URSolver.*PytorchSolver"): + AtomicActionAdapter(env) + + +@pytest.mark.parametrize("gripper_model", ["pgi", "robotiq"]) +def test_control_profiles_use_selected_gripper_joint_semantics( + gripper_model: str, +) -> None: + selected = get_gripper_profile(gripper_model) + hand_dof = len(selected.open_positions) + joint_ids = { + "physical_left_arm": [0, 1], + "physical_left_eef": list(range(2, 2 + hand_dof)), + "physical_right_arm": [2 + hand_dof, 3 + hand_dof], + "physical_right_eef": list(range(4 + hand_dof, 4 + 2 * hand_dof)), + } + robot = SimpleNamespace( + uid="profile_robot", + dof=4 + 2 * hand_dof, + control_parts=joint_ids, + get_joint_ids=lambda *, name: list(joint_ids[name]), + ) + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + robot=robot, + sim=SimpleNamespace(get_rigid_object=lambda _uid: None), + agent_gripper_model=gripper_model, + open_state=torch.tensor(selected.open_positions), + close_state=torch.tensor(selected.close_positions), + get_agent_arm_control_part=lambda is_left: ( + "physical_left_arm" if is_left else "physical_right_arm" + ), + get_agent_eef_control_part=lambda is_left: ( + "physical_left_eef" if is_left else "physical_right_eef" + ), + ) + + profiles = AtomicActionAdapter(env)._control_profiles() + + assert set(profiles) == {"physical_left_eef", "physical_right_eef"} + for command_profile in profiles.values(): + open_qpos = command_profile.commands["open"].resolve( + num_envs=1, + control_dof=hand_dof, + device="cpu", + ) + grasp_qpos = command_profile.commands["grasp"].resolve( + num_envs=1, + control_dof=hand_dof, + device="cpu", + ) + torch.testing.assert_close(open_qpos[0], torch.tensor(selected.open_positions)) + torch.testing.assert_close( + grasp_qpos[0], torch.tensor(selected.close_positions) + ) + + +def test_coordinated_grasp_generator_honors_ground_filter_policy() -> None: + adapter = AtomicActionAdapter(_planner_env()) + + generators = adapter._grasp_pose_generators(filter_ground_collision=False) + + assert generators["physical_left_eef"] is generators["physical_right_eef"] + assert ( + generators["physical_left_eef"].collision_cfg.filter_ground_collision is False + ) + + +def test_retreat_uses_row_local_motion_planner_reachability_search( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.05 + requested = reference.clone() + requested[:, 2, 3] = 1.35 + height_thresholds = torch.tensor([1.24, 1.00]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + height_reachable = target[:, 2, 3] <= height_thresholds + baseward_reachable = target[:, 1, 3] < -0.05 + success = height_reachable | baseward_reachable + terminal = target[:, 2, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=_planner_diagnostics(success), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "retreat_height": 0.30, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_reference_pose": reference, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert len(attempted_targets) > 1 + assert bool(outcome.success.all()) + selected_z = outcome.grounded.target.xpos[:, 2, 3] + assert selected_z.tolist() == pytest.approx([1.20, 1.35]) + assert outcome.grounded.target.xpos[:, 1, 3].tolist() == pytest.approx([0.0, -0.10]) + search = outcome.planner_trace["reachability_search"] + assert search["strategy"] == "bounded_motion_planner" + assert search["selected_target_z"].tolist() == pytest.approx([1.20, 1.35]) + assert len(search["attempts"]) == len(attempted_targets) + + +def test_lift_clear_reachability_search_uses_only_vertical_candidates( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.05 + requested = reference.clone() + requested[:, 2, 3] = 1.35 + height_thresholds = torch.tensor([1.24, 1.00]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + height_reachable = target[:, 2, 3] <= height_thresholds + baseward_reachable = target[:, 1, 3] < -0.05 + success = height_reachable | baseward_reachable + terminal = target[:, 2, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=_planner_diagnostics(success), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "retreat_height": 0.30, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_search_mode": "vertical_only", + "retreat_reference_pose": reference, + "minimum_retreat_height": 0.05, + "retreat_distance": 0.10, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert all( + torch.equal(target[:, 1, 3], reference[:, 1, 3]) for target in attempted_targets + ) + assert outcome.success.tolist() == [True, False] + assert outcome.grounded.target.xpos[:, 1, 3].tolist() == pytest.approx([0.0, 0.0]) + + +def test_retreat_after_lift_search_reduces_only_horizontal_distance( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, 2, 3] = 1.35 + requested = reference.clone() + requested[:, 1, 3] = -0.20 + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + success = target[:, 1, 3].abs() <= 0.10 + 1.0e-6 + if not bool(success.all()): + log_warning("Synthetic unreachable horizontal retreat candidate.") + terminal = target[:, 1, 3, None].repeat(1, 8) + positions = torch.stack((torch.zeros_like(terminal), terminal), dim=1) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=_planner_diagnostics(success), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=requested), + { + "sample_interval": 10, + "minimum_retreat_distance": 0.05, + "retreat_search_samples": 4, + }, + motion_policy={ + "collision_safety": "required", + "retreat_reachability_search": True, + "retreat_search_mode": "horizontal_only", + "retreat_reference_pose": reference, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + attempted_distances = torch.stack( + [ + torch.linalg.vector_norm(target[:, :2, 3] - reference[:, :2, 3], dim=1) + for target in attempted_targets + ] + ) + torch.testing.assert_close( + attempted_distances, + torch.tensor([[0.20, 0.20], [0.15, 0.15], [0.10, 0.10]]), + ) + assert all( + torch.equal(target[:, 2, 3], reference[:, 2, 3]) for target in attempted_targets + ) + assert bool(outcome.success.all()) + search = outcome.planner_trace["reachability_search"] + assert search["selected_candidates"] == ["distance_2", "distance_2"] + assert search["selected_target_distance"].tolist() == pytest.approx([0.10, 0.10]) + assert search["suppressed_warnings"] == 2 + + +def test_tool_down_reorientation_uses_waypoints_and_yaw_candidates( + monkeypatch: Any, +) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + reference = torch.eye(4).repeat(2, 1, 1) + reference[:, :3, :3] = _rotation_x(55.0) + reference[:, :3, 3] = torch.tensor([0.1, -0.2, 1.35]) + attempted_targets: list[torch.Tensor] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + target = invocation.goal.xpos.clone() + attempted_targets.append(target) + success = torch.full((2,), len(attempted_targets) >= 2, dtype=torch.bool) + if not bool(success.all()): + log_warning("Synthetic unreachable tool-down yaw candidate.") + positions = torch.zeros(2, 30, 8) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="move_end_effector", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=_planner_diagnostics(success), + expected_effects=StateDelta(), + ) + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveEndEffector", + "right_arm", + "arm", + EndEffectorPoseGoal(xpos=reference), + { + "sample_interval": 30, + "reorient_tool_down": True, + "reorient_reference_pose": reference, + "reorient_waypoint_count": 5, + "reorient_yaw_degrees": [0.0, 45.0, -45.0], + }, + motion_policy={ + "reorient_tool_down": True, + "reorient_reference_pose": reference, + }, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert len(attempted_targets) == 2 + for target in attempted_targets: + assert target.shape == (2, 5, 4, 4) + torch.testing.assert_close( + target[:, :, :3, 3], + reference[:, None, :3, 3].expand(-1, 5, -1), + ) + torch.testing.assert_close( + target[:, -1, :3, 2], + torch.tensor([0.0, 0.0, -1.0]).repeat(2, 1), + atol=1.0e-6, + rtol=1.0e-6, + ) + assert bool(outcome.success.all()) + assert outcome.grounded.motion_policy["reorient_selected_yaw_degrees"] == 45.0 + + +def test_curobo_generator_receives_generated_static_obstacles( + monkeypatch: Any, +) -> None: + table = object() + can = object() + captured: dict[str, Any] = {} + + def fake_motion_generator(*, cfg: Any) -> object: + captured["cfg"] = cfg + return object() + + monkeypatch.setattr(actions, "MotionGenerator", fake_motion_generator) + adapter = AtomicActionAdapter( + _planner_env(table=table, rigid_objects={"can": can}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["can"], + }, + ) + + generator = adapter._generator() + + assert generator is adapter._motion_generator + planner = captured["cfg"].planner_cfg + assert isinstance(planner, CuroboPlannerCfg) + assert planner.world.rigid_objects == {"table": table, "can": can} + assert planner.world.dynamic_obstacle_names == ["can"] + assert planner.world.obstacle_representation == "cuboid" + assert planner.world.collision_cache == {"cuboid": 8, "mesh": 2} + + +def test_curobo_generator_sizes_collision_cache_for_large_scene( + monkeypatch: Any, +) -> None: + rigid_objects = {f"object_{index:02d}": object() for index in range(13)} + captured: dict[str, Any] = {} + + def fake_motion_generator(*, cfg: Any) -> object: + captured["cfg"] = cfg + return object() + + monkeypatch.setattr(actions, "MotionGenerator", fake_motion_generator) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=rigid_objects), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(rigid_objects), + }, + ) + + adapter._generator() + + planner = captured["cfg"].planner_cfg + assert planner.world.collision_cache == {"cuboid": 13, "mesh": 2} + + +def test_dynamic_scene_separates_live_semantic_and_parked_collision_poses() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + actual[:, 2, 3] = torch.tensor([0.7, 0.8]) + entities = {uid: _PoseEntity(actual.clone()) for uid in ("target", "held", "other")} + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=entities), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(entities), + }, + ) + held_semantics = ObjectSemantics( + label="held", + entity_id="held", + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=held_semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + env_mask=torch.tensor([True, False]), + ) + state = ExecutionState( + last_qpos=torch.zeros(2, 8), + held_objects={"physical_left_arm": held}, + ) + grounded = GroundedAction( + "PickUp", + "right_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="target", + ) + + scene = adapter._scene_snapshot(grounded, state) + collision_poses = scene.collision_obstacle_poses( + batch_size=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + assert torch.equal(scene.entities["target"].pose, actual) + assert torch.equal(scene.entities["held"].pose, actual) + assert torch.equal(scene.entities["other"].pose, actual) + assert torch.equal( + collision_poses["target"][:, 2, 3], + actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET, + ) + assert collision_poses["held"][0, 2, 3] == ( + actual[0, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET + ) + assert collision_poses["held"][1, 2, 3] == actual[1, 2, 3] + assert torch.equal(collision_poses["other"], actual) + + +def test_released_object_returns_to_live_dynamic_collision_pose() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + actual[:, 0, 3] = torch.tensor([0.2, 0.4]) + entity = _PoseEntity(actual) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"released": entity}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["released"], + }, + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="released", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert torch.equal(scene.entities["released"].pose, actual) + + +def test_default_scene_provider_advances_only_after_material_change() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + entity = _PoseEntity(actual.clone()) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"can": entity}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["can"], + }, + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="can", + ) + state = ExecutionState(last_qpos=torch.zeros(2, 8)) + + first = adapter._scene_snapshot(grounded, state) + unchanged = adapter._scene_snapshot(grounded, state) + entity.pose[:, 0, 3] += 0.1 + changed = adapter._scene_snapshot(grounded, state) + + assert first.version == unchanged.version == 0 + assert changed.version == 1 + assert changed.collision_world_revisions(2) == (1, 1) + + +def test_external_scene_provider_is_used_by_planning_snapshot() -> None: + pose = torch.eye(4).repeat(2, 1, 1) + + class _Provider: + def snapshot(self, *, timestamp: float, env_ids: torch.Tensor) -> SceneSnapshot: + assert timestamp == 0.0 + assert torch.equal(env_ids, torch.tensor([0, 1])) + return SceneSnapshot( + timestamp=timestamp, + version=7, + entities={"can": actions.EntityState(pose)}, + ) + + adapter = AtomicActionAdapter( + _planner_env(), + scene_provider=_Provider(), + ) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + object_uid="can", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert scene.version == 7 + assert torch.equal(scene.entities["can"].pose, pose) + + +def test_start_session_delegates_to_shared_atomic_engine(monkeypatch: Any) -> None: + adapter = AtomicActionAdapter(_planner_env()) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ) + state = ExecutionState(last_qpos=torch.zeros(2, 8)) + marker = object() + captured: dict[str, Any] = {} + + monkeypatch.setattr(adapter, "_planning_context", lambda *_args: "context") + monkeypatch.setattr( + adapter, + "_invocation", + lambda *_args, **_kwargs: "invocation", + ) + + class _Engine: + def start(self, invocations: tuple[Any, ...], context: Any) -> object: + captured["invocations"] = invocations + captured["context"] = context + return marker + + monkeypatch.setattr(adapter, "_engine", lambda: _Engine()) + + result = adapter.start_session(grounded, state) + + assert result is marker + assert captured == {"invocations": ("invocation",), "context": "context"} + + +def test_retreat_parks_only_collision_poses_for_contact_objects() -> None: + actual = torch.eye(4).repeat(2, 1, 1) + entities = { + uid: _PoseEntity(actual.clone()) for uid in ("released", "container", "other") + } + adapter = AtomicActionAdapter( + _planner_env(rigid_objects=entities), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": list(entities), + }, + ) + grounded = GroundedAction( + "MoveEndEffector", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + motion_policy={ + "collision_exclusion_uids": ["released", "container"], + }, + object_uid="released", + ) + + scene = adapter._scene_snapshot( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + collision_poses = scene.collision_obstacle_poses( + batch_size=2, + device=torch.device("cpu"), + dtype=torch.float32, + ) + + parked_z = actual[:, 2, 3] + actions._COLLISION_PARKING_Z_OFFSET + assert torch.equal(scene.entities["released"].pose, actual) + assert torch.equal(scene.entities["container"].pose, actual) + assert torch.equal(scene.entities["other"].pose, actual) + assert torch.equal(collision_poses["released"][:, 2, 3], parked_z) + assert torch.equal(collision_poses["container"][:, 2, 3], parked_z) + assert torch.equal(collision_poses["other"], actual) + + +def test_action_outcome_commits_state_delta_only_for_verified_rows() -> None: + semantics = ObjectSemantics( + label="cube", + entity_id="cube", + geometry={}, + affordance=Affordance(), + ) + held = HeldObjectState( + semantics=semantics, + object_to_eef=torch.eye(4).repeat(2, 1, 1), + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + prior = ExecutionState(last_qpos=torch.zeros(2, 3)) + trajectory = torch.stack( + (torch.zeros(2, 3), torch.ones(2, 3)), + dim=1, + ) + delta = StateDelta(held_object_updates={"physical_left_arm": held}) + projected = ExecutionState.from_task_state( + delta.apply(prior.to_task_state(), torch.ones(2, dtype=torch.bool)), + last_qpos=trajectory[:, -1], + ) + grounded = GroundedAction( + "PickUp", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + ) + outcome = ActionOutcome( + trajectory=trajectory, + success=torch.ones(2, dtype=torch.bool), + next_state=projected, + grounded=grounded, + prior_state=prior, + expected_effects=delta, + ) + + committed = outcome.state_after(torch.tensor([True, False])) + + assert torch.equal(committed.last_qpos[0], torch.ones(3)) + assert torch.equal(committed.last_qpos[1], torch.zeros(3)) + committed_held = committed.get_held_object("physical_left_arm") + assert committed_held is not None + assert torch.equal(committed_held.env_mask, torch.tensor([True, False])) + + +def test_fallback_rows_keep_the_fallback_plan_effects(monkeypatch: Any) -> None: + env = _planner_env() + adapter = AtomicActionAdapter(env) + semantics = ObjectSemantics( + label="cube", + entity_id="cube", + geometry={}, + affordance=Affordance(), + ) + + def held_at(x: float) -> HeldObjectState: + relation = torch.eye(4).repeat(2, 1, 1) + relation[:, 0, 3] = x + return HeldObjectState( + semantics=semantics, + object_to_eef=relation, + grasp_xpos=torch.eye(4).repeat(2, 1, 1), + ) + + def action_plan( + success: torch.Tensor, + terminal: float, + held: HeldObjectState, + *, + messages: tuple[str, ...] = (), + ) -> ActionPlan: + positions = torch.full((2, 2, 8), terminal) + trajectory = TimedTrajectory.from_uniform_step( + positions, + env_ids=torch.arange(2), + step_dt=0.01, + ) + return ActionPlan( + skill_id="pick_up", + plan_success=success, + commands=_commands_for(trajectory), + joint_trajectory=trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=0, + planned_collision_world_revision=(0, 0), + diagnostics=PlannerDiagnostics( + backend="curobo", + messages=messages, + metadata={"marker": terminal}, + failure=( + None if bool(success.all()) else PlanningFailure("planning_failed") + ), + ), + expected_effects=StateDelta( + held_object_updates={"physical_left_arm": held} + ), + ) + + plans = iter( + ( + action_plan( + torch.tensor([True, False]), + 1.0, + held_at(1.0), + messages=("IK unreachable",), + ), + action_plan(torch.tensor([True, True]), 2.0, held_at(2.0)), + ) + ) + strategies: list[str] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + strategies.append(invocation.motion_policy.strategy) + return next(plans) + + monkeypatch.setattr( + adapter, + "_engine", + lambda: _FakeEngine(plan), + ) + grounded = GroundedAction( + "PickUp", + "left_arm", + "arm", + GraspGoal(semantics=semantics), + {}, + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert strategies == ["motion_gen", "ik_interp"] + assert torch.equal(outcome.success, torch.tensor([True, True])) + assert torch.equal(outcome.next_state.last_qpos[0], torch.ones(8)) + assert torch.equal(outcome.next_state.last_qpos[1], torch.full((8,), 2.0)) + held = outcome.next_state.get_held_object("physical_left_arm") + assert held is not None + assert held.object_to_eef[0, 0, 3] == 1.0 + assert held.object_to_eef[1, 0, 3] == 2.0 + assert torch.equal( + outcome.planner_trace["primary_success"], torch.tensor([True, False]) + ) + assert torch.equal( + outcome.planner_trace["fallback_attempted"], torch.tensor([False, True]) + ) + assert torch.equal( + outcome.planner_trace["fallback_used"], torch.tensor([False, True]) + ) + assert outcome.planner_trace["primary_action_diagnostics"]["marker"] == 1.0 + assert outcome.planner_trace["fallback_action_diagnostics"]["marker"] == 2.0 + assert outcome.planner_trace["gripper_model"] == "pgi" + assert outcome.planner_trace["requested_backend"] == "curobo" + assert outcome.planner_trace["effective_backend"] == "mixed" + assert outcome.planner_trace["effective_strategy"] == "mixed" + assert outcome.planner_trace["primary_effective_backend"] == "curobo" + assert bool(outcome.planner_trace["fallback_occurred"].any()) + assert outcome.planner_trace["planner_failure_reason"] == "IK unreachable" + assert outcome.planner_trace["collision_planning_capable"] is False + + +def test_collision_required_cleanup_does_not_use_unsafe_fallback( + monkeypatch: Any, +) -> None: + pose = torch.eye(4).repeat(2, 1, 1) + adapter = AtomicActionAdapter( + _planner_env(rigid_objects={"released": _PoseEntity(pose)}), + planner_policy={ + "dynamic_collision": True, + "dynamic_obstacle_uids": ["released"], + }, + ) + failed_trajectory = TimedTrajectory.from_uniform_step( + torch.zeros(2, 2, 8), + env_ids=torch.arange(2), + step_dt=0.01, + ) + failed_plan = ActionPlan( + skill_id="move_joints", + plan_success=torch.tensor([False, False]), + commands=_commands_for(failed_trajectory), + joint_trajectory=failed_trajectory, + recovery_policy=RecoveryPolicy(), + tracking_policy=TrackingPolicy.timed(), + planned_scene_version=1, + planned_collision_world_revision=(1, 1), + diagnostics=PlannerDiagnostics( + backend="fake", + failure=PlanningFailure("planning_failed"), + ), + expected_effects=StateDelta(), + ) + strategies: list[str] = [] + + def plan(invocation: Any, _context: Any) -> ActionPlan: + strategies.append(invocation.motion_policy.strategy) + assert invocation.motion_policy.dynamic_collision_mode.value == "required" + return failed_plan + + monkeypatch.setattr(adapter, "_engine", lambda: _FakeEngine(plan)) + grounded = GroundedAction( + "MoveJoints", + "left_arm", + "arm", + JointPositionGoal(target=torch.zeros(2, 2)), + {}, + motion_policy={"collision_safety": "required"}, + object_uid="released", + ) + + outcome = adapter.plan( + grounded, + ExecutionState(last_qpos=torch.zeros(2, 8)), + ) + + assert strategies == ["motion_gen"] + assert not bool(outcome.success.any()) + assert outcome.planner_trace["fallback_allowed"] is False + assert not bool(outcome.planner_trace["fallback_attempted"].any()) + assert not bool(outcome.planner_trace["fallback_used"].any()) + assert outcome.planner_trace["collision_obstacle_positions"]["released"].shape == ( + 2, + 3, + ) diff --git a/tests/gen_sim/action_engine/runtime/test_atomic_compat.py b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py new file mode 100644 index 000000000..54cdd37cc --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_atomic_compat.py @@ -0,0 +1,141 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace + +from embodichain.gen_sim.action_engine.runtime.actions import AtomicActionAdapter +from embodichain.gen_sim.action_engine.runtime.atomic_compat import ( + ActionEngineMoveJoints, + ActionEngineMoveJointsOptions, + ExactTargetMoveHeldObject, + ExactTargetMoveHeldObjectOptions, +) +from embodichain.lab.sim.atomic_actions import ( + MoveHeldObject, + MoveHeldObjectOptions, + MoveJoints, + MoveJointsOptions, + StateDelta, +) + + +def test_grounded_target_transport_uses_mainline_exact_target_contract() -> None: + action = ExactTargetMoveHeldObject() + assert type(action).__dict__["binding_contract"] is MoveHeldObject.binding_contract + assert action._plan.__func__ is MoveHeldObject._plan + assert "_apply_automatic_transport_rotation" not in type(action).__dict__ + + +def test_semantic_transport_config_has_no_task_facing_rotation_switch() -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + action = SimpleNamespace(cfg={}) + capability = SimpleNamespace( + config_type=MoveHeldObjectOptions, + target_materializer="semantic_held_object", + ) + + options = adapter._build_single_arm_config(action, capability) + + assert isinstance(options, ExactTargetMoveHeldObjectOptions) + assert not hasattr(options, "allow_automatic_transport_rotation") + + +def test_joint_config_materializes_single_release_only_when_requested() -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + capability = SimpleNamespace( + config_type=MoveJointsOptions, + target_materializer="joint_state", + ) + + release = adapter._build_single_arm_config( + SimpleNamespace(cfg={"single_release": True}), + capability, + ) + ordinary = adapter._build_single_arm_config( + SimpleNamespace(cfg={}), + capability, + ) + + assert isinstance(release, ActionEngineMoveJointsOptions) + assert release.single_release + assert isinstance(ordinary, ActionEngineMoveJointsOptions) + assert not ordinary.single_release + + +def test_single_release_binds_hand_motion_to_the_arm_held_state_key() -> None: + adapter = AtomicActionAdapter.__new__(AtomicActionAdapter) + adapter._parts = lambda _arm: ("physical_left_arm", "physical_left_hand", 2) + captured = {} + + class Engine: + def bind_control_parts(self, skill_id, endpoints, *, task_state_keys=None): + captured.update( + skill_id=skill_id, + endpoints=endpoints, + task_state_keys=task_state_keys, + ) + return object() + + adapter._binding( + SimpleNamespace( + arm="left_arm", + control="hand", + cfg={"single_release": True}, + ), + SimpleNamespace( + action_type=MoveJoints, + config_materializer="single_arm", + ), + engine=Engine(), + ) + + assert captured == { + "skill_id": "move_joints", + "endpoints": {"primary": {"motion": "physical_left_hand"}}, + "task_state_keys": {"primary": "physical_left_arm"}, + } + + +def test_single_release_plan_removes_only_the_bound_arm_attachment(monkeypatch) -> None: + @dataclass(frozen=True) + class Plan: + expected_effects: object + + monkeypatch.setattr( + MoveJoints, + "_plan", + lambda _self, _request, _context: Plan(expected_effects=object()), + ) + action = ActionEngineMoveJoints() + request = SimpleNamespace( + binding=SimpleNamespace( + endpoint=lambda _slot, _endpoint: SimpleNamespace(task_state_key="left_arm") + ), + skill_options=ActionEngineMoveJointsOptions(single_release=True), + ) + context = SimpleNamespace( + task=SimpleNamespace( + get_held_object=lambda key: object() if key == "left_arm" else None + ) + ) + + plan = action._plan(request, context) + + assert isinstance(plan.expected_effects, StateDelta) + assert dict(plan.expected_effects.held_object_updates) == {"left_arm": None} diff --git a/tests/gen_sim/action_engine/runtime/test_body_grasp.py b/tests/gen_sim/action_engine/runtime/test_body_grasp.py new file mode 100644 index 000000000..7c3ec8040 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_body_grasp.py @@ -0,0 +1,154 @@ +# ---------------------------------------------------------------------------- +# 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 elongated-object axis analysis and body-grasp filtering.""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime.body_grasp import ( + AxisAlignBodyGraspAdapter, + select_body_grasp_candidates, +) +from embodichain.gen_sim.action_engine.runtime.geometry_axes import ( + analyze_local_geometry_axes, +) +from embodichain.lab.sim.atomic_actions import ( + AxisAlignAffordance, + AxisAlignGoal, + ObjectSemantics, +) + + +def _box_vertices(extents: tuple[float, float, float]) -> torch.Tensor: + half = torch.tensor(extents) * 0.5 + return torch.tensor( + [ + [sx * half[0], sy * half[1], sz * half[2]] + for sx in (-1.0, 1.0) + for sy in (-1.0, 1.0) + for sz in (-1.0, 1.0) + ] + ) + + +def test_can_geometry_resolves_local_y_as_the_long_axis() -> None: + axes = analyze_local_geometry_axes(_box_vertices((0.0611, 0.1143, 0.0632))) + + assert axes.long_axis_index == 1 + assert axes.short_axis_index == 0 + torch.testing.assert_close(axes.long_axis, torch.tensor([0.0, 1.0, 0.0])) + assert axes.elongation_ratio == pytest.approx(0.1143 / 0.0632) + + +def test_axis_analysis_rejects_ambiguous_or_rotated_local_geometry() -> None: + with pytest.raises(ValueError, match="ambiguous"): + analyze_local_geometry_axes(_box_vertices((0.06, 0.06, 0.06))) + + vertices = _box_vertices((0.04, 0.12, 0.05)) + angle = math.radians(35.0) + rotation = torch.tensor( + [ + [math.cos(angle), -math.sin(angle), 0.0], + [math.sin(angle), math.cos(angle), 0.0], + [0.0, 0.0, 1.0], + ] + ) + with pytest.raises(ValueError, match="not aligned"): + analyze_local_geometry_axes(vertices @ rotation.T) + + +def test_body_grasp_rejects_caps_and_longitudinal_closing() -> None: + axes = analyze_local_geometry_axes(_box_vertices((0.06, 0.12, 0.06))) + candidates = torch.eye(4).repeat(1, 3, 1, 1) + candidates[0, 0, 1, 3] = 0.055 # End-cap candidate with the best raw cost. + candidates[0, 1, 1, 3] = 0.0 # Central radial body grasp. + candidates[0, 2, 1, 3] = 0.0 + candidates[0, 2, :3, 0] = torch.tensor([0.0, 1.0, 0.0]) + candidates[0, 2, :3, 1] = torch.tensor([1.0, 0.0, 0.0]) + costs = torch.tensor([[0.0, 0.2, 0.1]]) + + selected = select_body_grasp_candidates( + candidates, + costs, + torch.eye(4).unsqueeze(0), + axes, + ) + + assert selected.success.tolist() == [True] + assert selected.candidate_indices.tolist() == [1] + assert selected.body_candidate_counts.tolist() == [1] + assert selected.ranked_candidate_indices.tolist() == [[1]] + torch.testing.assert_close(selected.grasp_xpos[0], candidates[0, 1]) + + +def test_body_grasp_chooses_a_reachable_body_candidate() -> None: + axes = analyze_local_geometry_axes(_box_vertices((0.06, 0.12, 0.06))) + candidates = torch.eye(4).repeat(1, 2, 1, 1) + candidates[0, 1, 0, 3] = 0.01 + costs = torch.tensor([[0.0, 0.2]]) + + selected = select_body_grasp_candidates( + candidates, + costs, + torch.eye(4).unsqueeze(0), + axes, + feasible=torch.tensor([[False, True]]), + ) + + assert selected.candidate_indices.tolist() == [1] + assert selected.reachable_candidate_counts.tolist() == [1] + + +def test_axis_align_adapter_injects_the_selected_body_grasp_unchanged() -> None: + vertices = _box_vertices((0.06, 0.12, 0.06)) + triangles = torch.tensor([[0, 1, 2], [0, 2, 3]]) + goal = AxisAlignGoal( + semantics=ObjectSemantics( + label="can", + entity_id="can", + geometry={}, + affordance=AxisAlignAffordance( + mesh_vertices=vertices, + mesh_triangles=triangles, + internal_axis=torch.tensor([0.0, 1.0, 0.0]), + ), + ) + ) + candidate = torch.eye(4).unsqueeze(0) + + class Generator: + def get_valid_grasp_poses(self, **_kwargs): + return [(candidate, torch.tensor([0.1]))] + + adapted = AxisAlignBodyGraspAdapter().adapt( + goal, + object_pose=torch.eye(4).unsqueeze(0), + grasp_generator=Generator(), + approach_direction=torch.tensor([0.0, 0.0, -1.0]), + target_axis=torch.tensor([0.0, 0.0, 1.0]), + seed=7, + ) + + assert adapted.goal.grasp_xpos is not None + assert len(adapted.alternative_goals) == 1 + assert adapted.alternative_rank_indices == (0,) + explicit = adapted.goal.grasp_xpos + torch.testing.assert_close(explicit, candidate) diff --git a/tests/gen_sim/action_engine/runtime/test_coordinated_safety.py b/tests/gen_sim/action_engine/runtime/test_coordinated_safety.py new file mode 100644 index 000000000..b9a3eb5a6 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_coordinated_safety.py @@ -0,0 +1,185 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import pytest +import torch + +from embodichain.gen_sim.action_engine.runtime.coordinated_safety import ( + _canonicalize_parallel_jaw_poses, + _minimum_interarm_capsule_clearance, + _rank_non_crossing_grasp_pairs, + _segments_intersect_2d, + _trajectory_safety_report, +) + + +def _pose(x: float, y: float, z: float = 0.75) -> torch.Tensor: + pose = torch.eye(4, dtype=torch.float32) + pose[:3, 3] = torch.tensor([x, y, z]) + return pose + + +def test_parallel_jaw_pose_chooses_equivalent_half_turn_nearest_live_eef() -> None: + pose = torch.eye(4).unsqueeze(0) + pose[0, 0, 0] = -1.0 + pose[0, 1, 1] = -1.0 + + result = _canonicalize_parallel_jaw_poses(pose, torch.eye(4)) + + torch.testing.assert_close(result.poses[0], torch.eye(4)) + assert result.flipped.tolist() == [True] + assert result.selected_rotation_radians.tolist() == [0.0] + assert result.alternative_rotation_radians.tolist() == pytest.approx([torch.pi]) + + +def test_pair_ranking_rejects_reversal_overlap_and_xy_crossing() -> None: + left = torch.stack( + ( + _pose(0.0, -0.20), + _pose(1.0, 0.20), + _pose(0.0, 0.18), + ) + ) + right = torch.stack( + ( + _pose(0.0, 0.20), + _pose(-1.0, -0.20), + _pose(0.0, 0.19), + ) + ) + result = _rank_non_crossing_grasp_pairs( + left, + right, + left_costs=torch.tensor([0.0, 10.0, 10.0]), + right_costs=torch.tensor([0.0, 10.0, 10.0]), + left_rotation_costs=torch.zeros(3), + right_rotation_costs=torch.zeros(3), + left_base=_pose(-1.0, -0.30), + right_base=_pose(1.0, 0.30), + left_to_right_direction=torch.tensor([0.0, 1.0, 0.0]), + minimum_separation=0.08, + minimum_lateral_gap=0.05, + ) + + assert result.ranked_pairs[0] == (0, 0) + assert (1, 1) not in result.ranked_pairs # XY paths intersect. + assert (2, 2) not in result.ranked_pairs # Distinct poses are too close. + assert result.rejection_counts["reversed"] > 0 + assert result.rejection_counts["too_close"] > 0 + assert result.rejection_counts["path_crossing"] > 0 + + +def test_xy_route_intersection_includes_touching_and_collinear_overlap() -> None: + assert _segments_intersect_2d( + torch.tensor([0.0, 0.0, 0.0]), + torch.tensor([1.0, 0.0, 0.0]), + torch.tensor([0.5, 0.0, 0.0]), + torch.tensor([1.5, 0.0, 0.0]), + ) + assert _segments_intersect_2d( + torch.tensor([0.0, 0.0, 0.0]), + torch.tensor([1.0, 1.0, 0.0]), + torch.tensor([1.0, 1.0, 0.0]), + torch.tensor([2.0, 1.0, 0.0]), + ) + + +def test_capsule_clearance_covers_every_left_right_link_segment() -> None: + left = torch.tensor( + [[[[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]]], + dtype=torch.float32, + ) + crossing = torch.tensor( + [[[[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]]]], + dtype=torch.float32, + ) + clear = crossing.clone() + clear[..., 2] = 1.0 + + crossing_clearance = _minimum_interarm_capsule_clearance( + left, + crossing, + capsule_radius=0.05, + ) + safe_clearance = _minimum_interarm_capsule_clearance( + left, + clear, + capsule_radius=0.05, + ) + + torch.testing.assert_close(crossing_clearance, torch.tensor([[-0.10]])) + torch.testing.assert_close(safe_clearance, torch.tensor([[0.90]])) + + +def test_trajectory_safety_rejects_orientation_jump_order_and_capsule_collision() -> ( + None +): + left_qpos = torch.zeros(1, 3, 2) + right_qpos = torch.zeros(1, 3, 2) + left_qpos[:, 1, 0] = 0.40 + left_eef = torch.eye(4).repeat(1, 3, 1, 1) + right_eef = torch.eye(4).repeat(1, 3, 1, 1) + desired_left = left_eef.clone() + desired_right = right_eef.clone() + left_eef[:, 1, 0, 0] = -1.0 + left_eef[:, 1, 1, 1] = -1.0 + left_eef[:, :, 1, 3] = torch.tensor([-0.2, 0.2, -0.2]) + right_eef[:, :, 1, 3] = torch.tensor([0.2, -0.2, 0.2]) + left_links = torch.tensor( + [ + [ + [[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + [[-1.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + ] + ] + ) + right_links = torch.tensor( + [ + [ + [[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]], + [[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]], + [[0.0, -1.0, 0.0], [0.0, 1.0, 0.0]], + ] + ] + ) + + report = _trajectory_safety_report( + left_qpos=left_qpos, + right_qpos=right_qpos, + left_eef=left_eef, + right_eef=right_eef, + desired_left_eef=desired_left, + desired_right_eef=desired_right, + left_link_points=left_links, + right_link_points=right_links, + left_to_right_direction=torch.tensor([0.0, 1.0, 0.0]), + maximum_joint_step=0.25, + maximum_orientation_error=0.20, + minimum_lateral_gap=0.05, + capsule_radius=0.05, + minimum_capsule_clearance=0.0, + ) + + assert report.success.tolist() == [False] + assert report.failed_checks == { + "joint_step": [True], + "orientation": [True], + "lateral_order": [True], + "capsule_collision": [True], + } diff --git a/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py new file mode 100644 index 000000000..76f7949c2 --- /dev/null +++ b/tests/gen_sim/action_engine/runtime/test_grasp_diagnostics.py @@ -0,0 +1,227 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch + +from embodichain.gen_sim.action_engine.runtime.grasp_diagnostics import ( + _TracingAntipodalGraspPoseGenerator, +) +from embodichain.toolkits.graspkit.pg_grasp import AntipodalGraspPoseGenerator +from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg + + +class _CollisionChecker: + def __init__(self) -> None: + self.calls = 0 + + def query(self, *_args, **_kwargs): + self.calls += 1 + if self.calls == 1: + return torch.tensor([False, True, False]), torch.tensor([0.01, -0.002, 0.0]) + return torch.tensor([True, False]), torch.tensor([-0.004, 0.003]) + + +class _Backend: + device = torch.device("cpu") + _max_deviation_angle = torch.pi / 6 + _approach_direction_samples = 4 + + def __init__(self) -> None: + self.antipodal_pairs = torch.tensor( + [ + [[-0.08, -0.08, 0.0], [0.08, -0.08, 0.0]], + [[-0.08, -0.06, 0.0], [0.08, -0.06, 0.0]], + [[-0.08, 0.06, 0.0], [0.08, 0.06, 0.0]], + [[-0.08, 0.08, 0.0], [0.08, 0.08, 0.0]], + ], + dtype=torch.float32, + ) + self._collision_checker = _CollisionChecker() + + def get_dual_arm_valid_grasp_poses(self, **_kwargs): + pose = torch.eye(4).repeat(3, 1, 1) + left_colliding, _ = self._collision_checker.query(None, pose, torch.ones(3)) + right_colliding, _ = self._collision_checker.query( + None, pose[:2], torch.ones(2) + ) + return { + "left": { + "is_success": True, + "grasp_poses": pose[~left_colliding], + "open_lengths": torch.ones(2), + "total_cost": torch.tensor([0.1, 0.2]), + }, + "right": { + "is_success": True, + "grasp_poses": pose[:2][~right_colliding], + "open_lengths": torch.ones(1), + "total_cost": torch.tensor([0.3]), + }, + } + + +def test_dual_grasp_trace_separates_generation_angle_nms_and_collision( + monkeypatch, +) -> None: + generator = _TracingAntipodalGraspPoseGenerator( + ParallelJawGripperModelCfg(model_id="trace_test") + ) + backend = _Backend() + monkeypatch.setattr(generator, "_backend", lambda *_args: backend) + vertices = torch.tensor( + [ + [-0.1, -0.1, -0.02], + [0.1, -0.1, -0.02], + [0.1, 0.1, 0.02], + [-0.1, 0.1, 0.02], + ] + ) + + result = generator.get_dual_arm_valid_grasp_poses( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2], [0, 2, 3]]), + obj_poses=torch.eye(4).unsqueeze(0), + left_to_right_arm_direction=torch.tensor([0.0, 1.0, 0.0]), + approach_direction=torch.tensor([0.0, 0.0, -1.0]), + middle_empty_ratio=0.4, + ) + + assert result[0] is not None + trace = generator.last_dual_trace + assert trace is not None + assert trace["S1_grasp_pair_generation"]["antipodal_pair_count"] == 4 + assert trace["S2_approach_angle_filtering"] == { + "left_partition_pair_count": 2, + "right_partition_pair_count": 2, + "left_angle_valid_pair_count": 2, + "right_angle_valid_pair_count": 2, + } + assert trace["S3_nms"] == { + "left_candidate_count": 3, + "right_candidate_count": 2, + } + assert trace["S4_collision_filtering"]["left_candidate_count"] == 2 + assert trace["S4_collision_filtering"]["right_candidate_count"] == 1 + assert trace["S5_left_right_pairing"] == { + "left_final_count": 2, + "right_final_count": 1, + "paired": True, + } + assert generator.last_dual_trace is not trace + + +def test_generator_context_selects_non_crossing_pair_and_canonical_half_turn() -> None: + generator = _TracingAntipodalGraspPoseGenerator( + ParallelJawGripperModelCfg(model_id="pair_test") + ) + left_poses = torch.stack((_pose_with_y(-0.2), _pose_with_y(0.2))) + left_poses[0, 0, 0] = -1.0 + left_poses[0, 1, 1] = -1.0 + right_poses = torch.stack((_pose_with_y(0.2), _pose_with_y(-0.2))) + result = { + "left": { + "is_success": True, + "grasp_poses": left_poses, + "open_lengths": torch.ones(2), + "total_cost": torch.tensor([0.1, 0.0]), + }, + "right": { + "is_success": True, + "grasp_poses": right_poses, + "open_lengths": torch.ones(2), + "total_cost": torch.tensor([0.1, 0.0]), + }, + } + + with generator.dual_arm_selection_context( + left_eef=torch.eye(4).unsqueeze(0), + right_eef=torch.eye(4).unsqueeze(0), + left_base=_pose_with_y(-0.3).unsqueeze(0), + right_base=_pose_with_y(0.3).unsqueeze(0), + left_to_right_direction=torch.tensor([0.0, 1.0, 0.0]), + pair_rank=0, + minimum_separation=0.08, + minimum_lateral_gap=0.05, + ): + selected, trace = generator._select_pair(result, row_index=0) + + assert selected is not None + assert trace is not None + assert trace["selected"] is True + assert trace["selected_left_index"] == 0 + assert trace["selected_right_index"] == 0 + assert trace["selected_left_half_turn"] is True + torch.testing.assert_close( + selected["left"]["grasp_poses"][0, :3, :3], + torch.eye(3), + ) + + +def test_upright_context_rejects_end_clamps_and_ranks_mid_body_grasps( + monkeypatch, +) -> None: + generator = _TracingAntipodalGraspPoseGenerator( + ParallelJawGripperModelCfg(model_id="upright_test") + ) + poses = torch.eye(4).repeat(4, 1, 1) + poses[:, 1, 3] = torch.tensor([0.5, 0.05, 0.5, 0.7]) + poses[0, :3, 0] = torch.tensor([0.0, 1.0, 0.0]) + poses[0, :3, 1] = torch.tensor([-1.0, 0.0, 0.0]) + + monkeypatch.setattr( + AntipodalGraspPoseGenerator, + "get_valid_grasp_poses", + lambda _self, **_kwargs: [(poses, torch.tensor([0.0, 0.01, 0.2, 0.3]))], + ) + vertices = torch.tensor( + [ + [-0.1, 0.0, -0.1], + [0.1, 0.0, 0.1], + [-0.1, 1.0, 0.1], + [0.1, 1.0, -0.1], + ] + ) + + with generator.upright_selection_context(local_axis=torch.tensor([0.0, 1.0, 0.0])): + result = generator.get_valid_grasp_poses( + mesh_vertices=vertices, + mesh_triangles=torch.tensor([[0, 1, 2], [1, 2, 3]]), + obj_poses=torch.eye(4).unsqueeze(0), + approach_direction=torch.tensor([0.0, 0.0, -1.0]), + ) + + ranked_poses, ranked_costs = result[0] + assert ranked_poses[0, 1, 3] == 0.5 + assert torch.isfinite(ranked_costs[:3]).all() + assert torch.isinf(ranked_costs[-1]) + trace = generator.last_upright_trace + assert trace is not None + assert trace["local_axis"] == [0.0, 1.0, 0.0] + assert trace["candidate_count"] == 4 + assert trace["side_compatible_count"] == 3 + assert trace["central_band_count"] == 3 + assert trace["side_and_central_count"] == 2 + assert trace["retained_count"] == 3 + assert trace["best_candidate_axis_alignment"] == 0.0 + assert trace["best_candidate_axis_fraction"] == 0.5 + + +def _pose_with_y(y: float) -> torch.Tensor: + pose = torch.eye(4) + pose[1, 3] = y + return pose diff --git a/tests/gen_sim/action_engine/tasks/test_e3_pour.py b/tests/gen_sim/action_engine/tasks/test_e3_pour.py new file mode 100644 index 000000000..19c9e0f7e --- /dev/null +++ b/tests/gen_sim/action_engine/tasks/test_e3_pour.py @@ -0,0 +1,132 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Deterministic contracts for the E3 approximate pouring workflow.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from embodichain.gen_sim.action_engine.domain import validate_seed_graph +from embodichain.gen_sim.action_engine.runtime.predicates import evaluate_predicate +from embodichain.gen_sim.action_engine.tasks import instantiate_seed_graph + + +def _task() -> dict: + params = { + "source_role": "source", + "target_role": "target", + "required_arm": "right_arm", + } + return { + "schema_version": "action_engine_task_spec_v2", + "task_id": "e3_single", + "level": "L1", + "instruction": "Pour from the source container into the target container.", + "reasoning_type": "none", + "task_instances": [ + { + "id": "pour", + "task_type": "E3", + "params": params, + "depends_on": [], + "role": "primary", + } + ], + "success": {"type": "poured"}, + "oracle": {}, + "metadata": {}, + } + + +def _graph() -> dict: + return instantiate_seed_graph( + _task(), + {"source": "source_container", "target": "target_container"}, + ) + + +def test_single_arm_mode_preserves_the_historical_auto_actor_contract() -> None: + task = _task() + task["task_instances"][0]["params"].pop("required_arm") + + graph = instantiate_seed_graph( + task, + {"source": "source_container", "target": "target_container"}, + ) + + assert graph["task_groups"][0]["actor"] == {"mode": "auto"} + assert all(node["actor"] == {"mode": "auto"} for node in graph["nodes"]) + assert [node["atomic_action"] for node in graph["nodes"]] == [ + "PickUp", + "MoveHeldObject", + "Pour", + "MoveHeldObject", + "Place", + "MoveEndEffector", + "MoveJoints", + ] + + +def test_single_arm_recipe_binds_only_the_source_and_requested_arm() -> None: + graph = _graph() + + assert all(node["object_uid"] == "source_container" for node in graph["nodes"]) + assert all( + node["actor"] == {"mode": "required", "arm": "right_arm"} + for node in graph["nodes"] + ) + assert graph["task_groups"][0]["success"]["verification"] == "action_completion" + + +@pytest.mark.parametrize("field", ["pour_mode", "pouring_arm", "holding_arm"]) +def test_single_arm_recipe_rejects_legacy_dual_arm_fields(field: str) -> None: + task = _task() + task["task_instances"][0]["params"][field] = "dual_arm" + + with pytest.raises(ValueError, match="Dual-arm E3 is not supported"): + instantiate_seed_graph( + task, + {"source": "source_container", "target": "target_container"}, + ) + + +def test_seed_graph_rejects_legacy_dual_arm_e3_goal() -> None: + graph = _graph() + graph["task_groups"][0]["goal"]["pour_mode"] = "dual_arm" + + with pytest.raises(ValueError, match="unsupported dual-arm pour fields"): + validate_seed_graph(graph) + + +def test_approximate_poured_predicate_needs_no_scene_or_content_state() -> None: + env = SimpleNamespace( + num_envs=1, + device=torch.device("cpu"), + ) + + result = evaluate_predicate( + env, + { + "type": "poured", + "verification": "action_completion", + }, + ) + + assert result.tolist() == [True]