From e0fec8ab2e8a7d4819ec5921c964e28dd5f75227 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 19 Aug 2026 02:59:27 +0000 Subject: [PATCH] feat(benchmark): add atomic skill planner framework Add shared planner, robot, scenario, skill, and object extension points for physics-backed Atomic Action evaluation. Introduce a Franka + PGI cuRobo smoke suite with MoveEndEffector and PickUp, staged success metrics, execution timing, manifests, reports, documentation, and tests. --- scripts/benchmark/__main__.py | 4 +- .../motion_generation/BENCHMARK_DESIGN.md | 22 + scripts/benchmark/motion_generation/README.md | 48 +- .../motion_generation/aggregation.py | 249 +++- .../benchmark/motion_generation/artifacts.py | 30 +- scripts/benchmark/motion_generation/config.py | 14 + scripts/benchmark/motion_generation/models.py | 29 +- .../motion_generation/planners/base.py | 22 + .../motion_generation/planners/curobo.py | 6 +- .../benchmark/motion_generation/registry.py | 34 +- .../benchmark/motion_generation/reporting.py | 28 +- .../motion_generation/robots/__init__.py | 24 + .../motion_generation/robots/base.py | 48 + .../motion_generation/robots/franka.py | 68 + .../motion_generation/run_benchmark.py | 17 +- scripts/benchmark/motion_generation/runner.py | 149 ++- .../motion_generation/scenarios/__init__.py | 30 +- .../scenarios/atomic_objects.py | 194 +++ .../scenarios/atomic_task.py | 1176 +++++++++++++++++ .../motion_generation/scenarios/base.py | 103 +- .../motion_generation/scenarios/free_space.py | 1 + .../suites/atomic_franka_pgi_curobo.yaml | 102 ++ .../test_atomic_task_benchmark.py | 221 ++++ 23 files changed, 2527 insertions(+), 92 deletions(-) create mode 100644 scripts/benchmark/motion_generation/robots/__init__.py create mode 100644 scripts/benchmark/motion_generation/robots/base.py create mode 100644 scripts/benchmark/motion_generation/robots/franka.py create mode 100644 scripts/benchmark/motion_generation/scenarios/atomic_objects.py create mode 100644 scripts/benchmark/motion_generation/scenarios/atomic_task.py create mode 100644 scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml create mode 100644 tests/benchmark/motion_generation/test_atomic_task_benchmark.py diff --git a/scripts/benchmark/__main__.py b/scripts/benchmark/__main__.py index 226843aa4..68ccc98f4 100644 --- a/scripts/benchmark/__main__.py +++ b/scripts/benchmark/__main__.py @@ -51,7 +51,7 @@ def _run_rl_cli(_: argparse.Namespace) -> None: def _run_motion_generation_cli(args: argparse.Namespace) -> None: - """Run the free-space motion-generation benchmark.""" + """Run the shared planner and Atomic Task benchmark.""" from scripts.benchmark.motion_generation.run_benchmark import run_from_args run_from_args(args) @@ -118,7 +118,7 @@ def main(argv: Sequence[str] | None = None) -> None: motion_generation_parser = subparsers.add_parser( "motion-generation", - help="Benchmark free-space motion generation with cuRobo as baseline.", + help="Benchmark planners on fixed motion and Atomic Task cases.", ) add_parser_arguments(motion_generation_parser) motion_generation_parser.set_defaults(func=_run_motion_generation_cli) diff --git a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md index 66463bb2c..22cc90b01 100644 --- a/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md +++ b/scripts/benchmark/motion_generation/BENCHMARK_DESIGN.md @@ -15,6 +15,28 @@ The default comparison should be NMG versus cuRobo. IK plus interpolation and TOPPRA should remain optional diagnostic baselines rather than define the main leaderboard. +## Implemented vertical slice + +The first physics-backed slice is available as +`suites/atomic_franka_pgi_curobo.yaml`. It runs Franka + PGI with cuRobo only, +and covers `MoveEndEffector` plus antipodal-grasp `PickUp`. Both skills pin +`MotionPolicy(strategy="motion_gen", planner="curobo")` and compile through the +same `AtomicActionEngine`; scenario code never calls cuRobo directly. + +The shared runner now selects planners, scenarios, robots, Atomic Action case +providers, and object kinds through registries. Cases freeze the full robot +start state, explicit targets/grasp, object configuration, difficulty factors, +and independent sequential-IK evidence before measured planner calls. Reports +keep planning, kinematic motion validity, controller execution, and physical +task success separate while retaining exactly three tables. + +Run the slice with: + +```bash +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo --device cuda +``` + ## Motivation The existing NeuralPlanner benchmark provides useful latency, memory, rollout, diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index 9ec7ce42d..7c988e962 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -1,36 +1,58 @@ -# Motion Generation Benchmark +# Planner Motion Generation & Atomic Skill Benchmark -Free-space motion-generation suite with cuRobo as the default primary baseline. +Shared planner benchmark framework for fixed motion-generation cases and +physics-backed Atomic Actions. All Atomic Actions call the selected planner +through the adapter-owned `MotionGenerator`. Design background and roadmap: see [`BENCHMARK_DESIGN.md`](./BENCHMARK_DESIGN.md). ## Run ```bash -embodichain benchmark motion-generation --suite smoke -embodichain benchmark motion-generation --suite coverage -embodichain benchmark motion-generation --extra-baselines ik_interpolate toppra -embodichain benchmark motion-generation --path-shapes direct l_turn --start-state-bins nominal near_singularity +python -m scripts.benchmark.motion_generation.run_benchmark --suite smoke +python -m scripts.benchmark.motion_generation.run_benchmark --suite coverage +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo --device cuda +python -m scripts.benchmark.motion_generation.run_benchmark \ + --extra-baselines ik_interpolate toppra ``` -Artifacts land under `outputs/benchmarks/motion_generation//` +Artifacts land under `outputs/benchmarks///` (`resolved_suite.yaml`, `case_manifest.json`, `trials.jsonl`, `aggregates.json`, `report.md` with exactly three tables). ## Implemented -- Extensible planner/scenario registries and track-based suite YAML +- Extensible planner, scenario, robot, Atomic Action, and object registries - `free-space-common` track with fixed manifests and start-state bins +- `atomic-task` track with frozen robot/object/task manifests and common physics replay +- Initial Atomic Task slice: Franka + PGI, cuRobo, `MoveEndEffector`, and + antipodal-grasp `PickUp` on a declarative cube - Default matrix: cuRobo (`primary_baseline`); IK / TOPPRA optional diagnostics - NMG adapter stub (`candidate`, disabled until a checkpoint is ready) - Lifecycle timing: construct / prepare / cold / warm -- Ordered waypoint matching and external `motion_valid` (separate from - `PlanResult.success`) +- Distinct planning, motion-valid, execution, and physical task-success stages +- Planning latency, execution wall time, end-to-end time, nominal trajectory + duration, simulated task-completion time, controller tracking RMSE, and + task-specific object lift - One Markdown report: Time & Memory, Success & Other Metrics, Leaderboard -## Not implemented yet +## Extend + +- Planner: register a `PlannerAdapter`, expose its `MotionGenerator`, and + declare the `atomic_action` capability. +- Robot: register a `RobotProvider` and select it under `robot`; PickUp suites + also declare the gripper control part and open/grasp qpos under `gripper`. +- Object: add another `objects` entry for built-in `cube`/`mesh`, or register a + new object-kind factory. +- Atomic skill: implement and register an `AtomicSkillCaseProvider`; the runner, + artifact schema, aggregation, and report stay unchanged. + +## Current limits - Real NMG checkpoint adapter -- `collision-deployment` and `atomic-task` tracks -- Physics execution / task-success metrics +- `collision-deployment` and obstacle-aware common-input tracks +- Atomic Task execution is currently `B=1`; the supplied suite covers only + Franka + PGI and cuRobo +- Remaining skills: `MoveHeldObject`, `Place`, `Press`, and action chains - Latency-budget Pareto sweeps, confidence intervals, subprocess isolation diff --git a/scripts/benchmark/motion_generation/aggregation.py b/scripts/benchmark/motion_generation/aggregation.py index 25dbd45db..60eea4196 100644 --- a/scripts/benchmark/motion_generation/aggregation.py +++ b/scripts/benchmark/motion_generation/aggregation.py @@ -70,6 +70,51 @@ def _case_macro_rate( return sum(case_rates) / len(case_rates) +def _case_supports_attribute(case: BenchmarkCase, attribute: str) -> bool: + """Return whether a staged outcome applies to one case protocol.""" + if attribute == "task_success": + return case.primary_success == "task_success" + if attribute == "execution_success": + return case.primary_success in {"execution_success", "task_success"} + return True + + +def _case_macro_optional_rate( + measured: list[TrialRecord], + track_cases: list[BenchmarkCase], + attribute: str, +) -> float | None: + """Macro-average an applicable staged boolean, otherwise return ``None``.""" + applicable = [ + case for case in track_cases if _case_supports_attribute(case, attribute) + ] + if not applicable: + return None + return _case_macro_rate(measured, applicable, attribute) + + +def _case_macro_primary_rate( + measured: list[TrialRecord], track_cases: list[BenchmarkCase] +) -> float: + """Macro-average each case's explicitly declared primary success stage.""" + if not track_cases: + return 0.0 + outcomes_by_case: dict[str, list[CaseOutcome]] = defaultdict(list) + for record in measured: + outcomes_by_case[record.case_id].extend(record.outcomes) + case_rates: list[float] = [] + for case in track_cases: + outcomes = outcomes_by_case.get(case.case_id, []) + if not outcomes: + case_rates.append(0.0) + continue + case_rates.append( + sum(bool(getattr(outcome, case.primary_success)) for outcome in outcomes) + / len(outcomes) + ) + return sum(case_rates) / len(case_rates) + + def _case_macro_mean( measured: list[TrialRecord], track_cases: list[BenchmarkCase], @@ -183,15 +228,20 @@ def _performance_rows( cases: list[BenchmarkCase], ) -> list[dict[str, object]]: """Aggregate steady-state time and memory by track, algorithm, and input shape.""" - measured_groups: dict[tuple[str, str, int, int], list[TrialRecord]] = defaultdict( - list - ) + measured_groups: dict[ + tuple[str, str, str, str, str | None, str | None, int, int], + list[TrialRecord], + ] = defaultdict(list) for record in records: if record.phase is TrialPhase.MEASURED: measured_groups[ ( record.track, record.algorithm_id, + record.robot_id, + record.skill_id, + record.object_id, + record.task_difficulty, record.batch_size, record.waypoint_count, ) @@ -199,8 +249,29 @@ def _performance_rows( metadata_by_id = {item.algorithm_id: item for item in metadata} rows: list[dict[str, object]] = [] - for key in sorted(measured_groups): - track, algorithm_id, batch_size, waypoint_count = key + for key in sorted( + measured_groups, + key=lambda item: ( + item[0], + item[1], + item[2], + item[3], + item[4] or "", + item[5] or "", + item[6], + item[7], + ), + ): + ( + track, + algorithm_id, + robot_id, + skill_id, + object_id, + task_difficulty, + batch_size, + waypoint_count, + ) = key group = measured_groups[key] info = metadata_by_id[algorithm_id] costs = [record.cost_time_ms for record in group] @@ -210,6 +281,10 @@ def _performance_rows( "track": track, "algorithm": algorithm_id, "algorithm_role": info.algorithm_role.value, + "robot": robot_id, + "skill": skill_id, + "object": object_id, + "task_difficulty": task_difficulty, "batch_size": batch_size, "waypoint_count": waypoint_count, "num_trials": len(group), @@ -244,6 +319,23 @@ def _performance_rows( "cpu_delta_mb": _mean(record.cpu_delta_mb for record in group), "gpu_delta_mb": _mean(record.gpu_delta_mb for record in group), "peak_gpu_mb": _peak_gpu(group), + "execution_time_ms": _mean( + record.execution_time_ms for record in group + ), + "end_to_end_time_ms": _mean( + record.end_to_end_time_ms for record in group + ), + "trajectory_duration_s": _mean( + record.trajectory_duration_s for record in group + ), + "trajectory_waypoints": _mean( + ( + float(record.trajectory_waypoints) + if record.trajectory_waypoints is not None + else None + ) + for record in group + ), } ) @@ -257,6 +349,10 @@ def _performance_rows( "track": track, "algorithm": info.algorithm_id, "algorithm_role": info.algorithm_role.value, + "robot": None, + "skill": None, + "object": None, + "task_difficulty": None, "batch_size": None, "waypoint_count": None, "num_trials": 0, @@ -272,6 +368,10 @@ def _performance_rows( "cpu_delta_mb": None, "gpu_delta_mb": None, "peak_gpu_mb": None, + "execution_time_ms": None, + "end_to_end_time_ms": None, + "trajectory_duration_s": None, + "trajectory_waypoints": None, } ) return sorted( @@ -281,6 +381,8 @@ def _performance_rows( str(row["algorithm"]), int(row["batch_size"] or 0), int(row["waypoint_count"] or 0), + str(row["skill"] or ""), + str(row["object"] or ""), ), ) @@ -298,7 +400,21 @@ def _metric_rows( remain conditioned on externally motion-valid outcomes. """ measured_by_key: dict[ - tuple[str, str, str, int, int, str, str], list[TrialRecord] + tuple[ + str, + str, + str, + str, + str, + str | None, + str | None, + str, + int, + int, + str, + str, + ], + list[TrialRecord], ] = defaultdict(list) for record in records: if record.phase is not TrialPhase.MEASURED: @@ -307,6 +423,11 @@ def _metric_rows( record.track, record.algorithm_id, record.scenario_id, + record.robot_id, + record.skill_id, + record.object_id, + record.task_difficulty, + record.primary_success, record.batch_size, record.waypoint_count, record.path_shape, @@ -314,14 +435,46 @@ def _metric_rows( ) measured_by_key[key].append(record) - expected_by_group: Counter[tuple[str, str, int, int, str, str]] = Counter() - cases_by_group: dict[tuple[str, str, int, int, str, str], list[BenchmarkCase]] = ( - defaultdict(list) - ) + expected_by_group: Counter[ + tuple[ + str, + str, + str, + str, + str | None, + str | None, + str, + int, + int, + str, + str, + ] + ] = Counter() + cases_by_group: dict[ + tuple[ + str, + str, + str, + str, + str | None, + str | None, + str, + int, + int, + str, + str, + ], + list[BenchmarkCase], + ] = defaultdict(list) for case in cases: key = ( case.track, case.scenario_id, + case.robot_id, + case.skill_id, + case.object_id, + case.task_difficulty, + case.primary_success, case.batch_size, case.num_waypoints, case.path_shape, @@ -332,10 +485,30 @@ def _metric_rows( rows: list[dict[str, object]] = [] for info in metadata: - for group_key in sorted(expected_by_group): + for group_key in sorted( + expected_by_group, + key=lambda item: ( + item[0], + item[1], + item[2], + item[3], + item[4] or "", + item[5] or "", + item[6], + item[7], + item[8], + item[9], + item[10], + ), + ): ( track, scenario_id, + robot_id, + skill_id, + object_id, + task_difficulty, + primary_success, batch_size, waypoint_count, path_shape, @@ -347,6 +520,11 @@ def _metric_rows( track, info.algorithm_id, scenario_id, + robot_id, + skill_id, + object_id, + task_difficulty, + primary_success, batch_size, waypoint_count, path_shape, @@ -363,6 +541,11 @@ def _metric_rows( "scenario": scenario_id, "algorithm": info.algorithm_id, "algorithm_role": info.algorithm_role.value, + "robot": robot_id, + "skill": skill_id, + "object": object_id, + "task_difficulty": task_difficulty, + "primary_success": primary_success, "batch_size": batch_size, "waypoint_count": waypoint_count, "path_shape": path_shape, @@ -370,13 +553,19 @@ def _metric_rows( "cases": len(group_cases), "n_valid": len(valid_outcomes), "coverage_rate": min(1.0, len(outcomes) / max(expected, 1)), - # Free-space primary success is external motion validity. - "success_rate": _case_macro_rate( - measured, group_cases, "motion_valid" - ), + "success_rate": _case_macro_primary_rate(measured, group_cases), "planning_success_rate": _case_macro_rate( measured, group_cases, "planning_success" ), + "motion_valid_rate": _case_macro_rate( + measured, group_cases, "motion_valid" + ), + "execution_success_rate": _case_macro_optional_rate( + measured, group_cases, "execution_success" + ), + "task_success_rate": _case_macro_optional_rate( + measured, group_cases, "task_success" + ), "ordered_waypoint_success_rate": _case_macro_rate( measured, group_cases, "ordered_waypoints_reached" ), @@ -409,6 +598,25 @@ def _metric_rows( "path_efficiency": _mean( outcome.path_efficiency for outcome in valid_outcomes ), + "task_completion_time_s": _mean( + outcome.task_completion_time_s + for outcome in outcomes + if outcome.task_success + ), + "joint_tracking_rmse_rad": _mean( + outcome.joint_tracking_rmse_rad for outcome in outcomes + ), + "object_lift_delta_m": _mean( + outcome.object_lift_delta_m for outcome in outcomes + ), + "replan_count": _mean( + ( + float(outcome.replan_count) + if outcome.replan_count is not None + else None + ) + for outcome in outcomes + ), "top_failure": _top_failure(outcomes), } ) @@ -447,6 +655,11 @@ def _leaderboard_rows( coverage = min(1.0, len(outcomes) / max(expected_outcomes, 1)) motion_rate = _case_macro_rate(measured, track_cases, "motion_valid") planning_rate = _case_macro_rate(measured, track_cases, "planning_success") + execution_rate = _case_macro_optional_rate( + measured, track_cases, "execution_success" + ) + task_rate = _case_macro_optional_rate(measured, track_cases, "task_success") + primary_rate = _case_macro_primary_rate(measured, track_cases) latency_p95 = _case_macro_latency_p95(measured, track_cases) peak_gpu = _peak_gpu(measured) track_entries.append( @@ -458,11 +671,11 @@ def _leaderboard_rows( "planner_config_hash": info.config_hash[:12], "eligible": coverage >= 1.0 - 1.0e-12, "coverage_rate": coverage, - # free-space v1: primary_success == motion_valid - "overall_success_rate": motion_rate, + "overall_success_rate": primary_rate, "planning_success_rate": planning_rate, "motion_valid_rate": motion_rate, - "task_success_rate": None, + "execution_success_rate": execution_rate, + "task_success_rate": task_rate, "latency_p95_ms": latency_p95, "peak_gpu_mb": peak_gpu, } diff --git a/scripts/benchmark/motion_generation/artifacts.py b/scripts/benchmark/motion_generation/artifacts.py index 4544230f6..af8928b9d 100644 --- a/scripts/benchmark/motion_generation/artifacts.py +++ b/scripts/benchmark/motion_generation/artifacts.py @@ -124,21 +124,47 @@ def _case_to_dict(case: BenchmarkCase) -> dict[str, Any]: "num_waypoints": case.num_waypoints, "path_shape": case.path_shape, "start_state_bin": case.start_state_bin, + "robot_id": case.robot_id, + "skill_id": case.skill_id, + "object_id": case.object_id, + "task_difficulty": case.task_difficulty, + "primary_success": case.primary_success, "start_qpos": case.start_qpos.detach().cpu().tolist(), + "full_start_qpos": ( + None + if case.full_start_qpos is None + else case.full_start_qpos.detach().cpu().tolist() + ), "target_waypoints": case.target_waypoints.detach().cpu().tolist(), + "case_parameters": _to_json_value(case.case_parameters), "validity_evidence": { - "method": "reference_qpos_fk", + "method": ( + "reference_qpos_fk" + if case.skill_id == "N/A" + else "independent_sequential_ik" + ), "reference_qpos": case.reference_qpos.detach().cpu().tolist(), }, } +def _to_json_value(value: object) -> object: + """Recursively preserve tensors and numeric case configuration values.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().tolist() + if isinstance(value, dict): + return {str(key): _to_json_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_to_json_value(item) for item in value] + return value + + def write_case_manifest(path: str | Path, cases: list[BenchmarkCase]) -> Path: """Write the algorithm-independent case manifest.""" return write_json( path, { - "case_schema_version": 1, + "case_schema_version": 2, "cases": [_case_to_dict(case) for case in cases], }, ) diff --git a/scripts/benchmark/motion_generation/config.py b/scripts/benchmark/motion_generation/config.py index 61fb07019..a58321445 100644 --- a/scripts/benchmark/motion_generation/config.py +++ b/scripts/benchmark/motion_generation/config.py @@ -35,6 +35,7 @@ "FreeSpaceTrackCfg", "PlannerSpecCfg", "ProtocolCfg", + "RobotSpecCfg", "SuiteCfg", "TrackCfg", "load_suite", @@ -84,6 +85,15 @@ class ProtocolCfg: joint_limit_tolerance_rad: float = 1.0e-5 +@configclass +class RobotSpecCfg: + """Robot provider selected for every track in one suite run.""" + + id: str = "franka_panda" + provider: str = "franka_panda" + config: dict[str, Any] = {} + + @configclass class FreeSpaceTrackCfg: """Case matrix for the ``free-space-common`` track.""" @@ -114,6 +124,7 @@ class SuiteCfg: suite_version: str = "free_space_common_v2" profile: str = "smoke" planners: list[PlannerSpecCfg] = [] + robot: RobotSpecCfg = RobotSpecCfg() protocol: ProtocolCfg = ProtocolCfg() tracks: list[TrackCfg] = [] free_space: FreeSpaceTrackCfg = FreeSpaceTrackCfg() @@ -129,6 +140,7 @@ def from_dict(cls, data: dict[str, Any]) -> "SuiteCfg": suite_version=str(data.get("suite_version", "free_space_common_v2")), profile=str(data.get("profile", "smoke")), planners=planners, + robot=RobotSpecCfg(**data.get("robot", {})), protocol=ProtocolCfg(**data.get("protocol", {})), tracks=tracks, free_space=free_space, @@ -171,6 +183,8 @@ def validate_benchmark(self) -> None: "Every planner must define a non-empty id and adapter." ) AlgorithmRole(spec.role) + if not self.robot.id or not self.robot.provider: + raise ValueError("robot must define non-empty id and provider values.") if not self.tracks: raise ValueError("The benchmark suite must declare at least one track.") track_ids = [track.id for track in self.tracks] diff --git a/scripts/benchmark/motion_generation/models.py b/scripts/benchmark/motion_generation/models.py index 02868bc78..e670d564d 100644 --- a/scripts/benchmark/motion_generation/models.py +++ b/scripts/benchmark/motion_generation/models.py @@ -69,7 +69,12 @@ class PlannerMetadata: @dataclass(frozen=True) class BenchmarkCase: - """One env-batched free-space planning input frozen before execution.""" + """One env-batched planner input frozen before execution. + + The first fields retain the free-space case contract. The trailing fields + describe execution/task tracks without forcing planner adapters to know + about a particular robot, atomic skill, or object implementation. + """ suite_version: str track: str @@ -83,6 +88,13 @@ class BenchmarkCase: start_qpos: torch.Tensor target_waypoints: torch.Tensor reference_qpos: torch.Tensor + robot_id: str = "franka_panda" + skill_id: str = "N/A" + object_id: str | None = None + task_difficulty: str | None = None + primary_success: str = "motion_valid" + full_start_qpos: torch.Tensor | None = None + case_parameters: dict[str, object] = field(default_factory=dict) @dataclass(frozen=True) @@ -110,6 +122,12 @@ class CaseOutcome: path_efficiency: float | None failure_code: str | None = None planner_failure_code: str | None = None + execution_success: bool | None = None + task_success: bool | None = None + task_completion_time_s: float | None = None + joint_tracking_rmse_rad: float | None = None + object_lift_delta_m: float | None = None + replan_count: int | None = None @dataclass(frozen=True) @@ -138,6 +156,15 @@ class TrialRecord: cpu_delta_mb: float | None = None gpu_delta_mb: float | None = None peak_gpu_mb: float | None = None + robot_id: str = "franka_panda" + skill_id: str = "N/A" + object_id: str | None = None + task_difficulty: str | None = None + primary_success: str = "motion_valid" + execution_time_ms: float | None = None + end_to_end_time_ms: float | None = None + trajectory_duration_s: float | None = None + trajectory_waypoints: int | None = None metadata: dict[str, object] = field(default_factory=dict) outcomes: tuple[CaseOutcome, ...] = () diff --git a/scripts/benchmark/motion_generation/planners/base.py b/scripts/benchmark/motion_generation/planners/base.py index 895e5f11b..a8e85481b 100644 --- a/scripts/benchmark/motion_generation/planners/base.py +++ b/scripts/benchmark/motion_generation/planners/base.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from embodichain.lab.sim.objects import Robot + from embodichain.lab.sim.planners import MotionGenerator __all__ = ["PlannerAdapter", "PlannerContext"] @@ -43,6 +44,7 @@ class PlannerContext: control_part: str device: torch.device sample_interval: int + robot_id: str = "unknown" class PlannerAdapter(ABC): @@ -69,6 +71,7 @@ def metadata(self) -> PlannerMetadata: model_revision=str( self.spec.config.get("model_revision", self.model_revision) ), + supported_robots=(self.context.robot_id,), parameters=dict(self.spec.config), ) @@ -84,6 +87,25 @@ def prepare(self, case: BenchmarkCase) -> dict[str, object] | None: """Prepare a lazy backend, or return ``None`` when not applicable.""" return None + @property + def motion_policy_planner(self) -> str: + """Return the backend name pinned into Atomic Action motion policies.""" + return self.spec.adapter + + def require_motion_generator(self) -> "MotionGenerator": + """Return the adapter-owned MotionGenerator or fail clearly. + + Atomic-task scenarios use this boundary instead of reaching into a + backend implementation. Every planner that opts into the + ``atomic_action`` capability must expose its generator here. + """ + motion_generator = getattr(self, "motion_generator", None) + if motion_generator is None: + raise RuntimeError( + f"Planner adapter {self.spec.id!r} does not expose a MotionGenerator." + ) + return motion_generator + @abstractmethod def plan(self, case: BenchmarkCase) -> PlanResult: """Plan one env-batched benchmark case.""" diff --git a/scripts/benchmark/motion_generation/planners/curobo.py b/scripts/benchmark/motion_generation/planners/curobo.py index 2b4e34716..153d903a1 100644 --- a/scripts/benchmark/motion_generation/planners/curobo.py +++ b/scripts/benchmark/motion_generation/planners/curobo.py @@ -46,7 +46,9 @@ class CuroboAdapter(PlannerAdapter): """Run cuRobo with a frozen, empty-world operational configuration.""" - capabilities = frozenset({"eef_waypoint", "batched", "empty_world"}) + capabilities = frozenset( + {"eef_waypoint", "batched", "empty_world", "atomic_action"} + ) model_revision = "curobo-v2" separate_prepare = True @@ -69,7 +71,7 @@ def build(self) -> None: auto_values = dict(values.get("auto_gen", {})) if bool(world_values.get("multi_env", False)): raise ValueError( - "free-space-common requires one shared empty cuRobo world " + "The current cuRobo benchmark adapter requires one shared empty world " "with world.multi_env=false." ) world = CuroboWorldCfg( diff --git a/scripts/benchmark/motion_generation/registry.py b/scripts/benchmark/motion_generation/registry.py index d10efb3fe..7611b015e 100644 --- a/scripts/benchmark/motion_generation/registry.py +++ b/scripts/benchmark/motion_generation/registry.py @@ -21,22 +21,27 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from .config import PlannerSpecCfg + from .config import PlannerSpecCfg, RobotSpecCfg from .planners.base import PlannerAdapter, PlannerContext + from .robots.base import RobotProvider from .scenarios.base import ScenarioProvider __all__ = [ "create_planner_adapter", "create_scenario_provider", + "create_robot_provider", "planner_adapter_names", "register_planner_adapter", "register_scenario_provider", + "register_robot_provider", + "robot_provider_names", "scenario_provider_names", "unregister_planner_adapter", ] _PLANNER_ADAPTERS: dict[str, type["PlannerAdapter"]] = {} _SCENARIO_PROVIDERS: dict[str, type["ScenarioProvider"]] = {} +_ROBOT_PROVIDERS: dict[str, type["RobotProvider"]] = {} def register_planner_adapter(name: str, adapter_cls: type["PlannerAdapter"]) -> None: @@ -73,6 +78,33 @@ def create_planner_adapter( return adapter_cls(spec=spec, context=context) +def register_robot_provider(name: str, provider_cls: type["RobotProvider"]) -> None: + """Register one robot provider under a stable suite name.""" + if not name: + raise ValueError("Robot provider name must not be empty.") + previous = _ROBOT_PROVIDERS.get(name) + if previous is not None and previous is not provider_cls: + raise ValueError(f"Robot provider {name!r} is already registered.") + _ROBOT_PROVIDERS[name] = provider_cls + + +def robot_provider_names() -> tuple[str, ...]: + """Return registered robot-provider names in deterministic order.""" + return tuple(sorted(_ROBOT_PROVIDERS)) + + +def create_robot_provider(spec: "RobotSpecCfg") -> "RobotProvider": + """Construct the robot provider selected by a suite.""" + try: + provider_cls = _ROBOT_PROVIDERS[spec.provider] + except KeyError as exc: + raise ValueError( + f"Unknown robot provider {spec.provider!r}; " + f"registered providers: {robot_provider_names()}." + ) from exc + return provider_cls(spec) + + def register_scenario_provider( name: str, provider_cls: type["ScenarioProvider"] ) -> None: diff --git a/scripts/benchmark/motion_generation/reporting.py b/scripts/benchmark/motion_generation/reporting.py index e5d26af63..49fc50f75 100644 --- a/scripts/benchmark/motion_generation/reporting.py +++ b/scripts/benchmark/motion_generation/reporting.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Render the free-space benchmark as exactly three Markdown tables.""" +"""Render planner and Atomic Task tracks as exactly three Markdown tables.""" from __future__ import annotations @@ -30,6 +30,10 @@ "track", "algorithm", "algorithm_role", + "robot", + "skill", + "object", + "task_difficulty", "batch_size", "waypoint_count", "num_trials", @@ -45,6 +49,10 @@ "cpu_delta_mb", "gpu_delta_mb", "peak_gpu_mb", + "execution_time_ms", + "end_to_end_time_ms", + "trajectory_duration_s", + "trajectory_waypoints", ) METRIC_COLUMNS = ( @@ -52,6 +60,11 @@ "scenario", "algorithm", "algorithm_role", + "robot", + "skill", + "object", + "task_difficulty", + "primary_success", "batch_size", "waypoint_count", "path_shape", @@ -61,6 +74,9 @@ "coverage_rate", "success_rate", "planning_success_rate", + "motion_valid_rate", + "execution_success_rate", + "task_success_rate", "ordered_waypoint_success_rate", "waypoint_completion_rate", "final_pos_err_mm", @@ -71,6 +87,10 @@ "joint_path_length_rad", "cartesian_path_length_m", "path_efficiency", + "task_completion_time_s", + "joint_tracking_rmse_rad", + "object_lift_delta_m", + "replan_count", "top_failure", ) @@ -86,6 +106,7 @@ "overall_success_rate", "planning_success_rate", "motion_valid_rate", + "execution_success_rate", "task_success_rate", "latency_p95_ms", "peak_gpu_mb", @@ -132,7 +153,7 @@ def write_markdown_report( output = Path(path) output.parent.mkdir(parents=True, exist_ok=True) lines = [ - "# Motion Generation Benchmark Report", + "# Planner Motion Generation & Atomic Skill Benchmark Report", "", f"Generated at: {datetime.now(timezone.utc).isoformat(timespec='seconds')}", "", @@ -158,7 +179,8 @@ def write_markdown_report( "## Success & Other Metrics", "", "Continuous error/path columns are conditioned on `motion_valid` " - "outcomes; use `n_valid` as the denominator before comparing them.", + "outcomes; use `n_valid` as the denominator. Atomic Task " + "`success_rate` follows the case-owned `primary_success` stage.", "", ] ) diff --git a/scripts/benchmark/motion_generation/robots/__init__.py b/scripts/benchmark/motion_generation/robots/__init__.py new file mode 100644 index 000000000..417f45245 --- /dev/null +++ b/scripts/benchmark/motion_generation/robots/__init__.py @@ -0,0 +1,24 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Built-in robot providers for motion-generation benchmarks.""" + +from __future__ import annotations + +from .base import RobotProvider +from .franka import FrankaPandaProvider, FrankaPgiProvider + +__all__ = ["FrankaPandaProvider", "FrankaPgiProvider", "RobotProvider"] diff --git a/scripts/benchmark/motion_generation/robots/base.py b/scripts/benchmark/motion_generation/robots/base.py new file mode 100644 index 000000000..5d7e1786f --- /dev/null +++ b/scripts/benchmark/motion_generation/robots/base.py @@ -0,0 +1,48 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Robot-provider contract for planner benchmark suites.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.cfg import RobotCfg + from embodichain.lab.sim.objects import Robot + + from ..config import RobotSpecCfg + +__all__ = ["RobotProvider"] + + +class RobotProvider(ABC): + """Build one benchmark embodiment behind a stable suite identifier.""" + + control_part: str = "arm" + + def __init__(self, spec: "RobotSpecCfg") -> None: + self.spec = spec + + @abstractmethod + def build_cfg(self) -> "RobotCfg": + """Build the robot configuration without mutating a simulation.""" + + def add_robot(self, simulation: "SimulationManager") -> "Robot": + """Add the configured robot to a simulation.""" + return simulation.add_robot(cfg=self.build_cfg()) diff --git a/scripts/benchmark/motion_generation/robots/franka.py b/scripts/benchmark/motion_generation/robots/franka.py new file mode 100644 index 000000000..e7cf6280c --- /dev/null +++ b/scripts/benchmark/motion_generation/robots/franka.py @@ -0,0 +1,68 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Built-in Franka robot providers.""" + +from __future__ import annotations + +from embodichain.lab.sim.cfg import RobotCfg +from embodichain.lab.sim.robots import FrankaPandaCfg +from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg +from scripts.tutorials.atomic_action.tutorial_utils import ( + create_franka_panda_robot_cfg, +) + +from ..registry import register_robot_provider +from .base import RobotProvider + +__all__ = ["FrankaPandaProvider", "FrankaPgiProvider"] + + +class FrankaPandaProvider(RobotProvider): + """Build the stock Franka Panda used by the free-space track.""" + + def build_cfg(self) -> RobotCfg: + """Build a stock Panda while applying optional suite overrides.""" + values = { + "uid": "benchmark_franka_panda", + "robot_type": "panda", + **dict(self.spec.config), + } + return FrankaPandaCfg.from_dict(values) + + +class FrankaPgiProvider(RobotProvider): + """Build a Franka arm assembled with the tutorial PGI gripper. + + The configuration mirrors the Franka compatibility in + ``scripts/tutorials/atomic_action``: an arm-only Panda URDF, the shared + ``DH_PGI_140_80`` component, a 180-degree base rotation, and the PGI TCP. + """ + + def build_cfg(self) -> RobotCfg: + """Build the Franka + PGI benchmark configuration.""" + cfg = create_franka_panda_robot_cfg() + return merge_robot_cfg( + cfg, + { + "uid": "benchmark_franka_pgi", + **dict(self.spec.config), + }, + ) + + +register_robot_provider("franka_panda", FrankaPandaProvider) +register_robot_provider("franka_pgi", FrankaPgiProvider) diff --git a/scripts/benchmark/motion_generation/run_benchmark.py b/scripts/benchmark/motion_generation/run_benchmark.py index 4a9ff9ee6..d34a4d00e 100644 --- a/scripts/benchmark/motion_generation/run_benchmark.py +++ b/scripts/benchmark/motion_generation/run_benchmark.py @@ -14,13 +14,15 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run the extensible free-space motion-generation benchmark. +"""Run the extensible planner motion-generation benchmark. cuRobo is the default primary baseline. IK interpolation and TOPPRA are optional diagnostic baselines. NMG remains an explicitly configurable, unsupported adapter stub until its production checkpoint contract is ready. -Run: ``embodichain benchmark motion-generation --suite smoke`` +Run: ``python -m scripts.benchmark.motion_generation.run_benchmark --suite +smoke`` or select the Franka + PGI Atomic Task slice with +``--suite atomic_franka_pgi_curobo``. """ from __future__ import annotations @@ -43,11 +45,14 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: - """Add free-space benchmark options to an existing argument parser.""" + """Add planner benchmark options to an existing argument parser.""" parser.add_argument( "--suite", default="smoke", - help="Suite short name (smoke/coverage) or an explicit YAML path.", + help=( + "Suite short name (smoke/coverage/atomic_franka_pgi_curobo) " + "or an explicit YAML path." + ), ) parser.add_argument( "--algorithms", @@ -214,7 +219,7 @@ def run_all_benchmarks( nmg_rot_eps: float | None = None, output_root: str | Path = "outputs/benchmarks", ) -> BenchmarkRunResult: - """Resolve configuration and run all selected free-space benchmarks.""" + """Resolve configuration and run all selected benchmark tracks.""" from .runner import BenchmarkRunner suite = load_suite(suite_name) @@ -274,7 +279,7 @@ def run_from_args(args: argparse.Namespace) -> BenchmarkRunResult: def _parse_args() -> argparse.Namespace: """Parse standalone module arguments using the unified option schema.""" parser = argparse.ArgumentParser( - description="Benchmark motion generation on fixed free-space cases." + description="Benchmark planners on fixed motion and Atomic Task cases." ) add_parser_arguments(parser) return parser.parse_args() diff --git a/scripts/benchmark/motion_generation/runner.py b/scripts/benchmark/motion_generation/runner.py index 6f6badc87..fc660fe9e 100644 --- a/scripts/benchmark/motion_generation/runner.py +++ b/scripts/benchmark/motion_generation/runner.py @@ -25,10 +25,9 @@ import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.planners.utils import PlanResult -from embodichain.lab.sim.robots import FrankaPandaCfg from . import planners as _builtin_planners # noqa: F401 - registry side effects +from . import robots as _builtin_robots # noqa: F401 - registry side effects from . import scenarios as _builtin_scenarios # noqa: F401 - registry side effects from .aggregation import aggregate_results from .artifacts import ( @@ -40,8 +39,7 @@ write_resolved_suite, ) from .config import PlannerSpecCfg, SuiteCfg -from .metrics import compute_case_outcomes, timed_call -from .metrics.trajectory import make_failure_outcomes +from .metrics import timed_call from .models import ( BenchmarkCase, PlannerMetadata, @@ -49,8 +47,14 @@ TrialRecord, ) from .planners.base import PlannerAdapter, PlannerContext -from .registry import create_planner_adapter, create_scenario_provider +from .registry import ( + create_planner_adapter, + create_robot_provider, + create_scenario_provider, +) from .reporting import write_markdown_report +from .scenarios.base import ScenarioEvaluation, ScenarioProvider +from .scenarios.free_space import FreeSpaceScenario if TYPE_CHECKING: from collections.abc import Callable @@ -60,8 +64,6 @@ __all__ = ["BenchmarkRunResult", "BenchmarkRunner", "resolve_device"] _T = TypeVar("_T") -_CONTROL_PART = "arm" -_ROBOT_UID = "benchmark_franka_panda" @dataclass(frozen=True) @@ -109,13 +111,15 @@ def __init__( self.device = resolve_device(device) self.headless = headless self.output_root = Path(output_root) + self.robot_provider = create_robot_provider(suite.robot) + self.control_part = self.robot_provider.control_part self.records: list[TrialRecord] = [] self.cases: list[BenchmarkCase] = [] self.metadata: dict[str, PlannerMetadata] = {} self.notes: list[str] = [] def _create_simulation(self, batch_size: int) -> tuple[SimulationManager, "Robot"]: - """Create one isolated Franka simulator for a fixed batch size.""" + """Create one isolated suite-selected robot for a fixed batch size.""" sim = SimulationManager( SimulationManagerCfg( headless=self.headless, @@ -124,22 +128,10 @@ def _create_simulation(self, batch_size: int) -> tuple[SimulationManager, "Robot arena_space=2.0, ) ) - robot = sim.add_robot( - cfg=FrankaPandaCfg.from_dict({"uid": _ROBOT_UID, "robot_type": "panda"}) - ) + robot = self.robot_provider.add_robot(sim) sim.update(step=1) return sim, robot - @staticmethod - def _set_case_start( - sim: SimulationManager, robot: "Robot", case: BenchmarkCase - ) -> None: - """Restore current and target robot state outside the timed region.""" - robot.set_qpos(case.start_qpos, name=_CONTROL_PART, target=False) - robot.set_qpos(case.start_qpos, name=_CONTROL_PART, target=True) - robot.clear_dynamics() - sim.update(step=1) - def _append(self, writer: TrialJsonlWriter, record: TrialRecord) -> None: """Retain and immediately persist one raw record.""" self.records.append(record) @@ -169,6 +161,11 @@ def _base_record( "waypoint_count": case.num_waypoints, "path_shape": case.path_shape, "start_state_bin": case.start_state_bin, + "robot_id": case.robot_id, + "skill_id": case.skill_id, + "object_id": case.object_id, + "task_difficulty": case.task_difficulty, + "primary_success": case.primary_success, "phase": phase, } @@ -237,49 +234,50 @@ def _run_plan_call( adapter: PlannerAdapter, metadata: PlannerMetadata, case: BenchmarkCase, + provider: ScenarioProvider, phase: TrialPhase, repeat: int, ) -> None: """Time one plan, validate outside timing, and persist the record.""" - self._set_case_start(sim, robot, case) - measured = timed_call(lambda: _capture(lambda: adapter.plan(case))) + provider.reset_case(sim, robot, case, self.control_part) + measured = timed_call( + lambda: _capture(lambda: provider.plan_case(adapter, case)) + ) result, error = measured.result failure_code = None failure_message = None status = "ok" + evaluation: ScenarioEvaluation | None = None if error is not None: status = "error" failure_code = "planner_exception" failure_message = str(error) - outcomes = make_failure_outcomes(case.batch_size, failure_code) - elif not isinstance(result, PlanResult): + outcomes = provider.failure_outcomes(case, failure_code) + elif (contract_error := provider.plan_contract_error(result)) is not None: status = "error" failure_code = "planner_contract_error" - failure_message = f"Expected PlanResult, got {type(result).__name__}." - outcomes = make_failure_outcomes(case.batch_size, failure_code) + failure_message = contract_error + outcomes = provider.failure_outcomes(case, failure_code) elif phase in (TrialPhase.WARMUP, TrialPhase.COLD): # Cold/warmup timing must not pay for FK validation that is unused # by aggregation. outcomes = () else: try: - outcomes = compute_case_outcomes( + evaluation = provider.evaluate_case( result, case, robot, - _CONTROL_PART, - validation_samples=self.suite.protocol.validation_samples, - position_threshold_m=self.suite.protocol.position_threshold_m, - rotation_threshold_rad=self.suite.protocol.rotation_threshold_rad, - joint_limit_tolerance_rad=( - self.suite.protocol.joint_limit_tolerance_rad - ), + self.control_part, + self.suite, + planning_time_ms=measured.cost_time_ms, ) + outcomes = evaluation.outcomes except Exception as exc: # noqa: BLE001 - metric failure is recorded status = "error" failure_code = "metric_evaluation_error" failure_message = str(exc) - outcomes = make_failure_outcomes(case.batch_size, failure_code) + outcomes = provider.failure_outcomes(case, failure_code) self._append( writer, @@ -292,13 +290,26 @@ def _run_plan_call( cpu_delta_mb=measured.cpu_delta_mb, gpu_delta_mb=measured.gpu_delta_mb, peak_gpu_mb=measured.peak_gpu_mb, + execution_time_ms=( + None if evaluation is None else evaluation.execution_time_ms + ), + end_to_end_time_ms=( + None if evaluation is None else evaluation.end_to_end_time_ms + ), + trajectory_duration_s=( + None if evaluation is None else evaluation.trajectory_duration_s + ), + trajectory_waypoints=( + None if evaluation is None else evaluation.trajectory_waypoints + ), + metadata={} if evaluation is None else evaluation.metadata, outcomes=outcomes, ), ) if phase is not TrialPhase.WARMUP: print( f" {metadata.algorithm_id:<16} B={case.batch_size:>3d} " - f"W={case.num_waypoints} {case.path_shape:<16} " + f"W={case.num_waypoints} {case.skill_id:<20} " f"{phase.value:<8} {measured.cost_time_ms:>10.3f} ms " f"status={status}" ) @@ -311,13 +322,16 @@ def _run_adapter( spec: PlannerSpecCfg, cases: list[BenchmarkCase], required_capabilities: frozenset[str], + provider: ScenarioProvider | None = None, ) -> None: """Execute one adapter over every case for a fixed simulator batch.""" + provider = provider or FreeSpaceScenario() context = PlannerContext( robot=robot, - control_part=_CONTROL_PART, + control_part=self.control_part, device=self.device, sample_interval=self.suite.protocol.sample_interval, + robot_id=self.suite.robot.id, ) adapter = create_planner_adapter(spec, context) metadata = adapter.metadata @@ -354,6 +368,7 @@ def _run_adapter( if build_error is not None: adapter.close() return + scenario_prepared = False try: if adapter.separate_prepare: _, prepare_error = self._record_timed_lifecycle( @@ -366,6 +381,25 @@ def _run_adapter( if prepare_error is not None: return + try: + provider.prepare_planner(adapter, first_case) + scenario_prepared = True + except Exception as exc: # noqa: BLE001 - recorded benchmark failure + self._append( + writer, + TrialRecord( + **self._base_record(metadata, first_case, TrialPhase.PREPARE), + status="error", + failure_code="scenario_prepare_error", + failure_message=str(exc), + ), + ) + self.notes.append( + f"{metadata.algorithm_id} scenario prepare failed for " + f"B={first_case.batch_size}: {exc}" + ) + return + self._run_plan_call( writer, sim, @@ -373,6 +407,7 @@ def _run_adapter( adapter, metadata, first_case, + provider, TrialPhase.COLD, repeat=-1, ) @@ -385,6 +420,7 @@ def _run_adapter( adapter, metadata, case, + provider, TrialPhase.WARMUP, repeat=warmup_index, ) @@ -396,10 +432,13 @@ def _run_adapter( adapter, metadata, case, + provider, TrialPhase.MEASURED, repeat=repeat, ) finally: + if scenario_prepared: + provider.close_planner(adapter) adapter.close() def run(self) -> BenchmarkRunResult: @@ -423,10 +462,19 @@ def run(self) -> BenchmarkRunResult: provider = create_scenario_provider(track.scenario) for batch_size in provider.batch_sizes(self.suite, track): sim: SimulationManager | None = None + runtime_configured = False try: sim, robot = self._create_simulation(batch_size) + provider.configure_runtime( + sim, + robot, + self.suite, + track, + self.control_part, + ) + runtime_configured = True cases = provider.generate_cases( - self.suite, track, robot, _CONTROL_PART, batch_size + self.suite, track, robot, self.control_part, batch_size ) self.cases.extend(cases) for spec in self.planner_specs: @@ -437,8 +485,11 @@ def run(self) -> BenchmarkRunResult: spec, cases, provider.required_capabilities, + provider, ) finally: + if runtime_configured: + provider.close_runtime() if sim is not None: # Benchmarks must aggregate and report after simulator # teardown; the SimulationManager default exits the whole @@ -459,6 +510,24 @@ def run(self) -> BenchmarkRunResult: self.suite.protocol.measured_trials, ) write_json(run_dir / "aggregates.json", aggregates) + enabled_scenarios = {track.scenario for track in enabled_tracks} + track_notes: list[str] = [] + if "free_space" in enabled_scenarios: + track_notes.append( + "Collision, dynamic, execution, and task metrics are N/A in " + "free-space-common v1." + ) + if "atomic_task" in enabled_scenarios: + track_notes.extend( + [ + "Atomic Task cost_time_ms measures AtomicActionEngine.compile only; " + "execution_time_ms is common physics-replay wall time and " + "end_to_end_time_ms is their sum.", + "trajectory_duration_s is planner-native nominal duration; " + "task_completion_time_s is simulated replay time through the " + "stability hold for successful tasks.", + ] + ) report_path = write_markdown_report( run_dir / "report.md", self.suite, @@ -473,7 +542,7 @@ def run(self) -> BenchmarkRunResult: "path lengths, path_efficiency) average only motion_valid outcomes " "(success-conditioned / survivor-biased). Always read them with n_valid; " "a high path_efficiency on n_valid=2 is not comparable to n_valid=200.", - "Collision, dynamic, execution, and task metrics are N/A in free-space-common v1.", + *track_notes, "Leaderboard and Success-table boolean rates " "(overall_success_rate / success_rate / motion_valid_rate / " "planning_success_rate / ordered_waypoint_success_rate) are macro averages " diff --git a/scripts/benchmark/motion_generation/scenarios/__init__.py b/scripts/benchmark/motion_generation/scenarios/__init__.py index e959d74df..34b4328a3 100644 --- a/scripts/benchmark/motion_generation/scenarios/__init__.py +++ b/scripts/benchmark/motion_generation/scenarios/__init__.py @@ -18,7 +18,33 @@ from __future__ import annotations -from .base import ScenarioProvider +from .atomic_objects import ( + AtomicObjectHandle, + atomic_object_kind_names, + create_atomic_object, + register_atomic_object_kind, +) +from .atomic_task import ( + AtomicSkillCaseProvider, + AtomicTaskScenario, + atomic_skill_provider_names, + create_atomic_skill_provider, + register_atomic_skill_provider, +) +from .base import ScenarioEvaluation, ScenarioProvider from .free_space import FreeSpaceScenario -__all__ = ["FreeSpaceScenario", "ScenarioProvider"] +__all__ = [ + "AtomicObjectHandle", + "AtomicSkillCaseProvider", + "AtomicTaskScenario", + "atomic_object_kind_names", + "atomic_skill_provider_names", + "create_atomic_object", + "create_atomic_skill_provider", + "FreeSpaceScenario", + "register_atomic_object_kind", + "register_atomic_skill_provider", + "ScenarioEvaluation", + "ScenarioProvider", +] diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_objects.py b/scripts/benchmark/motion_generation/scenarios/atomic_objects.py new file mode 100644 index 000000000..cb706a361 --- /dev/null +++ b/scripts/benchmark/motion_generation/scenarios/atomic_objects.py @@ -0,0 +1,194 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Config-driven rigid objects for Atomic Task benchmark cases.""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import torch + +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import RigidObject + +__all__ = [ + "AtomicObjectHandle", + "atomic_object_kind_names", + "create_atomic_object", + "register_atomic_object_kind", +] + +AtomicShapeFactory = Callable[[Mapping[str, object]], object] +_ATOMIC_OBJECT_KINDS: dict[str, AtomicShapeFactory] = {} + + +@dataclass +class AtomicObjectHandle: + """Simulation object plus its algorithm-independent frozen state.""" + + object_id: str + kind: str + config: dict[str, object] + entity: "RigidObject" + initial_pose: torch.Tensor + + def reset(self) -> None: + """Restore the frozen initial pose and clear residual dynamics.""" + self.entity.set_local_pose(self.initial_pose) + self.entity.clear_dynamics() + + def park(self, index: int) -> None: + """Move an inactive object outside every benchmark workspace.""" + pose = self.initial_pose.clone() + pose[:, 0, 3] = 8.0 + float(index) + pose[:, 1, 3] = 8.0 + pose[:, 2, 3] = 1.0 + self.entity.set_local_pose(pose) + self.entity.clear_dynamics() + + +def register_atomic_object_kind(name: str, factory: AtomicShapeFactory) -> None: + """Register a config-only object shape factory.""" + if not name: + raise ValueError("Atomic object kind must not be empty.") + previous = _ATOMIC_OBJECT_KINDS.get(name) + if previous is not None and previous is not factory: + raise ValueError(f"Atomic object kind {name!r} is already registered.") + _ATOMIC_OBJECT_KINDS[name] = factory + + +def atomic_object_kind_names() -> tuple[str, ...]: + """Return registered object kinds in deterministic order.""" + return tuple(sorted(_ATOMIC_OBJECT_KINDS)) + + +def _vector( + value: object, *, name: str, length: int, default: Sequence[float] +) -> list[float]: + """Validate and normalize a numeric vector from YAML configuration.""" + resolved = default if value is None else value + if not isinstance(resolved, Sequence) or isinstance(resolved, (str, bytes)): + raise TypeError(f"{name} must be a sequence of {length} numbers.") + result = [float(item) for item in resolved] + if len(result) != length or not all(math.isfinite(item) for item in result): + raise ValueError(f"{name} must contain {length} finite values.") + return result + + +def _cube_shape(config: Mapping[str, object]) -> CubeCfg: + """Build a cube shape from its declarative size.""" + size = _vector( + config.get("size"), name="cube.size", length=3, default=(0.05, 0.05, 0.05) + ) + if any(value <= 0.0 for value in size): + raise ValueError("cube.size values must be greater than zero.") + return CubeCfg(size=size) + + +def _mesh_shape(config: Mapping[str, object]) -> MeshCfg: + """Build a mesh shape from an absolute or EmbodiChain data path.""" + asset_path = config.get("asset_path") + if not isinstance(asset_path, str) or not asset_path: + raise ValueError("mesh.asset_path must be a non-empty string.") + resolved = Path(asset_path) + return MeshCfg( + fpath=str(resolved if resolved.is_absolute() else get_data_path(asset_path)) + ) + + +def create_atomic_object( + simulation: "SimulationManager", object_config: Mapping[str, object] +) -> AtomicObjectHandle: + """Create one config-driven object and freeze its settled initial pose.""" + object_id = object_config.get("id") + kind = object_config.get("kind") + if not isinstance(object_id, str) or not object_id: + raise ValueError("Every atomic object must define a non-empty id.") + if not isinstance(kind, str) or not kind: + raise ValueError(f"Atomic object {object_id!r} must define a kind.") + try: + factory = _ATOMIC_OBJECT_KINDS[kind] + except KeyError as exc: + raise ValueError( + f"Unknown atomic object kind {kind!r}; registered kinds: " + f"{atomic_object_kind_names()}." + ) from exc + + config = dict(object_config) + position = _vector( + config.get("position"), + name=f"objects[{object_id}].position", + length=3, + default=(-0.42, -0.08, 0.05), + ) + rotation = _vector( + config.get("rotation_deg"), + name=f"objects[{object_id}].rotation_deg", + length=3, + default=(0.0, 0.0, 0.0), + ) + scale = _vector( + config.get("scale"), + name=f"objects[{object_id}].scale", + length=3, + default=(1.0, 1.0, 1.0), + ) + entity = simulation.add_rigid_object( + cfg=RigidObjectCfg( + uid=f"atomic_benchmark_{object_id}", + shape=factory(config), + attrs=RigidBodyAttributesCfg( + mass=float(config.get("mass", 0.05)), + dynamic_friction=float(config.get("dynamic_friction", 0.97)), + static_friction=float(config.get("static_friction", 0.99)), + restitution=float(config.get("restitution", 0.0)), + contact_offset=float(config.get("contact_offset", 0.003)), + rest_offset=float(config.get("rest_offset", 0.001)), + linear_damping=float(config.get("linear_damping", 0.7)), + angular_damping=float(config.get("angular_damping", 0.7)), + min_position_iters=int(config.get("min_position_iters", 32)), + min_velocity_iters=int(config.get("min_velocity_iters", 8)), + ), + max_convex_hull_num=int(config.get("max_convex_hull_num", 16)), + init_pos=position, + init_rot=rotation, + body_scale=scale, + use_usd_properties=bool(config.get("use_usd_properties", False)), + ) + ) + simulation.update(step=int(config.get("settle_steps", 10))) + entity.clear_dynamics() + return AtomicObjectHandle( + object_id=object_id, + kind=kind, + config=config, + entity=entity, + initial_pose=entity.get_local_pose(to_matrix=True).clone(), + ) + + +register_atomic_object_kind("cube", _cube_shape) +register_atomic_object_kind("mesh", _mesh_shape) diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py new file mode 100644 index 000000000..2a77500a0 --- /dev/null +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -0,0 +1,1176 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Physics-backed Atomic Task track shared by planner adapters.""" + +from __future__ import annotations + +import math +import time +from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +import torch + +from embodichain.lab.sim.atomic_actions import ( + ActionBinding, + ActionInvocation, + Affordance, + AtomicActionEngine, + ControlPartCommandProfile, + EndEffectorPoseGoal, + GraspGoal, + MotionPolicy, + ObjectSemantics, + PickUpOptions, +) +from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory +from embodichain.lab.sim.planners.utils import PlanResult + +from ..config import SuiteCfg, TrackCfg +from ..metrics.trajectory import compute_case_outcomes +from ..models import BenchmarkCase, CaseOutcome +from ..registry import register_scenario_provider +from .atomic_objects import AtomicObjectHandle, create_atomic_object +from .base import ScenarioEvaluation, ScenarioProvider + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + from embodichain.lab.sim.objects import Robot + + from ..planners.base import PlannerAdapter + +__all__ = [ + "AtomicSkillCaseProvider", + "AtomicTaskScenario", + "atomic_skill_provider_names", + "create_atomic_skill_provider", + "register_atomic_skill_provider", +] + +_TOP_DOWN_ROTATION = ( + (-0.0539, -0.9985, -0.0022), + (-0.9977, 0.0540, -0.0401), + (0.0401, 0.0000, -0.9992), +) +_TASK_DIFFICULTIES = {"simple", "medium", "hard"} + + +@dataclass(frozen=True) +class _ExecutionObservation: + """Common-execution measurements used by skill-specific task rules.""" + + execution_success: torch.Tensor + final_tcp_pose: torch.Tensor + joint_tracking_rmse_rad: torch.Tensor + execution_time_ms: float + task_completion_time_s: float + object_lift_delta_m: torch.Tensor | None = None + + +class AtomicSkillCaseProvider(ABC): + """Generate and ground one Atomic Action without planner-specific logic.""" + + skill_id: str + + @abstractmethod + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + """Generate one frozen case and independent IK validity evidence.""" + + @abstractmethod + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + """Ground the case into one planner-independent action invocation.""" + + def object_id(self, case: BenchmarkCase) -> str | None: + """Return the manipulated object identifier when one exists.""" + return case.object_id + + def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: + """Return the first lift waypoint that should release object dynamics.""" + return None + + @abstractmethod + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + """Return per-environment task success and its stable failure code.""" + + +AtomicSkillProviderType = type[AtomicSkillCaseProvider] +_ATOMIC_SKILL_PROVIDERS: dict[str, AtomicSkillProviderType] = {} + + +def register_atomic_skill_provider( + skill_id: str, provider_type: AtomicSkillProviderType +) -> None: + """Register one Atomic Action case provider.""" + if not skill_id: + raise ValueError("Atomic skill id must not be empty.") + previous = _ATOMIC_SKILL_PROVIDERS.get(skill_id) + if previous is not None and previous is not provider_type: + raise ValueError(f"Atomic skill provider {skill_id!r} is already registered.") + _ATOMIC_SKILL_PROVIDERS[skill_id] = provider_type + + +def atomic_skill_provider_names() -> tuple[str, ...]: + """Return registered Atomic Action case providers.""" + return tuple(sorted(_ATOMIC_SKILL_PROVIDERS)) + + +def create_atomic_skill_provider(skill_id: str) -> AtomicSkillCaseProvider: + """Construct one registered Atomic Action case provider.""" + try: + provider_type = _ATOMIC_SKILL_PROVIDERS[skill_id] + except KeyError as exc: + raise ValueError( + f"Unknown atomic skill {skill_id!r}; registered skills: " + f"{atomic_skill_provider_names()}." + ) from exc + return provider_type() + + +def _float_vector( + value: object, + *, + name: str, + length: int, + default: Sequence[float] | None = None, +) -> list[float]: + """Resolve a finite numeric vector from a YAML-compatible value.""" + resolved = default if value is None else value + if not isinstance(resolved, Sequence) or isinstance(resolved, (str, bytes)): + raise TypeError(f"{name} must be a sequence of {length} numbers.") + result = [float(item) for item in resolved] + if len(result) != length or not all(math.isfinite(item) for item in result): + raise ValueError(f"{name} must contain {length} finite values.") + return result + + +def _case_name(config: Mapping[str, object]) -> str: + """Return and validate a stable case name.""" + name = config.get("name") + if not isinstance(name, str) or not name: + raise ValueError("Every atomic skill case must define a non-empty name.") + return name + + +def _difficulty(config: Mapping[str, object]) -> str: + """Resolve the explicit, frozen Atomic Task difficulty label.""" + difficulty = str(config.get("task_difficulty", "simple")) + if difficulty not in _TASK_DIFFICULTIES: + raise ValueError( + f"task_difficulty must be one of {sorted(_TASK_DIFFICULTIES)}." + ) + return difficulty + + +class _MoveEndEffectorCases(AtomicSkillCaseProvider): + """Deterministic robot-relative MoveEndEffector cases.""" + + skill_id = "move_end_effector" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + scenario.restore_base_robot() + raw_offsets = config.get("target_offsets_m") + if not isinstance(raw_offsets, Sequence) or isinstance( + raw_offsets, (str, bytes) + ): + raise TypeError("target_offsets_m must be a non-empty list of xyz vectors.") + offsets = [ + _float_vector(value, name="target_offsets_m", length=3) + for value in raw_offsets + ] + if not offsets: + raise ValueError("target_offsets_m must not be empty.") + + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + start_pose = scenario.robot.compute_fk( + start_qpos, name=scenario.control_part, to_matrix=True + ) + targets = start_pose[:, None].repeat(1, len(offsets), 1, 1) + targets[:, :, :3, 3] += torch.tensor( + offsets, dtype=targets.dtype, device=targets.device + )[None] + references = scenario.solve_reference_qpos(start_qpos, targets) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=len(offsets), + path_shape="robot_relative_waypoints", + start_state_bin="pre_action", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=scenario.robot.get_qpos().clone(), + case_parameters={ + "sample_count": int(config.get("sample_count", 80)), + "target_offsets_m": offsets, + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + return ActionInvocation( + skill_id=self.skill_id, + goal=EndEffectorPoseGoal(case.target_waypoints), + binding=ActionBinding(manipulators={"primary": scenario.control_part}), + motion_policy=MotionPolicy( + planner=adapter.motion_policy_planner, + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + del compiled + target = case.target_waypoints[:, -1] + translation = torch.linalg.vector_norm( + observation.final_tcp_pose[:, :3, 3] - target[:, :3, 3], dim=-1 + ) + relative = ( + target[:, :3, :3].transpose(-1, -2) @ observation.final_tcp_pose[:, :3, :3] + ) + trace = torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) + rotation = torch.arccos(torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0)) + motion_valid = torch.tensor( + [item.motion_valid for item in motion_outcomes], + dtype=torch.bool, + device=translation.device, + ) + success = ( + observation.execution_success + & motion_valid + & (translation <= scenario.suite.protocol.position_threshold_m) + & (rotation <= scenario.suite.protocol.rotation_threshold_rad) + ) + return success, "task_goal_miss" + + +class _PickUpCases(AtomicSkillCaseProvider): + """Explicit-grasp PickUp cases that isolate motion-planner performance.""" + + skill_id = "pick_up" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + object_id = config.get("object") + if not isinstance(object_id, str) or not object_id: + raise ValueError("PickUp cases must reference a non-empty object id.") + handle = scenario.activate_object(object_id) + scenario.restore_base_robot() + + object_pose = handle.entity.get_local_pose(to_matrix=True).clone() + arm_start = scenario.robot.get_qpos(name=scenario.control_part) + pre_pick_pose = scenario.robot.compute_fk( + arm_start, name=scenario.control_part, to_matrix=True + ).clone() + pre_pick_pose[:, :2, 3] = object_pose[:, :2, 3] + pre_pick_pose[:, 2, 3] = float(config.get("pre_pick_height_m", 0.36)) + success, pre_pick_qpos = scenario.robot.compute_ik( + pose=pre_pick_pose, + joint_seed=arm_start, + name=scenario.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + raise RuntimeError( + f"Independent IK rejected PickUp case {_case_name(config)!r}." + ) + scenario.set_robot_start(pre_pick_qpos, open_gripper=True) + + approach = torch.tensor( + _float_vector( + config.get("approach_direction"), + name="approach_direction", + length=3, + default=(0.0, 0.0, -1.0), + ), + dtype=object_pose.dtype, + device=object_pose.device, + ) + approach_norm = torch.linalg.vector_norm(approach) + if float(approach_norm.item()) <= 1.0e-6: + raise ValueError("approach_direction must be non-zero.") + approach = approach / approach_norm + pre_grasp_distance = float(config.get("pre_grasp_distance_m", 0.15)) + lift_height = float(config.get("lift_height_m", 0.16)) + + grasp_source = str(config.get("grasp_source", "fixed")) + if grasp_source == "antipodal": + grasp_pose = scenario.resolve_antipodal_grasp( + handle, + object_pose, + approach, + seed=seed, + start_qpos=scenario.robot.get_qpos(name=scenario.control_part), + pre_grasp_distance=pre_grasp_distance, + lift_height=lift_height, + n_sample=int(config.get("grasp_sample_count", 10_000)), + max_candidates=int(config.get("grasp_max_candidates", 128)), + alignment_max_angle_deg=float( + config.get("grasp_alignment_max_angle_deg", 10.0) + ), + ) + elif grasp_source == "fixed": + grasp_pose = object_pose.clone() + rotation_value = config.get("grasp_rotation", _TOP_DOWN_ROTATION) + if not isinstance(rotation_value, Sequence) or len(rotation_value) != 3: + raise ValueError("grasp_rotation must be a 3x3 matrix.") + rotation = torch.tensor( + rotation_value, dtype=grasp_pose.dtype, device=grasp_pose.device + ) + if rotation.shape != (3, 3): + raise ValueError("grasp_rotation must be a 3x3 matrix.") + grasp_pose[:, :3, :3] = rotation + else: + raise ValueError("grasp_source must be 'fixed' or 'antipodal'.") + grasp_offset = torch.tensor( + _float_vector( + config.get("grasp_offset_m"), + name="grasp_offset_m", + length=3, + default=(0.0, 0.0, 0.0), + ), + dtype=grasp_pose.dtype, + device=grasp_pose.device, + ) + grasp_pose[:, :3, 3] += grasp_offset + pre_grasp = grasp_pose.clone() + pre_grasp[:, :3, 3] -= approach * pre_grasp_distance + lift = grasp_pose.clone() + lift[:, 2, 3] += lift_height + targets = torch.stack([pre_grasp, grasp_pose, lift], dim=1) + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + references = scenario.solve_reference_qpos(start_qpos, targets) + name = _case_name(config) + return BenchmarkCase( + suite_version=suite.suite_version, + track=track.id, + scenario_id=self.skill_id, + case_id=f"{track.id}:{self.skill_id}:{name}:s{seed}", + seed=seed, + batch_size=batch_size, + num_waypoints=3, + path_shape="approach_grasp_lift", + start_state_bin="pre_pick", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + object_id=object_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=scenario.robot.get_qpos().clone(), + case_parameters={ + "sample_count": int(config.get("sample_count", 120)), + "grasp_source": grasp_source, + "hand_interp_steps": int(config.get("hand_interp_steps", 12)), + "approach_direction": approach.detach().cpu().tolist(), + "pre_grasp_distance_m": pre_grasp_distance, + "lift_height_m": lift_height, + "minimum_object_lift_m": float( + config.get("minimum_object_lift_m", 0.04) + ), + "grasp_pose": grasp_pose.detach().cpu().tolist(), + "object_initial_pose": object_pose.detach().cpu().tolist(), + "object_config": dict(handle.config), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + handle = scenario.object_handle(case.object_id) + if scenario.end_effector_part is None: + raise RuntimeError( + "PickUp requires a configured end-effector control part." + ) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + properties={"benchmark_object_id": handle.object_id}, + label=handle.object_id, + entity=handle.entity, + ) + return ActionInvocation( + skill_id=self.skill_id, + goal=GraspGoal( + semantics=semantics, + grasp_xpos=case.target_waypoints[:, 1], + ), + binding=ActionBinding( + manipulators={"primary": scenario.control_part}, + end_effectors={"primary": scenario.end_effector_part}, + ), + motion_policy=MotionPolicy( + planner=adapter.motion_policy_planner, + strategy="motion_gen", + sample_count=int(case.case_parameters["sample_count"]), + ), + skill_options=PickUpOptions( + approach_direction=torch.tensor( + case.case_parameters["approach_direction"], + dtype=torch.float32, + device=scenario.robot.device, + ), + pre_grasp_distance=float(case.case_parameters["pre_grasp_distance_m"]), + lift_height=float(case.case_parameters["lift_height_m"]), + hand_interp_steps=int(case.case_parameters["hand_interp_steps"]), + ), + ) + + def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: + """Return the lift boundary emitted by PickUp.""" + return compiled.segment(0, "lift").start + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + held_created = ( + compiled.projected_context.get_held_object(scenario.control_part) + is not None + ) + lift = observation.object_lift_delta_m + if lift is None: + lift = torch.full( + (case.batch_size,), + -torch.inf, + device=observation.execution_success.device, + ) + motion_valid = torch.tensor( + [item.motion_valid for item in motion_outcomes], + dtype=torch.bool, + device=lift.device, + ) + success = ( + observation.execution_success + & motion_valid + & held_created + & (lift >= float(case.case_parameters["minimum_object_lift_m"])) + ) + return success, "object_not_grasped" + + +class AtomicTaskScenario(ScenarioProvider): + """Run fixed Atomic Actions through an adapter-owned MotionGenerator.""" + + required_capabilities = frozenset({"eef_waypoint", "atomic_action"}) + + def __init__(self) -> None: + self.simulation: SimulationManager | None = None + self.robot: Robot | None = None + self.suite: SuiteCfg | None = None + self.track: TrackCfg | None = None + self.control_part = "arm" + self.end_effector_part: str | None = None + self._objects: dict[str, AtomicObjectHandle] = {} + self._case_providers: dict[str, AtomicSkillCaseProvider] = {} + self._base_full_qpos: torch.Tensor | None = None + self._gripper_open: torch.Tensor | None = None + self._gripper_grasp: torch.Tensor | None = None + self._engine: AtomicActionEngine | None = None + + def batch_sizes(self, suite: SuiteCfg, track: TrackCfg) -> list[int]: + """Return the explicitly configured physical-execution batch sizes.""" + del suite + values = [int(value) for value in track.config.get("batch_sizes", [1])] + if values != [1]: + raise ValueError( + "The initial Atomic Task implementation supports batch_sizes: [1] only." + ) + return values + + def configure_runtime( + self, + simulation: SimulationManager, + robot: Robot, + suite: SuiteCfg, + track: TrackCfg, + control_part: str, + ) -> None: + """Create the declarative object pool and cache robot command states.""" + self.simulation = simulation + self.robot = robot + self.suite = suite + self.track = track + self.control_part = control_part + self._base_full_qpos = robot.get_qpos().clone() + + object_values = track.config.get("objects", []) + if not isinstance(object_values, list): + raise TypeError("atomic-task objects must be a list of mappings.") + for value in object_values: + if not isinstance(value, Mapping): + raise TypeError("Every atomic-task object must be a mapping.") + handle = create_atomic_object(simulation, value) + if handle.object_id in self._objects: + raise ValueError(f"Duplicate atomic object id {handle.object_id!r}.") + self._objects[handle.object_id] = handle + for index, handle in enumerate(self._objects.values()): + handle.park(index) + simulation.update(step=1) + + if any( + isinstance(value, Mapping) and value.get("id") == "pick_up" + for value in track.config.get("skills", []) + ): + gripper_value = track.config.get("gripper") + if not isinstance(gripper_value, Mapping): + raise ValueError("Atomic PickUp tracks must define a gripper mapping.") + end_effector_part = gripper_value.get("control_part") + if not isinstance(end_effector_part, str) or not end_effector_part: + raise ValueError("gripper.control_part must be a non-empty string.") + self.end_effector_part = end_effector_part + limits = robot.get_qpos_limits(name=end_effector_part)[0].to( + device=robot.device, dtype=torch.float32 + ) + dofs = limits.shape[0] + open_qpos = _float_vector( + gripper_value.get("open_qpos"), + name="gripper.open_qpos", + length=dofs, + default=limits[:, 0].detach().cpu().tolist(), + ) + grasp_qpos = _float_vector( + gripper_value.get("grasp_qpos"), + name="gripper.grasp_qpos", + length=dofs, + ) + self._gripper_open = torch.tensor( + open_qpos, dtype=limits.dtype, device=limits.device + ) + self._gripper_grasp = torch.tensor( + grasp_qpos, dtype=limits.dtype, device=limits.device + ) + if bool( + ( + (self._gripper_open < limits[:, 0]) + | (self._gripper_open > limits[:, 1]) + | (self._gripper_grasp < limits[:, 0]) + | (self._gripper_grasp > limits[:, 1]) + ) + .any() + .item() + ): + raise ValueError( + "gripper open/grasp qpos must lie within joint limits." + ) + + def generate_cases( + self, + suite: SuiteCfg, + track: TrackCfg, + robot: Robot, + control_part: str, + batch_size: int, + ) -> list[BenchmarkCase]: + """Generate the algorithm-independent Atomic Task case manifest.""" + if robot is not self.robot or control_part != self.control_part: + raise RuntimeError( + "AtomicTaskScenario must be configured before generation." + ) + skill_values = track.config.get("skills", []) + if not isinstance(skill_values, list) or not skill_values: + raise ValueError("atomic-task skills must be a non-empty list.") + seeds = [int(value) for value in track.config.get("seeds", [11])] + if not seeds: + raise ValueError("atomic-task seeds must not be empty.") + + cases: list[BenchmarkCase] = [] + for skill_value in skill_values: + if not isinstance(skill_value, Mapping): + raise TypeError("Every atomic-task skill entry must be a mapping.") + skill_id = skill_value.get("id") + if not isinstance(skill_id, str) or not skill_id: + raise ValueError("Every atomic-task skill entry must define an id.") + raw_cases = skill_value.get("cases", []) + if not isinstance(raw_cases, list) or not raw_cases: + raise ValueError(f"Atomic skill {skill_id!r} needs at least one case.") + defaults = { + key: value + for key, value in skill_value.items() + if key not in {"id", "cases"} + } + for raw_case in raw_cases: + if not isinstance(raw_case, Mapping): + raise TypeError("Every atomic skill case must be a mapping.") + config = {**defaults, **dict(raw_case)} + for seed in seeds: + provider = create_atomic_skill_provider(skill_id) + case = provider.generate_case( + self, + suite, + track, + config, + seed=seed, + batch_size=batch_size, + ) + if case.case_id in self._case_providers: + raise ValueError( + f"Duplicate Atomic Task case {case.case_id!r}." + ) + self._case_providers[case.case_id] = provider + cases.append(case) + self.restore_base_robot() + for index, handle in enumerate(self._objects.values()): + handle.park(index) + return cases + + def prepare_planner( + self, adapter: PlannerAdapter, first_case: BenchmarkCase + ) -> None: + """Bind one AtomicActionEngine to the adapter-owned MotionGenerator.""" + del first_case + control_profiles = None + if any( + provider.skill_id == "pick_up" for provider in self._case_providers.values() + ): + if ( + self.end_effector_part is None + or self._gripper_open is None + or self._gripper_grasp is None + ): + raise RuntimeError("Configured gripper command states are unavailable.") + control_profiles = { + self.end_effector_part: ControlPartCommandProfile.joint_positions( + open=self._gripper_open, + grasp=self._gripper_grasp, + ) + } + self._engine = AtomicActionEngine( + motion_generator=adapter.require_motion_generator(), + control_profiles=control_profiles, + ) + + def close_planner(self, adapter: PlannerAdapter) -> None: + """Drop the engine before its adapter closes the shared generator.""" + del adapter + self._engine = None + + def reset_case( + self, + simulation: SimulationManager, + robot: Robot, + case: BenchmarkCase, + control_part: str, + ) -> None: + """Restore full robot and object state before every planner call.""" + del control_part + if case.full_start_qpos is None: + raise ValueError("Atomic Task cases require full_start_qpos.") + for target in (False, True): + robot.set_qpos(case.full_start_qpos, target=target) + robot.clear_dynamics() + active_id = case.object_id + for index, handle in enumerate(self._objects.values()): + if handle.object_id == active_id: + handle.reset() + else: + handle.park(index) + simulation.update(step=2) + + def plan_case(self, adapter: PlannerAdapter, case: BenchmarkCase) -> object: + """Compile one Atomic Action with an explicitly pinned motion backend.""" + if self._engine is None: + raise RuntimeError("Atomic Task planner resources were not prepared.") + provider = self._case_providers[case.case_id] + invocation = provider.build_invocation(self, case, adapter) + policy = invocation.motion_policy + if ( + policy.strategy != "motion_gen" + or policy.planner != adapter.motion_policy_planner + ): + raise RuntimeError( + "Atomic Task invocations must pin the adapter motion backend." + ) + return self._engine.compile((invocation,)) + + def plan_contract_error(self, result: object) -> str | None: + """Accept compiled Atomic Action trajectories instead of raw plans.""" + if isinstance(result, CompiledTrajectory): + return None + return f"Expected CompiledTrajectory, got {type(result).__name__}." + + def failure_outcomes( + self, case: BenchmarkCase, failure_code: str + ) -> tuple[CaseOutcome, ...]: + """Mark execution/task stages false after runner-level failures.""" + return tuple( + replace( + outcome, + execution_success=False, + task_success=False, + replan_count=0, + ) + for outcome in super().failure_outcomes(case, failure_code) + ) + + def evaluate_case( + self, + result: object, + case: BenchmarkCase, + robot: Robot, + control_part: str, + suite: SuiteCfg, + *, + planning_time_ms: float, + ) -> ScenarioEvaluation: + """Validate motion, replay physics, and evaluate physical task success.""" + if not isinstance(result, CompiledTrajectory): + raise TypeError(self.plan_contract_error(result)) + arm_joint_ids = list(robot.get_joint_ids(name=control_part)) + arm_positions = result.trajectory.positions[:, :, arm_joint_ids] + motion_outcomes = compute_case_outcomes( + PlanResult(success=result.plan_success, positions=arm_positions), + case, + robot, + control_part, + validation_samples=suite.protocol.validation_samples, + position_threshold_m=suite.protocol.position_threshold_m, + rotation_threshold_rad=suite.protocol.rotation_threshold_rad, + joint_limit_tolerance_rad=suite.protocol.joint_limit_tolerance_rad, + ) + provider = self._case_providers[case.case_id] + observation = self._execute(result, case, provider) + if observation is None: + execution_success = torch.zeros( + case.batch_size, dtype=torch.bool, device=robot.device + ) + task_success = execution_success.clone() + execution_time_ms = None + tracking = torch.full((case.batch_size,), torch.nan, device=robot.device) + object_lift = None + task_failure_code = "task_goal_miss" + else: + execution_success = observation.execution_success + task_success, task_failure_code = provider.task_result( + self, case, result, observation, motion_outcomes + ) + execution_time_ms = observation.execution_time_ms + tracking = observation.joint_tracking_rmse_rad + object_lift = observation.object_lift_delta_m + + durations = result.trajectory.duration.detach().to("cpu") + outcomes: list[CaseOutcome] = [] + for index, outcome in enumerate(motion_outcomes): + executed = bool(execution_success[index].item()) + task_done = bool(task_success[index].item()) + if outcome.failure_code is not None: + failure_code = outcome.failure_code + elif not executed: + failure_code = "controller_tracking_failure" + elif not task_done: + failure_code = task_failure_code + else: + failure_code = None + tracking_value = float(tracking[index].item()) + outcomes.append( + replace( + outcome, + execution_success=executed, + task_success=task_done, + task_completion_time_s=( + observation.task_completion_time_s + if observation is not None and task_done + else None + ), + joint_tracking_rmse_rad=( + tracking_value if math.isfinite(tracking_value) else None + ), + object_lift_delta_m=( + None + if object_lift is None + else float(object_lift[index].item()) + ), + replan_count=0, + failure_code=failure_code, + ) + ) + trajectory_duration = ( + float(durations.mean().item()) if durations.numel() else None + ) + return ScenarioEvaluation( + outcomes=tuple(outcomes), + execution_time_ms=execution_time_ms, + end_to_end_time_ms=( + planning_time_ms + execution_time_ms + if execution_time_ms is not None + else None + ), + trajectory_duration_s=trajectory_duration, + trajectory_waypoints=result.trajectory.waypoint_count, + metadata={ + "timing_scope": "atomic_action_compile", + "motion_policy_strategy": "motion_gen", + "physics_validation": "common_joint_target_replay", + "constraint_information": ( + "empty external world; manipulated target excluded from " + "cuRobo collision obstacles" + ), + }, + ) + + def _execute( + self, + compiled: CompiledTrajectory, + case: BenchmarkCase, + provider: AtomicSkillCaseProvider, + ) -> _ExecutionObservation | None: + """Replay a successful full-robot trajectory under common physics.""" + if self.simulation is None or self.robot is None or self.track is None: + raise RuntimeError("Atomic Task runtime is not configured.") + trajectory = compiled.trajectory.positions + if ( + trajectory.shape[1] == 0 + or not bool(compiled.plan_success.all().item()) + or not bool(torch.isfinite(trajectory).all().item()) + ): + return None + object_handle = ( + None if case.object_id is None else self.object_handle(case.object_id) + ) + initial_object_position = ( + None + if object_handle is None + else object_handle.entity.get_local_pose(to_matrix=True)[:, :3, 3].clone() + ) + physics = dict(self.track.config.get("physics", {})) + steps_per_waypoint = int(physics.get("steps_per_waypoint", 4)) + hold_steps = int(physics.get("hold_steps", 80)) + hold_sim_steps = int(physics.get("hold_sim_steps", 2)) + tracking_tolerance = float(physics.get("joint_tracking_tolerance_rad", 0.05)) + if steps_per_waypoint < 1 or hold_steps < 0 or hold_sim_steps < 1: + raise ValueError("Atomic Task physics step counts are invalid.") + + lift_start = provider.lift_segment_start(compiled) + dynamics_cleared = False + squared_tracking_error = torch.zeros( + case.batch_size, dtype=trajectory.dtype, device=trajectory.device + ) + tracking_value_count = 0 + self._synchronize() + started = time.perf_counter() + for waypoint_index in range(trajectory.shape[1]): + positions = trajectory[:, waypoint_index] + self.robot.set_qpos(positions, target=True) + self.simulation.update(step=steps_per_waypoint) + observed = self.robot.get_qpos() + squared_tracking_error += ((observed - positions) ** 2).sum(dim=1) + tracking_value_count += observed.shape[1] + if ( + object_handle is not None + and lift_start is not None + and not dynamics_cleared + and waypoint_index + 1 >= lift_start + ): + object_handle.entity.clear_dynamics() + dynamics_cleared = True + final_command = trajectory[:, -1] + for _ in range(hold_steps): + self.robot.set_qpos(final_command, target=True) + self.simulation.update(step=hold_sim_steps) + observed = self.robot.get_qpos() + squared_tracking_error += ((observed - final_command) ** 2).sum(dim=1) + tracking_value_count += observed.shape[1] + self._synchronize() + elapsed_ms = (time.perf_counter() - started) * 1000.0 + + observed = self.robot.get_qpos() + tracking = torch.sqrt(squared_tracking_error / max(tracking_value_count, 1)) + simulated_execution_time = ( + trajectory.shape[1] * steps_per_waypoint + hold_steps * hold_sim_steps + ) * float(self.simulation.sim_config.physics_dt) + final_tcp = self.robot.compute_fk( + self.robot.get_qpos(name=self.control_part), + name=self.control_part, + to_matrix=True, + ) + object_lift = None + if object_handle is not None and initial_object_position is not None: + final_object_position = object_handle.entity.get_local_pose(to_matrix=True)[ + :, :3, 3 + ] + object_lift = final_object_position[:, 2] - initial_object_position[:, 2] + return _ExecutionObservation( + execution_success=torch.isfinite(observed).all(dim=1) + & (tracking <= tracking_tolerance), + final_tcp_pose=final_tcp, + joint_tracking_rmse_rad=tracking, + execution_time_ms=elapsed_ms, + task_completion_time_s=simulated_execution_time, + object_lift_delta_m=object_lift, + ) + + def solve_reference_qpos( + self, start_qpos: torch.Tensor, target_waypoints: torch.Tensor + ) -> torch.Tensor: + """Build independent sequential-IK validity evidence for a case.""" + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + seed = start_qpos + references: list[torch.Tensor] = [] + for index in range(target_waypoints.shape[1]): + success, seed = self.robot.compute_ik( + pose=target_waypoints[:, index], + joint_seed=seed, + name=self.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + raise RuntimeError( + f"Independent IK rejected atomic target waypoint {index}." + ) + references.append(seed.clone()) + return torch.stack(references, dim=1) + + def resolve_antipodal_grasp( + self, + handle: AtomicObjectHandle, + object_pose: torch.Tensor, + approach_direction: torch.Tensor, + *, + seed: int, + start_qpos: torch.Tensor, + pre_grasp_distance: float, + lift_height: float, + n_sample: int, + max_candidates: int, + alignment_max_angle_deg: float, + ) -> torch.Tensor: + """Freeze one geometry-aware, independently reachable PGI grasp pose. + + Grasp sampling and sequential IK screening happen once while the case + manifest is built, before any planner adapter is evaluated. Every + planner therefore receives the same explicit grasp pose and planning + latency excludes grasp generation. + """ + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + if n_sample < 1 or max_candidates < 1: + raise ValueError( + "Antipodal grasp sample/candidate counts must be positive." + ) + if not 0.0 < alignment_max_angle_deg <= 90.0: + raise ValueError("grasp_alignment_max_angle_deg must be in (0, 90].") + from scripts.tutorials.atomic_action.tutorial_utils import ( + create_antipodal_semantics, + ) + + fork_devices = ( + [] + if self.robot.device.type != "cuda" + else [self.robot.device.index or torch.cuda.current_device()] + ) + with torch.random.fork_rng(devices=fork_devices): + torch.manual_seed(seed) + semantics = create_antipodal_semantics( + handle.entity, + label=handle.object_id, + n_sample=n_sample, + force_reannotate=False, + ) + candidates, costs = semantics.affordance.get_valid_grasp_poses( + obj_poses=object_pose, + approach_direction=approach_direction, + )[0] + if candidates.shape[0] == 0: + raise RuntimeError( + f"No antipodal grasp candidates were found for {handle.object_id!r}." + ) + finite_indices = torch.nonzero(torch.isfinite(costs), as_tuple=False).flatten() + if finite_indices.numel() == 0: + raise RuntimeError( + f"No valid antipodal grasp candidates were found for {handle.object_id!r}." + ) + ranked = finite_indices[torch.argsort(costs[finite_indices])] + ranked = ranked[:max_candidates].detach().to("cpu").tolist() + minimum_alignment = math.cos(math.radians(alignment_max_angle_deg)) + for candidate_index in ranked: + candidate = candidates[candidate_index].to( + device=self.robot.device, dtype=torch.float32 + ) + if ( + float(torch.dot(candidate[:3, 2], approach_direction).item()) + < minimum_alignment + ): + continue + mirrored = candidate.clone() + mirrored[:3, 0] = -mirrored[:3, 0] + mirrored[:3, 1] = -mirrored[:3, 1] + for variant in (candidate, mirrored): + grasp = variant.unsqueeze(0) + pre_grasp = grasp.clone() + pre_grasp[:, :3, 3] -= approach_direction * pre_grasp_distance + lift = grasp.clone() + lift[:, 2, 3] += lift_height + seed = start_qpos + feasible = True + for pose in (pre_grasp, grasp, lift): + success, seed = self.robot.compute_ik( + pose=pose, + joint_seed=seed, + name=self.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + feasible = False + break + if feasible: + return grasp + raise RuntimeError( + f"No independently reachable antipodal grasp remained for " + f"{handle.object_id!r} after screening {len(ranked)} candidates." + ) + + def restore_base_robot(self) -> None: + """Restore the robot state captured before scenario case generation.""" + if self.robot is None or self._base_full_qpos is None: + raise RuntimeError("Atomic Task runtime is not configured.") + for target in (False, True): + self.robot.set_qpos(self._base_full_qpos, target=target) + self.robot.clear_dynamics() + + def set_robot_start( + self, manipulator_qpos: torch.Tensor, *, open_gripper: bool + ) -> None: + """Install a manipulator start and optional open-gripper command.""" + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + for target in (False, True): + self.robot.set_qpos(manipulator_qpos, name=self.control_part, target=target) + if open_gripper: + if self.end_effector_part is None or self._gripper_open is None: + raise RuntimeError("Open-gripper state is unavailable.") + command = self._gripper_open.unsqueeze(0).expand( + manipulator_qpos.shape[0], -1 + ) + self.robot.set_qpos(command, name=self.end_effector_part, target=target) + self.robot.clear_dynamics() + + def activate_object(self, object_id: str) -> AtomicObjectHandle: + """Reset one case object and park every other configured object.""" + handle = self.object_handle(object_id) + for index, candidate in enumerate(self._objects.values()): + if candidate is handle: + candidate.reset() + else: + candidate.park(index) + if self.simulation is not None: + self.simulation.update(step=2) + return handle + + def object_handle(self, object_id: str | None) -> AtomicObjectHandle: + """Resolve a configured object identifier with an actionable error.""" + if object_id is None: + raise ValueError("This Atomic Task case has no object id.") + try: + return self._objects[object_id] + except KeyError as exc: + raise ValueError( + f"Unknown atomic object {object_id!r}; configured objects: " + f"{sorted(self._objects)}." + ) from exc + + def close_runtime(self) -> None: + """Release Python references before SimulationManager teardown.""" + self._engine = None + self._case_providers.clear() + self._objects.clear() + self._base_full_qpos = None + self.end_effector_part = None + self._gripper_open = None + self._gripper_grasp = None + self.simulation = None + self.robot = None + self.suite = None + self.track = None + + @staticmethod + def _synchronize() -> None: + """Synchronize CUDA around physical wall-time measurement.""" + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +register_atomic_skill_provider("move_end_effector", _MoveEndEffectorCases) +register_atomic_skill_provider("pick_up", _PickUpCases) +register_scenario_provider("atomic_task", AtomicTaskScenario) diff --git a/scripts/benchmark/motion_generation/scenarios/base.py b/scripts/benchmark/motion_generation/scenarios/base.py index 11435cb71..30dc3c73f 100644 --- a/scripts/benchmark/motion_generation/scenarios/base.py +++ b/scripts/benchmark/motion_generation/scenarios/base.py @@ -19,19 +19,39 @@ from __future__ import annotations from abc import ABC, abstractmethod +from dataclasses import dataclass, field from typing import TYPE_CHECKING +from embodichain.lab.sim.planners.utils import PlanResult + +from ..metrics.trajectory import compute_case_outcomes, make_failure_outcomes +from ..models import CaseOutcome + if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.objects import Robot from ..config import SuiteCfg, TrackCfg from ..models import BenchmarkCase + from ..planners.base import PlannerAdapter + +__all__ = ["ScenarioEvaluation", "ScenarioProvider"] + -__all__ = ["ScenarioProvider"] +@dataclass(frozen=True) +class ScenarioEvaluation: + """Outcomes and execution metrics produced outside planner timing.""" + + outcomes: tuple[CaseOutcome, ...] + execution_time_ms: float | None = None + end_to_end_time_ms: float | None = None + trajectory_duration_s: float | None = None + trajectory_waypoints: int | None = None + metadata: dict[str, object] = field(default_factory=dict) class ScenarioProvider(ABC): - """Generate fixed cases for one registered scenario kind.""" + """Generate, plan, execute, and evaluate one registered scenario kind.""" required_capabilities: frozenset[str] = frozenset() @@ -49,3 +69,82 @@ def generate_cases( batch_size: int, ) -> list["BenchmarkCase"]: """Build the frozen case manifest for one batch size.""" + + def configure_runtime( + self, + simulation: "SimulationManager", + robot: "Robot", + suite: "SuiteCfg", + track: "TrackCfg", + control_part: str, + ) -> None: + """Create scenario-owned runtime entities before case generation.""" + + def close_runtime(self) -> None: + """Release references to scenario-owned simulation entities.""" + + def prepare_planner( + self, adapter: "PlannerAdapter", first_case: "BenchmarkCase" + ) -> None: + """Bind scenario resources to a built planner outside trial timing.""" + + def close_planner(self, adapter: "PlannerAdapter") -> None: + """Release scenario resources that retain a planner adapter.""" + + def reset_case( + self, + simulation: "SimulationManager", + robot: "Robot", + case: "BenchmarkCase", + control_part: str, + ) -> None: + """Restore the frozen robot start state before a planning call.""" + if case.full_start_qpos is not None: + for target in (False, True): + robot.set_qpos(case.full_start_qpos, target=target) + else: + for target in (False, True): + robot.set_qpos(case.start_qpos, name=control_part, target=target) + robot.clear_dynamics() + simulation.update(step=1) + + def plan_case(self, adapter: "PlannerAdapter", case: "BenchmarkCase") -> object: + """Plan one case through the selected adapter.""" + return adapter.plan(case) + + def plan_contract_error(self, result: object) -> str | None: + """Return a diagnostic when a planner artifact violates this scenario.""" + if isinstance(result, PlanResult): + return None + return f"Expected PlanResult, got {type(result).__name__}." + + def failure_outcomes( + self, case: "BenchmarkCase", failure_code: str + ) -> tuple[CaseOutcome, ...]: + """Create scenario-appropriate outcomes after a runner-level failure.""" + return make_failure_outcomes(case.batch_size, failure_code) + + def evaluate_case( + self, + result: object, + case: "BenchmarkCase", + robot: "Robot", + control_part: str, + suite: "SuiteCfg", + *, + planning_time_ms: float, + ) -> ScenarioEvaluation: + """Externally validate a planner-only trajectory.""" + if not isinstance(result, PlanResult): + raise TypeError(self.plan_contract_error(result)) + outcomes = compute_case_outcomes( + result, + case, + robot, + control_part, + validation_samples=suite.protocol.validation_samples, + position_threshold_m=suite.protocol.position_threshold_m, + rotation_threshold_rad=suite.protocol.rotation_threshold_rad, + joint_limit_tolerance_rad=suite.protocol.joint_limit_tolerance_rad, + ) + return ScenarioEvaluation(outcomes=outcomes) diff --git a/scripts/benchmark/motion_generation/scenarios/free_space.py b/scripts/benchmark/motion_generation/scenarios/free_space.py index ce2b7fff4..f299d7fe6 100644 --- a/scripts/benchmark/motion_generation/scenarios/free_space.py +++ b/scripts/benchmark/motion_generation/scenarios/free_space.py @@ -199,6 +199,7 @@ def _build_case( start_qpos=start_qpos, target_waypoints=target_waypoints, reference_qpos=reference_qpos, + robot_id=suite.robot.id, ) diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml new file mode 100644 index 000000000..2bfd23ef4 --- /dev/null +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml @@ -0,0 +1,102 @@ +schema_version: 1 +name: atomic_skill_franka_pgi_curobo +suite_version: atomic_franka_pgi_curobo_smoke_v1 +profile: smoke + +robot: + id: franka_pgi + provider: franka_pgi + config: {} + +planners: + - id: curobo + adapter: curobo + role: primary_baseline + enabled: true + config: + max_attempts: 5 + max_planning_time: null + interpolation_dt: 0.025 + collision_activation_distance: 0.01 + use_cuda_graph: true + cuda_graph_fallback: true + warmup_iterations: 1 + preserve_plan_samples: false + world: + obstacle_representation: sphere + multi_env: false + auto_gen: + fit_type: voxel + sphere_density: 0.1 + collision_sphere_buffer: 0.0 + +protocol: + warmup_trials: 1 + measured_trials: 1 + sample_interval: 80 + validation_samples: 128 + position_threshold_m: 0.01 + rotation_threshold_rad: 0.1 + joint_limit_tolerance_rad: 0.00001 + +tracks: + - id: atomic-task + scenario: atomic_task + enabled: true + config: + batch_sizes: [1] + seeds: [11] + physics: + steps_per_waypoint: 4 + hold_steps: 80 + hold_sim_steps: 2 + joint_tracking_tolerance_rad: 0.05 + gripper: + control_part: hand + open_qpos: [0.0] + grasp_qpos: [0.024] + objects: + - id: cube_50mm + kind: cube + size: [0.05, 0.05, 0.05] + position: [-0.42, -0.08, 0.05] + mass: 0.05 + dynamic_friction: 0.97 + static_friction: 0.99 + contact_offset: 0.003 + rest_offset: 0.001 + settle_steps: 10 + skills: + - id: move_end_effector + sample_count: 80 + cases: + - name: relative_two_waypoint + task_difficulty: simple + target_offsets_m: + - [-0.08, -0.08, 0.08] + - [0.04, 0.12, 0.04] + difficulty_factors: + waypoint_count: 2 + maximum_translation_m: 0.13 + obstacles: 0 + - id: pick_up + sample_count: 120 + grasp_source: antipodal + grasp_sample_count: 10000 + grasp_max_candidates: 128 + grasp_alignment_max_angle_deg: 10.0 + hand_interp_steps: 12 + pre_grasp_distance_m: 0.15 + lift_height_m: 0.16 + minimum_object_lift_m: 0.04 + approach_direction: [0.0, 0.0, -1.0] + pre_pick_height_m: 0.36 + cases: + - name: cube_top_center + object: cube_50mm + task_difficulty: simple + grasp_offset_m: [0.0, 0.0, 0.0] + difficulty_factors: + approach: top + pre_grasp_clearance_m: 0.15 + required_lift_m: 0.04 diff --git a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py new file mode 100644 index 000000000..9defa233e --- /dev/null +++ b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py @@ -0,0 +1,221 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Pure-logic tests for the planner-oriented Atomic Task benchmark track.""" + +from __future__ import annotations + +import json +from unittest.mock import Mock + +import pytest +import torch + +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.models import ( + AlgorithmRole, + BenchmarkCase, + CaseOutcome, + PlannerMetadata, + TrialPhase, + TrialRecord, +) +from scripts.benchmark.motion_generation.registry import create_robot_provider +from scripts.benchmark.motion_generation import robots as _robots # noqa: F401 +from scripts.benchmark.motion_generation.scenarios.atomic_objects import ( + atomic_object_kind_names, + create_atomic_object, +) +from scripts.benchmark.motion_generation.scenarios.atomic_task import ( + atomic_skill_provider_names, + create_atomic_skill_provider, +) + + +def _atomic_case() -> BenchmarkCase: + target = torch.eye(4).reshape(1, 1, 4, 4) + return BenchmarkCase( + suite_version="atomic_test_v1", + track="atomic-task", + scenario_id="move_end_effector", + case_id="atomic-task:move_end_effector:simple:s11", + seed=11, + batch_size=1, + num_waypoints=1, + path_shape="robot_relative_waypoints", + start_state_bin="pre_action", + start_qpos=torch.zeros(1, 7), + target_waypoints=target, + reference_qpos=torch.zeros(1, 1, 7), + robot_id="franka_pgi", + skill_id="move_end_effector", + task_difficulty="simple", + primary_success="task_success", + full_start_qpos=torch.zeros(1, 9), + case_parameters={"sample_count": 80, "target_offsets_m": [[0.1, 0.0, 0.0]]}, + ) + + +def _atomic_outcome() -> CaseOutcome: + return CaseOutcome( + env_index=0, + planning_success=True, + finite=True, + ordered_waypoints_reached=True, + motion_valid=True, + completed_waypoint_ratio=1.0, + final_translation_err_mm=1.0, + final_rotation_err_deg=1.0, + waypoint_translation_err_mm_mean=1.0, + waypoint_translation_err_mm_p95=1.0, + waypoint_translation_err_mm_max=1.0, + waypoint_rotation_err_deg_mean=1.0, + waypoint_rotation_err_deg_p95=1.0, + waypoint_rotation_err_deg_max=1.0, + joint_limit_violation=False, + max_normalized_joint_violation=0.0, + joint_path_length_rad=0.2, + cartesian_path_length_m=0.1, + path_efficiency=1.0, + execution_success=True, + task_success=True, + task_completion_time_s=1.5, + joint_tracking_rmse_rad=0.002, + replan_count=0, + ) + + +def test_atomic_suite_is_franka_pgi_and_curobo_only(): + suite = load_suite("atomic_franka_pgi_curobo") + + assert suite.robot.id == "franka_pgi" + assert suite.robot.provider == "franka_pgi" + assert [spec.id for spec in suite.planners if spec.enabled] == ["curobo"] + assert [(track.id, track.scenario) for track in suite.enabled_tracks()] == [ + ("atomic-task", "atomic_task") + ] + skills = suite.enabled_tracks()[0].config["skills"] + assert [item["id"] for item in skills] == ["move_end_effector", "pick_up"] + gripper = suite.enabled_tracks()[0].config["gripper"] + assert gripper == { + "control_part": "hand", + "open_qpos": [0.0], + "grasp_qpos": [0.024], + } + + +def test_franka_pgi_robot_provider_exposes_arm_hand_and_tcp(): + suite = load_suite("atomic_franka_pgi_curobo") + cfg = create_robot_provider(suite.robot).build_cfg() + + assert cfg.uid == "benchmark_franka_pgi" + assert len(cfg.control_parts["arm"]) == 7 + assert cfg.control_parts["hand"] == ["gripper_finger1_joint_1"] + assert cfg.solver_cfg["arm"].end_link_name == "fr3_link8" + assert cfg.solver_cfg["arm"].tcp[2][3] == pytest.approx(0.15) + assert len(cfg.init_qpos) == 9 + + +def test_atomic_invocation_pins_motion_generator_and_selected_planner(): + provider = create_atomic_skill_provider("move_end_effector") + invocation = provider.build_invocation( + Mock(control_part="manipulator"), + _atomic_case(), + Mock(motion_policy_planner="curobo"), + ) + + assert invocation.motion_policy.strategy == "motion_gen" + assert invocation.motion_policy.planner == "curobo" + assert invocation.skill_id == "move_end_effector" + assert invocation.binding.manipulators == {"primary": "manipulator"} + + +def test_atomic_skill_and_object_extensions_are_registry_driven(): + assert atomic_skill_provider_names() == ("move_end_effector", "pick_up") + assert atomic_object_kind_names() == ("cube", "mesh") + with pytest.raises(ValueError, match="Unknown atomic object kind"): + create_atomic_object(Mock(), {"id": "new_object", "kind": "not_registered"}) + + +def test_atomic_primary_success_and_execution_efficiency_aggregate(): + case = _atomic_case() + metadata = [ + PlannerMetadata( + algorithm_id="curobo", + algorithm_role=AlgorithmRole.PRIMARY_BASELINE, + adapter="curobo", + config_hash="abc", + capabilities=frozenset({"eef_waypoint", "atomic_action"}), + supported_robots=("franka_pgi",), + ) + ] + record = TrialRecord( + suite_version=case.suite_version, + track=case.track, + scenario_id=case.scenario_id, + case_id=case.case_id, + algorithm_id="curobo", + algorithm_role=AlgorithmRole.PRIMARY_BASELINE, + model_revision="curobo-v2", + planner_config_hash="abc", + seed=case.seed, + repeat=0, + batch_size=1, + waypoint_count=1, + path_shape=case.path_shape, + start_state_bin=case.start_state_bin, + phase=TrialPhase.MEASURED, + cost_time_ms=20.0, + robot_id=case.robot_id, + skill_id=case.skill_id, + task_difficulty=case.task_difficulty, + primary_success=case.primary_success, + execution_time_ms=30.0, + end_to_end_time_ms=50.0, + trajectory_duration_s=1.5, + trajectory_waypoints=80, + outcomes=(_atomic_outcome(),), + ) + + aggregates = aggregate_results([record], metadata, [case], measured_trials=1) + metrics = aggregates["success_and_metrics"][0] + performance = aggregates["time_and_memory"][0] + leaderboard = aggregates["leaderboard"][0] + + assert metrics["primary_success"] == "task_success" + assert metrics["success_rate"] == pytest.approx(1.0) + assert metrics["execution_success_rate"] == pytest.approx(1.0) + assert metrics["task_success_rate"] == pytest.approx(1.0) + assert performance["execution_time_ms"] == pytest.approx(30.0) + assert performance["end_to_end_time_ms"] == pytest.approx(50.0) + assert leaderboard["overall_success_rate"] == pytest.approx(1.0) + assert leaderboard["task_success_rate"] == pytest.approx(1.0) + + +def test_atomic_case_manifest_retains_robot_skill_object_and_parameters(tmp_path): + case = _atomic_case() + path = write_case_manifest(tmp_path / "case_manifest.json", [case]) + payload = json.loads(path.read_text(encoding="utf-8")) + serialized = payload["cases"][0] + + assert payload["case_schema_version"] == 2 + assert serialized["robot_id"] == "franka_pgi" + assert serialized["skill_id"] == "move_end_effector" + assert serialized["primary_success"] == "task_success" + assert serialized["case_parameters"]["sample_count"] == 80 + assert serialized["validity_evidence"]["method"] == "independent_sequential_ik"