diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index fc72b4869..1b1e45d26 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -239,9 +239,11 @@ population; later populations extend the manifest. ### Provider and evaluator interfaces `atomic_skill.py` owns `AtomicSkillCaseProvider`; the previous import from -`scenarios.atomic_task` remains available. Existing providers continue to use -`task_result()`. A migrated provider overrides `physical_evaluator()` to return -an implementation of `contracts.PhysicalStageEvaluator`: +`scenarios.atomic_task` remains available. All 15 registered Atomic Skills use +v1 physical evaluators. Terminal skills score settled replay state; staged +skills declare compiled-coordinate replay windows and score the resulting raw +simulator traces. Providers override `physical_evaluator()` to return an +implementation of `contracts.PhysicalStageEvaluator`: ```python def evaluate( @@ -255,13 +257,14 @@ def evaluate( The evaluator receives frozen task inputs, measured replay observations, and motion-validation results. It receives neither the scenario/engine nor a -`CompiledTrajectory` or `projected_context`. `ExecutionObservation` currently -contains the existing measured replay summary, plus the target joint at the end -of the declared articulation segment, the held object's peak signed rotation -in a provider-declared window, and each configured gripper's final joint -positions. Missing optional measurements -mean unavailable evidence; the follow-up evaluation work must extend observation -collection before claiming grasp/release/retention or intermediate-stage success. +`CompiledTrajectory` or `projected_context`. A provider instead implements +`replay_observation_windows()` while it still has the compiled segment map. Each +window selects raw TCP pose, active-object pose, named scene-object poses, +articulation-joint, or gripper-joint channels and may include terminal settle. +`ExecutionObservation` carries the sampled traces plus configured gripper +commands; missing optional measurements remain unavailable evidence rather than +successful effects. Push contact uses an explicitly labeled measured object- +motion proxy; it is not reported as contact-sensor evidence. ### Atomic Action runtime evidence @@ -301,9 +304,9 @@ variance, retention, and eligibility. - **Embodiment and case integration:** build against `AtomicSkillCaseProvider`, reuse official component loading, and keep case generation independent of planner outcomes. Do not duplicate the runner lifecycle or artifact writers. -- **Physical evaluation:** extend `ExecutionObservation` with measured segment - and contact evidence, implement evaluators and counterexample tests, then opt - providers into `physical_evaluator()` individually. +- **Physical evaluation:** add provider-declared replay windows and measured + contact evidence when a new stage needs them, implement counterexample tests, + then opt the provider into `physical_evaluator()` individually. - **Framework and statistics:** add domain generators, physical observation collection, and richer confidence intervals on top of the current domain identity and aggregation contracts. diff --git a/scripts/benchmark/motion_generation/atomic_skill.py b/scripts/benchmark/motion_generation/atomic_skill.py index ffe5a0ba6..3dfe22ff9 100644 --- a/scripts/benchmark/motion_generation/atomic_skill.py +++ b/scripts/benchmark/motion_generation/atomic_skill.py @@ -31,7 +31,11 @@ ) from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory from .config import SuiteCfg, TrackCfg - from .contracts import ExecutionObservation, PhysicalStageEvaluator + from .contracts import ( + ExecutionObservation, + PhysicalStageEvaluator, + ReplayObservationWindow, + ) from .models import BenchmarkCase, CaseOutcome from .planners.base import PlannerAdapter from .scenarios.atomic_task import AtomicTaskScenario @@ -143,6 +147,18 @@ def object_rotation_segment( del compiled return None + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Declare raw simulator traces required by the physical evaluator. + + Providers resolve semantic action segments here. The replay adapter + stores only measured tensors, so the evaluator remains independent of + compiled planning state and the live simulator. + """ + del compiled, case + return () + def physical_evaluator(self) -> PhysicalStageEvaluator | None: """Return a measured-stage evaluator, or use legacy task_result. diff --git a/scripts/benchmark/motion_generation/contracts.py b/scripts/benchmark/motion_generation/contracts.py index 3757073f5..16f274e1d 100644 --- a/scripts/benchmark/motion_generation/contracts.py +++ b/scripts/benchmark/motion_generation/contracts.py @@ -21,6 +21,7 @@ import math from collections.abc import Mapping from dataclasses import dataclass, field +from types import MappingProxyType from typing import TYPE_CHECKING, Literal, Protocol import torch @@ -34,11 +35,159 @@ "ExecutionObservation", "PhysicalEvaluation", "PhysicalStageEvaluator", + "ReplayObservationTrace", + "ReplayObservationWindow", "StageOutcome", "validate_physical_evaluations", ] +_REPLAY_CHANNELS = frozenset( + { + "tcp_pose", + "object_pose", + "scene_object_pose", + "articulation_joint", + "gripper_qpos", + } +) + + +@dataclass(frozen=True) +class ReplayObservationWindow: + """One compiled-coordinate window requested by a physical evaluator. + + ``start`` and ``stop`` are half-open waypoint indices. A terminal-settle + request uses ``start == stop`` at the trajectory boundary together with + ``include_settle=True``. Providers resolve semantic segment names while + they still have the compiled trajectory; evaluators receive only the + measured trace identified by ``observation_id``. + """ + + observation_id: str + start: int + stop: int + channels: frozenset[str] + include_settle: bool = False + scene_object_ids: tuple[str, ...] = () + + def __post_init__(self) -> None: + if not isinstance(self.observation_id, str) or not self.observation_id.strip(): + raise ValueError("observation_id must be a non-empty string.") + if type(self.start) is not int or type(self.stop) is not int: + raise TypeError("Replay observation bounds must be integers.") + if self.start < 0 or self.stop < self.start: + raise ValueError( + "Replay observation bounds must be ordered and non-negative." + ) + channels = frozenset(self.channels) + if not channels or not channels <= _REPLAY_CHANNELS: + raise ValueError( + "Replay observation channels must be a non-empty subset of " + f"{sorted(_REPLAY_CHANNELS)}." + ) + if self.start == self.stop and not self.include_settle: + raise ValueError("An empty waypoint window must include terminal settle.") + scene_object_ids = tuple(self.scene_object_ids) + if any( + not isinstance(object_id, str) or not object_id.strip() + for object_id in scene_object_ids + ): + raise ValueError("scene_object_ids must contain non-empty strings.") + if len(scene_object_ids) != len(set(scene_object_ids)): + raise ValueError("scene_object_ids must be unique.") + if ("scene_object_pose" in channels) != bool(scene_object_ids): + raise ValueError( + "scene_object_pose windows must declare scene_object_ids, and " + "scene_object_ids require the scene_object_pose channel." + ) + object.__setattr__(self, "channels", channels) + object.__setattr__(self, "scene_object_ids", scene_object_ids) + + +@dataclass(frozen=True) +class ReplayObservationTrace: + """Raw synchronized simulator samples for one requested replay window. + + Every tensor starts with ``(sample, environment)``. Optional channels are + absent evidence; gripper tensors additionally retain their joint axis. + """ + + tcp_pose: torch.Tensor | None = None + object_pose: torch.Tensor | None = None + articulation_joint: torch.Tensor | None = None + gripper_qpos: Mapping[str, torch.Tensor] = field(default_factory=dict) + scene_object_poses: Mapping[str, torch.Tensor] = field(default_factory=dict) + + def __post_init__(self) -> None: + tensors = [ + value + for value in (self.tcp_pose, self.object_pose, self.articulation_joint) + if value is not None + ] + grippers = dict(self.gripper_qpos) + scene_objects = dict(self.scene_object_poses) + tensors.extend(grippers.values()) + tensors.extend(scene_objects.values()) + if not tensors: + raise ValueError( + "A replay observation trace requires at least one channel." + ) + if any(not isinstance(value, torch.Tensor) for value in tensors): + raise TypeError("Replay observation channels must be torch tensors.") + shape = tensors[0].shape[:2] + if len(shape) != 2 or shape[0] < 1 or shape[1] < 1: + raise ValueError( + "Replay observation channels must start with non-empty sample and " + "environment dimensions." + ) + if any(value.dim() < 2 or value.shape[:2] != shape for value in tensors): + raise ValueError( + "Replay observation channels must share sample and environment axes." + ) + if self.tcp_pose is not None and self.tcp_pose.shape[2:] != (4, 4): + raise ValueError("tcp_pose trace must have shape (S, B, 4, 4).") + if self.object_pose is not None and self.object_pose.shape[2:] != (4, 4): + raise ValueError("object_pose trace must have shape (S, B, 4, 4).") + if self.articulation_joint is not None and self.articulation_joint.dim() != 2: + raise ValueError("articulation_joint trace must have shape (S, B).") + if any(not part for part in grippers): + raise ValueError("Gripper trace keys must be non-empty strings.") + if any(not object_id for object_id in scene_objects): + raise ValueError("Scene-object trace keys must be non-empty strings.") + if any(value.shape[2:] != (4, 4) for value in scene_objects.values()): + raise ValueError("scene_object_poses traces must have shape (S, B, 4, 4).") + object.__setattr__( + self, "tcp_pose", None if self.tcp_pose is None else self.tcp_pose.clone() + ) + object.__setattr__( + self, + "object_pose", + None if self.object_pose is None else self.object_pose.clone(), + ) + object.__setattr__( + self, + "articulation_joint", + ( + None + if self.articulation_joint is None + else self.articulation_joint.clone() + ), + ) + object.__setattr__( + self, + "gripper_qpos", + MappingProxyType({part: value.clone() for part, value in grippers.items()}), + ) + object.__setattr__( + self, + "scene_object_poses", + MappingProxyType( + {object_id: value.clone() for object_id, value in scene_objects.items()} + ), + ) + + @dataclass(frozen=True) class EvaluationDomain: """Versioned evaluation population; classification does not generate cases.""" @@ -199,9 +348,9 @@ def failure_code(self) -> str | None: class ExecutionObservation: """Measured execution summary shared with physical evaluators. - Optional values are unavailable evidence, not successful effects. Segment - sampling and skill-specific contact evidence can extend this contract in - the physical-evaluation follow-up; no projected context is carried here. + Optional values and missing replay channels are unavailable evidence, not + successful effects. Provider-declared traces carry raw segment samples; + no compiled trajectory or projected task context is carried here. """ execution_success: torch.Tensor @@ -219,6 +368,9 @@ class ExecutionObservation: object_rotation_peak_signed_rad: torch.Tensor | None = None scene_object_final_poses: Mapping[str, torch.Tensor] = field(default_factory=dict) final_gripper_qpos: Mapping[str, torch.Tensor] = field(default_factory=dict) + gripper_open_qpos: Mapping[str, torch.Tensor] = field(default_factory=dict) + gripper_grasp_qpos: Mapping[str, torch.Tensor] = field(default_factory=dict) + replay_traces: Mapping[str, ReplayObservationTrace] = field(default_factory=dict) class PhysicalStageEvaluator(Protocol): diff --git a/scripts/benchmark/motion_generation/evaluation.py b/scripts/benchmark/motion_generation/evaluation.py index 900756a56..3b8518c4a 100644 --- a/scripts/benchmark/motion_generation/evaluation.py +++ b/scripts/benchmark/motion_generation/evaluation.py @@ -22,7 +22,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Literal @@ -37,6 +37,7 @@ "Criterion", "StageCriterion", "StageEvidence", + "evaluate_skill_stage_batch", "evaluate_skill_stages", "validate_standard_evaluation", ] @@ -190,6 +191,7 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: Criterion( "tcp_position_error_m", "le", TOLERANCES["primitive_position_m"] ), + _rotation("tcp_rotation_error_rad"), ), ), "move_joints": ( @@ -210,6 +212,15 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: "actuation_end", "task_goal_miss", _rotation("axis_angle_error_rad"), + Criterion("lift_delta_m", "ge", TOLERANCES["pick_min_lift_m"]), + _NO_DROP, + ), + _stage( + "settled", + "terminal_settle", + "task_goal_miss", + _rotation("axis_angle_error_rad"), + Criterion("lift_delta_m", "ge", TOLERANCES["pick_min_lift_m"]), _NO_DROP, ), _HELD, @@ -223,6 +234,15 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: _rotation(), _NO_DROP, ), + _stage( + "settled", + "terminal_settle", + "task_goal_miss", + _position(), + _rotation(), + _NO_DROP, + ), + _HELD, ), "place": ( _RELEASE, @@ -230,7 +250,8 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: "placed", "terminal_settle", "task_goal_miss", - _position("xy_error_m"), + _position(), + _rotation("orientation_error_rad"), _NO_DROP, ), _STABLE, @@ -243,6 +264,14 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: _rotation("peak_signed_axis_angle_error_rad"), _NO_DROP, ), + _stage( + "settled", + "terminal_settle", + "task_goal_miss", + Criterion("lift_delta_m", "ge", TOLERANCES["pick_min_lift_m"]), + _NO_DROP, + ), + _HELD, ), "push_object": ( _stage( @@ -263,6 +292,7 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: "terminal_settle", "task_goal_miss", _bool("on_support_plane"), + _position("in_plane_error_m"), _rotation("tilt_error_rad"), ), ), @@ -328,6 +358,14 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: _rotation(), _NO_DROP, ), + _stage( + "settled", + "terminal_settle", + "task_goal_miss", + _position(), + _rotation(), + _NO_DROP, + ), _HELD, ), "coordinated_placement": ( @@ -343,10 +381,19 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: "released", "release_end", "release_failure", - _bool("left_release_established"), - _bool("right_release_established"), + _bool("placing_release_established"), + _bool("support_retained"), + ), + _stage( + "stable", + "settle_interval", + "task_goal_miss", + _position("relative_position_error_m"), + _rotation("relative_rotation_error_rad"), + _position("position_drift_m"), + _rotation("rotation_drift_rad"), + _NO_DROP, ), - _STABLE, ), "hand_over": ( _GRASP, @@ -365,6 +412,12 @@ def _rotation(name: str = "rotation_error_rad") -> Criterion: _bool("receiver_retained"), _NO_DROP, ), + _stage( + "receiver_released", + "terminal_settle", + "release_failure", + _bool("release_established"), + ), _stage( "placed", "terminal_settle", @@ -444,6 +497,30 @@ def evaluate_skill_stages( return PhysicalEvaluation(env_index, tuple(results), STANDARD_VERSION) +def evaluate_skill_stage_batch( + skill_id: str, + evidence_by_env: Sequence[Mapping[str, StageEvidence]], +) -> tuple[PhysicalEvaluation, ...]: + """Evaluate row-local stage evidence in stable environment order. + + This is the bridge between batched replay observations and the standard's + deliberately row-local evaluator. Observation adapters remain responsible + for sampling physical state and naming the measurement method; this helper + only assigns environment indices and applies the frozen criteria. + + Args: + skill_id: Canonical snake-case Atomic Skill identifier. + evidence_by_env: One stage-evidence mapping per environment row. + + Returns: + Physical evaluations ordered by environment index. + """ + return tuple( + evaluate_skill_stages(skill_id, evidence, env_index=env_index) + for env_index, evidence in enumerate(evidence_by_env) + ) + + def validate_standard_evaluation(skill_id: str, result: PhysicalEvaluation) -> None: """Reject a claimed v1 result inconsistent with the executable criteria. diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_dual_arm_providers.py b/scripts/benchmark/motion_generation/scenarios/atomic_dual_arm_providers.py index cf0c85abf..666012e96 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_dual_arm_providers.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_dual_arm_providers.py @@ -60,6 +60,8 @@ from ..atomic_skill import AtomicSkillCaseProvider from ..config import SuiteCfg, TrackCfg +from ..contracts import PhysicalEvaluation, ReplayObservationWindow +from ..evaluation import StageEvidence, evaluate_skill_stage_batch from ..models import BenchmarkCase, CaseOutcome from .atomic_task import ( _ExecutionObservation, @@ -69,10 +71,16 @@ _case_pose, _difficulty, _float_vector, + _gripper_holding, + _gripper_released, + _merge_stage_evidence_batches, + _minimum_height_above_initial_support, _motion_valid_mask, _randomized_vector_batch, _randomization_parameters, + _replay_trace, _snapshot_case_qpos, + _stage_evidence_batch, _with_plan_options, register_atomic_skill_provider, ) @@ -93,6 +101,130 @@ _GRIPPER_HOLD_GAP = 0.001 +class _HandOverEvaluator: + """Evaluate measured pickup, transfer, handover, and placement stages.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + parameters = case.case_parameters + parts = dict(parameters["control_parts"]) + source = str(parts["source"]["grasp"]) + destination = str(parts["destination"]["grasp"]) + pickup = _replay_trace(observation, "handover_pickup") + transfer = _replay_trace(observation, "handover_transfer") + exchange = _replay_trace(observation, "handover_exchange") + settle = _replay_trace(observation, "handover_settle") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if pickup is not None: + grasped = _gripper_holding( + pickup, observation, source, gap=_GRIPPER_HOLD_GAP + ) + if grasped is not None: + batches.append( + _stage_evidence_batch( + stage_id="grasped", + window="grasp_end", + measurements={"grasp_established": grasped}, + method="source_gripper_joint_state", + source="measured_proxy", + ) + ) + if transfer is not None: + source_held = _gripper_holding( + transfer, observation, source, gap=_GRIPPER_HOLD_GAP + ) + destination_held = _gripper_holding( + transfer, observation, destination, gap=_GRIPPER_HOLD_GAP + ) + min_height = _minimum_height_above_initial_support(transfer, case) + if ( + source_held is not None + and destination_held is not None + and min_height is not None + ): + batches.append( + _stage_evidence_batch( + stage_id="transferred", + window="transfer_end", + measurements={ + "both_grasps_established": source_held & destination_held, + "min_height_above_support_m": min_height, + }, + method="dual_gripper_and_object_trace", + source="measured_proxy", + ) + ) + if exchange is not None: + giver_released = _gripper_released(exchange, observation, source) + receiver_retained = _gripper_holding( + exchange, observation, destination, gap=_GRIPPER_HOLD_GAP + ) + min_height = _minimum_height_above_initial_support(exchange, case) + if ( + giver_released is not None + and receiver_retained is not None + and min_height is not None + ): + batches.append( + _stage_evidence_batch( + stage_id="handed_over", + window="handover_end", + measurements={ + "giver_released": giver_released, + "receiver_retained": receiver_retained, + "min_height_above_support_m": min_height, + }, + method="dual_gripper_and_object_trace", + source="measured_proxy", + ) + ) + if settle is not None and settle.object_pose is not None: + target = _case_pose( + case, "target_object_pose", device=settle.object_pose.device + ).to(dtype=settle.object_pose.dtype) + final = settle.object_pose[-1] + xy_error = torch.linalg.vector_norm( + final[:, :2, 3] - target[:, :2, 3], dim=-1 + ) + signed_z_error = final[:, 2, 3] - target[:, 2, 3] + min_height = _minimum_height_above_initial_support(settle, case) + if min_height is not None: + receiver_released = _gripper_released(settle, observation, destination) + if receiver_released is not None: + batches.append( + _stage_evidence_batch( + stage_id="receiver_released", + window="terminal_settle", + measurements={"release_established": receiver_released}, + method="receiver_gripper_joint_state", + source="measured_proxy", + ) + ) + batches.append( + _stage_evidence_batch( + stage_id="placed", + window="terminal_settle", + measurements={ + "xy_error_m": xy_error, + "signed_z_error_m": signed_z_error, + "min_height_above_support_m": min_height, + }, + method="terminal_object_pose_trace", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("hand_over", evidence) + + def _rotation_error(actual: torch.Tensor, target: torch.Tensor) -> torch.Tensor: """Return the angle of the rotation between two ``(B, 4, 4)`` poses.""" relative = target[:, :3, :3].transpose(-1, -2) @ actual[:, :3, :3] @@ -133,6 +265,253 @@ def _grippers_holding( return held +class _CoordinatedPickmentEvaluator: + """Evaluate dual grasp, lift, synchronized move, and retention.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + parts = dict(case.case_parameters["control_parts"]) + left = str(parts["left"]["grasp"]) + right = str(parts["right"]["grasp"]) + grasp = _replay_trace(observation, "dual_pick_grasp") + lift = _replay_trace(observation, "dual_pick_lift") + move = _replay_trace(observation, "dual_pick_move") + settle = _replay_trace(observation, "dual_pick_settle") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if grasp is not None: + left_held = _gripper_holding( + grasp, observation, left, gap=_GRIPPER_HOLD_GAP + ) + right_held = _gripper_holding( + grasp, observation, right, gap=_GRIPPER_HOLD_GAP + ) + if left_held is not None and right_held is not None: + batches.append( + _stage_evidence_batch( + stage_id="dual_grasped", + window="grasp_end", + measurements={ + "left_grasp_established": left_held, + "right_grasp_established": right_held, + }, + method="dual_gripper_joint_state", + source="measured_proxy", + ) + ) + if lift is not None and lift.object_pose is not None: + initial = torch.as_tensor( + case.case_parameters["object_initial_pose"], + dtype=lift.object_pose.dtype, + device=lift.object_pose.device, + ) + min_height = _minimum_height_above_initial_support(lift, case) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="lifted", + window="lift_end", + measurements={ + "lift_delta_m": lift.object_pose[-1, :, 2, 3] + - initial[:, 2, 3], + "min_height_above_support_m": min_height, + }, + method="object_lift_trace", + ) + ) + if move is not None and move.object_pose is not None: + final = move.object_pose[-1] + target = _case_pose(case, "target_object_pose", device=final.device).to( + dtype=final.dtype + ) + min_height = _minimum_height_above_initial_support(move, case) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="moved", + window="actuation_end", + measurements={ + "position_error_m": torch.linalg.vector_norm( + final[:, :3, 3] - target[:, :3, 3], dim=-1 + ), + "rotation_error_rad": _rotation_error(final, target), + "min_height_above_support_m": min_height, + }, + method="object_move_trace", + ) + ) + if settle is not None and settle.object_pose is not None: + left_held = _gripper_holding( + settle, observation, left, gap=_GRIPPER_HOLD_GAP + ) + right_held = _gripper_holding( + settle, observation, right, gap=_GRIPPER_HOLD_GAP + ) + min_height = _minimum_height_above_initial_support(settle, case) + final = settle.object_pose[-1] + target = _case_pose(case, "target_object_pose", device=final.device).to( + dtype=final.dtype + ) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="settled", + window="terminal_settle", + measurements={ + "position_error_m": torch.linalg.vector_norm( + final[:, :3, 3] - target[:, :3, 3], dim=-1 + ), + "rotation_error_rad": _rotation_error(final, target), + "min_height_above_support_m": min_height, + }, + method="terminal_object_pose_trace", + ) + ) + if left_held is not None and right_held is not None: + batches.append( + _stage_evidence_batch( + stage_id="held", + window="terminal_settle", + measurements={ + "object_retained": left_held & right_held, + "min_height_above_support_m": min_height, + }, + method="dual_gripper_and_object_settle_trace", + source="measured_proxy", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("coordinated_pickment", evidence) + + +class _CoordinatedPlacementEvaluator: + """Evaluate relative alignment, role-aware release, and settled stability.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + parameters = case.case_parameters + parts = dict(parameters["control_parts"]) + placing_part = str(parts["placing"]["grasp"]) + support_part = str(parts["support"]["grasp"]) + support_id = str(parameters["support_object_id"]) + target = _case_pose( + case, + "placing_relative_target_pose", + device=observation.execution_success.device, + ) + aligned = _replay_trace(observation, "coordinated_place_alignment") + release = _replay_trace(observation, "coordinated_place_release") + settle = _replay_trace(observation, "coordinated_place_settle") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if ( + aligned is not None + and aligned.object_pose is not None + and support_id in aligned.scene_object_poses + ): + placing = aligned.object_pose[-1] + support = aligned.scene_object_poses[support_id][-1] + relative = torch.bmm(pose_inv(support), placing) + local_target = target.to(device=relative.device, dtype=relative.dtype) + min_height = _minimum_height_above_initial_support(aligned, case) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="aligned", + window="actuation_end", + measurements={ + "relative_position_error_m": torch.linalg.vector_norm( + relative[:, :3, 3] - local_target[:, :3, 3], dim=-1 + ), + "relative_rotation_error_rad": _rotation_error( + relative, local_target + ), + "min_height_above_support_m": min_height, + }, + method="relative_object_pose_trace", + ) + ) + if release is not None: + placing_released = _gripper_released(release, observation, placing_part) + support_retained = _gripper_holding( + release, observation, support_part, gap=_GRIPPER_HOLD_GAP + ) + if placing_released is not None and support_retained is not None: + batches.append( + _stage_evidence_batch( + stage_id="released", + window="release_end", + measurements={ + "placing_release_established": placing_released, + "support_retained": support_retained, + }, + method="role_gripper_joint_state", + source="measured_proxy", + ) + ) + if ( + settle is not None + and settle.object_pose is not None + and support_id in settle.scene_object_poses + and settle.object_pose.shape[0] >= 2 + ): + relative = torch.matmul( + pose_inv(settle.scene_object_poses[support_id]), settle.object_pose + ) + relative_start = relative[:1].expand_as(relative) + final_relative = relative[-1] + local_target = target.to( + device=final_relative.device, dtype=final_relative.dtype + ) + min_height = _minimum_height_above_initial_support(settle, case) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="stable", + window="settle_interval", + measurements={ + "relative_position_error_m": torch.linalg.vector_norm( + final_relative[:, :3, 3] - local_target[:, :3, 3], + dim=-1, + ), + "relative_rotation_error_rad": _rotation_error( + final_relative, local_target + ), + "position_drift_m": torch.linalg.vector_norm( + relative[:, :, :3, 3] - relative_start[:, :, :3, 3], + dim=-1, + ).amax(dim=0), + "rotation_drift_rad": _rotation_error( + relative.flatten(0, 1), + relative_start.flatten(0, 1), + ) + .reshape(relative.shape[:2]) + .amax(dim=0), + "min_height_above_support_m": min_height, + }, + method="relative_object_settle_drift_trace", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("coordinated_placement", evidence) + + # --------------------------------------------------------------------------- # Reference compile # --------------------------------------------------------------------------- @@ -651,6 +1030,48 @@ def frozen_grasps( ).items() } + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect each physical handover phase and terminal settle.""" + del case + pickup = compiled.segment(0, "pickup_close") + transfer_start = compiled.segment(0, "pickup_transport") + transfer_end = compiled.segment(0, "receive_close") + exchange = compiled.segment(0, "handover_release") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "handover_pickup", + pickup.start, + pickup.stop, + frozenset({"gripper_qpos"}), + ), + ReplayObservationWindow( + "handover_transfer", + transfer_start.start, + transfer_end.stop, + frozenset({"object_pose", "gripper_qpos"}), + ), + ReplayObservationWindow( + "handover_exchange", + exchange.start, + exchange.stop, + frozenset({"object_pose", "gripper_qpos"}), + ), + ReplayObservationWindow( + "handover_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose", "gripper_qpos"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _HandOverEvaluator: + """Return the v1 HandOver evaluator.""" + return _HandOverEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -861,6 +1282,47 @@ def build_invocation( skill_options=_pickment_options(case.case_parameters, device), ) + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect dual-grasp, lift, move, and terminal retention evidence.""" + del case + close = compiled.segment(0, "close") + lift = compiled.segment(0, "lift") + move = compiled.segment(0, "move") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "dual_pick_grasp", + close.start, + close.stop, + frozenset({"gripper_qpos"}), + ), + ReplayObservationWindow( + "dual_pick_lift", + lift.start, + lift.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "dual_pick_move", + move.start, + move.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "dual_pick_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose", "gripper_qpos"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _CoordinatedPickmentEvaluator: + """Return the v1 CoordinatedPickment evaluator.""" + return _CoordinatedPickmentEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -1292,6 +1754,43 @@ def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: """Release residual contact impulses once the first PickUp lifts.""" return compiled.segment(0, "lift").start + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect relative alignment, release, and stability evidence.""" + support_id = str(case.case_parameters["support_object_id"]) + alignment = compiled.segment(2, "approach") + release = compiled.segment(2, "release") + waypoint_count = compiled.trajectory.waypoint_count + pose_channels = frozenset({"object_pose", "scene_object_pose"}) + return ( + ReplayObservationWindow( + "coordinated_place_alignment", + alignment.start, + alignment.stop, + pose_channels, + scene_object_ids=(support_id,), + ), + ReplayObservationWindow( + "coordinated_place_release", + release.start, + release.stop, + frozenset({"gripper_qpos"}), + ), + ReplayObservationWindow( + "coordinated_place_settle", + waypoint_count, + waypoint_count, + pose_channels, + include_settle=True, + scene_object_ids=(support_id,), + ), + ) + + def physical_evaluator(self) -> _CoordinatedPlacementEvaluator: + """Return the v1 CoordinatedPlacement evaluator.""" + return _CoordinatedPlacementEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_skill_providers.py b/scripts/benchmark/motion_generation/scenarios/atomic_skill_providers.py index 3493542d1..3659e69c0 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_skill_providers.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_skill_providers.py @@ -56,6 +56,8 @@ from ..atomic_skill import AtomicSkillCaseProvider from ..config import SuiteCfg, TrackCfg +from ..contracts import PhysicalEvaluation, ReplayObservationWindow +from ..evaluation import StageEvidence, evaluate_skill_stage_batch from ..models import BenchmarkCase, CaseOutcome from .atomic_task import ( _TOP_DOWN_ROTATION, @@ -66,11 +68,18 @@ _case_pose, _difficulty, _float_vector, + _gripper_holding, + _gripper_released, + _merge_stage_evidence_batches, + _minimum_height_above_initial_support, _motion_valid_mask, _object_pose_success, _randomization_parameters, _randomized_vector_batch, + _replay_trace, + _signed_rotation_about_axis, _snapshot_case_qpos, + _stage_evidence_batch, register_atomic_skill_provider, ) @@ -85,6 +94,265 @@ _TASK_POSITION_TOLERANCE_M = 0.01 +class _OpenDoorEvaluator: + """Evaluate hinge engagement, opening, and measured gripper release.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + actuation = _replay_trace(observation, "door_actuation") + release = _replay_trace(observation, "door_release") + if actuation is None or actuation.articulation_joint is None: + return evaluate_skill_stage_batch( + "open_door", tuple({} for _ in range(case.batch_size)) + ) + joint = actuation.articulation_joint + displacement = joint[-1] - joint[0] + grasped = displacement.abs() > float( + case.case_parameters.get("contact_minimum_joint_delta", 1.0e-5) + ) + opened_error = ( + joint[-1] - float(case.case_parameters["target_joint_position"]) + ).abs() + batches: list[Sequence[Mapping[str, StageEvidence]]] = [ + _stage_evidence_batch( + stage_id="grasped", + window="grasp_end", + measurements={"grasp_established": grasped}, + method="hinge_motion_proxy", + source="measured_proxy", + ), + _stage_evidence_batch( + stage_id="opened", + window="actuation_end", + measurements={"hinge_angle_error_rad": opened_error}, + method="hinge_segment_end", + ), + ] + if release is not None: + part = next(iter(observation.gripper_open_qpos), "") + released = _gripper_released(release, observation, part) + if released is not None: + batches.append( + _stage_evidence_batch( + stage_id="released", + window="release_end", + measurements={"release_established": released}, + method="gripper_joint_state", + source="measured_proxy", + ) + ) + return evaluate_skill_stage_batch( + "open_door", _merge_stage_evidence_batches(*batches) + ) + + +class _PourEvaluator: + """Evaluate the pour rotation plus settled lift and retention.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + actuation = _replay_trace(observation, "pour_actuation") + settle = _replay_trace(observation, "pour_settle") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if actuation is not None and actuation.object_pose is not None: + poses = actuation.object_pose + axis = torch.as_tensor( + case.case_parameters["object_rotation_axis"], + dtype=poses.dtype, + device=poses.device, + ) + axis = axis / torch.linalg.vector_norm(axis) + reference = poses[0, :, :3, :3] + world_axis = reference @ axis + angles = torch.stack( + [ + _signed_rotation_about_axis(reference, pose[:, :3, :3], world_axis) + for pose in poses + ] + ) + sign = float(case.case_parameters.get("object_rotation_sign", 1.0)) + peak = (angles * sign).amax(dim=0) + error = (peak - abs(float(case.case_parameters["rotate_angle_rad"]))).abs() + min_height = _minimum_height_above_initial_support(actuation, case) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="poured", + window="actuation_interval", + measurements={ + "peak_signed_axis_angle_error_rad": error, + "min_height_above_support_m": min_height, + }, + method="object_pose_trace", + ) + ) + if settle is not None and settle.object_pose is not None: + part = next(iter(observation.gripper_open_qpos), "") + retained = _gripper_holding(settle, observation, part) + min_height = _minimum_height_above_initial_support(settle, case) + initial = torch.as_tensor( + case.case_parameters["object_initial_pose"], + dtype=settle.object_pose.dtype, + device=settle.object_pose.device, + ) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="settled", + window="terminal_settle", + measurements={ + "lift_delta_m": ( + settle.object_pose[-1, :, 2, 3] - initial[:, 2, 3] + ), + "min_height_above_support_m": min_height, + }, + method="terminal_object_pose_trace", + ) + ) + if retained is not None: + batches.append( + _stage_evidence_batch( + stage_id="held", + window="terminal_settle", + measurements={ + "object_retained": retained, + "min_height_above_support_m": min_height, + }, + method="gripper_and_object_settle_trace", + source="measured_proxy", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("pour", evidence) + + +class _AxisAlignEvaluator: + """Evaluate grasp proxy, aligned object axis, and terminal retention.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + actuation = _replay_trace(observation, "axis_align_actuation") + settle = _replay_trace(observation, "axis_align_settle") + if actuation is None or actuation.object_pose is None: + return evaluate_skill_stage_batch( + "axis_align", tuple({} for _ in range(case.batch_size)) + ) + poses = actuation.object_pose + internal = torch.as_tensor( + case.case_parameters["internal_axis"], + dtype=poses.dtype, + device=poses.device, + ) + target = torch.as_tensor( + case.case_parameters["target_axis"], + dtype=poses.dtype, + device=poses.device, + ) + initial_axis = torch.nn.functional.normalize( + poses[0, :, :3, :3] @ internal, dim=-1 + ) + final_axis = torch.nn.functional.normalize( + poses[-1, :, :3, :3] @ internal, dim=-1 + ) + moved = torch.arccos( + torch.clamp((initial_axis * final_axis).sum(dim=-1), -1.0, 1.0) + ) + aligned_error = torch.arccos(torch.clamp(final_axis @ target, -1.0, 1.0)) + actuation_height = _minimum_height_above_initial_support(actuation, case) + if actuation_height is None: + return evaluate_skill_stage_batch( + "axis_align", tuple({} for _ in range(case.batch_size)) + ) + initial = torch.as_tensor( + case.case_parameters["object_initial_pose"], + dtype=poses.dtype, + device=poses.device, + ) + batches: list[Sequence[Mapping[str, StageEvidence]]] = [ + _stage_evidence_batch( + stage_id="grasped", + window="grasp_end", + measurements={"grasp_established": moved > 1.0e-5}, + method="object_axis_motion_proxy", + source="measured_proxy", + ), + _stage_evidence_batch( + stage_id="aligned", + window="actuation_end", + measurements={ + "axis_angle_error_rad": aligned_error, + "lift_delta_m": poses[-1, :, 2, 3] - initial[:, 2, 3], + "min_height_above_support_m": actuation_height, + }, + method="object_pose_trace", + ), + ] + if settle is not None: + min_height = _minimum_height_above_initial_support(settle, case) + part = next(iter(observation.gripper_open_qpos), "") + retained = _gripper_holding(settle, observation, part) + if settle.object_pose is not None and min_height is not None: + final_pose = settle.object_pose[-1] + final_axis = torch.nn.functional.normalize( + final_pose[:, :3, :3] @ internal, dim=-1 + ) + terminal_error = torch.arccos( + torch.clamp(final_axis @ target, -1.0, 1.0) + ) + settle_initial = initial.to( + device=final_pose.device, dtype=final_pose.dtype + ) + batches.append( + _stage_evidence_batch( + stage_id="settled", + window="terminal_settle", + measurements={ + "axis_angle_error_rad": terminal_error, + "lift_delta_m": ( + final_pose[:, 2, 3] - settle_initial[:, 2, 3] + ), + "min_height_above_support_m": min_height, + }, + method="terminal_object_pose_trace", + ) + ) + if retained is not None: + batches.append( + _stage_evidence_batch( + stage_id="held", + window="terminal_settle", + measurements={ + "object_retained": retained, + "min_height_above_support_m": min_height, + }, + method="gripper_and_object_settle_trace", + source="measured_proxy", + ) + ) + return evaluate_skill_stage_batch( + "axis_align", _merge_stage_evidence_batches(*batches) + ) + + def _unit_vector(value: object, *, name: str, default: Sequence[float]) -> list[float]: """Resolve a non-zero 3-vector and normalize it.""" vector = _float_vector(value, name=name, length=3, default=default) @@ -626,6 +894,32 @@ def articulation_effect_segment(self, case: BenchmarkCase) -> str: del case return "open" + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect hinge actuation and gripper release samples.""" + del case + actuation = compiled.segment(0, "open") + release = compiled.segment(0, "release") + return ( + ReplayObservationWindow( + "door_actuation", + actuation.start, + actuation.stop, + frozenset({"articulation_joint"}), + ), + ReplayObservationWindow( + "door_release", + release.start, + release.stop, + frozenset({"gripper_qpos"}), + ), + ) + + def physical_evaluator(self) -> _OpenDoorEvaluator: + """Return the v1 OpenDoor evaluator.""" + return _OpenDoorEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -891,6 +1185,33 @@ def object_rotation_segment(self, compiled: CompiledTrajectory) -> tuple[int, in segment = compiled.segment(1, "pour") return segment.start, segment.stop + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect pour motion plus terminal lift and retention evidence.""" + del case + segment = compiled.segment(1, "pour") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "pour_actuation", + segment.start, + segment.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "pour_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose", "gripper_qpos"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _PourEvaluator: + """Return the v1 Pour evaluator.""" + return _PourEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -1218,6 +1539,33 @@ def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: """Release residual contact impulses once the held object starts moving.""" return compiled.segment(0, "manipulate").start + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect alignment motion and the terminal settle interval.""" + del case + segment = compiled.segment(0, "manipulate") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "axis_align_actuation", + segment.start, + segment.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "axis_align_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose", "gripper_qpos"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _AxisAlignEvaluator: + """Return the v1 AxisAlign evaluator.""" + return _AxisAlignEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -1261,6 +1609,109 @@ def task_result( # --------------------------------------------------------------------------- +class _PushObjectEvaluator: + """Evaluate contact proxy, planar displacement, and settled support pose.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + contact = _replay_trace(observation, "push_contact") + actuation = _replay_trace(observation, "push_actuation") + settle = _replay_trace(observation, "push_settle") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if contact is not None and contact.object_pose is not None: + contact_motion = torch.linalg.vector_norm( + contact.object_pose[-1, :, :3, 3] - contact.object_pose[0, :, :3, 3], + dim=-1, + ) + batches.append( + _stage_evidence_batch( + stage_id="contacted", + window="contact_end", + measurements={ + "contact_established": contact_motion + > float( + case.case_parameters.get( + "contact_motion_threshold_m", 1.0e-5 + ) + ) + }, + method="object_motion_during_contact", + source="measured_proxy", + ) + ) + if actuation is not None and actuation.object_pose is not None: + poses = actuation.object_pose + initial = torch.as_tensor( + case.case_parameters["object_initial_pose"], + dtype=poses.dtype, + device=poses.device, + ) + target = _case_pose(case, "target_object_pose", device=poses.device).to( + dtype=poses.dtype + ) + normal = torch.nn.functional.normalize(target[:, :3, 2], dim=-1) + desired = target[:, :3, 3] - initial[:, :3, 3] + desired -= (desired * normal).sum(dim=-1, keepdim=True) * normal + direction = torch.nn.functional.normalize(desired, dim=-1) + actual = poses[-1, :, :3, 3] - initial[:, :3, 3] + residual = target[:, :3, 3] - poses[-1, :, :3, 3] + residual -= (residual * normal).sum(dim=-1, keepdim=True) * normal + batches.append( + _stage_evidence_batch( + stage_id="pushed", + window="actuation_end", + measurements={ + "signed_displacement_m": (actual * direction).sum(dim=-1), + "in_plane_error_m": torch.linalg.vector_norm(residual, dim=-1), + }, + method="object_push_segment_end", + ) + ) + if settle is not None and settle.object_pose is not None: + final = settle.object_pose[-1] + target = _case_pose(case, "target_object_pose", device=final.device).to( + dtype=final.dtype + ) + normal = torch.nn.functional.normalize(target[:, :3, 2], dim=-1) + final_normal = torch.nn.functional.normalize(final[:, :3, 2], dim=-1) + support_error = ( + ((final[:, :3, 3] - target[:, :3, 3]) * normal).sum(dim=-1).abs() + ) + residual = target[:, :3, 3] - final[:, :3, 3] + residual -= (residual * normal).sum(dim=-1, keepdim=True) * normal + batches.append( + _stage_evidence_batch( + stage_id="on_support_plane", + window="terminal_settle", + measurements={ + "on_support_plane": support_error + <= float( + case.case_parameters.get( + "support_plane_tolerance_m", + _TASK_POSITION_TOLERANCE_M, + ) + ), + "in_plane_error_m": torch.linalg.vector_norm(residual, dim=-1), + "tilt_error_rad": torch.arccos( + torch.clamp((final_normal * normal).sum(dim=-1), -1.0, 1.0) + ), + }, + method="terminal_object_support_pose", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("push_object", evidence) + + def _push_poses( object_pose: torch.Tensor, target_pose: torch.Tensor, @@ -1425,6 +1876,12 @@ def generate_case( "object_tilt_threshold_rad", _TASK_ROTATION_TOLERANCE_RAD ) ), + "contact_motion_threshold_m": float( + config.get("contact_motion_threshold_m", 1.0e-5) + ), + "support_plane_tolerance_m": float( + config.get("support_plane_tolerance_m", _TASK_POSITION_TOLERANCE_M) + ), "task_success_scope": "physical_object_pose_after_push", "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), @@ -1471,6 +1928,40 @@ def build_invocation( ), ) + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect contact, push, and terminal support-plane evidence.""" + del case + contact = compiled.segment(0, "contact") + push = compiled.segment(0, "push") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "push_contact", + contact.start, + contact.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "push_actuation", + push.start, + push.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "push_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _PushObjectEvaluator: + """Return the v1 PushObject evaluator.""" + return _PushObjectEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py index 825999be2..7c2e29b44 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_task.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -21,7 +21,7 @@ import math import time from collections.abc import Callable, Mapping, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -69,8 +69,12 @@ from ..atomic_skill import AtomicSkillCaseProvider from ..contracts import ( ExecutionObservation as _ExecutionObservation, + PhysicalEvaluation, + ReplayObservationTrace, + ReplayObservationWindow, validate_physical_evaluations, ) +from ..evaluation import StageEvidence, evaluate_skill_stage_batch from ..config import ( SuiteCfg, TrackCfg, @@ -140,6 +144,38 @@ class _ReplayMetrics: articulation_joint_peak_signed_delta: torch.Tensor | None = None articulation_joint_effect_end: torch.Tensor | None = None object_rotation_peak_signed_rad: torch.Tensor | None = None + replay_traces: Mapping[str, ReplayObservationTrace] = field(default_factory=dict) + + +@dataclass +class _ReplayTraceAccumulator: + """Mutable samples for one validated replay observation window.""" + + window: ReplayObservationWindow + tcp_pose: list[torch.Tensor] = field(default_factory=list) + object_pose: list[torch.Tensor] = field(default_factory=list) + articulation_joint: list[torch.Tensor] = field(default_factory=list) + gripper_qpos: dict[str, list[torch.Tensor]] = field(default_factory=dict) + scene_object_poses: dict[str, list[torch.Tensor]] = field(default_factory=dict) + + def finish(self) -> ReplayObservationTrace: + """Stack synchronized samples into the public immutable trace.""" + return ReplayObservationTrace( + tcp_pose=torch.stack(self.tcp_pose) if self.tcp_pose else None, + object_pose=torch.stack(self.object_pose) if self.object_pose else None, + articulation_joint=( + torch.stack(self.articulation_joint) + if self.articulation_joint + else None + ), + gripper_qpos={ + part: torch.stack(values) for part, values in self.gripper_qpos.items() + }, + scene_object_poses={ + object_id: torch.stack(values) + for object_id, values in self.scene_object_poses.items() + }, + ) @dataclass(frozen=True) @@ -767,6 +803,615 @@ def _executed_final_pose_errors( return translation, rotation +def _stage_evidence_batch( + *, + stage_id: str, + window: str, + measurements: Mapping[str, torch.Tensor], + method: str, + source: Literal["simulation_state", "contact_sensor", "measured_proxy"] = ( + "simulation_state" + ), +) -> tuple[Mapping[str, StageEvidence], ...]: + """Convert measured batch scalars into row-local standard evidence.""" + if not measurements: + raise ValueError("Stage evidence requires at least one measurement.") + values = tuple(measurements.values()) + if any(value.dim() != 1 for value in values): + raise ValueError("Stage measurement values must be one-dimensional batches.") + batch_size = values[0].shape[0] + if any(value.shape[0] != batch_size for value in values): + raise ValueError("Stage measurement batches must have the same size.") + + def scalar(value: torch.Tensor, index: int) -> float | bool: + item = value[index].item() + return bool(item) if value.dtype == torch.bool else float(item) + + rows: list[Mapping[str, StageEvidence]] = [] + for env_index in range(batch_size): + row_measurements = { + name: scalar(value, env_index) for name, value in measurements.items() + } + if any( + type(value) is not bool and not math.isfinite(value) + for value in row_measurements.values() + ): + rows.append({}) + continue + rows.append( + { + stage_id: StageEvidence( + window=window, + source=source, + measurements=row_measurements, + method=method, + ) + } + ) + return tuple(rows) + + +def _merge_stage_evidence_batches( + *batches: Sequence[Mapping[str, StageEvidence]], +) -> tuple[Mapping[str, StageEvidence], ...]: + """Merge stage batches while preserving stable environment order.""" + if not batches: + return () + batch_size = len(batches[0]) + if any(len(batch) != batch_size for batch in batches): + raise ValueError("Stage evidence batches must have the same environment size.") + merged: list[dict[str, StageEvidence]] = [] + for env_index in range(batch_size): + row: dict[str, StageEvidence] = {} + for batch in batches: + overlap = set(row) & set(batch[env_index]) + if overlap: + raise ValueError(f"Duplicate stage evidence: {sorted(overlap)}.") + row.update(batch[env_index]) + merged.append(row) + return tuple(merged) + + +def _replay_trace( + observation: _ExecutionObservation, observation_id: str +) -> ReplayObservationTrace | None: + """Return one measured replay trace without fabricating missing evidence.""" + return observation.replay_traces.get(observation_id) + + +def _minimum_height_above_initial_support( + trace: ReplayObservationTrace, + case: BenchmarkCase, +) -> torch.Tensor | None: + """Return the minimum object-reference height above its frozen start plane.""" + if trace.object_pose is None: + return None + initial = torch.as_tensor( + case.case_parameters["object_initial_pose"], + dtype=trace.object_pose.dtype, + device=trace.object_pose.device, + ) + return trace.object_pose[:, :, 2, 3].min(dim=0).values - initial[:, 2, 3] + + +def _gripper_released( + trace: ReplayObservationTrace, + observation: _ExecutionObservation, + part: str, +) -> torch.Tensor | None: + """Return whether the final window sample is nearer open than grasp.""" + qpos = trace.gripper_qpos.get(part) + open_qpos = observation.gripper_open_qpos.get(part) + grasp_qpos = observation.gripper_grasp_qpos.get(part) + if qpos is None or open_qpos is None or grasp_qpos is None: + return None + final = qpos[-1] + open_qpos = open_qpos.to(device=final.device, dtype=final.dtype) + grasp_qpos = grasp_qpos.to(device=final.device, dtype=final.dtype) + return torch.linalg.vector_norm( + final - open_qpos, dim=-1 + ) <= torch.linalg.vector_norm(final - grasp_qpos, dim=-1) + + +def _gripper_holding( + trace: ReplayObservationTrace, + observation: _ExecutionObservation, + part: str, + *, + gap: float = 1.0e-3, +) -> torch.Tensor | None: + """Return whether the final sample stopped between open and grasp commands.""" + qpos = trace.gripper_qpos.get(part) + open_qpos = observation.gripper_open_qpos.get(part) + grasp_qpos = observation.gripper_grasp_qpos.get(part) + if qpos is None or open_qpos is None or grasp_qpos is None: + return None + final = qpos[-1] + open_qpos = open_qpos.to(device=final.device, dtype=final.dtype) + grasp_qpos = grasp_qpos.to(device=final.device, dtype=final.dtype) + direction = torch.sign(grasp_qpos - open_qpos) + closed = (final - open_qpos) * direction + remaining = (grasp_qpos - final) * direction + return ((closed >= gap) & (remaining >= gap)).all(dim=-1) + + +def _pose_rotation_error(actual: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Return the geodesic rotation error for matching pose batches.""" + relative = target[..., :3, :3].transpose(-1, -2) @ actual[..., :3, :3] + trace = torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) + return torch.arccos(torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0)) + + +class _PickUpEvaluator: + """Evaluate grasp, lift, and terminal retention from replayed state.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + grasp = _replay_trace(observation, "pickup_grasp") + lift = _replay_trace(observation, "pickup_lift") + settle = _replay_trace(observation, "pickup_settle") + part = next(iter(observation.gripper_open_qpos), "") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if grasp is not None: + established = _gripper_holding(grasp, observation, part) + if established is not None: + batches.append( + _stage_evidence_batch( + stage_id="grasped", + window="grasp_end", + measurements={"grasp_established": established}, + method="gripper_joint_state", + source="measured_proxy", + ) + ) + if lift is not None and lift.object_pose is not None: + initial = torch.as_tensor( + case.case_parameters["object_initial_pose"], + dtype=lift.object_pose.dtype, + device=lift.object_pose.device, + ) + min_height = _minimum_height_above_initial_support(lift, case) + if min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="lifted", + window="lift_end", + measurements={ + "lift_delta_m": lift.object_pose[-1, :, 2, 3] + - initial[:, 2, 3], + "min_height_above_support_m": min_height, + }, + method="object_lift_trace", + ) + ) + if settle is not None: + retained = _gripper_holding(settle, observation, part) + min_height = _minimum_height_above_initial_support(settle, case) + if retained is not None and min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="held", + window="terminal_settle", + measurements={ + "object_retained": retained, + "min_height_above_support_m": min_height, + }, + method="gripper_and_object_settle_trace", + source="measured_proxy", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("pick_up", evidence) + + +class _MoveHeldObjectEvaluator: + """Evaluate transport plus the held object's settled pose and retention.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + transport = _replay_trace(observation, "held_object_transport") + settle = _replay_trace(observation, "held_object_settle") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + for stage_id, window, trace in ( + ("transported", "actuation_end", transport), + ("settled", "terminal_settle", settle), + ): + if trace is None or trace.object_pose is None: + continue + final = trace.object_pose[-1] + target = _case_pose( + case, "target_object_pose", device=trace.object_pose.device + ).to(dtype=trace.object_pose.dtype) + min_height = _minimum_height_above_initial_support(trace, case) + if min_height is None: + continue + measurements: dict[str, torch.Tensor] = { + "position_error_m": torch.linalg.vector_norm( + final[:, :3, 3] - target[:, :3, 3], dim=-1 + ), + "rotation_error_rad": _pose_rotation_error(final, target), + "min_height_above_support_m": min_height, + } + batches.append( + _stage_evidence_batch( + stage_id=stage_id, + window=window, + measurements=measurements, + method=( + "object_transport_trace" + if stage_id == "transported" + else "terminal_object_pose_trace" + ), + ) + ) + if settle is not None: + part = next(iter(observation.gripper_open_qpos), "") + retained = _gripper_holding(settle, observation, part) + min_height = _minimum_height_above_initial_support(settle, case) + if retained is not None and min_height is not None: + batches.append( + _stage_evidence_batch( + stage_id="held", + window="terminal_settle", + measurements={ + "object_retained": retained, + "min_height_above_support_m": min_height, + }, + method="gripper_and_object_settle_trace", + source="measured_proxy", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("move_held_object", evidence) + + +class _PlaceEvaluator: + """Evaluate release, terminal placement, and post-release stability.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + release = _replay_trace(observation, "place_release") + settle = _replay_trace(observation, "place_settle") + part = next(iter(observation.gripper_open_qpos), "") + batches: list[Sequence[Mapping[str, StageEvidence]]] = [] + if release is not None: + released = _gripper_released(release, observation, part) + if released is not None: + batches.append( + _stage_evidence_batch( + stage_id="released", + window="release_end", + measurements={"release_established": released}, + method="gripper_joint_state", + source="measured_proxy", + ) + ) + if settle is not None and settle.object_pose is not None: + poses = settle.object_pose + target = _case_pose(case, "target_object_pose", device=poses.device).to( + dtype=poses.dtype + ) + min_height = _minimum_height_above_initial_support(settle, case) + if min_height is not None: + final = poses[-1] + if "object_rotation_threshold_rad" in case.case_parameters: + orientation_error = _pose_rotation_error(final, target) + else: + orientation_error = torch.arccos( + torch.clamp( + (final[:, :3, 2] * target[:, :3, 2]).sum(dim=-1), + -1.0, + 1.0, + ) + ) + batches.append( + _stage_evidence_batch( + stage_id="placed", + window="terminal_settle", + measurements={ + "position_error_m": torch.linalg.vector_norm( + final[:, :3, 3] - target[:, :3, 3], dim=-1 + ), + "orientation_error_rad": orientation_error, + "min_height_above_support_m": min_height, + }, + method="terminal_object_pose_trace", + ) + ) + if poses.shape[0] >= 2: + position_drift = torch.linalg.vector_norm( + poses[:, :, :3, 3] - poses[:1, :, :3, 3], dim=-1 + ).amax(dim=0) + rotation_drift = _pose_rotation_error( + poses, poses[:1].expand_as(poses) + ).amax(dim=0) + batches.append( + _stage_evidence_batch( + stage_id="stable", + window="settle_interval", + measurements={ + "position_drift_m": position_drift, + "rotation_drift_rad": rotation_drift, + "min_height_above_support_m": min_height, + }, + method="object_settle_drift_trace", + ) + ) + evidence = ( + _merge_stage_evidence_batches(*batches) + if batches + else tuple({} for _ in range(case.batch_size)) + ) + return evaluate_skill_stage_batch("place", evidence) + + +class _PressEvaluator: + """Evaluate contact and signed button stroke inside the press segment.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + trace = _replay_trace(observation, "press_actuation") + if trace is None or trace.articulation_joint is None: + return evaluate_skill_stage_batch( + "press", tuple({} for _ in range(case.batch_size)) + ) + values = trace.articulation_joint + sign = float(case.case_parameters.get("effect_joint_sign", 1.0)) + peak = ((values - values[:1]) * sign).amax(dim=0) + stroke = float(case.case_parameters["articulation_joint_stroke"]) + contacted = peak > float( + case.case_parameters.get("contact_minimum_joint_delta", 1.0e-5) + ) + evidence = _merge_stage_evidence_batches( + _stage_evidence_batch( + stage_id="contacted", + window="actuation_interval", + measurements={"contact_established": contacted}, + method="target_joint_motion_proxy", + source="measured_proxy", + ), + _stage_evidence_batch( + stage_id="pressed", + window="actuation_interval", + measurements={"peak_signed_stroke_ratio": peak / stroke}, + method="target_joint_trace", + ), + ) + return evaluate_skill_stage_batch("press", evidence) + + +class _SlideEvaluator: + """Evaluate grasp proxy, drawer travel, and release at their phase windows.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + actuation = _replay_trace(observation, "slide_actuation") + release = _replay_trace(observation, "slide_release") + if actuation is None or actuation.articulation_joint is None: + return evaluate_skill_stage_batch( + "slide", tuple({} for _ in range(case.batch_size)) + ) + joint = actuation.articulation_joint + displacement = joint[-1] - joint[0] + effect_sign = float(case.case_parameters.get("effect_joint_sign", 1.0)) + requested_displacement = effect_sign * float( + case.case_parameters["translation_distance_m"] + ) + grasped = displacement * effect_sign > float( + case.case_parameters.get("contact_minimum_joint_delta", 1.0e-5) + ) + travel_error = (displacement - requested_displacement).abs() + batches: list[Sequence[Mapping[str, StageEvidence]]] = [ + _stage_evidence_batch( + stage_id="grasped", + window="grasp_end", + measurements={"grasp_established": grasped}, + method="target_joint_motion_proxy", + source="measured_proxy", + ), + _stage_evidence_batch( + stage_id="slid", + window="actuation_end", + measurements={"travel_error_m": travel_error}, + method="target_joint_segment_end", + ), + ] + if release is not None: + part = next(iter(observation.gripper_open_qpos), "") + released = _gripper_released(release, observation, part) + if released is not None: + batches.append( + _stage_evidence_batch( + stage_id="released", + window="release_end", + measurements={"release_established": released}, + method="gripper_joint_state", + source="measured_proxy", + ) + ) + return evaluate_skill_stage_batch( + "slide", _merge_stage_evidence_batches(*batches) + ) + + +class _TwistEvaluator: + """Evaluate executed TCP twist rather than substituting knob displacement.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + actuation = _replay_trace(observation, "twist_actuation") + release = _replay_trace(observation, "twist_release") + if ( + actuation is None + or actuation.articulation_joint is None + or actuation.tcp_pose is None + ): + return evaluate_skill_stage_batch( + "twist", tuple({} for _ in range(case.batch_size)) + ) + joint = actuation.articulation_joint + grasped = (joint[-1] - joint[0]).abs() > float( + case.case_parameters.get("contact_minimum_joint_delta", 1.0e-5) + ) + tcp = actuation.tcp_pose + target_pose = _case_pose(case, "twist_target_pose", device=tcp.device).to( + dtype=tcp.dtype + ) + axis_local = torch.as_tensor( + case.case_parameters["twist_axis"], dtype=tcp.dtype, device=tcp.device + ) + axis_local = axis_local / torch.linalg.vector_norm(axis_local) + world_axis = target_pose[:, :3, :3] @ axis_local + executed = _signed_rotation_about_axis( + tcp[0, :, :3, :3], tcp[-1, :, :3, :3], world_axis + ) + error = (executed - float(case.case_parameters["twist_angle_rad"])).abs() + batches: list[Sequence[Mapping[str, StageEvidence]]] = [ + _stage_evidence_batch( + stage_id="grasped", + window="grasp_end", + measurements={"grasp_established": grasped}, + method="target_joint_motion_proxy", + source="measured_proxy", + ), + _stage_evidence_batch( + stage_id="twisted", + window="actuation_end", + measurements={"tcp_axis_angle_error_rad": error}, + method="tcp_segment_rotation", + ), + ] + if release is not None: + part = next(iter(observation.gripper_open_qpos), "") + released = _gripper_released(release, observation, part) + if released is not None: + batches.append( + _stage_evidence_batch( + stage_id="released", + window="release_end", + measurements={"release_established": released}, + method="gripper_joint_state", + source="measured_proxy", + ) + ) + return evaluate_skill_stage_batch( + "twist", _merge_stage_evidence_batches(*batches) + ) + + +class _MoveEndEffectorEvaluator: + """Score the physically settled TCP position without planner state.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + final_tcp_pose = observation.final_tcp_pose + target = case.target_waypoints[:, -1].to( + device=final_tcp_pose.device, + dtype=final_tcp_pose.dtype, + ) + if final_tcp_pose.shape != target.shape: + raise ValueError( + "Final TCP observations must match the batched terminal target pose." + ) + position_error, rotation_error = _executed_final_pose_errors( + case, final_tcp_pose + ) + return evaluate_skill_stage_batch( + "move_end_effector", + _stage_evidence_batch( + stage_id="reached", + window="terminal_settle", + measurements={ + "tcp_position_error_m": position_error, + "tcp_rotation_error_rad": rotation_error, + }, + method="terminal_tcp_pose", + ), + ) + + +class _MoveJointsEvaluator: + """Score the physically settled arm joints without planner state.""" + + def evaluate( + self, + case: BenchmarkCase, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[PhysicalEvaluation, ...]: + del motion_outcomes + final_arm_qpos = observation.final_arm_qpos + if final_arm_qpos is None: + return evaluate_skill_stage_batch( + "move_joints", tuple({} for _ in range(case.batch_size)) + ) + target = torch.as_tensor( + case.case_parameters["target_qpos"], + dtype=final_arm_qpos.dtype, + device=final_arm_qpos.device, + ) + if target.dim() != 3 or target.shape[0] != case.batch_size: + raise ValueError("MoveJoints target_qpos must have shape (B, W, D).") + target = target[:, -1] + if final_arm_qpos.shape != target.shape: + raise ValueError( + "Final arm observations must match the batched terminal joint target." + ) + error = torch.linalg.vector_norm(final_arm_qpos - target, dim=-1) + return evaluate_skill_stage_batch( + "move_joints", + _stage_evidence_batch( + stage_id="reached", + window="terminal_settle", + measurements={"joint_error_l2_rad": error}, + method="terminal_arm_qpos", + ), + ) + + def _grasp_sample_count(case: BenchmarkCase) -> int: """Return the grasp sampler size a self-sampling case was frozen with.""" return int(case.case_parameters.get("grasp_sample_count", 10_000)) @@ -1424,6 +2069,10 @@ def task_result( "task_goal_miss", ) + def physical_evaluator(self) -> _MoveEndEffectorEvaluator: + """Evaluate the terminal measured TCP position under the v1 standard.""" + return _MoveEndEffectorEvaluator() + class _PickUpCases(AtomicSkillCaseProvider): """Explicit-grasp PickUp cases that isolate motion-planner performance.""" @@ -1655,6 +2304,40 @@ def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: """Return the lift boundary emitted by PickUp.""" return compiled.segment(0, "lift").start + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect grasp, lift, and terminal retention evidence.""" + del case + close = compiled.segment(0, "close") + lift = compiled.segment(0, "lift") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "pickup_grasp", + close.start, + close.stop, + frozenset({"gripper_qpos"}), + ), + ReplayObservationWindow( + "pickup_lift", + lift.start, + lift.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "pickup_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose", "gripper_qpos"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _PickUpEvaluator: + """Return the v1 PickUp evaluator.""" + return _PickUpEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -1831,6 +2514,10 @@ def task_result( ) return success, "joint_goal_miss" + def physical_evaluator(self) -> _MoveJointsEvaluator: + """Evaluate the terminal measured joint state under the v1 standard.""" + return _MoveJointsEvaluator() + class _HeldObjectCases(AtomicSkillCaseProvider): """Shared isolated held-object precondition for transport and placement.""" @@ -2156,6 +2843,33 @@ def build_invocation( skill_options=MoveHeldObjectOptions(), ) + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect transport plus settled pose and grasp state.""" + del case + transport = compiled.segment(1, "transport") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "held_object_transport", + transport.start, + transport.stop, + frozenset({"object_pose"}), + ), + ReplayObservationWindow( + "held_object_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose", "gripper_qpos"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _MoveHeldObjectEvaluator: + """Return the v1 MoveHeldObject evaluator.""" + return _MoveHeldObjectEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -2342,6 +3056,33 @@ def build_invocation( ), ) + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect release and terminal stability evidence.""" + del case + release = compiled.segment(1, "release") + waypoint_count = compiled.trajectory.waypoint_count + return ( + ReplayObservationWindow( + "place_release", + release.start, + release.stop, + frozenset({"gripper_qpos"}), + ), + ReplayObservationWindow( + "place_settle", + waypoint_count, + waypoint_count, + frozenset({"object_pose"}), + include_settle=True, + ), + ) + + def physical_evaluator(self) -> _PlaceEvaluator: + """Return the v1 Place evaluator.""" + return _PlaceEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -2404,6 +3145,14 @@ def generate_case( start_qpos, name=scenario.control_part, to_matrix=True ) articulation_initial_qpos = handle.entity.get_qpos().clone() + joint_index = handle.entity.joint_names.index(target_joint) + joint_limits = handle.entity.get_qpos_limits(joint_ids=[joint_index])[0, 0] + articulation_joint_stroke = float((joint_limits[1] - joint_limits[0]).item()) + if ( + not math.isfinite(articulation_joint_stroke) + or articulation_joint_stroke <= 0 + ): + raise ValueError("Press target joint must expose a finite positive stroke.") target_pose = handle.link_pose(target_link).to(device=start_pose.device) geometry = scenario.sample_articulation_geometry( handle, @@ -2498,6 +3247,10 @@ def generate_case( "target_joint": target_joint, "target_entity_id": f"{handle.entity.uid}:{target_link}", "effect_joint_sign": float(config.get("effect_joint_sign", 1.0)), + "articulation_joint_stroke": articulation_joint_stroke, + "contact_minimum_joint_delta": float( + config.get("contact_minimum_joint_delta", 1.0e-5) + ), "minimum_articulation_joint_delta": float( config.get("minimum_joint_delta_m", 0.004) ), @@ -2563,6 +3316,25 @@ def articulation_effect_segment(self, case: BenchmarkCase) -> str: del case return "press" + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect the target-joint trace only during commanded pressing.""" + del case + segment = compiled.segment(0, "press") + return ( + ReplayObservationWindow( + "press_actuation", + segment.start, + segment.stop, + frozenset({"articulation_joint"}), + ), + ) + + def physical_evaluator(self) -> _PressEvaluator: + """Return the v1 press-stage evaluator.""" + return _PressEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -2822,6 +3594,33 @@ def articulation_effect_segment(self, case: BenchmarkCase) -> str: """Score displacement while the primitive pulls or pushes the drawer.""" return _slide_direction(case.case_parameters["direction"]) + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect the commanded slide and release windows.""" + actuation = compiled.segment( + 0, _slide_direction(case.case_parameters["direction"]) + ) + release = compiled.segment(0, "open") + return ( + ReplayObservationWindow( + "slide_actuation", + actuation.start, + actuation.stop, + frozenset({"articulation_joint"}), + ), + ReplayObservationWindow( + "slide_release", + release.start, + release.stop, + frozenset({"gripper_qpos"}), + ), + ) + + def physical_evaluator(self) -> _SlideEvaluator: + """Return the v1 slide-stage evaluator.""" + return _SlideEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -3061,6 +3860,32 @@ def articulation_effect_segment(self, case: BenchmarkCase) -> str: del case return "twist" + def replay_observation_windows( + self, compiled: CompiledTrajectory, case: BenchmarkCase + ) -> tuple[ReplayObservationWindow, ...]: + """Collect executed TCP/joint traces and the release boundary.""" + del case + actuation = compiled.segment(0, "twist") + release = compiled.segment(0, "open") + return ( + ReplayObservationWindow( + "twist_actuation", + actuation.start, + actuation.stop, + frozenset({"tcp_pose", "articulation_joint"}), + ), + ReplayObservationWindow( + "twist_release", + release.start, + release.stop, + frozenset({"gripper_qpos"}), + ), + ) + + def physical_evaluator(self) -> _TwistEvaluator: + """Return the v1 twist-stage evaluator.""" + return _TwistEvaluator() + def task_result( self, scenario: "AtomicTaskScenario", @@ -4143,6 +4968,15 @@ def _execute( part: self.robot.get_qpos(name=part).clone() for part in self._configured_grippers() }, + gripper_open_qpos={ + part: states[0].clone() + for part, states in self._configured_grippers().items() + }, + gripper_grasp_qpos={ + part: states[1].clone() + for part, states in self._configured_grippers().items() + }, + replay_traces=replay_metrics.replay_traces, ) def _physics_settings(self) -> _PhysicsReplaySettings: @@ -4183,6 +5017,78 @@ def _is_recordable(compiled: CompiledTrajectory) -> bool: trajectory = compiled.trajectory.positions return trajectory.shape[1] > 0 and bool(torch.isfinite(trajectory).all().item()) + def _capture_replay_samples( + self, + accumulators: Sequence[_ReplayTraceAccumulator], + *, + object_handle: AtomicObjectHandle | None, + scene_object_handles: Mapping[str, AtomicObjectHandle], + articulation_handle: AtomicArticulationHandle | None, + articulation_joint_name: str | None, + ) -> None: + """Append one synchronized simulator sample to active windows.""" + if not accumulators: + return + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + channels = frozenset( + channel + for accumulator in accumulators + for channel in accumulator.window.channels + ) + tcp_pose = None + object_pose = None + articulation_joint = None + grippers: dict[str, torch.Tensor] = {} + scene_object_poses: dict[str, torch.Tensor] = {} + if "tcp_pose" in channels: + tcp_pose = self.robot.compute_fk( + self.robot.get_qpos(name=self.control_part), + name=self.control_part, + to_matrix=True, + ).clone() + if "object_pose" in channels and object_handle is not None: + object_pose = object_handle.entity.get_local_pose(to_matrix=True).clone() + if ( + "articulation_joint" in channels + and articulation_handle is not None + and articulation_joint_name is not None + ): + articulation_joint = articulation_handle.joint_qpos(articulation_joint_name) + if "gripper_qpos" in channels: + grippers = { + part: self.robot.get_qpos(name=part).clone() + for part in self._configured_grippers() + } + if "scene_object_pose" in channels: + requested_ids = { + object_id + for accumulator in accumulators + for object_id in accumulator.window.scene_object_ids + } + scene_object_poses = { + object_id: scene_object_handles[object_id] + .entity.get_local_pose(to_matrix=True) + .clone() + for object_id in requested_ids + } + for accumulator in accumulators: + requested = accumulator.window.channels + if "tcp_pose" in requested and tcp_pose is not None: + accumulator.tcp_pose.append(tcp_pose) + if "object_pose" in requested and object_pose is not None: + accumulator.object_pose.append(object_pose) + if "articulation_joint" in requested and articulation_joint is not None: + accumulator.articulation_joint.append(articulation_joint) + if "gripper_qpos" in requested: + for part, qpos in grippers.items(): + accumulator.gripper_qpos.setdefault(part, []).append(qpos) + if "scene_object_pose" in requested: + for object_id in accumulator.window.scene_object_ids: + accumulator.scene_object_poses.setdefault(object_id, []).append( + scene_object_poses[object_id] + ) + def _replay_physics( self, compiled: CompiledTrajectory, @@ -4291,7 +5197,58 @@ def _replay_physics( # then observe through release, retract, and the final hold. effect_start, effect_stop = segment.start, trajectory.shape[1] effect_end_index = segment.stop - 1 + observation_windows = ( + () + if provider is None or not collect_metrics + else provider.replay_observation_windows(compiled, case) + ) + if len({window.observation_id for window in observation_windows}) != len( + observation_windows + ): + raise ValueError("Replay observation IDs must be unique within one case.") + waypoint_count = trajectory.shape[1] + for window in observation_windows: + if window.stop > waypoint_count: + raise ValueError( + f"Replay observation {window.observation_id!r} exceeds the " + "compiled trajectory." + ) + if window.include_settle and window.stop != waypoint_count: + raise ValueError( + "Only a window ending at the trajectory boundary may include " + "terminal settle." + ) + trace_accumulators = { + window.observation_id: _ReplayTraceAccumulator(window) + for window in observation_windows + } + requested_scene_object_ids = { + object_id + for window in observation_windows + for object_id in window.scene_object_ids + } + missing_scene_object_ids = requested_scene_object_ids - set(self._objects) + if missing_scene_object_ids: + raise ValueError( + "Replay observations reference unknown scene objects: " + f"{sorted(missing_scene_object_ids)}." + ) + scene_object_handles = { + object_id: self._objects[object_id] + for object_id in requested_scene_object_ids + } for waypoint_index in range(trajectory.shape[1]): + self._capture_replay_samples( + [ + trace_accumulators[window.observation_id] + for window in observation_windows + if window.start == waypoint_index + ], + object_handle=object_handle, + scene_object_handles=scene_object_handles, + articulation_handle=articulation_handle, + articulation_joint_name=articulation_joint_name, + ) if rotation_window is not None and waypoint_index == rotation_window[0]: assert object_handle is not None and rotation_axis_local is not None rotation_reference = ( @@ -4341,6 +5298,17 @@ def _replay_physics( articulation_peak, (current_joint - articulation_initial) * effect_sign, ) + self._capture_replay_samples( + [ + trace_accumulators[window.observation_id] + for window in observation_windows + if window.start <= waypoint_index < window.stop + ], + object_handle=object_handle, + scene_object_handles=scene_object_handles, + articulation_handle=articulation_handle, + articulation_joint_name=articulation_joint_name, + ) if ( articulation_handle is not None and articulation_joint_name is not None @@ -4362,6 +5330,18 @@ def _replay_physics( object_handle.entity.clear_dynamics() dynamics_cleared = True final_command = trajectory[:, -1] + settling = [ + trace_accumulators[window.observation_id] + for window in observation_windows + if window.include_settle + ] + self._capture_replay_samples( + settling, + object_handle=object_handle, + scene_object_handles=scene_object_handles, + articulation_handle=articulation_handle, + articulation_joint_name=articulation_joint_name, + ) for _ in range(settings.hold_steps): self.robot.set_qpos(final_command, target=True) self.simulation.update(step=settings.hold_sim_steps) @@ -4380,6 +5360,13 @@ def _replay_physics( observed = self.robot.get_qpos() squared_tracking_error += ((observed - final_command) ** 2).sum(dim=1) tracking_value_count += observed.shape[1] + self._capture_replay_samples( + settling, + object_handle=object_handle, + scene_object_handles=scene_object_handles, + articulation_handle=articulation_handle, + articulation_joint_name=articulation_joint_name, + ) if collect_metrics: if squared_tracking_error is None: raise RuntimeError("Timed replay lost its tracking accumulator.") @@ -4397,6 +5384,17 @@ def _replay_physics( object_rotation_peak_signed_rad=( None if rotation_peak is None else rotation_peak * rotation_sign ), + replay_traces={ + observation_id: accumulator.finish() + for observation_id, accumulator in trace_accumulators.items() + if ( + accumulator.tcp_pose + or accumulator.object_pose + or accumulator.articulation_joint + or accumulator.gripper_qpos + or accumulator.scene_object_poses + ) + }, ) return None diff --git a/tests/benchmark/motion_generation/test_atomic_dual_arm_providers.py b/tests/benchmark/motion_generation/test_atomic_dual_arm_providers.py index c09de3643..b6fa019ff 100644 --- a/tests/benchmark/motion_generation/test_atomic_dual_arm_providers.py +++ b/tests/benchmark/motion_generation/test_atomic_dual_arm_providers.py @@ -33,6 +33,10 @@ ) from embodichain.utils.math import axis_angle_to_rotation_matrix from scripts.benchmark.motion_generation.config import load_suite +from scripts.benchmark.motion_generation.contracts import ( + ReplayObservationTrace, + ReplayObservationWindow, +) from scripts.benchmark.motion_generation.models import BenchmarkCase, CaseOutcome from scripts.benchmark.motion_generation import robots as _robots # noqa: F401 from scripts.benchmark.motion_generation.registry import create_robot_provider @@ -116,6 +120,391 @@ def _observation(**values: object) -> _ExecutionObservation: ) +def test_hand_over_physical_evaluator_uses_phase_local_observations() -> None: + target = _pose(0.2, -0.1, 0.1).unsqueeze(0) + initial = _pose(0.0).unsqueeze(0) + object_trace = torch.stack((target, target)) + case = _case( + "hand_over", + control_parts={ + "source": {"motion": "left_arm", "grasp": "left_hand"}, + "destination": {"motion": "right_arm", "grasp": "right_hand"}, + }, + object_initial_pose=initial.tolist(), + target_object_pose=target.tolist(), + ) + holding = torch.full((2, 1, 1), 0.5) + released = torch.zeros(2, 1, 1) + observation = _observation( + gripper_open_qpos={ + "left_hand": torch.tensor([0.0]), + "right_hand": torch.tensor([0.0]), + }, + gripper_grasp_qpos={ + "left_hand": torch.tensor([1.0]), + "right_hand": torch.tensor([1.0]), + }, + replay_traces={ + "handover_pickup": ReplayObservationTrace( + gripper_qpos={"left_hand": holding, "right_hand": released} + ), + "handover_transfer": ReplayObservationTrace( + object_pose=object_trace, + gripper_qpos={"left_hand": holding, "right_hand": holding}, + ), + "handover_exchange": ReplayObservationTrace( + object_pose=object_trace, + gripper_qpos={"left_hand": released, "right_hand": holding}, + ), + "handover_settle": ReplayObservationTrace( + object_pose=object_trace, + gripper_qpos={"left_hand": released, "right_hand": released}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("hand_over").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_outcome(),))[0] + + assert result.task_success + assert [stage.stage_id for stage in result.stages] == [ + "grasped", + "transferred", + "handed_over", + "receiver_released", + "placed", + ] + + +def test_hand_over_requires_receiver_release_at_terminal_target() -> None: + target = _pose(0.2, -0.1, 0.1).unsqueeze(0) + initial = _pose(0.0).unsqueeze(0) + object_trace = torch.stack((target, target)) + holding = torch.full((2, 1, 1), 0.5) + receiver_holding = torch.full((2, 1, 1), 0.75) + released = torch.zeros(2, 1, 1) + case = _case( + "hand_over", + control_parts={ + "source": {"motion": "left_arm", "grasp": "left_hand"}, + "destination": {"motion": "right_arm", "grasp": "right_hand"}, + }, + object_initial_pose=initial.tolist(), + target_object_pose=target.tolist(), + ) + observation = _observation( + gripper_open_qpos={ + "left_hand": torch.tensor([0.0]), + "right_hand": torch.tensor([0.0]), + }, + gripper_grasp_qpos={ + "left_hand": torch.tensor([1.0]), + "right_hand": torch.tensor([1.0]), + }, + replay_traces={ + "handover_pickup": ReplayObservationTrace( + gripper_qpos={"left_hand": holding, "right_hand": released} + ), + "handover_transfer": ReplayObservationTrace( + object_pose=object_trace, + gripper_qpos={"left_hand": holding, "right_hand": holding}, + ), + "handover_exchange": ReplayObservationTrace( + object_pose=object_trace, + gripper_qpos={"left_hand": released, "right_hand": holding}, + ), + "handover_settle": ReplayObservationTrace( + object_pose=object_trace, + gripper_qpos={ + "left_hand": released, + "right_hand": receiver_holding, + }, + ), + }, + ) + + evaluator = create_atomic_skill_provider("hand_over").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "receiver_released" + + +def test_coordinated_pickment_evaluator_uses_all_four_physical_stages() -> None: + initial = _pose(0.0).unsqueeze(0) + target = _pose(0.1, 0.0, 0.1).unsqueeze(0) + lift = torch.stack((initial, _pose(0.0, 0.0, 0.05).unsqueeze(0))) + move = torch.stack((lift[-1], target)) + holding = torch.full((2, 1, 1), 0.5) + case = _case( + "coordinated_pickment", + control_parts={ + "left": {"motion": "left_arm", "grasp": "left_hand"}, + "right": {"motion": "right_arm", "grasp": "right_hand"}, + }, + object_initial_pose=initial.tolist(), + target_object_pose=target.tolist(), + ) + observation = _observation( + gripper_open_qpos={ + "left_hand": torch.tensor([0.0]), + "right_hand": torch.tensor([0.0]), + }, + gripper_grasp_qpos={ + "left_hand": torch.tensor([1.0]), + "right_hand": torch.tensor([1.0]), + }, + replay_traces={ + "dual_pick_grasp": ReplayObservationTrace( + gripper_qpos={"left_hand": holding, "right_hand": holding} + ), + "dual_pick_lift": ReplayObservationTrace(object_pose=lift), + "dual_pick_move": ReplayObservationTrace(object_pose=move), + "dual_pick_settle": ReplayObservationTrace( + object_pose=target.unsqueeze(0).repeat(2, 1, 1, 1), + gripper_qpos={"left_hand": holding, "right_hand": holding}, + ), + }, + ) + + evaluator = create_atomic_skill_provider( + "coordinated_pickment" + ).physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_outcome(),))[0] + + assert result.task_success + assert [stage.stage_id for stage in result.stages] == [ + "dual_grasped", + "lifted", + "moved", + "settled", + "held", + ] + + +def test_coordinated_pickment_rechecks_target_after_settle() -> None: + initial = _pose(0.0).unsqueeze(0) + target = _pose(0.1, 0.0, 0.1).unsqueeze(0) + lift = torch.stack((initial, _pose(0.0, 0.0, 0.05).unsqueeze(0))) + move = torch.stack((lift[-1], target)) + settled_wrong = _pose(0.4, 0.0, 0.1).reshape(1, 1, 4, 4).repeat(2, 1, 1, 1) + holding = torch.full((2, 1, 1), 0.5) + case = _case( + "coordinated_pickment", + control_parts={ + "left": {"motion": "left_arm", "grasp": "left_hand"}, + "right": {"motion": "right_arm", "grasp": "right_hand"}, + }, + object_initial_pose=initial.tolist(), + target_object_pose=target.tolist(), + ) + observation = _observation( + gripper_open_qpos={ + "left_hand": torch.tensor([0.0]), + "right_hand": torch.tensor([0.0]), + }, + gripper_grasp_qpos={ + "left_hand": torch.tensor([1.0]), + "right_hand": torch.tensor([1.0]), + }, + replay_traces={ + "dual_pick_grasp": ReplayObservationTrace( + gripper_qpos={"left_hand": holding, "right_hand": holding} + ), + "dual_pick_lift": ReplayObservationTrace(object_pose=lift), + "dual_pick_move": ReplayObservationTrace(object_pose=move), + "dual_pick_settle": ReplayObservationTrace( + object_pose=settled_wrong, + gripper_qpos={"left_hand": holding, "right_hand": holding}, + ), + }, + ) + + evaluator = create_atomic_skill_provider( + "coordinated_pickment" + ).physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "settled" + + +def test_coordinated_placement_evaluator_uses_named_support_pose_trace() -> None: + pose = _pose(0.0).unsqueeze(0) + pose_trace = pose.unsqueeze(0).repeat(2, 1, 1, 1) + placing_released = torch.zeros(2, 1, 1) + support_held = torch.full((2, 1, 1), 0.5) + case = _case( + "coordinated_placement", + control_parts={ + "placing": {"motion": "left_arm", "grasp": "left_hand"}, + "support": {"motion": "right_arm", "grasp": "right_hand"}, + }, + support_object_id="support", + object_initial_pose=pose.tolist(), + placing_relative_target_pose=pose.tolist(), + ) + observation = _observation( + gripper_open_qpos={ + "left_hand": torch.tensor([0.0]), + "right_hand": torch.tensor([0.0]), + }, + gripper_grasp_qpos={ + "left_hand": torch.tensor([1.0]), + "right_hand": torch.tensor([1.0]), + }, + replay_traces={ + "coordinated_place_alignment": ReplayObservationTrace( + object_pose=pose_trace, + scene_object_poses={"support": pose_trace}, + ), + "coordinated_place_release": ReplayObservationTrace( + gripper_qpos={ + "left_hand": placing_released, + "right_hand": support_held, + } + ), + "coordinated_place_settle": ReplayObservationTrace( + object_pose=pose_trace, + scene_object_poses={"support": pose_trace}, + ), + }, + ) + + evaluator = create_atomic_skill_provider( + "coordinated_placement" + ).physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_outcome(),))[0] + + assert result.task_success + assert result.stages[1].measurements == { + "placing_release_established": True, + "support_retained": True, + } + + +def test_coordinated_placement_rechecks_relative_target_after_release() -> None: + target = _pose(0.0).unsqueeze(0) + aligned_trace = target.unsqueeze(0).repeat(2, 1, 1, 1) + support_trace = target.unsqueeze(0).repeat(2, 1, 1, 1) + settled_wrong = _pose(0.5).reshape(1, 1, 4, 4).repeat(2, 1, 1, 1) + placing_released = torch.zeros(2, 1, 1) + support_held = torch.full((2, 1, 1), 0.5) + case = _case( + "coordinated_placement", + control_parts={ + "placing": {"motion": "left_arm", "grasp": "left_hand"}, + "support": {"motion": "right_arm", "grasp": "right_hand"}, + }, + support_object_id="support", + object_initial_pose=target.tolist(), + placing_relative_target_pose=target.tolist(), + ) + observation = _observation( + gripper_open_qpos={ + "left_hand": torch.tensor([0.0]), + "right_hand": torch.tensor([0.0]), + }, + gripper_grasp_qpos={ + "left_hand": torch.tensor([1.0]), + "right_hand": torch.tensor([1.0]), + }, + replay_traces={ + "coordinated_place_alignment": ReplayObservationTrace( + object_pose=aligned_trace, + scene_object_poses={"support": support_trace}, + ), + "coordinated_place_release": ReplayObservationTrace( + gripper_qpos={ + "left_hand": placing_released, + "right_hand": support_held, + } + ), + "coordinated_place_settle": ReplayObservationTrace( + object_pose=settled_wrong, + scene_object_poses={"support": support_trace}, + ), + }, + ) + + evaluator = create_atomic_skill_provider( + "coordinated_placement" + ).physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "stable" + + +def test_replay_collects_named_scene_object_poses_in_sync() -> None: + state = {"step": 0} + scenario = AtomicTaskScenario() + scenario.simulation = Mock() + scenario.simulation.update.side_effect = lambda step: state.update( + step=state["step"] + step + ) + scenario.robot = Mock() + scenario.robot.get_qpos.side_effect = lambda: torch.zeros(1, 6) + scenario.track = Mock( + config={ + "physics": { + "steps_per_waypoint": 1, + "hold_steps": 0, + "hold_sim_steps": 1, + } + } + ) + + def moving_pose(to_matrix: bool) -> torch.Tensor: + del to_matrix + pose = torch.eye(4).unsqueeze(0) + pose[:, 0, 3] = 0.01 * state["step"] + return pose + + placing = Mock() + placing.entity.get_local_pose.side_effect = moving_pose + support = Mock() + support.entity.get_local_pose.return_value = _pose(0.2).unsqueeze(0) + scenario._objects = {"placing": placing, "support": support} + compiled = SimpleNamespace( + trajectory=SimpleNamespace(positions=torch.zeros(1, 3, 6)), + plan_success=torch.tensor([True]), + ) + provider = Mock() + provider.lift_segment_start.return_value = None + provider.object_rotation_segment.return_value = None + provider.replay_observation_windows.return_value = ( + ReplayObservationWindow( + "relative", + 0, + 2, + frozenset({"object_pose", "scene_object_pose"}), + scene_object_ids=("support",), + ), + ) + case = replace( + _case("coordinated_placement", support_object_id="support"), + object_id="placing", + ) + + metrics = scenario._replay_physics(compiled, case, provider, collect_metrics=True) + + assert metrics is not None + trace = metrics.replay_traces["relative"] + assert trace.object_pose is not None + assert trace.object_pose.shape == (3, 1, 4, 4) + assert trace.scene_object_poses["support"].shape == (3, 1, 4, 4) + torch.testing.assert_close( + trace.scene_object_poses["support"][:, 0, 0, 3], torch.full((3,), 0.2) + ) + + def _two_arm_robot() -> Mock: """A 6-joint robot whose arm FK maps the first three joints to x, y, z.""" robot = Mock(device=CPU) diff --git a/tests/benchmark/motion_generation/test_atomic_skill_providers.py b/tests/benchmark/motion_generation/test_atomic_skill_providers.py index 996d60b45..c58651ab7 100644 --- a/tests/benchmark/motion_generation/test_atomic_skill_providers.py +++ b/tests/benchmark/motion_generation/test_atomic_skill_providers.py @@ -39,6 +39,10 @@ from embodichain.lab.sim.motion.planners.utils import TrajectorySampleMethod from embodichain.utils.math import axis_angle_to_rotation_matrix, pose_inv from scripts.benchmark.motion_generation.config import PlannerSpecCfg, load_suite +from scripts.benchmark.motion_generation.contracts import ( + ReplayObservationTrace, + ReplayObservationWindow, +) from scripts.benchmark.motion_generation.models import ( AlgorithmRole, BenchmarkCase, @@ -127,6 +131,252 @@ def _rotation_z(angle: float) -> torch.Tensor: return axis_angle_to_rotation_matrix(torch.tensor([0.0, 0.0, angle])) +def test_open_door_physical_evaluator_uses_hinge_and_release_windows() -> None: + case = _case("open_door", 2, target_joint_position=0.5) + observation = _observation( + 2, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "door_actuation": ReplayObservationTrace( + articulation_joint=torch.tensor([[0.0, 0.0], [0.5, 0.3]]) + ), + "door_release": ReplayObservationTrace( + gripper_qpos={"hand": torch.tensor([[[1.0], [1.0]], [[0.0], [0.0]]])} + ), + }, + ) + + evaluator = create_atomic_skill_provider("open_door").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, _outcomes(2)) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_stage == "opened" + + +def test_pour_physical_evaluator_uses_object_pose_interval() -> None: + initial = torch.eye(4).repeat(2, 1, 1) + poses = torch.eye(4).repeat(3, 2, 1, 1) + poses[:, :, 2, 3] = 0.1 + poses[1, :, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[0.0, 0.0, 0.5], [0.0, 0.0, 0.2]]) + ) + case = _case( + "pour", + 2, + object_initial_pose=initial.tolist(), + object_rotation_axis=[0.0, 0.0, 1.0], + object_rotation_sign=1.0, + rotate_angle_rad=0.5, + ) + settle = poses[-1:].repeat(2, 1, 1, 1) + observation = _observation( + 2, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "pour_actuation": ReplayObservationTrace(object_pose=poses), + "pour_settle": ReplayObservationTrace( + object_pose=settle, + gripper_qpos={"hand": torch.full((2, 2, 1), 0.5)}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("pour").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, _outcomes(2)) + + assert [result.task_success for result in results] == [True, False] + assert results[0].stages[0].measurements[ + "peak_signed_axis_angle_error_rad" + ] == pytest.approx(0.0, abs=1.0e-6) + + +def test_pour_requires_terminal_lift_and_retention() -> None: + initial = torch.eye(4).repeat(2, 1, 1) + actuation = torch.eye(4).repeat(3, 2, 1, 1) + actuation[:, :, 2, 3] = 0.1 + actuation[1, :, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[0.0, 0.0, 0.5], [0.0, 0.0, 0.5]]) + ) + settle = torch.eye(4).repeat(2, 2, 1, 1) + settle[:, 1, 2, 3] = 0.1 + terminal_gripper = torch.full((2, 2, 1), 0.5) + terminal_gripper[:, 1] = 1.0 + case = _case( + "pour", + 2, + object_initial_pose=initial.tolist(), + object_rotation_axis=[0.0, 0.0, 1.0], + object_rotation_sign=1.0, + rotate_angle_rad=0.5, + ) + observation = _observation( + 2, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "pour_actuation": ReplayObservationTrace(object_pose=actuation), + "pour_settle": ReplayObservationTrace( + object_pose=settle, + gripper_qpos={"hand": terminal_gripper}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("pour").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, _outcomes(2)) + + assert [result.failure_stage for result in results] == ["settled", "held"] + + +def test_axis_align_physical_evaluator_keeps_terminal_retention_separate() -> None: + initial = torch.eye(4).repeat(2, 1, 1) + actuation = torch.eye(4).repeat(2, 2, 1, 1) + actuation[:, :, 2, 3] = 0.1 + actuation[-1, :, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[0.0, 0.0, math.pi / 2], [0.0, 0.0, math.pi / 4]]) + ) + settle = actuation[-1:].clone().repeat(2, 1, 1, 1) + case = _case( + "axis_align", + 2, + object_initial_pose=initial.tolist(), + internal_axis=[1.0, 0.0, 0.0], + target_axis=[0.0, 1.0, 0.0], + ) + observation = _observation( + 2, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "axis_align_actuation": ReplayObservationTrace(object_pose=actuation), + "axis_align_settle": ReplayObservationTrace( + object_pose=settle, + gripper_qpos={"hand": torch.full((2, 2, 1), 0.5)}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("axis_align").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, _outcomes(2)) + + assert [result.task_success for result in results] == [True, False] + assert results[0].stages[-1].stage_id == "held" + assert results[1].failure_stage == "aligned" + + +def test_axis_align_requires_lift_and_terminal_alignment() -> None: + initial = torch.eye(4).repeat(2, 1, 1) + actuation = torch.eye(4).repeat(2, 2, 1, 1) + actuation[-1, :, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[0.0, 0.0, math.pi / 2], [0.0, 0.0, math.pi / 2]]) + ) + actuation[:, 1, 2, 3] = 0.1 + settle = actuation[-1:].repeat(2, 1, 1, 1) + settle[:, 1, :3, :3] = torch.eye(3) + case = _case( + "axis_align", + 2, + object_initial_pose=initial.tolist(), + internal_axis=[1.0, 0.0, 0.0], + target_axis=[0.0, 1.0, 0.0], + ) + observation = _observation( + 2, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "axis_align_actuation": ReplayObservationTrace(object_pose=actuation), + "axis_align_settle": ReplayObservationTrace( + object_pose=settle, + gripper_qpos={"hand": torch.full((2, 2, 1), 0.5)}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("axis_align").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, _outcomes(2)) + + assert [result.failure_stage for result in results] == ["aligned", "settled"] + + +def test_push_object_physical_evaluator_uses_contact_and_support_windows() -> None: + initial = torch.eye(4).repeat(2, 1, 1) + target = initial.clone() + target[:, 1, 3] = 0.1 + contact = initial.unsqueeze(0).repeat(2, 1, 1, 1) + contact[-1, 0, 1, 3] = 0.001 + actuation = initial.unsqueeze(0).repeat(2, 1, 1, 1) + actuation[-1] = target + settle = target.unsqueeze(0).repeat(2, 1, 1, 1) + case = _case( + "push_object", + 2, + object_initial_pose=initial.tolist(), + target_object_pose=target.tolist(), + contact_motion_threshold_m=1.0e-5, + support_plane_tolerance_m=0.01, + ) + observation = _observation( + 2, + replay_traces={ + "push_contact": ReplayObservationTrace(object_pose=contact), + "push_actuation": ReplayObservationTrace(object_pose=actuation), + "push_settle": ReplayObservationTrace(object_pose=settle), + }, + ) + + evaluator = create_atomic_skill_provider("push_object").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, _outcomes(2)) + + assert [result.task_success for result in results] == [True, False] + assert results[0].stages[0].evidence_source == ( + "measured_proxy:object_motion_during_contact" + ) + assert results[1].failure_stage == "contacted" + + +def test_push_object_rechecks_in_plane_goal_after_settle() -> None: + initial = torch.eye(4).reshape(1, 4, 4) + target = initial.clone() + target[:, 1, 3] = 0.1 + contact = initial.unsqueeze(0).repeat(2, 1, 1, 1) + contact[-1, :, 1, 3] = 0.001 + actuation = torch.stack((initial, target)) + settle = target.unsqueeze(0).repeat(2, 1, 1, 1) + settle[:, :, 0, 3] = 0.05 + case = _case( + "push_object", + 1, + object_initial_pose=initial.tolist(), + target_object_pose=target.tolist(), + contact_motion_threshold_m=1.0e-5, + support_plane_tolerance_m=0.01, + ) + observation = _observation( + 1, + replay_traces={ + "push_contact": ReplayObservationTrace(object_pose=contact), + "push_actuation": ReplayObservationTrace(object_pose=actuation), + "push_settle": ReplayObservationTrace(object_pose=settle), + }, + ) + + evaluator = create_atomic_skill_provider("push_object").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, _outcomes(1))[0] + + assert not result.task_success + assert result.failure_stage == "on_support_plane" + + def _door_affordance() -> OpenDoorAffordance: """A handle at local +X around a local +Z hinge through the origin.""" return OpenDoorAffordance( @@ -614,6 +864,9 @@ def object_pose(to_matrix: bool) -> torch.Tensor: provider = Mock() provider.lift_segment_start.return_value = None provider.object_rotation_segment.return_value = (1, 3) + provider.replay_observation_windows.return_value = ( + ReplayObservationWindow("pour_trace", 1, 3, frozenset({"object_pose"})), + ) case = _case( "pour", 1, @@ -627,6 +880,7 @@ def object_pose(to_matrix: bool) -> torch.Tensor: assert metrics.object_rotation_peak_signed_rad.item() == pytest.approx( -0.2, abs=1e-5 ) + assert metrics.replay_traces["pour_trace"].object_pose.shape == (3, 1, 4, 4) state["step"] = 0 articulation = Mock() @@ -637,6 +891,7 @@ def object_pose(to_matrix: bool) -> torch.Tensor: scenario._articulations = {"microwave": articulation} compiled.segment = Mock(return_value=SimpleNamespace(start=1, stop=3)) provider.articulation_effect_segment.return_value = "open" + provider.replay_observation_windows.return_value = () door_case = replace( _case("open_door", 1, target_joint="door_hinge"), object_id="microwave" ) diff --git a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py index 3825cf43e..c8e17040b 100644 --- a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py +++ b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py @@ -20,8 +20,10 @@ import argparse import json +import math from dataclasses import replace from pathlib import Path +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -29,9 +31,11 @@ from embodichain.lab.sim.atomic_actions import ArticulationAffordanceGeometry from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory, TimedTrajectory +from embodichain.utils.math import axis_angle_to_rotation_matrix from scripts.benchmark.motion_generation.aggregation import aggregate_results from scripts.benchmark.motion_generation.artifacts import write_case_manifest from scripts.benchmark.motion_generation.config import load_suite +from scripts.benchmark.motion_generation.contracts import ReplayObservationTrace from scripts.benchmark.motion_generation.models import ( AlgorithmRole, BenchmarkCase, @@ -138,6 +142,665 @@ def test_atomic_case_identity_includes_simulator_batch_size() -> None: assert case_b1 != case_b8 +@pytest.mark.parametrize( + ("skill_id", "parameters", "expected_ids"), + [ + ( + "pick_up", + {}, + ("pickup_grasp", "pickup_lift", "pickup_settle"), + ), + ( + "move_held_object", + {}, + ("held_object_transport", "held_object_settle"), + ), + ("place", {}, ("place_release", "place_settle")), + ("press", {}, ("press_actuation",)), + ("slide", {"direction": "pull"}, ("slide_actuation", "slide_release")), + ("twist", {}, ("twist_actuation", "twist_release")), + ("open_door", {}, ("door_actuation", "door_release")), + ("pour", {}, ("pour_actuation", "pour_settle")), + ( + "axis_align", + {}, + ("axis_align_actuation", "axis_align_settle"), + ), + ( + "hand_over", + {}, + ( + "handover_pickup", + "handover_transfer", + "handover_exchange", + "handover_settle", + ), + ), + ( + "push_object", + {}, + ("push_contact", "push_actuation", "push_settle"), + ), + ( + "coordinated_pickment", + {}, + ( + "dual_pick_grasp", + "dual_pick_lift", + "dual_pick_move", + "dual_pick_settle", + ), + ), + ( + "coordinated_placement", + {"support_object_id": "support"}, + ( + "coordinated_place_alignment", + "coordinated_place_release", + "coordinated_place_settle", + ), + ), + ], +) +def test_migrated_skill_providers_declare_replay_observation_windows( + skill_id: str, + parameters: dict[str, object], + expected_ids: tuple[str, ...], +) -> None: + segments = { + name: SimpleNamespace(start=index * 5, stop=(index + 1) * 5) + for index, name in enumerate( + ( + "press", + "pull", + "open", + "twist", + "release", + "pour", + "manipulate", + "pickup_close", + "pickup_transport", + "receive_close", + "handover_release", + "close", + "lift", + "transport", + "contact", + "push", + "move", + "approach", + ) + ) + } + compiled = Mock(trajectory=SimpleNamespace(waypoint_count=60)) + compiled.segment.side_effect = lambda action, name: segments[name] + case = replace( + _atomic_case(), + scenario_id=skill_id, + skill_id=skill_id, + case_parameters=parameters, + ) + + windows = create_atomic_skill_provider(skill_id).replay_observation_windows( + compiled, case + ) + + assert tuple(window.observation_id for window in windows) == expected_ids + + +def test_every_atomic_skill_provider_uses_the_versioned_physical_evaluator() -> None: + assert all( + create_atomic_skill_provider(skill_id).physical_evaluator() is not None + for skill_id in atomic_skill_provider_names() + ) + + +def test_move_end_effector_provider_uses_settled_tcp_standard_evidence() -> None: + case = replace( + _atomic_case(), + batch_size=2, + target_waypoints=torch.eye(4).repeat(2, 1, 1, 1), + ) + final_tcp = torch.eye(4).repeat(2, 1, 1) + final_tcp[:, 0, 3] = torch.tensor([0.002, 0.004]) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=final_tcp, + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + ) + outcomes = (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)) + + evaluator = create_atomic_skill_provider("move_end_effector").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, outcomes) + + assert [result.task_success for result in results] == [True, False] + assert results[0].protocol_version == "atomic-skill-v1" + assert results[0].stages[0].evidence_source == ( + "simulation_state:terminal_tcp_pose" + ) + + +def test_move_end_effector_rejects_wrong_terminal_orientation() -> None: + case = replace( + _atomic_case(), + target_waypoints=torch.eye(4).reshape(1, 1, 4, 4), + ) + final_tcp = torch.eye(4).reshape(1, 4, 4) + final_tcp[:, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[0.0, 0.0, math.pi]]) + ) + observation = _ExecutionObservation( + execution_success=torch.ones(1, dtype=torch.bool), + final_tcp_pose=final_tcp, + joint_tracking_rmse_rad=torch.zeros(1), + execution_time_ms=1.0, + task_completion_time_s=0.5, + ) + + evaluator = create_atomic_skill_provider("move_end_effector").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_atomic_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "reached" + + +def test_move_end_effector_keeps_non_finite_failures_row_local() -> None: + case = replace( + _atomic_case(), + batch_size=2, + target_waypoints=torch.eye(4).repeat(2, 1, 1, 1), + ) + final_tcp = torch.eye(4).repeat(2, 1, 1) + final_tcp[1, 0, 3] = torch.nan + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=final_tcp, + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + ) + + evaluator = create_atomic_skill_provider("move_end_effector").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate( + case, + observation, + (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)), + ) + + assert results[0].task_success + assert results[1].failure_code == "evidence_unavailable" + + +def test_move_joints_provider_uses_l2_settled_joint_error() -> None: + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="move_joints", + skill_id="move_joints", + case_parameters={ + "sample_count": 80, + "target_qpos": [[[0.0, 0.0]], [[0.0, 0.0]]], + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + final_arm_qpos=torch.tensor([[0.003, 0.003], [0.004, 0.004]]), + ) + outcomes = (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)) + + evaluator = create_atomic_skill_provider("move_joints").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, outcomes) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_code == "task_goal_miss" + assert results[0].stages[0].measurements["joint_error_l2_rad"] == pytest.approx( + 2**0.5 * 0.003 + ) + + +def test_move_joints_missing_terminal_observation_is_unavailable() -> None: + case = replace( + _atomic_case(), + scenario_id="move_joints", + skill_id="move_joints", + case_parameters={"target_qpos": [[[0.0, 0.0]]]}, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(1, dtype=torch.bool), + final_tcp_pose=torch.eye(4).reshape(1, 4, 4), + joint_tracking_rmse_rad=torch.zeros(1), + execution_time_ms=1.0, + task_completion_time_s=0.5, + ) + + evaluator = create_atomic_skill_provider("move_joints").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_atomic_outcome(),))[0] + + assert result.stages[0].status == "unavailable" + assert result.failure_code == "evidence_unavailable" + + +def test_pick_up_provider_uses_grasp_lift_and_retention_windows() -> None: + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="pick_up", + skill_id="pick_up", + case_parameters={"object_initial_pose": torch.eye(4).repeat(2, 1, 1).tolist()}, + ) + lift = torch.eye(4).repeat(2, 2, 1, 1) + lift[-1, :, 2, 3] = torch.tensor([0.05, 0.03]) + held = torch.full((2, 2, 1), 0.5) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "pickup_grasp": ReplayObservationTrace(gripper_qpos={"hand": held}), + "pickup_lift": ReplayObservationTrace(object_pose=lift), + "pickup_settle": ReplayObservationTrace( + object_pose=lift[-1:].repeat(2, 1, 1, 1), + gripper_qpos={"hand": held}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("pick_up").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate( + case, + observation, + (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)), + ) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_stage == "lifted" + assert results[0].stages[0].evidence_source == ( + "measured_proxy:gripper_joint_state" + ) + + +def test_move_held_object_provider_uses_measured_transport_pose() -> None: + target = torch.eye(4).repeat(2, 1, 1) + target[:, 2, 3] = 0.1 + transport = target.unsqueeze(0).repeat(2, 1, 1, 1) + poses = target.unsqueeze(0).repeat(2, 1, 1, 1) + poses[-1, :, 0, 3] = torch.tensor([0.005, 0.02]) + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="move_held_object", + skill_id="move_held_object", + case_parameters={ + "object_initial_pose": torch.eye(4).repeat(2, 1, 1).tolist(), + "target_object_pose": target.tolist(), + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "held_object_transport": ReplayObservationTrace(object_pose=transport), + "held_object_settle": ReplayObservationTrace( + object_pose=poses, + gripper_qpos={"hand": torch.full((2, 2, 1), 0.5)}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("move_held_object").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate( + case, + observation, + (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)), + ) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_stage == "settled" + + +def test_move_held_object_requires_terminal_retention() -> None: + target = torch.eye(4).reshape(1, 4, 4) + target[:, 2, 3] = 0.1 + poses = target.unsqueeze(0).repeat(2, 1, 1, 1) + case = replace( + _atomic_case(), + scenario_id="move_held_object", + skill_id="move_held_object", + case_parameters={ + "object_initial_pose": torch.eye(4).reshape(1, 4, 4).tolist(), + "target_object_pose": target.tolist(), + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(1, dtype=torch.bool), + final_tcp_pose=torch.eye(4).reshape(1, 4, 4), + joint_tracking_rmse_rad=torch.zeros(1), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "held_object_transport": ReplayObservationTrace(object_pose=poses), + "held_object_settle": ReplayObservationTrace( + object_pose=poses, + gripper_qpos={"hand": torch.ones(2, 1, 1)}, + ), + }, + ) + + evaluator = create_atomic_skill_provider("move_held_object").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_atomic_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "held" + assert result.failure_code == "object_dropped" + + +def test_place_provider_separates_release_goal_and_stability() -> None: + target = torch.eye(4).repeat(2, 1, 1) + target[:, 2, 3] = 0.02 + settle = target.unsqueeze(0).repeat(3, 1, 1, 1) + settle[-1, :, 0, 3] = torch.tensor([0.005, 0.02]) + released = torch.zeros(2, 2, 1) + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="place", + skill_id="place", + case_parameters={ + "object_initial_pose": torch.eye(4).repeat(2, 1, 1).tolist(), + "target_object_pose": target.tolist(), + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "place_release": ReplayObservationTrace(gripper_qpos={"hand": released}), + "place_settle": ReplayObservationTrace(object_pose=settle), + }, + ) + + evaluator = create_atomic_skill_provider("place").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate( + case, + observation, + (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)), + ) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_stage == "placed" + + +def test_place_rejects_wrong_terminal_height_and_orientation() -> None: + target = torch.eye(4).repeat(2, 1, 1) + settle = target.unsqueeze(0).repeat(2, 1, 1, 1) + settle[:, 0, 2, 3] = 0.1 + settle[:, 1, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[math.pi / 2, 0.0, 0.0]]) + ) + released = torch.zeros(2, 2, 1) + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="place", + skill_id="place", + case_parameters={ + "object_initial_pose": torch.eye(4).repeat(2, 1, 1).tolist(), + "target_object_pose": target.tolist(), + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "place_release": ReplayObservationTrace(gripper_qpos={"hand": released}), + "place_settle": ReplayObservationTrace(object_pose=settle), + }, + ) + + evaluator = create_atomic_skill_provider("place").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate( + case, + observation, + (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)), + ) + + assert [result.task_success for result in results] == [False, False] + assert [result.failure_stage for result in results] == ["placed", "placed"] + + +def test_press_provider_scores_only_the_requested_actuation_trace() -> None: + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="press", + skill_id="press", + case_parameters={ + "articulation_joint_stroke": 0.006, + "effect_joint_sign": 1.0, + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + replay_traces={ + "press_actuation": ReplayObservationTrace( + articulation_joint=torch.tensor( + [[0.0, 0.0], [0.003, 0.003], [0.005, 0.004]] + ) + ) + }, + ) + outcomes = (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)) + + evaluator = create_atomic_skill_provider("press").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, outcomes) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_stage == "pressed" + + +def test_slide_provider_uses_actuation_and_release_windows() -> None: + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="slide", + skill_id="slide", + case_parameters={ + "translation_distance_m": 0.1, + "direction": "pull", + "effect_joint_sign": 1.0, + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "slide_actuation": ReplayObservationTrace( + articulation_joint=torch.tensor([[0.0, 0.0], [0.1, 0.07]]) + ), + "slide_release": ReplayObservationTrace( + gripper_qpos={"hand": torch.tensor([[[1.0], [1.0]], [[0.0], [0.0]]])} + ), + }, + ) + outcomes = (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)) + + evaluator = create_atomic_skill_provider("slide").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, outcomes) + + assert [result.task_success for result in results] == [True, False] + assert results[0].stages[-1].stage_id == "released" + assert results[1].failure_stage == "slid" + + +def test_slide_rejects_opposite_joint_direction() -> None: + case = replace( + _atomic_case(), + scenario_id="slide", + skill_id="slide", + case_parameters={ + "translation_distance_m": 0.1, + "direction": "pull", + "effect_joint_sign": 1.0, + }, + ) + observation = _ExecutionObservation( + execution_success=torch.ones(1, dtype=torch.bool), + final_tcp_pose=torch.eye(4).reshape(1, 4, 4), + joint_tracking_rmse_rad=torch.zeros(1), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "slide_actuation": ReplayObservationTrace( + articulation_joint=torch.tensor([[0.0], [-0.1]]) + ), + "slide_release": ReplayObservationTrace( + gripper_qpos={"hand": torch.tensor([[[1.0]], [[0.0]]])} + ), + }, + ) + + evaluator = create_atomic_skill_provider("slide").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_atomic_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "grasped" + + +def test_twist_provider_scores_executed_tcp_rotation_not_knob_angle() -> None: + case = replace( + _atomic_case(), + batch_size=2, + scenario_id="twist", + skill_id="twist", + case_parameters={ + "twist_target_pose": torch.eye(4).repeat(2, 1, 1).tolist(), + "twist_axis": [0.0, 0.0, 1.0], + "twist_angle_rad": 0.5, + }, + ) + tcp = torch.eye(4).repeat(2, 2, 1, 1) + tcp[-1, :, :3, :3] = axis_angle_to_rotation_matrix( + torch.tensor([[0.0, 0.0, 0.5], [0.0, 0.0, 0.2]]) + ) + observation = _ExecutionObservation( + execution_success=torch.ones(2, dtype=torch.bool), + final_tcp_pose=torch.eye(4).repeat(2, 1, 1), + joint_tracking_rmse_rad=torch.zeros(2), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "twist_actuation": ReplayObservationTrace( + tcp_pose=tcp, + articulation_joint=torch.tensor([[0.0, 0.0], [0.1, 0.1]]), + ), + "twist_release": ReplayObservationTrace( + gripper_qpos={"hand": torch.tensor([[[1.0], [1.0]], [[0.0], [0.0]]])} + ), + }, + ) + outcomes = (_atomic_outcome(), replace(_atomic_outcome(), env_index=1)) + + evaluator = create_atomic_skill_provider("twist").physical_evaluator() + assert evaluator is not None + results = evaluator.evaluate(case, observation, outcomes) + + assert [result.task_success for result in results] == [True, False] + assert results[1].failure_stage == "twisted" + + +def test_twist_rejects_opposite_tcp_rotation() -> None: + case = replace( + _atomic_case(), + scenario_id="twist", + skill_id="twist", + case_parameters={ + "twist_target_pose": torch.eye(4).reshape(1, 4, 4).tolist(), + "twist_axis": [0.0, 0.0, 1.0], + "twist_angle_rad": -0.5, + }, + ) + tcp = torch.eye(4).repeat(2, 1, 1, 1) + tcp[-1, :, :3, :3] = axis_angle_to_rotation_matrix(torch.tensor([[0.0, 0.0, 0.5]])) + observation = _ExecutionObservation( + execution_success=torch.ones(1, dtype=torch.bool), + final_tcp_pose=torch.eye(4).reshape(1, 4, 4), + joint_tracking_rmse_rad=torch.zeros(1), + execution_time_ms=1.0, + task_completion_time_s=0.5, + gripper_open_qpos={"hand": torch.tensor([0.0])}, + gripper_grasp_qpos={"hand": torch.tensor([1.0])}, + replay_traces={ + "twist_actuation": ReplayObservationTrace( + tcp_pose=tcp, + articulation_joint=torch.tensor([[0.0], [0.1]]), + ), + "twist_release": ReplayObservationTrace( + gripper_qpos={"hand": torch.tensor([[[1.0]], [[0.0]]])} + ), + }, + ) + + evaluator = create_atomic_skill_provider("twist").physical_evaluator() + assert evaluator is not None + result = evaluator.evaluate(case, observation, (_atomic_outcome(),))[0] + + assert not result.task_success + assert result.failure_stage == "twisted" + + def _fake_articulation_geometry(*, joint_type: str) -> dict[str, torch.Tensor]: """Build non-degenerate link-local geometry for articulation providers.""" target_points = torch.tensor( @@ -569,6 +1232,8 @@ def test_new_atomic_skill_cases_freeze_reference_waypoints( robot.compute_fk.return_value = torch.eye(4).repeat(batch_size, 1, 1) scenario = Mock(robot=robot, control_part="arm") entity = Mock(uid=f"benchmark_{config['articulation']}") + entity.joint_names = [config["target_joint"]] + entity.get_qpos_limits.return_value = torch.tensor([[[0.0, 0.006]]]) articulation_dof = 4 if config["articulation"] == "microwave" else 1 articulation_pose = torch.eye(4).repeat(batch_size, 1, 1) articulation_qpos = torch.zeros(batch_size, articulation_dof) diff --git a/tests/benchmark/motion_generation/test_contracts.py b/tests/benchmark/motion_generation/test_contracts.py index 6e48dd17e..53cb7d2a4 100644 --- a/tests/benchmark/motion_generation/test_contracts.py +++ b/tests/benchmark/motion_generation/test_contracts.py @@ -301,3 +301,54 @@ def test_batch_results_require_the_same_stage_order() -> None: ) with pytest.raises(ValueError): contracts.validate_physical_evaluations(rows, batch_size=2) + + +def test_replay_observation_window_accepts_segments_and_terminal_settle() -> None: + segment = contracts.ReplayObservationWindow( + "actuation", 3, 7, frozenset({"object_pose", "gripper_qpos"}) + ) + settle = contracts.ReplayObservationWindow( + "settle", + 9, + 9, + frozenset({"object_pose"}), + include_settle=True, + ) + scene = contracts.ReplayObservationWindow( + "relative", + 4, + 8, + frozenset({"object_pose", "scene_object_pose"}), + scene_object_ids=("support",), + ) + + assert segment.stop == 7 + assert settle.include_settle + assert scene.scene_object_ids == ("support",) + with pytest.raises(ValueError, match="empty waypoint window"): + contracts.ReplayObservationWindow("invalid", 2, 2, frozenset({"tcp_pose"})) + with pytest.raises(ValueError, match="channels"): + contracts.ReplayObservationWindow( + "invalid", 0, 1, frozenset({"projected_effect"}) + ) + with pytest.raises(ValueError, match="scene_object_pose"): + contracts.ReplayObservationWindow( + "invalid", 0, 1, frozenset({"scene_object_pose"}) + ) + + +def test_replay_observation_trace_requires_synchronized_batch_axes() -> None: + trace = contracts.ReplayObservationTrace( + object_pose=torch.eye(4).repeat(5, 2, 1, 1), + gripper_qpos={"hand": torch.zeros(5, 2, 2)}, + scene_object_poses={"support": torch.eye(4).repeat(5, 2, 1, 1)}, + ) + + assert trace.object_pose is not None + assert trace.object_pose.shape == (5, 2, 4, 4) + assert trace.scene_object_poses["support"].shape == (5, 2, 4, 4) + with pytest.raises(ValueError, match="share sample and environment axes"): + contracts.ReplayObservationTrace( + object_pose=torch.eye(4).repeat(5, 2, 1, 1), + gripper_qpos={"hand": torch.zeros(4, 2, 2)}, + ) diff --git a/tests/benchmark/motion_generation/test_evaluation.py b/tests/benchmark/motion_generation/test_evaluation.py index bca7cdb41..706a6ff10 100644 --- a/tests/benchmark/motion_generation/test_evaluation.py +++ b/tests/benchmark/motion_generation/test_evaluation.py @@ -37,6 +37,7 @@ STANDARD_VERSION, SKILL_CRITERIA, StageEvidence, + evaluate_skill_stage_batch, evaluate_skill_stages, validate_standard_evaluation, ) @@ -59,7 +60,13 @@ def _evidence(window: str, **measurements: float | bool) -> StageEvidence: def _primitive(error: float, env_index: int = 0) -> PhysicalEvaluation: return evaluate_skill_stages( "move_end_effector", - {"reached": _evidence("terminal_settle", tcp_position_error_m=error)}, + { + "reached": _evidence( + "terminal_settle", + tcp_position_error_m=error, + tcp_rotation_error_rad=0.0, + ) + }, env_index=env_index, ) @@ -166,6 +173,33 @@ def test_joint_criterion_uses_l2_measurement_not_max_joint_error() -> None: assert result.failure_code == "task_goal_miss" +def test_batched_stage_evidence_preserves_environment_order() -> None: + results = evaluate_skill_stage_batch( + "move_end_effector", + ( + { + "reached": _evidence( + "terminal_settle", + tcp_position_error_m=0.002, + tcp_rotation_error_rad=0.0, + ) + }, + { + "reached": _evidence( + "terminal_settle", + tcp_position_error_m=0.004, + tcp_rotation_error_rad=0.0, + ) + }, + ), + ) + + assert [result.env_index for result in results] == [0, 1] + assert [result.task_success for result in results] == [True, False] + for result in results: + validate_standard_evaluation("move_end_effector", result) + + def test_lift_then_drop_fails_closing_stage() -> None: result = evaluate_skill_stages( "pick_up", @@ -193,7 +227,10 @@ def test_later_observations_do_not_repair_failed_release() -> None: { "released": _evidence("release_end", release_established=False), "placed": _evidence( - "terminal_settle", xy_error_m=0.0, min_height_above_support_m=0.0 + "terminal_settle", + position_error_m=0.0, + orientation_error_rad=0.0, + min_height_above_support_m=0.0, ), "stable": _evidence( "settle_interval", @@ -286,6 +323,7 @@ def test_handover_signed_height_bounds(height: float, passed: bool) -> None: receiver_retained=True, min_height_above_support_m=0.0, ), + "receiver_released": _evidence("terminal_settle", release_established=True), "placed": _evidence( "terminal_settle", xy_error_m=0.03, @@ -466,12 +504,13 @@ def test_unknown_skill_and_projected_evidence_are_rejected() -> None: ) -def test_catalog_preserves_unimplemented_skills_as_declarations() -> None: +def test_catalog_preserves_full_skill_stage_declarations() -> None: assert len(SKILL_CRITERIA) == 15 assert tuple(s.stage_id for s in SKILL_CRITERIA["coordinated_pickment"]) == ( "dual_grasped", "lifted", "moved", + "settled", "held", )