From e0fec8ab2e8a7d4819ec5921c964e28dd5f75227 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Wed, 19 Aug 2026 02:59:27 +0000 Subject: [PATCH 01/11] 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" From 2a1f06f3757f8114aecb8aeeeb666ed0b5e9ec15 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Thu, 20 Aug 2026 12:23:28 +0800 Subject: [PATCH 02/11] Add record camera --- scripts/benchmark/motion_generation/README.md | 7 +- .../motion_generation/run_benchmark.py | 34 +++ scripts/benchmark/motion_generation/runner.py | 58 ++++- .../scenarios/atomic_task.py | 192 +++++++++++++---- .../motion_generation/scenarios/base.py | 28 +++ scripts/benchmark/motion_generation/video.py | 198 ++++++++++++++++++ .../test_atomic_task_benchmark.py | 195 +++++++++++++++++ 7 files changed, 665 insertions(+), 47 deletions(-) create mode 100644 scripts/benchmark/motion_generation/video.py diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index 7c988e962..77b87224a 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -13,13 +13,16 @@ 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 \ + --suite atomic_franka_pgi_curobo --device cuda --record-video python -m scripts.benchmark.motion_generation.run_benchmark \ --extra-baselines ik_interpolate toppra ``` Artifacts land under `outputs/benchmarks///` (`resolved_suite.yaml`, `case_manifest.json`, `trials.jsonl`, `aggregates.json`, -`report.md` with exactly three tables). +`report.md` with exactly three tables). Atomic Task videos, when enabled, land +in that run's `videos/` directory. ## Implemented @@ -36,6 +39,8 @@ Artifacts land under `outputs/benchmarks///` duration, simulated task-completion time, controller tracking RMSE, and task-specific object lift - One Markdown report: Time & Memory, Success & Other Metrics, Leaderboard +- Optional Atomic Task headless replay videos after measured evaluation + (`--record-video`, `--record-failed-video`) ## Extend diff --git a/scripts/benchmark/motion_generation/run_benchmark.py b/scripts/benchmark/motion_generation/run_benchmark.py index d34a4d00e..50eb8fac8 100644 --- a/scripts/benchmark/motion_generation/run_benchmark.py +++ b/scripts/benchmark/motion_generation/run_benchmark.py @@ -33,6 +33,7 @@ from typing import TYPE_CHECKING from .config import PlannerSpecCfg, SuiteCfg, load_suite +from .video import VideoRecordCfg, video_cfg_from_args if TYPE_CHECKING: from .runner import BenchmarkRunResult @@ -110,6 +111,36 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--no-headless", action="store_false", dest="headless", help="Open a viewer." ) + parser.add_argument( + "--record-video", + action="store_true", + help="Record Atomic Task measured physics-replay videos after evaluation.", + ) + parser.add_argument( + "--record-failed-video", + action="store_true", + help="With --record-video, also record failed cases as static debug scenes.", + ) + parser.add_argument( + "--video-case-limit", + type=int, + default=0, + help="Maximum recorded videos. Use 0 to record every selected case.", + ) + parser.add_argument( + "--video-dir", + default=None, + help="Override video directory. Default is /videos.", + ) + parser.add_argument("--video-fps", type=int, default=20) + parser.add_argument("--video-width", type=int, default=640) + parser.add_argument("--video-height", type=int, default=480) + parser.add_argument( + "--video-max-memory", + type=int, + default=2048, + help="Maximum recorder frame-buffer memory in MB.", + ) def _planner_by_id(suite: SuiteCfg, planner_id: str) -> PlannerSpecCfg: @@ -218,6 +249,7 @@ def run_all_benchmarks( nmg_pos_eps: float | None = None, nmg_rot_eps: float | None = None, output_root: str | Path = "outputs/benchmarks", + video: VideoRecordCfg | None = None, ) -> BenchmarkRunResult: """Resolve configuration and run all selected benchmark tracks.""" from .runner import BenchmarkRunner @@ -247,6 +279,7 @@ def run_all_benchmarks( device=sim_device, headless=headless, output_root=output_root, + video=video, ).run() @@ -273,6 +306,7 @@ def run_from_args(args: argparse.Namespace) -> BenchmarkRunResult: nmg_pos_eps=args.nmg_pos_eps, nmg_rot_eps=args.nmg_rot_eps, output_root=args.output_root, + video=video_cfg_from_args(args), ) diff --git a/scripts/benchmark/motion_generation/runner.py b/scripts/benchmark/motion_generation/runner.py index fc660fe9e..c1ce0dd70 100644 --- a/scripts/benchmark/motion_generation/runner.py +++ b/scripts/benchmark/motion_generation/runner.py @@ -55,6 +55,11 @@ from .reporting import write_markdown_report from .scenarios.base import ScenarioEvaluation, ScenarioProvider from .scenarios.free_space import FreeSpaceScenario +from .video import ( + VideoRecordCfg, + should_record_case, + summarize_video_recording, +) if TYPE_CHECKING: from collections.abc import Callable @@ -105,18 +110,22 @@ def __init__( device: str = "auto", headless: bool = True, output_root: str | Path = "outputs/benchmarks", + video: VideoRecordCfg | None = None, ) -> None: self.suite = suite self.planner_specs = planner_specs self.device = resolve_device(device) self.headless = headless self.output_root = Path(output_root) + self.video = VideoRecordCfg() if video is None else video 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] = [] + self._run_dir: Path | None = None + self._video_paths: list[str] = [] def _create_simulation(self, batch_size: int) -> tuple[SimulationManager, "Robot"]: """Create one isolated suite-selected robot for a fixed batch size.""" @@ -279,6 +288,13 @@ def _run_plan_call( failure_message = str(exc) outcomes = provider.failure_outcomes(case, failure_code) + record_metadata = {} if evaluation is None else dict(evaluation.metadata) + video_path = self._maybe_record_replay( + provider, result, case, evaluation, metadata.algorithm_id, phase + ) + if video_path is not None: + record_metadata["video_path"] = str(video_path) + self._append( writer, TrialRecord( @@ -302,7 +318,7 @@ def _run_plan_call( trajectory_waypoints=( None if evaluation is None else evaluation.trajectory_waypoints ), - metadata={} if evaluation is None else evaluation.metadata, + metadata=record_metadata, outcomes=outcomes, ), ) @@ -314,6 +330,44 @@ def _run_plan_call( f"status={status}" ) + def _maybe_record_replay( + self, + provider: ScenarioProvider, + result: object, + case: BenchmarkCase, + evaluation: ScenarioEvaluation | None, + algorithm_id: str, + phase: TrialPhase, + ) -> Path | None: + """Record one measured Atomic Task replay outside planner timing.""" + if phase is not TrialPhase.MEASURED or not self.video.enabled: + return None + success = bool( + evaluation is not None + and evaluation.outcomes + and all(bool(outcome.task_success) for outcome in evaluation.outcomes) + ) + if not should_record_case(self.video, len(self._video_paths), success): + return None + if self._run_dir is None: + raise RuntimeError("Benchmark run directory is not initialized.") + output_dir = ( + self.video.output_dir + if self.video.output_dir is not None + else self._run_dir / "videos" + ) + video_path = provider.record_replay( + result, + case, + evaluation, + output_dir=output_dir, + algorithm_id=algorithm_id, + video=self.video, + ) + if video_path is not None: + self._video_paths.append(str(video_path)) + return video_path + def _run_adapter( self, writer: TrialJsonlWriter, @@ -444,6 +498,7 @@ def _run_adapter( def run(self) -> BenchmarkRunResult: """Run the suite and write all required artifacts.""" run_dir = create_run_directory(self.output_root, self.suite.name) + self._run_dir = run_dir write_resolved_suite(run_dir / "resolved_suite.yaml", self.suite) write_json(run_dir / "environment.json", environment_metadata()) writer = TrialJsonlWriter(run_dir / "trials.jsonl") @@ -528,6 +583,7 @@ def run(self) -> BenchmarkRunResult: "stability hold for successful tasks.", ] ) + track_notes.extend(summarize_video_recording(self.video, self._video_paths)) report_path = write_markdown_report( run_dir / "report.md", self.suite, diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py index 2a77500a0..2bc1e66c4 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_task.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -23,6 +23,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from dataclasses import dataclass, replace +from pathlib import Path from typing import TYPE_CHECKING import torch @@ -46,6 +47,7 @@ from ..metrics.trajectory import compute_case_outcomes from ..models import BenchmarkCase, CaseOutcome from ..registry import register_scenario_provider +from ..video import VideoRecordCfg, build_video_path, record_with_window from .atomic_objects import AtomicObjectHandle, create_atomic_object from .base import ScenarioEvaluation, ScenarioProvider @@ -83,6 +85,16 @@ class _ExecutionObservation: object_lift_delta_m: torch.Tensor | None = None +@dataclass(frozen=True) +class _PhysicsReplaySettings: + """Shared physics-step counts for timed evaluation and video replay.""" + + steps_per_waypoint: int + hold_steps: int + hold_sim_steps: int + joint_tracking_tolerance_rad: float + + class AtomicSkillCaseProvider(ABC): """Generate and ground one Atomic Action without planner-specific logic.""" @@ -891,6 +903,36 @@ def evaluate_case( }, ) + def record_replay( + self, + result: object, + case: BenchmarkCase, + evaluation: ScenarioEvaluation | None, + *, + output_dir: Path, + algorithm_id: str, + video: VideoRecordCfg, + ) -> Path | None: + """Reset the case and record a second untimed physics replay.""" + del evaluation + if self.simulation is None or self.robot is None: + return None + self.reset_case(self.simulation, self.robot, case, self.control_part) + compiled = result if isinstance(result, CompiledTrajectory) else None + replayable = compiled is not None and self._is_replayable(compiled) + provider = self._case_providers.get(case.case_id) + video_path = build_video_path( + output_dir, algorithm_id, case.skill_id, case.case_id + ) + + def _replay() -> None: + if replayable and compiled is not None: + self._replay_physics(compiled, case, provider, collect_metrics=False) + else: + self._hold_static() + + return record_with_window(self.simulation, video, video_path, _replay) + def _execute( self, compiled: CompiledTrajectory, @@ -900,12 +942,7 @@ def _execute( """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()) - ): + if not self._is_replayable(compiled): return None object_handle = ( None if case.object_id is None else self.object_handle(case.object_id) @@ -915,51 +952,23 @@ def _execute( 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 + settings = self._physics_settings() 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] + tracking_parts = self._replay_physics( + compiled, case, provider, collect_metrics=True + ) self._synchronize() elapsed_ms = (time.perf_counter() - started) * 1000.0 - + if tracking_parts is None: + raise RuntimeError("Timed Atomic Task replay did not return metrics.") + squared_tracking_error, tracking_value_count = tracking_parts observed = self.robot.get_qpos() tracking = torch.sqrt(squared_tracking_error / max(tracking_value_count, 1)) + trajectory = compiled.trajectory.positions simulated_execution_time = ( - trajectory.shape[1] * steps_per_waypoint + hold_steps * hold_sim_steps + trajectory.shape[1] * settings.steps_per_waypoint + + settings.hold_steps * settings.hold_sim_steps ) * float(self.simulation.sim_config.physics_dt) final_tcp = self.robot.compute_fk( self.robot.get_qpos(name=self.control_part), @@ -974,7 +983,7 @@ def _execute( object_lift = final_object_position[:, 2] - initial_object_position[:, 2] return _ExecutionObservation( execution_success=torch.isfinite(observed).all(dim=1) - & (tracking <= tracking_tolerance), + & (tracking <= settings.joint_tracking_tolerance_rad), final_tcp_pose=final_tcp, joint_tracking_rmse_rad=tracking, execution_time_ms=elapsed_ms, @@ -982,6 +991,99 @@ def _execute( object_lift_delta_m=object_lift, ) + def _physics_settings(self) -> _PhysicsReplaySettings: + """Resolve and validate the common physics-replay step counts.""" + if self.track is None: + raise RuntimeError("Atomic Task runtime is not configured.") + physics = dict(self.track.config.get("physics", {})) + settings = _PhysicsReplaySettings( + 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)), + joint_tracking_tolerance_rad=float( + physics.get("joint_tracking_tolerance_rad", 0.05) + ), + ) + if ( + settings.steps_per_waypoint < 1 + or settings.hold_steps < 0 + or settings.hold_sim_steps < 1 + ): + raise ValueError("Atomic Task physics step counts are invalid.") + return settings + + @staticmethod + def _is_replayable(compiled: CompiledTrajectory) -> bool: + """Return whether a compiled trajectory can be physically replayed.""" + trajectory = compiled.trajectory.positions + return ( + trajectory.shape[1] > 0 + and bool(compiled.plan_success.all().item()) + and bool(torch.isfinite(trajectory).all().item()) + ) + + def _replay_physics( + self, + compiled: CompiledTrajectory, + case: BenchmarkCase, + provider: AtomicSkillCaseProvider | None, + *, + collect_metrics: bool, + ) -> tuple[torch.Tensor, int] | None: + """Replay one compiled trajectory with the evaluation physics contract.""" + if self.simulation is None or self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + trajectory = compiled.trajectory.positions + object_handle = ( + None if case.object_id is None else self.object_handle(case.object_id) + ) + settings = self._physics_settings() + lift_start = None if provider is None else provider.lift_segment_start(compiled) + dynamics_cleared = False + squared_tracking_error: torch.Tensor | None = None + tracking_value_count = 0 + if collect_metrics: + squared_tracking_error = torch.zeros( + case.batch_size, dtype=trajectory.dtype, device=trajectory.device + ) + for waypoint_index in range(trajectory.shape[1]): + positions = trajectory[:, waypoint_index] + self.robot.set_qpos(positions, target=True) + self.simulation.update(step=settings.steps_per_waypoint) + if squared_tracking_error is not None: + 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(settings.hold_steps): + self.robot.set_qpos(final_command, target=True) + self.simulation.update(step=settings.hold_sim_steps) + if squared_tracking_error is not None: + observed = self.robot.get_qpos() + squared_tracking_error += ((observed - final_command) ** 2).sum(dim=1) + tracking_value_count += observed.shape[1] + if collect_metrics: + if squared_tracking_error is None: + raise RuntimeError("Timed replay lost its tracking accumulator.") + return squared_tracking_error, tracking_value_count + return None + + def _hold_static(self) -> None: + """Hold the current scene so a failed-case debug video has frames.""" + if self.simulation is None: + raise RuntimeError("Atomic Task runtime is not configured.") + settings = self._physics_settings() + for _ in range(max(settings.hold_steps, 1)): + self.simulation.update(step=settings.hold_sim_steps) + def solve_reference_qpos( self, start_qpos: torch.Tensor, target_waypoints: torch.Tensor ) -> torch.Tensor: diff --git a/scripts/benchmark/motion_generation/scenarios/base.py b/scripts/benchmark/motion_generation/scenarios/base.py index 30dc3c73f..4be5b3c72 100644 --- a/scripts/benchmark/motion_generation/scenarios/base.py +++ b/scripts/benchmark/motion_generation/scenarios/base.py @@ -20,6 +20,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field +from pathlib import Path from typing import TYPE_CHECKING from embodichain.lab.sim.planners.utils import PlanResult @@ -34,6 +35,7 @@ from ..config import SuiteCfg, TrackCfg from ..models import BenchmarkCase from ..planners.base import PlannerAdapter + from ..video import VideoRecordCfg __all__ = ["ScenarioEvaluation", "ScenarioProvider"] @@ -148,3 +150,29 @@ def evaluate_case( joint_limit_tolerance_rad=suite.protocol.joint_limit_tolerance_rad, ) return ScenarioEvaluation(outcomes=outcomes) + + def record_replay( + self, + result: object, + case: "BenchmarkCase", + evaluation: ScenarioEvaluation | None, + *, + output_dir: Path, + algorithm_id: str, + video: "VideoRecordCfg", + ) -> Path | None: + """Optionally record a second, untimed replay. Default is a no-op. + + Args: + result: Planner or compiled-action artifact from the measured trial. + case: Frozen case identity used for the output filename. + evaluation: Timed evaluation, or ``None`` after a runner-level failure. + output_dir: Directory that should receive the mp4. + algorithm_id: Planner id used in the filename. + video: Recording policy and encoder settings. + + Returns: + Path to a saved video, or ``None`` when this scenario does not record. + """ + del result, case, evaluation, output_dir, algorithm_id, video + return None diff --git a/scripts/benchmark/motion_generation/video.py b/scripts/benchmark/motion_generation/video.py new file mode 100644 index 000000000..30c140b12 --- /dev/null +++ b/scripts/benchmark/motion_generation/video.py @@ -0,0 +1,198 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Optional Atomic Task replay recording helpers. + +Recording is a second, untimed physics pass. Failures never change trial +success and never add a fourth Markdown table. +""" + +from __future__ import annotations + +import argparse +import re +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from embodichain.lab.sim import SimulationManager + +__all__ = [ + "DEFAULT_VIDEO_LOOK_AT", + "VideoRecordCfg", + "build_video_path", + "record_with_window", + "should_record_case", + "summarize_video_recording", + "video_cfg_from_args", +] + +DEFAULT_VIDEO_FPS = 20 +DEFAULT_VIDEO_MAX_MEMORY_MB = 2048 +DEFAULT_VIDEO_WIDTH = 640 +DEFAULT_VIDEO_HEIGHT = 480 +DEFAULT_VIDEO_CASE_LIMIT = 0 +DEFAULT_VIDEO_LOOK_AT = ( + (-1.25, -1.15, 0.95), + (-0.25, -0.02, 0.25), + (0.0, 0.0, 1.0), +) + +LookAt = tuple[Sequence[float], Sequence[float], Sequence[float]] + + +@dataclass(frozen=True) +class VideoRecordCfg: + """CLI-resolved recording policy for Atomic Task measured replays.""" + + enabled: bool = False + record_failed: bool = False + case_limit: int = DEFAULT_VIDEO_CASE_LIMIT + fps: int = DEFAULT_VIDEO_FPS + width: int = DEFAULT_VIDEO_WIDTH + height: int = DEFAULT_VIDEO_HEIGHT + max_memory_mb: int = DEFAULT_VIDEO_MAX_MEMORY_MB + look_at: LookAt = DEFAULT_VIDEO_LOOK_AT + output_dir: Path | None = None + + def __post_init__(self) -> None: + if self.case_limit < 0: + raise ValueError("video case_limit must be non-negative.") + if self.fps <= 0: + raise ValueError("video fps must be positive.") + if self.width <= 0 or self.height <= 0: + raise ValueError("video width and height must be positive.") + if self.max_memory_mb <= 0: + raise ValueError("video max_memory_mb must be positive.") + + +def should_record_case(cfg: VideoRecordCfg, recorded_count: int, success: bool) -> bool: + """Return whether one measured Atomic Task case should emit a video.""" + if not cfg.enabled: + return False + if not success and not cfg.record_failed: + return False + return cfg.case_limit == 0 or recorded_count < cfg.case_limit + + +def build_video_path( + output_dir: Path, + algorithm_id: str, + skill_id: str, + case_id: str, +) -> Path: + """Build ``{algorithm}_{skill}_{case}.mp4`` under the run videos directory.""" + + def _sanitize(value: str) -> str: + return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_") or "unknown" + + output_dir.mkdir(parents=True, exist_ok=True) + filename = ( + f"{_sanitize(algorithm_id)}_{_sanitize(skill_id)}_{_sanitize(case_id)}.mp4" + ) + return output_dir / filename + + +def record_with_window( + sim: "SimulationManager", + cfg: VideoRecordCfg, + video_path: Path, + replay_fn: Callable[[], None], +) -> Path | None: + """Record one headless replay. Failures print a warning and return ``None``.""" + try: + original_width = sim.sim_config.width + original_height = sim.sim_config.height + recording_started = False + try: + sim.sim_config.width = cfg.width + sim.sim_config.height = cfg.height + recording_started = sim.start_window_record( + save_path=str(video_path), + fps=cfg.fps, + max_memory=cfg.max_memory_mb, + look_at=cfg.look_at, + use_sim_time=True, + ) + finally: + sim.sim_config.width = original_width + sim.sim_config.height = original_height + if not recording_started: + return None + + stop_success = False + try: + replay_fn() + finally: + if sim.is_window_recording(): + stop_success = sim.stop_window_record() + sim.wait_window_record_saves() + return video_path if stop_success else None + except Exception as exc: # noqa: BLE001 - recording must not change success + try: + if sim.is_window_recording(): + sim.stop_window_record() + sim.wait_window_record_saves() + except Exception: + pass + print( + "Warning: failed to record Atomic Task replay video " + f"{video_path}: {type(exc).__name__}: {exc}" + ) + return None + + +def summarize_video_recording( + cfg: VideoRecordCfg, video_paths: Sequence[str] +) -> list[str]: + """Return report notes describing recording coverage without extra tables.""" + if not cfg.enabled: + return ["Video policy: disabled."] + if cfg.record_failed: + policy = ( + "records Atomic Task measured success replays and failed-case " + "static scenes when capture is available." + ) + else: + policy = ( + "records Atomic Task measured success replays only; failed " + "cases are reported in the tables but do not emit videos." + ) + rendered = ", ".join(video_paths) if video_paths else "none" + return [ + f"Video policy: {policy}", + f"videos={len(video_paths)}", + f"Replay videos: {rendered}", + ] + + +def video_cfg_from_args(args: argparse.Namespace) -> VideoRecordCfg: + """Build recording config from the motion-generation CLI namespace.""" + output_dir = getattr(args, "video_dir", None) + return VideoRecordCfg( + enabled=bool(getattr(args, "record_video", False)), + record_failed=bool(getattr(args, "record_failed_video", False)), + case_limit=int(getattr(args, "video_case_limit", DEFAULT_VIDEO_CASE_LIMIT)), + fps=int(getattr(args, "video_fps", DEFAULT_VIDEO_FPS)), + width=int(getattr(args, "video_width", DEFAULT_VIDEO_WIDTH)), + height=int(getattr(args, "video_height", DEFAULT_VIDEO_HEIGHT)), + max_memory_mb=int( + getattr(args, "video_max_memory", DEFAULT_VIDEO_MAX_MEMORY_MB) + ), + output_dir=None if output_dir in (None, "") else Path(output_dir), + ) diff --git a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py index 9defa233e..abe0276e0 100644 --- a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py +++ b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py @@ -18,7 +18,9 @@ from __future__ import annotations +import argparse import json +from pathlib import Path from unittest.mock import Mock import pytest @@ -36,15 +38,27 @@ TrialRecord, ) from scripts.benchmark.motion_generation.registry import create_robot_provider +from scripts.benchmark.motion_generation.runner import BenchmarkRunner 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 ( + AtomicTaskScenario, atomic_skill_provider_names, create_atomic_skill_provider, ) +from scripts.benchmark.motion_generation.scenarios.base import ScenarioProvider +from scripts.benchmark.motion_generation.scenarios.free_space import FreeSpaceScenario +from scripts.benchmark.motion_generation.video import ( + VideoRecordCfg, + build_video_path, + record_with_window, + should_record_case, + summarize_video_recording, + video_cfg_from_args, +) def _atomic_case() -> BenchmarkCase: @@ -219,3 +233,184 @@ def test_atomic_case_manifest_retains_robot_skill_object_and_parameters(tmp_path assert serialized["primary_success"] == "task_success" assert serialized["case_parameters"]["sample_count"] == 80 assert serialized["validity_evidence"]["method"] == "independent_sequential_ik" + + +def test_should_record_case_respects_enable_failure_and_limit(): + disabled = VideoRecordCfg() + enabled = VideoRecordCfg(enabled=True) + with_failed = VideoRecordCfg(enabled=True, record_failed=True) + limited = VideoRecordCfg(enabled=True, case_limit=1) + + assert should_record_case(disabled, 0, True) is False + assert should_record_case(enabled, 0, True) is True + assert should_record_case(enabled, 0, False) is False + assert should_record_case(with_failed, 0, False) is True + assert should_record_case(limited, 0, True) is True + assert should_record_case(limited, 1, True) is False + + +def test_build_video_path_sanitizes_case_id(tmp_path): + path = build_video_path( + tmp_path, + "curobo", + "pick_up", + "atomic-task:pick_up:cube_top_center:s11", + ) + + assert path.parent == tmp_path + assert path.name == "curobo_pick_up_atomic-task_pick_up_cube_top_center_s11.mp4" + assert tmp_path.is_dir() + + +def test_default_scenario_record_replay_is_noop(tmp_path): + provider = FreeSpaceScenario() + path = provider.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True), + ) + + assert path is None + assert isinstance(provider, ScenarioProvider) + + +def test_atomic_record_replay_without_runtime_returns_none(tmp_path): + scenario = AtomicTaskScenario() + path = scenario.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True), + ) + assert path is None + + +def test_record_with_window_swallows_exceptions_and_does_not_raise(tmp_path): + sim = Mock() + sim.sim_config.width = 64 + sim.sim_config.height = 64 + sim.start_window_record.side_effect = RuntimeError("recorder failed") + sim.is_window_recording.return_value = False + + path = record_with_window( + sim, + VideoRecordCfg(enabled=True), + tmp_path / "failed.mp4", + lambda: None, + ) + + assert path is None + sim.wait_window_record_saves.assert_called() + + +def test_atomic_record_replay_swallows_recorder_errors(tmp_path): + scenario = AtomicTaskScenario() + sim = Mock() + sim.sim_config.width = 64 + sim.sim_config.height = 64 + sim.start_window_record.side_effect = RuntimeError("boom") + sim.is_window_recording.return_value = False + scenario.simulation = sim + scenario.robot = Mock() + scenario.reset_case = Mock() + + path = scenario.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True, record_failed=True), + ) + + assert path is None + scenario.reset_case.assert_called_once() + + +def test_atomic_failed_plan_records_static_hold(tmp_path): + scenario = AtomicTaskScenario() + sim = Mock() + sim.sim_config.width = 64 + sim.sim_config.height = 64 + sim.start_window_record.return_value = True + sim.is_window_recording.return_value = True + sim.stop_window_record.return_value = True + scenario.simulation = sim + scenario.robot = Mock() + scenario.track = Mock( + config={ + "physics": {"hold_steps": 2, "hold_sim_steps": 1, "steps_per_waypoint": 4} + } + ) + scenario.reset_case = Mock() + + path = scenario.record_replay( + None, + _atomic_case(), + None, + output_dir=tmp_path, + algorithm_id="curobo", + video=VideoRecordCfg(enabled=True, record_failed=True), + ) + + assert path is not None + assert path.name.startswith("curobo_move_end_effector_") + assert sim.update.call_count == 2 + + +def test_video_cfg_from_args_and_summary_notes(): + args = argparse.Namespace( + record_video=True, + record_failed_video=False, + video_case_limit=0, + video_fps=20, + video_width=640, + video_height=480, + video_max_memory=2048, + video_dir=None, + ) + cfg = video_cfg_from_args(args) + notes = summarize_video_recording(cfg, ()) + + assert cfg.enabled is True + assert cfg.record_failed is False + assert cfg.output_dir is None + assert "Video policy: disabled." not in notes + assert "videos=0" in notes + assert summarize_video_recording(VideoRecordCfg(), ()) == [ + "Video policy: disabled." + ] + with pytest.raises(ValueError, match="case_limit"): + VideoRecordCfg(enabled=True, case_limit=-1) + + +def test_runner_skips_video_outside_measured_phase(tmp_path): + suite = load_suite("atomic_franka_pgi_curobo") + specs = [spec for spec in suite.planners if spec.enabled] + runner = BenchmarkRunner( + suite, + specs, + device="cpu", + output_root=tmp_path, + video=VideoRecordCfg(enabled=True, record_failed=True), + ) + runner._run_dir = tmp_path + provider = Mock() + provider.record_replay.return_value = tmp_path / "should_not_write.mp4" + + path = runner._maybe_record_replay( + provider, + None, + _atomic_case(), + None, + "curobo", + TrialPhase.WARMUP, + ) + + assert path is None + provider.record_replay.assert_not_called() From f6122f1640e47cea694d1534aced9edbdcc859b3 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Sat, 22 Aug 2026 22:16:25 +0800 Subject: [PATCH 03/11] Add ONNX NMG planner and Atomic benchmark integration --- .../topics/motion-planning/motion-planning.md | 9 +- docs/source/overview/sim/planners/index.rst | 4 +- .../overview/sim/planners/neural_planner.md | 32 +- .../lab/sim/planners/neural_planner.py | 472 +++++++++--------- examples/sim/planners/neural_planner.py | 20 +- pyproject.toml | 3 + scripts/benchmark/motion_generation/config.py | 4 +- .../motion_generation/planners/__init__.py | 4 +- .../motion_generation/planners/neural.py | 55 -- .../motion_generation/planners/nmg_onnx.py | 117 +++++ .../motion_generation/run_benchmark.py | 19 +- .../suites/atomic_franka_pgi_curobo.yaml | 24 + .../motion_generation/suites/coverage.yaml | 14 +- .../motion_generation/suites/smoke.yaml | 14 +- tests/sim/planners/test_neural_batched.py | 24 +- tests/sim/planners/test_neural_planner.py | 198 +++++--- 16 files changed, 591 insertions(+), 422 deletions(-) delete mode 100644 scripts/benchmark/motion_generation/planners/neural.py create mode 100644 scripts/benchmark/motion_generation/planners/nmg_onnx.py diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index b5edd4a8b..877fdb9b4 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -9,7 +9,6 @@ | TOPPRA planner | `embodichain/lab/sim/planners/toppra_planner.py` → `ToppraPlanner`, `ToppraPlannerCfg`, `ToppraPlanOptions` | | Neural planner | `embodichain/lab/sim/planners/neural_planner.py` → `NeuralPlanner`, `NeuralPlannerCfg`, `NeuralPlanOptions` | | cuRobo planner | `embodichain/lab/sim/planners/curobo/curobo_planner.py` → `CuroboPlanner`, `CuroboPlannerCfg`, `CuroboWorldCfg`, `CuroboPlanOptions` | -| Planner assets | `embodichain/data/assets/planner_assets.py` → `download_neural_planner_checkpoint()` | | Motion generator | `embodichain/lab/sim/planners/motion_generator.py` → `MotionGenerator`, `MotionGenCfg`, `MotionGenOptions` | | Planner utilities & data types | `embodichain/lab/sim/planners/utils.py` → `PlanState`, `PlanResult`, `MoveType`, `MovePart`, `TrajectorySampleMethod`, `interpolate_xpos_batched` | @@ -40,7 +39,7 @@ Config hierarchy: ``` BasePlannerCfg robot_uid (MISSING), planner_type ├─ ToppraPlannerCfg planner_type = "toppra", max_workers, mp_context - └─ NeuralPlannerCfg planner_type = "neural", checkpoint_path (MISSING) + └─ NeuralPlannerCfg planner_type = "neural", onnx_model_path (MISSING) MotionGenCfg planner_cfg (MISSING — must be a BasePlannerCfg subclass) @@ -89,11 +88,11 @@ Worker details: Learning-based EEF waypoint planner. Franka Panda only. -- Checkpoint: `download_neural_planner_checkpoint()` from HuggingFace (gated, needs `HF_TOKEN`) +- Runtime: standalone `.onnx` policy with normalization embedded; install the `nmg` extra - Use via `MotionGenerator` with `planner_type="neural"` and `plan_opts=NeuralPlanOptions(...)` - Input: `EEF_MOVE` `PlanState` list with batched `xpos:(B, 4, 4)` -- Key cfg: `checkpoint_path` (from download), `control_part` -- Natively batched: transformer forward, reach checks, and convergence holds all operate on `(B, ...)`. +- Key cfg: `onnx_model_path`, `control_part`, `num_waypoints`, `policy_frame_from_world`, `runtime_tcp_from_policy_tcp` +- The current exported policy has fixed batch size one; dynamic-batch ONNX is not yet supported. ### CuroboPlanner collision worlds diff --git a/docs/source/overview/sim/planners/index.rst b/docs/source/overview/sim/planners/index.rst index 3ee955363..9ea901a09 100644 --- a/docs/source/overview/sim/planners/index.rst +++ b/docs/source/overview/sim/planners/index.rst @@ -27,8 +27,8 @@ The `embodichain` project provides a unified interface for robot trajectory plan These tools can be used to generate smooth and dynamically feasible robot trajectories. Install NVIDIA's CUDA-matched cuRobo source package separately when collision-aware planning against an explicit cuRobo world is required. -Use NeuralPlanner (experimental) when you have a trained APG checkpoint and need -learned EEF waypoint rollout on Franka Panda. +Use NeuralPlanner (experimental) when you have a standalone NMG ONNX policy and +need learned EEF waypoint rollout on Franka Panda. See also -------- diff --git a/docs/source/overview/sim/planners/neural_planner.md b/docs/source/overview/sim/planners/neural_planner.md index 632e2c488..7df57e4a9 100644 --- a/docs/source/overview/sim/planners/neural_planner.md +++ b/docs/source/overview/sim/planners/neural_planner.md @@ -3,21 +3,24 @@ ````{admonition} Experimental :class: warning -`NeuralPlanner` is an **experimental** feature. The API, checkpoint format, +`NeuralPlanner` is an **experimental** feature. The API, ONNX model contract, and default parameters may change without a deprecation cycle. It is currently only validated on the **Franka Panda** robot. ```` `NeuralPlanner` is a learning-based EEF waypoint planner. It rolls out a -trained APG checkpoint through `MotionGenerator` to reach Cartesian targets. +standalone NMG ONNX policy through `MotionGenerator` to reach Cartesian targets. +The ONNX graph must include raw-observation normalization. ## Configuration -Pre-trained checkpoints are hosted on HuggingFace and can be downloaded with -`download_neural_planner_checkpoint()` (requires `HF_TOKEN` environment variable). +Install the optional runtime and export the trained policy to ONNX before use: + +```bash +pip install -e '.[nmg]' +``` ```python -from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint from embodichain.lab.sim.planners import ( MotionGenCfg, MotionGenOptions, @@ -28,13 +31,13 @@ from embodichain.lab.sim.planners import ( ) from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions -checkpoint_path = download_neural_planner_checkpoint() +onnx_model_path = "/path/to/best_mean.onnx" motion_generator = MotionGenerator( cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid=robot.uid, - checkpoint_path=checkpoint_path, + onnx_model_path=onnx_model_path, control_part="main_arm", ) ) @@ -57,7 +60,18 @@ result = motion_generator.generate( ## Example ```bash -python examples/sim/planners/neural_planner.py --headless --device cuda +python examples/sim/planners/neural_planner.py \ + --headless --device cuda \ + --onnx-model-path /path/to/best_mean.onnx ``` -The example downloads the checkpoint automatically on first run. +The current exported policy has a fixed batch size of one. If the runtime robot +base frame or TCP differs from training, configure `policy_frame_from_world` and +`runtime_tcp_from_policy_tcp` as explicit homogeneous transforms. The conversion +is + +```text +policy_T_policy_tcp = policy_T_world + @ world_T_runtime_tcp + @ runtime_tcp_T_policy_tcp +``` diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index ed1083f3b..506767790 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -17,10 +17,8 @@ from dataclasses import MISSING from pathlib import Path - import numpy as np import torch -import torch.nn as nn from embodichain.lab.sim.planners.base_planner import ( BasePlanner, @@ -40,169 +38,99 @@ ] -def _safe_torch_load(path: Path, map_location: torch.device) -> dict: - """Load a PyTorch checkpoint with safe deserialization when possible. - - Attempts ``weights_only=True`` first. If that fails (e.g. on older PyTorch - versions or checkpoints with unsupported pickle objects), falls back to - ``weights_only=False`` and logs a warning. +class _OnnxPolicy: + """Small ONNX Runtime wrapper that keeps NMG model code out of EmbodiChain.""" - Args: - path: Path to the checkpoint file. - map_location: Device to map tensors to. + def __init__(self, path: Path, providers: list[str] | None = None) -> None: + try: + import onnxruntime as ort + except ImportError as exc: + raise ImportError( + "NeuralPlanner requires onnxruntime. Install EmbodiChain with " + "the 'nmg' optional dependency." + ) from exc - Returns: - The loaded checkpoint dictionary. - """ - try: - return torch.load(path, map_location=map_location, weights_only=True) - except (TypeError, RuntimeError, AttributeError) as exc: - logger.log_warning( - f"Failed to load checkpoint with weights_only=True from {path}: {exc}. " - "Falling back to weights_only=False." + options = ort.SessionOptions() + options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL + self.session = ort.InferenceSession( + str(path), sess_options=options, providers=providers ) - return torch.load(path, map_location=map_location, weights_only=False) - - -def _layer_init(layer: nn.Linear, std: float = np.sqrt(2), bias_const: float = 0.0): - torch.nn.init.orthogonal_(layer.weight, std) - torch.nn.init.constant_(layer.bias, bias_const) - return layer - - -class _RunningObsNormalizer: - def __init__(self, mean: torch.Tensor, var: torch.Tensor): - self.mean = mean - self.var = var - - def normalize(self, obs: torch.Tensor) -> torch.Tensor: - return (obs - self.mean) / (self.var.sqrt() + 1e-8) - - -class _WaypointTransformerActor(nn.Module): - """APG waypoint actor runtime copied in lightweight form for inference.""" - - def __init__( - self, - obs_dim: int, - action_dim: int, - num_waypoints: int, - use_relative_obs: bool = True, - hidden_dim: int = 256, - transformer_nhead: int = 4, - transformer_num_layers: int = 2, - transformer_ff_dim: int | None = None, - ): - super().__init__() - self.num_waypoints = int(num_waypoints) - self.use_relative_obs = bool(use_relative_obs) - if int(action_dim) != 7: + inputs = self.session.get_inputs() + outputs = self.session.get_outputs() + if len(inputs) != 1 or len(outputs) != 1: raise ValueError( - "Waypoint transformer checkpoints currently assume a 7-DoF arm. " - f"Got action_dim={action_dim}." + "NMG ONNX policy must expose exactly one input and one output; " + f"got {len(inputs)} inputs and {len(outputs)} outputs." ) - self.state_dim = 7 + 7 + 7 + (7 if self.use_relative_obs else 0) - self.waypoint_token_dim = 3 + 4 + 1 + 1 - - expected_obs_dim = 7 + 7 + self.num_waypoints * 3 - expected_obs_dim += self.num_waypoints * 4 + self.num_waypoints * 2 + 7 - if self.use_relative_obs: - expected_obs_dim += 7 - if int(obs_dim) != expected_obs_dim: + self.input_name = inputs[0].name + self.output_name = outputs[0].name + input_shape = inputs[0].shape + output_shape = outputs[0].shape + if len(input_shape) != 2 or not isinstance(input_shape[1], int): raise ValueError( - "Waypoint transformer expected obs_dim " - f"{expected_obs_dim}, got {obs_dim}." + f"Expected ONNX input shape [batch, obs], got {input_shape}." ) + if len(output_shape) != 2 or output_shape[1] != 7: + raise ValueError( + f"Expected ONNX output shape [batch, 7], got {output_shape}." + ) + self.obs_dim = int(input_shape[1]) + self.fixed_batch_size = ( + int(input_shape[0]) if isinstance(input_shape[0], int) else None + ) - if hidden_dim % transformer_nhead != 0: - raise ValueError("hidden_dim must be divisible by transformer_nhead") + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + if self.fixed_batch_size is not None and obs.shape[0] != self.fixed_batch_size: + raise ValueError( + f"ONNX policy has fixed batch size {self.fixed_batch_size}, " + f"but received {obs.shape[0]}." + ) + obs_np = np.ascontiguousarray(obs.detach().cpu().numpy(), dtype=np.float32) + action = self.session.run([self.output_name], {self.input_name: obs_np})[0] + return torch.as_tensor(action, dtype=torch.float32, device=obs.device) - ff_dim = transformer_ff_dim or hidden_dim * 4 - self.state_proj = _layer_init(nn.Linear(self.state_dim, hidden_dim)) - self.waypoint_proj = _layer_init(nn.Linear(self.waypoint_token_dim, hidden_dim)) - self.state_type_embedding = nn.Parameter(torch.zeros(1, 1, hidden_dim)) - self.waypoint_type_embedding = nn.Parameter(torch.zeros(1, 1, hidden_dim)) - self.waypoint_index_embedding = nn.Parameter( - torch.zeros(1, self.num_waypoints, hidden_dim) - ) - encoder_layer = nn.TransformerEncoderLayer( - d_model=hidden_dim, - nhead=transformer_nhead, - dim_feedforward=ff_dim, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=False, - ) - self.encoder = nn.TransformerEncoder( - encoder_layer, num_layers=transformer_num_layers - ) - self.actor_head = nn.Sequential( - nn.LayerNorm(hidden_dim), - _layer_init(nn.Linear(hidden_dim, hidden_dim)), - nn.Tanh(), - _layer_init(nn.Linear(hidden_dim, action_dim), std=0.01), - ) +def _waypoint_obs_dim(num_waypoints: int, use_relative_obs: bool) -> int: + """Return the unified NMG constraint-observation width.""" + n = int(num_waypoints) + dim = 7 + 7 + n * (3 + 4 + 7 + 5) + 7 + n + if use_relative_obs: + dim += 7 + n * (3 + 4 + 7) + return dim - def _parse_obs(self, x: torch.Tensor): - n = self.num_waypoints - cursor = 0 - joint = x[:, cursor : cursor + 7] - cursor += 7 - eef_pose = x[:, cursor : cursor + 7] - cursor += 7 - waypoint_pos = x[:, cursor : cursor + 3 * n].reshape(-1, n, 3) - cursor += 3 * n - waypoint_quat = x[:, cursor : cursor + 4 * n].reshape(-1, n, 4) - cursor += 4 * n - active_onehot = x[:, cursor : cursor + n].reshape(-1, n, 1) - cursor += n - valid_mask = x[:, cursor : cursor + n].reshape(-1, n, 1) - cursor += n - last_action = x[:, cursor : cursor + 7] - cursor += 7 - - state_parts = [joint, eef_pose, last_action] - if self.use_relative_obs: - state_parts.append(x[:, cursor : cursor + 7]) - - state = torch.cat(state_parts, dim=-1) - waypoint_tokens = torch.cat( - [waypoint_pos, waypoint_quat, active_onehot, valid_mask], dim=-1 - ) - return state, waypoint_tokens, valid_mask.squeeze(-1) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - state, waypoint_tokens, valid_mask = self._parse_obs(x) - state_token = self.state_proj(state).unsqueeze(1) + self.state_type_embedding - waypoint_tokens = ( - self.waypoint_proj(waypoint_tokens) - + self.waypoint_type_embedding - + self.waypoint_index_embedding - ) - tokens = torch.cat([state_token, waypoint_tokens], dim=1) - state_padding = torch.zeros( - valid_mask.shape[0], 1, dtype=torch.bool, device=valid_mask.device - ) - waypoint_padding = valid_mask < 0.5 - padding_mask = torch.cat([state_padding, waypoint_padding], dim=1) - encoded = self.encoder(tokens, src_key_padding_mask=padding_mask) - return self.actor_head(encoded[:, 0]) + +def _quat_inverse_xyzw(q: torch.Tensor) -> torch.Tensor: + """Return the inverse of a unit quaternion stored as xyzw.""" + return torch.cat([-q[..., :3], q[..., 3:4]], dim=-1) + + +def _quat_mul_xyzw(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Multiply quaternion tensors stored as xyzw.""" + ax, ay, az, aw = a.unbind(dim=-1) + bx, by, bz, bw = b.unbind(dim=-1) + return torch.stack( + [ + aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + aw * bw - ax * bx - ay * by - az * bz, + ], + dim=-1, + ) @configclass class NeuralPlannerCfg(BasePlannerCfg): planner_type: str = "neural" - checkpoint_path: str = MISSING - """Path to an APG waypoint checkpoint (.pt), e.g. from ``download_neural_planner_checkpoint()``.""" + onnx_model_path: str = MISSING + """Path to a standalone NMG ONNX policy, including observation normalization.""" control_part: str | None = None """Robot control part used for FK and qpos, e.g. 'left_arm'.""" - max_steps: int | None = None - """Maximum rollout steps. If None, uses checkpoint max_episode_steps.""" + max_steps: int = 240 + """Maximum closed-loop kinematic rollout steps.""" action_scale: float = 0.2 """Delta joint scaling factor in radians.""" @@ -210,11 +138,34 @@ class NeuralPlannerCfg(BasePlannerCfg): num_arm_joints: int = 7 """Number of arm joints controlled by the APG policy.""" - pos_eps: float | None = None - """Waypoint position threshold. If None, uses checkpoint waypoint_pos_threshold.""" + num_waypoints: int = 8 + """Number of constraint slots encoded by the ONNX policy.""" + + use_relative_obs: bool = True + """Whether the exported policy uses the unified relative-observation blocks.""" + + intermediate_orientation: bool = True + """Whether every Cartesian waypoint requires its orientation constraint.""" + + pos_eps: float = 0.01 + """Waypoint position threshold in metres.""" + + rot_eps: float = 0.1 + """Waypoint rotation threshold in radians.""" - rot_eps: float | None = None - """Waypoint rotation threshold. If None, uses checkpoint waypoint_rot_threshold.""" + onnx_providers: list[str] | None = None + """Optional ONNX Runtime execution-provider priority list.""" + + policy_frame_from_world: list[list[float]] | None = None + """Optional left transform ``policy_T_world`` for runtime poses.""" + + runtime_tcp_from_policy_tcp: list[list[float]] | None = None + """Optional right transform ``runtime_tcp_T_policy_tcp`` for runtime poses. + + Together, pose conversion is + ``policy_T_policy_tcp = policy_T_world @ world_T_runtime_tcp + @ runtime_tcp_T_policy_tcp``. + """ dt: float = 0.01 """Nominal timestep reported in PlanResult.""" @@ -230,8 +181,8 @@ class NeuralPlanOptions(PlanOptions): class NeuralPlanner(BasePlanner): r"""Neural motion planner based on an APG waypoint transformer policy. - The planner loads a checkpoint containing a waypoint-conditioned actor and - rolls it out in closed loop to drive the arm toward a sequence of + The planner loads a standalone ONNX waypoint policy and rolls it out in + closed loop to drive the arm toward a sequence of end-effector waypoints. Velocities and accelerations in the returned :class:`PlanResult` are estimated via finite differences over the generated position trajectory. @@ -240,9 +191,9 @@ class NeuralPlanner(BasePlanner): cfg: Configuration for the neural planner. Raises: - ValueError: If ``checkpoint_path`` is missing or invalid. - FileNotFoundError: If the checkpoint file does not exist. - KeyError: If the checkpoint is missing required keys. + ValueError: If ``onnx_model_path`` is missing or invalid. + FileNotFoundError: If the ONNX model file does not exist. + ImportError: If ONNX Runtime is not installed. """ supported_move_types = frozenset({MoveType.EEF_MOVE}) @@ -251,9 +202,9 @@ def __init__(self, cfg: NeuralPlannerCfg): super().__init__(cfg) self.cfg: NeuralPlannerCfg = cfg - if cfg.checkpoint_path is MISSING or not str(cfg.checkpoint_path): - logger.log_error("checkpoint_path is required", ValueError) - self._load_checkpoint(Path(cfg.checkpoint_path)) + if cfg.onnx_model_path is MISSING or not str(cfg.onnx_model_path): + logger.log_error("onnx_model_path is required", ValueError) + self._load_onnx_model(Path(cfg.onnx_model_path)) def default_plan_options(self) -> NeuralPlanOptions: return NeuralPlanOptions() @@ -274,92 +225,57 @@ def with_motion_context( options.start_qpos = start_qpos return options - def _load_checkpoint(self, checkpoint_path: Path) -> None: - if not checkpoint_path.exists(): - logger.log_error( - f"Checkpoint not found: {checkpoint_path}", FileNotFoundError - ) - - ckpt = _safe_torch_load(checkpoint_path, map_location=self.device) - if "agent" not in ckpt: - raise KeyError( - f"Checkpoint at '{checkpoint_path}' is missing 'agent'. " - f"Available keys: {list(ckpt.keys())}." - ) - if "obs_normalizer" not in ckpt: - raise KeyError( - f"Checkpoint at '{checkpoint_path}' is missing 'obs_normalizer'. " - f"Available keys: {list(ckpt.keys())}." + def _load_onnx_model(self, model_path: Path) -> None: + if not model_path.exists(): + logger.log_error(f"ONNX policy not found: {model_path}", FileNotFoundError) + if model_path.suffix.lower() != ".onnx": + raise ValueError( + "NeuralPlanner only accepts standalone .onnx policies; " + f"got {model_path}." ) - for subkey in ("mean", "var"): - if subkey not in ckpt["obs_normalizer"]: - raise KeyError( - f"Checkpoint obs_normalizer is missing '{subkey}'. " - f"Available: {list(ckpt['obs_normalizer'].keys())}." - ) - self._ckpt_args = ckpt.get("args", {}) - self._num_waypoints = int(self._ckpt_args.get("waypoint_max", 1)) - self._use_relative_obs = bool( - self._ckpt_args.get("waypoint_use_relative_obs", True) - ) - self._policy_arch = self._ckpt_args.get("policy_arch") - if self._policy_arch != "transformer": + self._num_waypoints = int(self.cfg.num_waypoints) + self._use_relative_obs = bool(self.cfg.use_relative_obs) + self._action_dim = int(self.cfg.num_arm_joints) + if self._action_dim != 7: raise ValueError( - "NeuralPlanner only supports transformer waypoint checkpoints. " - f"Got policy_arch={self._policy_arch!r}." + f"NMG ONNX policy controls 7 arm joints, got {self._action_dim}." ) - self._hidden_dim = int(self._ckpt_args.get("hidden_dim", 256)) - self._action_dim = int(self.cfg.num_arm_joints) - self._obs_dim = int(ckpt["obs_normalizer"]["mean"].numel()) - self._max_steps = int( - self.cfg.max_steps or self._ckpt_args.get("max_episode_steps", 30) - ) - self._pos_eps = float( - self.cfg.pos_eps - if self.cfg.pos_eps is not None - else self._ckpt_args.get("waypoint_pos_threshold", 0.05) + self._max_steps = int(self.cfg.max_steps) + self._pos_eps = float(self.cfg.pos_eps) + self._rot_eps = float(self.cfg.rot_eps) + self._intermediate_orientation = bool(self.cfg.intermediate_orientation) + self._policy = _OnnxPolicy(model_path, self.cfg.onnx_providers) + self._obs_dim = self._policy.obs_dim + self._policy_frame_from_world = self._as_transform( + self.cfg.policy_frame_from_world, "policy_frame_from_world" ) - self._rot_eps = float( - self.cfg.rot_eps - if self.cfg.rot_eps is not None - else self._ckpt_args.get("waypoint_rot_threshold", 0.3) + self._runtime_tcp_from_policy_tcp = self._as_transform( + self.cfg.runtime_tcp_from_policy_tcp, "runtime_tcp_from_policy_tcp" ) - self._intermediate_orientation = bool( - self._ckpt_args.get("waypoint_intermediate_orientation", True) + expected_obs_dim = _waypoint_obs_dim( + self._num_waypoints, self._use_relative_obs ) + if self._obs_dim != expected_obs_dim: + raise ValueError( + f"ONNX input has obs dim {self._obs_dim}, but the configured " + f"unified constraint layout requires {expected_obs_dim}." + ) - self._normalizer = _RunningObsNormalizer( - ckpt["obs_normalizer"]["mean"].to(self.device), - ckpt["obs_normalizer"]["var"].to(self.device), - ) - self._actor = self._build_actor().to(self.device) - - state_dict = { - k.replace("actor_mean.", ""): v.to(self.device) - for k, v in ckpt["agent"].items() - if k.startswith("actor_mean.") - } - if not state_dict: - raise KeyError("Checkpoint agent has no actor_mean.* weights.") - self._actor.load_state_dict(state_dict) - self._actor.eval() - - def _build_actor(self) -> nn.Module: - return _WaypointTransformerActor( - obs_dim=self._obs_dim, - action_dim=self._action_dim, - num_waypoints=self._num_waypoints, - use_relative_obs=self._use_relative_obs, - hidden_dim=self._hidden_dim, - transformer_nhead=int(self._ckpt_args.get("transformer_nhead", 4)), - transformer_num_layers=int( - self._ckpt_args.get("transformer_num_layers", 2) - ), - transformer_ff_dim=( - int(self._ckpt_args.get("transformer_ff_dim", 0)) or None - ), - ) + def _as_transform(self, value: list[list[float]] | None, name: str) -> torch.Tensor: + """Convert an optional homogeneous-transform config to a device tensor.""" + if value is None: + return torch.eye(4, dtype=torch.float32, device=self.device) + transform = torch.as_tensor(value, dtype=torch.float32, device=self.device) + if transform.shape != (4, 4): + raise ValueError(f"{name} must have shape (4, 4), got {transform.shape}.") + return transform + + def _to_policy_frame(self, xpos: torch.Tensor) -> torch.Tensor: + """Map runtime-world TCP poses to the NMG training base and TCP frame.""" + left = self._policy_frame_from_world.expand(xpos.shape[0], -1, -1) + right = self._runtime_tcp_from_policy_tcp.expand(xpos.shape[0], -1, -1) + return torch.bmm(torch.bmm(left, xpos), right) @validate_plan_options(options_cls=NeuralPlanOptions) def plan( @@ -432,7 +348,7 @@ def plan( active_idx, last_action, ) - action = self._actor(self._normalizer.normalize(obs)).clamp(-1.0, 1.0) + action = self._policy(obs).clamp(-1.0, 1.0) # Hold converged envs: zero their action so qpos does not drift. # `converged` reflects state up to the end of the previous step, so # once an env converged at the end of step N its action is masked @@ -494,7 +410,7 @@ def _parse_waypoints( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: if len(target_states) > self._num_waypoints: logger.log_error( - f"Received {len(target_states)} waypoints, but checkpoint supports " + f"Received {len(target_states)} waypoints, but the ONNX policy supports " f"at most {self._num_waypoints}.", ValueError, ) @@ -511,9 +427,10 @@ def _parse_waypoints( xpos = torch.as_tensor(target.xpos, dtype=torch.float32, device=self.device) if xpos.dim() == 2: xpos = xpos.unsqueeze(0) - waypoint_pos[:, idx] = xpos[:, :3, 3] + policy_xpos = self._to_policy_frame(xpos) + waypoint_pos[:, idx] = policy_xpos[:, :3, 3] waypoint_quat[:, idx] = convert_quat( - quat_from_matrix(xpos[:, :3, :3]), to="xyzw" + quat_from_matrix(policy_xpos[:, :3, :3]), to="xyzw" ) valid_mask[:, idx] = 1.0 return waypoint_pos, waypoint_quat, valid_mask, len(target_states) @@ -539,9 +456,9 @@ def _fk_matrix(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: return self.robot.compute_fk(qpos=qpos, name=control_part, to_matrix=True) def _fk_pose_xyzw(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: - fk = self.robot.compute_fk(qpos=qpos, name=control_part, to_matrix=False) - pos = fk[:, :3] - quat_xyzw = convert_quat(fk[:, 3:7], to="xyzw") + fk = self._to_policy_frame(self._fk_matrix(qpos, control_part)) + pos = fk[:, :3, 3] + quat_xyzw = convert_quat(quat_from_matrix(fk[:, :3, :3]), to="xyzw") return torch.cat([pos, quat_xyzw], dim=-1) def _build_obs( @@ -558,22 +475,85 @@ def _build_obs( active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) active_onehot = torch.zeros(b, self._num_waypoints, device=self.device) active_onehot.scatter_(1, active_idx_clamped.unsqueeze(1), 1.0) + pos_mask = valid_mask.clone() + rot_mask = valid_mask.clone() + if not self._intermediate_orientation: + waypoint_ids = torch.arange(self._num_waypoints, device=self.device) + episode_k = valid_mask.sum(dim=-1).long() + rot_mask = (waypoint_ids.unsqueeze(0) == (episode_k - 1).unsqueeze(1)).to( + valid_mask.dtype + ) + joint_mask = torch.zeros_like(valid_mask) + waypoint_joint = torch.zeros( + b, + self._num_waypoints, + self._action_dim, + dtype=joint_pos.dtype, + device=self.device, + ) + pos_block = waypoint_pos * pos_mask.unsqueeze(-1) + identity = torch.tensor( + [0.0, 0.0, 0.0, 1.0], + dtype=waypoint_quat.dtype, + device=self.device, + ) + quat_block = torch.where( + rot_mask.unsqueeze(-1) > 0.5, + waypoint_quat, + identity.view(1, 1, 4), + ) obs_parts = [ joint_pos, ee_pose, - waypoint_pos.reshape(b, self._num_waypoints * 3), - waypoint_quat.reshape(b, self._num_waypoints * 4), + pos_block.reshape(b, self._num_waypoints * 3), + quat_block.reshape(b, self._num_waypoints * 4), + waypoint_joint.reshape(b, self._num_waypoints * self._action_dim), active_onehot, valid_mask, + pos_mask, + rot_mask, + joint_mask, last_action, ] if self._use_relative_obs: idx = torch.arange(b, device=self.device) - active_pos = waypoint_pos[idx, active_idx_clamped] - active_quat = waypoint_quat[idx, active_idx_clamped] + active_pos = pos_block[idx, active_idx_clamped] + active_quat = quat_block[idx, active_idx_clamped] + inv_eef = _quat_inverse_xyzw(ee_pose[:, 3:7]) + active_rel_quat = _quat_mul_xyzw(active_quat, inv_eef) obs_parts.append( - torch.cat([active_pos - ee_pose[:, :3], active_quat], dim=-1) + torch.cat([active_pos - ee_pose[:, :3], active_rel_quat], dim=-1) + ) + rel_pos = (pos_block - ee_pose[:, None, :3]) * pos_mask.unsqueeze(-1) + rel_quat = _quat_mul_xyzw( + quat_block, + inv_eef[:, None, :].expand_as(quat_block), + ) + rel_quat = torch.where( + rot_mask.unsqueeze(-1) > 0.5, + rel_quat, + identity.view(1, 1, 4), ) + joint_err = torch.zeros_like(waypoint_joint) + obs_parts.extend( + [ + rel_pos.reshape(b, self._num_waypoints * 3), + rel_quat.reshape(b, self._num_waypoints * 4), + joint_err.reshape(b, self._num_waypoints * self._action_dim), + ] + ) + waypoint_type = torch.zeros( + b, + self._num_waypoints, + dtype=joint_pos.dtype, + device=self.device, + ) + waypoint_type = torch.where( + (valid_mask > 0.5) & (rot_mask < 0.5), + torch.ones_like(waypoint_type), + waypoint_type, + ) + obs_parts.append(waypoint_type) obs = torch.cat(obs_parts, dim=-1) if obs.shape[-1] != self._obs_dim: raise ValueError( diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/planners/neural_planner.py index 9ce5a18db..e326933cd 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/planners/neural_planner.py @@ -14,12 +14,12 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run the env-batched NeuralPlanner waypoint example. +"""Run the ONNX NeuralPlanner waypoint example. From the repository root:: - python examples/sim/planners/neural_planner.py --headless - python examples/sim/planners/neural_planner.py --headless --device cuda:1 + python examples/sim/planners/neural_planner.py --headless \ + --onnx-model-path /path/to/policy.onnx """ from __future__ import annotations @@ -30,7 +30,6 @@ import numpy as np import torch -from embodichain.data.assets.planner_assets import download_neural_planner_checkpoint from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args @@ -53,6 +52,11 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="NeuralPlanner waypoint example") add_env_launcher_args_to_parser(parser) parser.set_defaults(device=default_device, arena_space=2.0) + parser.add_argument( + "--onnx-model-path", + required=True, + help="Path to a standalone NMG ONNX policy.", + ) parser.add_argument( "--num-waypoints", type=int, @@ -188,15 +192,15 @@ def play_trajectory( def main() -> None: args = parse_args() - if args.num_envs < 1: - raise ValueError("--num_envs must be at least 1.") + if args.num_envs != 1: + raise ValueError("The current exported NMG ONNX policy requires --num-envs 1.") if args.num_waypoints < 1: raise ValueError("--num-waypoints must be at least 1.") if args.step_repeat < 1: raise ValueError("--step-repeat must be at least 1.") if args.hold_steps < 0: raise ValueError("--hold-steps must be non-negative.") - checkpoint_path = download_neural_planner_checkpoint() + onnx_model_path = args.onnx_model_path sim_device = _resolve_device(args.device, args.gpu_id) resolved_device = torch.device(sim_device) @@ -246,7 +250,7 @@ def main() -> None: cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid=robot.uid, - checkpoint_path=checkpoint_path, + onnx_model_path=onnx_model_path, control_part=arm_name, ) ) diff --git a/pyproject.toml b/pyproject.toml index ae17152e6..d3d5ececb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,9 @@ dependencies = [ ] [project.optional-dependencies] +nmg = [ + "onnxruntime>=1.20", +] gensim = [ "bpy", "pyrender==0.1.45", diff --git a/scripts/benchmark/motion_generation/config.py b/scripts/benchmark/motion_generation/config.py index a58321445..b5f99dcdd 100644 --- a/scripts/benchmark/motion_generation/config.py +++ b/scripts/benchmark/motion_generation/config.py @@ -213,9 +213,9 @@ def validate_benchmark(self) -> None: _validate_free_space(self.free_space) nmg = next((spec for spec in self.planners if spec.id == "nmg"), None) if nmg is not None: - if float(nmg.config.get("pos_eps", 0.05)) <= 0.0: + if float(nmg.config.get("pos_eps", 0.01)) <= 0.0: raise ValueError("NMG pos_eps must be > 0.") - if float(nmg.config.get("rot_eps", 0.3)) <= 0.0: + if float(nmg.config.get("rot_eps", 0.1)) <= 0.0: raise ValueError("NMG rot_eps must be > 0.") diff --git a/scripts/benchmark/motion_generation/planners/__init__.py b/scripts/benchmark/motion_generation/planners/__init__.py index 65b3edabc..27722579f 100644 --- a/scripts/benchmark/motion_generation/planners/__init__.py +++ b/scripts/benchmark/motion_generation/planners/__init__.py @@ -21,13 +21,13 @@ from .base import PlannerAdapter, PlannerContext from .curobo import CuroboAdapter from .ik_interpolate import IkInterpolateAdapter -from .neural import NeuralAdapterStub +from .nmg_onnx import NmgOnnxAdapter from .toppra import ToppraAdapter __all__ = [ "CuroboAdapter", "IkInterpolateAdapter", - "NeuralAdapterStub", + "NmgOnnxAdapter", "PlannerAdapter", "PlannerContext", "ToppraAdapter", diff --git a/scripts/benchmark/motion_generation/planners/neural.py b/scripts/benchmark/motion_generation/planners/neural.py deleted file mode 100644 index 20af60396..000000000 --- a/scripts/benchmark/motion_generation/planners/neural.py +++ /dev/null @@ -1,55 +0,0 @@ -# ---------------------------------------------------------------------------- -# 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. -# ---------------------------------------------------------------------------- - -"""Capability-aware NMG adapter stub reserved for a future checkpoint.""" - -from __future__ import annotations - -from embodichain.lab.sim.planners import PlanResult - -from ..models import BenchmarkCase -from ..registry import register_planner_adapter -from .base import PlannerAdapter - -__all__ = ["NeuralAdapterStub"] - - -class NeuralAdapterStub(PlannerAdapter): - """Expose configurable NMG precision without initializing an unavailable model.""" - - capabilities = frozenset({"eef_waypoint", "batched", "empty_world"}) - model_revision = "not-ready" - - def availability(self) -> tuple[bool, str | None]: - """Mark the placeholder unsupported until the checkpoint contract lands.""" - pos_eps = float(self.spec.config.get("pos_eps", 0.05)) - rot_eps = float(self.spec.config.get("rot_eps", 0.3)) - return ( - False, - "NMG adapter is a stub pending the production checkpoint; " - f"configured pos_eps={pos_eps} m, rot_eps={rot_eps} rad.", - ) - - def build(self) -> None: - """Reject accidental construction of the explicit placeholder.""" - raise RuntimeError("The NMG adapter is not implemented yet.") - - def plan(self, case: BenchmarkCase) -> PlanResult: # noqa: ARG002 - """Reject accidental execution of the explicit placeholder.""" - raise RuntimeError("The NMG adapter is not implemented yet.") - - -register_planner_adapter("neural_stub", NeuralAdapterStub) diff --git a/scripts/benchmark/motion_generation/planners/nmg_onnx.py b/scripts/benchmark/motion_generation/planners/nmg_onnx.py new file mode 100644 index 000000000..c4d263a4d --- /dev/null +++ b/scripts/benchmark/motion_generation/planners/nmg_onnx.py @@ -0,0 +1,117 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""NMG ONNX benchmark adapter.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +from embodichain.lab.sim.planners import ( + MotionGenCfg, + MotionGenOptions, + MotionGenerator, + NeuralPlanOptions, + NeuralPlannerCfg, + PlanResult, + PlanState, +) + +from ..config import PlannerSpecCfg +from ..models import BenchmarkCase +from ..registry import register_planner_adapter +from .base import PlannerAdapter, PlannerContext + +__all__ = ["NmgOnnxAdapter"] + + +class NmgOnnxAdapter(PlannerAdapter): + """Run a standalone NMG ONNX policy through :class:`NeuralPlanner`.""" + + capabilities = frozenset({"eef_waypoint", "empty_world", "atomic_action"}) + model_revision = "nmg-onnx-v1" + + def __init__(self, spec: PlannerSpecCfg, context: PlannerContext) -> None: + super().__init__(spec, context) + self.motion_generator: MotionGenerator | None = None + + def availability(self) -> tuple[bool, str | None]: + """Require ONNX Runtime and a resolved standalone policy path.""" + if importlib.util.find_spec("onnxruntime") is None: + return False, "onnxruntime is not installed." + model_path = self.spec.config.get("onnx_model_path") + if not model_path: + return False, "NMG requires config.onnx_model_path or --nmg-onnx-path." + path = Path(str(model_path)).expanduser() + if not path.is_file(): + return False, f"NMG ONNX policy does not exist: {path}." + if path.suffix.lower() != ".onnx": + return False, f"NMG requires a standalone .onnx policy, got {path}." + return True, None + + @property + def motion_policy_planner(self) -> str: + """Select EmbodiChain's neural MotionGenerator backend.""" + return "neural" + + def build(self) -> None: + """Construct the ONNX-backed MotionGenerator.""" + values = self.spec.config + model_path = str(Path(str(values["onnx_model_path"])).expanduser()) + providers = values.get("onnx_providers") + planner_cfg = NeuralPlannerCfg( + robot_uid=self.context.robot.uid, + onnx_model_path=model_path, + control_part=self.context.control_part, + max_steps=int(values.get("max_steps", 240)), + action_scale=float(values.get("action_scale", 0.2)), + num_arm_joints=int(values.get("num_arm_joints", 7)), + num_waypoints=int(values.get("num_waypoints", 8)), + use_relative_obs=bool(values.get("use_relative_obs", True)), + intermediate_orientation=bool(values.get("intermediate_orientation", True)), + pos_eps=float(values.get("pos_eps", 0.01)), + rot_eps=float(values.get("rot_eps", 0.1)), + onnx_providers=list(providers) if providers is not None else None, + policy_frame_from_world=values.get("policy_frame_from_world"), + runtime_tcp_from_policy_tcp=values.get("runtime_tcp_from_policy_tcp"), + dt=float(values.get("dt", 0.01)), + ) + self.motion_generator = MotionGenerator(MotionGenCfg(planner_cfg=planner_cfg)) + + def plan(self, case: BenchmarkCase) -> PlanResult: + """Plan all ordered Cartesian constraints in one closed-loop rollout.""" + if self.motion_generator is None: + raise RuntimeError("NMG adapter must be built before plan().") + targets = [ + PlanState.from_xpos(case.target_waypoints[:, index]) + for index in range(case.num_waypoints) + ] + return self.motion_generator.generate( + targets, + MotionGenOptions( + start_qpos=case.start_qpos, + control_part=self.context.control_part, + plan_opts=NeuralPlanOptions(), + ), + ) + + def close(self) -> None: + """Release the adapter-owned MotionGenerator.""" + self._close_motion_generator() + + +register_planner_adapter("nmg_onnx", NmgOnnxAdapter) diff --git a/scripts/benchmark/motion_generation/run_benchmark.py b/scripts/benchmark/motion_generation/run_benchmark.py index 50eb8fac8..49003d8a8 100644 --- a/scripts/benchmark/motion_generation/run_benchmark.py +++ b/scripts/benchmark/motion_generation/run_benchmark.py @@ -17,8 +17,7 @@ """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. +optional diagnostic baselines. NMG is loaded from a standalone ONNX policy. Run: ``python -m scripts.benchmark.motion_generation.run_benchmark --suite smoke`` or select the Franka + PGI Atomic Task slice with @@ -98,9 +97,9 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: help="NMG internal waypoint rotation threshold in radians.", ) parser.add_argument( - "--checkpoint-path", + "--nmg-onnx-path", default=None, - help="Reserved NMG checkpoint path; the current NMG adapter remains a stub.", + help="Path to the standalone NMG ONNX policy.", ) parser.add_argument( "--output-root", default="outputs/benchmarks", help="Artifact root directory." @@ -190,7 +189,7 @@ def _apply_overrides( rotation_threshold_rad: float | None = None, nmg_pos_eps: float | None = None, nmg_rot_eps: float | None = None, - checkpoint_path: str | None = None, + nmg_onnx_path: str | None = None, ) -> None: """Apply explicit CLI/programmatic overrides to a loaded suite.""" if batch_sizes is not None: @@ -222,8 +221,8 @@ def _apply_overrides( nmg.config["pos_eps"] = nmg_pos_eps if nmg_rot_eps is not None: nmg.config["rot_eps"] = nmg_rot_eps - if checkpoint_path is not None: - nmg.config["checkpoint_path"] = str(Path(checkpoint_path)) + if nmg_onnx_path is not None: + nmg.config["onnx_model_path"] = str(Path(nmg_onnx_path)) suite.validate_benchmark() @@ -231,7 +230,7 @@ def run_all_benchmarks( num_waypoints_list: list[int] | None = None, sim_device: str = "auto", headless: bool = True, - checkpoint_path: str | None = None, + nmg_onnx_path: str | None = None, *, suite_name: str = "smoke", algorithms: list[str] | None = None, @@ -270,7 +269,7 @@ def run_all_benchmarks( rotation_threshold_rad=rotation_threshold_rad, nmg_pos_eps=nmg_pos_eps, nmg_rot_eps=nmg_rot_eps, - checkpoint_path=checkpoint_path, + nmg_onnx_path=nmg_onnx_path, ) specs = _resolve_planners(suite, algorithms, list(extra_baselines or [])) return BenchmarkRunner( @@ -289,7 +288,7 @@ def run_from_args(args: argparse.Namespace) -> BenchmarkRunResult: num_waypoints_list=args.num_waypoints, sim_device=args.device, headless=args.headless, - checkpoint_path=args.checkpoint_path, + nmg_onnx_path=args.nmg_onnx_path, suite_name=args.suite, algorithms=args.algorithms, extra_baselines=args.extra_baselines, diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml index 2bfd23ef4..35ce3e94f 100644 --- a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml @@ -9,6 +9,30 @@ robot: config: {} planners: + - id: nmg + adapter: nmg_onnx + role: candidate + enabled: false + config: + model_revision: nmg-cartesian-300m + onnx_model_path: null + num_waypoints: 8 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 + policy_frame_from_world: + - [-1.0, 0.0, 0.0, 0.0] + - [0.0, -1.0, 0.0, 0.0] + - [0.0, 0.0, 1.0, 0.0] + - [0.0, 0.0, 0.0, 1.0] + runtime_tcp_from_policy_tcp: + - [0.70710678, 0.70710678, 0.0, 0.0] + - [-0.70710678, 0.70710678, 0.0, 0.0] + - [0.0, 0.0, 1.0, -0.0466] + - [0.0, 0.0, 0.0, 1.0] - id: curobo adapter: curobo role: primary_baseline diff --git a/scripts/benchmark/motion_generation/suites/coverage.yaml b/scripts/benchmark/motion_generation/suites/coverage.yaml index ebcd8faa9..f29b77940 100644 --- a/scripts/benchmark/motion_generation/suites/coverage.yaml +++ b/scripts/benchmark/motion_generation/suites/coverage.yaml @@ -37,13 +37,19 @@ planners: velocity: 0.2 acceleration: 0.5 - id: nmg - adapter: neural_stub + adapter: nmg_onnx role: candidate enabled: false config: - model_revision: not-ready - pos_eps: 0.05 - rot_eps: 0.3 + model_revision: nmg-onnx-v1 + onnx_model_path: null + num_waypoints: 8 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 protocol: warmup_trials: 3 diff --git a/scripts/benchmark/motion_generation/suites/smoke.yaml b/scripts/benchmark/motion_generation/suites/smoke.yaml index eafbbc5a6..2a29fe2b0 100644 --- a/scripts/benchmark/motion_generation/suites/smoke.yaml +++ b/scripts/benchmark/motion_generation/suites/smoke.yaml @@ -37,13 +37,19 @@ planners: velocity: 0.2 acceleration: 0.5 - id: nmg - adapter: neural_stub + adapter: nmg_onnx role: candidate enabled: false config: - model_revision: not-ready - pos_eps: 0.05 - rot_eps: 0.3 + model_revision: nmg-onnx-v1 + onnx_model_path: null + num_waypoints: 8 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 protocol: warmup_trials: 1 diff --git a/tests/sim/planners/test_neural_batched.py b/tests/sim/planners/test_neural_batched.py index 78055af88..c1e2acc73 100644 --- a/tests/sim/planners/test_neural_batched.py +++ b/tests/sim/planners/test_neural_batched.py @@ -16,7 +16,11 @@ from __future__ import annotations import torch -import pytest + + +def _set_identity_policy_frames(planner) -> None: + planner._policy_frame_from_world = torch.eye(4) + planner._runtime_tcp_from_policy_tcp = torch.eye(4) class TestNeuralParseWaypoints: @@ -28,6 +32,7 @@ def test_parse_waypoints_batched(self): planner = NeuralPlanner.__new__(NeuralPlanner) planner.device = torch.device("cpu") planner._num_waypoints = 4 + _set_identity_policy_frames(planner) B = 3 states = [ PlanState.from_xpos( @@ -60,7 +65,8 @@ def test_plan_returns_batched_success(self, monkeypatch): planner._rot_eps = 1e9 planner._intermediate_orientation = True planner._use_relative_obs = False - planner._obs_dim = 57 # 7+7+4*3+4*4+4+4+7 + planner._obs_dim = 101 + _set_identity_policy_frames(planner) planner.cfg = type( "c", (), @@ -72,9 +78,8 @@ def test_plan_returns_batched_success(self, monkeypatch): }, )() - # stub actor: returns zeros so qpos never changes but eps is huge -> reached - planner._actor = lambda obs: torch.zeros(obs.shape[0], 7) - planner._normalizer = type("n", (), {"normalize": lambda self, o: o})() + # Stub ONNX policy: qpos stays fixed but the huge eps marks each goal reached. + planner._policy = lambda obs: torch.zeros(obs.shape[0], 7) # stub robot FK + limits class _Robot: @@ -127,8 +132,8 @@ def test_converged_env_holds_qpos(self): planner._rot_eps = 1e9 planner._intermediate_orientation = True planner._use_relative_obs = False - # joint(7)+ee(7)+wp_pos(4*3)+wp_quat(4*4)+onehot(4)+valid(4)+last_act(7)=57 - planner._obs_dim = 57 + planner._obs_dim = 101 + _set_identity_policy_frames(planner) planner.cfg = type( "c", (), @@ -140,9 +145,8 @@ def test_converged_env_holds_qpos(self): }, )() - # actor: non-trivial action so qpos would drift if not masked - planner._actor = lambda obs: torch.ones(obs.shape[0], 7) * 0.5 - planner._normalizer = type("n", (), {"normalize": lambda self, o: o})() + # Non-trivial policy action so qpos would drift if not masked. + planner._policy = lambda obs: torch.ones(obs.shape[0], 7) * 0.5 class _Robot: num_instances = 3 diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index c1e2d2786..998c5f46e 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -27,52 +27,38 @@ NeuralPlannerCfg, PlanState, ) -from embodichain.lab.sim.planners.neural_planner import ( - NeuralPlanOptions, - _WaypointTransformerActor, -) +from embodichain.lab.sim.planners.neural_planner import NeuralPlanOptions +from embodichain.lab.sim.planners import neural_planner as neural_planner_module from embodichain.lab.sim.sim_manager import SimulationManager NUM_ARM_JOINTS = 7 -NUM_WAYPOINTS = 3 -OBS_DIM = 28 + 9 * NUM_WAYPOINTS -HIDDEN_DIM = 32 - - -def _create_fake_checkpoint(tmp_path) -> str: - actor = _WaypointTransformerActor( - obs_dim=OBS_DIM, - action_dim=NUM_ARM_JOINTS, - num_waypoints=NUM_WAYPOINTS, - use_relative_obs=True, - hidden_dim=HIDDEN_DIM, - transformer_nhead=4, - transformer_num_layers=1, - ) - checkpoint = { - "agent": {f"actor_mean.{k}": v for k, v in actor.state_dict().items()}, - "obs_normalizer": { - "mean": torch.zeros(OBS_DIM), - "var": torch.ones(OBS_DIM), - "count": 1.0, - }, - "args": { - "policy_arch": "transformer", - "hidden_dim": HIDDEN_DIM, - "transformer_nhead": 4, - "transformer_num_layers": 1, - "transformer_ff_dim": 0, - "waypoint_max": NUM_WAYPOINTS, - "waypoint_use_relative_obs": True, - "waypoint_intermediate_orientation": True, - "max_episode_steps": 3, - "waypoint_pos_threshold": 0.05, - "waypoint_rot_threshold": 0.3, - }, - } - checkpoint_path = tmp_path / "fake_neural_planner.pt" - torch.save(checkpoint, checkpoint_path) - return str(checkpoint_path) +NUM_WAYPOINTS = 8 +OBS_DIM = 300 + + +def _create_fake_onnx_model(tmp_path) -> str: + model_path = tmp_path / "fake_neural_planner.onnx" + model_path.write_bytes(b"fake-onnx") + return str(model_path) + + +class FakeOnnxPolicy: + obs_dim = OBS_DIM + fixed_batch_size = 1 + + def __init__(self, path, providers=None): + self.path = path + self.providers = providers + self.last_obs = None + + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + self.last_obs = obs.clone() + return torch.zeros(obs.shape[0], NUM_ARM_JOINTS, device=obs.device) + + +@pytest.fixture(autouse=True) +def _mock_onnx_policy(monkeypatch): + monkeypatch.setattr(neural_planner_module, "_OnnxPolicy", FakeOnnxPolicy) class FakeRobot: @@ -115,8 +101,8 @@ def test_neural_planner_is_registered(): assert MotionGenerator._support_planner_dict["neural"][1] is NeuralPlannerCfg -def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) +def test_neural_planner_generate_with_fake_onnx_model(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -126,7 +112,7 @@ def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -152,7 +138,7 @@ def test_neural_planner_generate_with_fake_checkpoint(tmp_path, monkeypatch): def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -162,7 +148,7 @@ def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -185,7 +171,7 @@ def test_neural_planner_uses_plan_opts_start_qpos(tmp_path, monkeypatch): def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -195,7 +181,7 @@ def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -216,7 +202,7 @@ def test_neural_planner_rejects_short_start_qpos(tmp_path, monkeypatch): def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -226,7 +212,7 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -252,6 +238,83 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat assert torch.isfinite(result.accelerations).all() +def test_neural_planner_builds_unified_300d_cartesian_observation( + tmp_path, monkeypatch +): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + joint = torch.zeros(1, 7) + eef = torch.tensor([[0.1, 0.2, 0.3, 0.0, 0.0, 0.0, 1.0]]) + waypoint_pos = torch.zeros(1, 8, 3) + waypoint_quat = torch.zeros(1, 8, 4) + waypoint_quat[..., 3] = 1.0 + valid = torch.zeros(1, 8) + valid[:, :2] = 1.0 + obs = planner._build_obs( + joint, + eef, + waypoint_pos, + waypoint_quat, + valid, + torch.zeros(1, dtype=torch.long), + torch.zeros(1, 7), + ) + + assert obs.shape == (1, 300) + # Unified layout semantic blocks: active, valid, pos, rot, joint masks. + semantic_start = 7 + 7 + 8 * (3 + 4 + 7) + expected_active = torch.zeros_like(valid) + expected_active[:, 0] = 1.0 + assert torch.equal(obs[:, semantic_start : semantic_start + 8], expected_active) + assert torch.equal(obs[:, semantic_start + 8 : semantic_start + 16], valid) + assert torch.equal(obs[:, semantic_start + 16 : semantic_start + 24], valid) + assert torch.equal(obs[:, semantic_start + 24 : semantic_start + 32], valid) + assert torch.count_nonzero(obs[:, semantic_start + 32 : semantic_start + 40]) == 0 + + +def test_neural_planner_applies_policy_frame_and_tcp_transforms(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + left = [ + [-1.0, 0.0, 0.0, 0.0], + [0.0, -1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + right = [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, -0.0466], + [0.0, 0.0, 0.0, 1.0], + ] + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + policy_frame_from_world=left, + runtime_tcp_from_policy_tcp=right, + ) + ) + pose = torch.eye(4).unsqueeze(0) + pose[:, :3, 3] = torch.tensor([[-0.5, 0.1, 0.4]]) + transformed = planner._to_policy_frame(pose) + + assert torch.allclose(transformed[0, :3, 3], torch.tensor([0.5, -0.1, 0.3534])) + + def test_neural_planner_finite_diff_helper(): """Finite-difference estimates match a known polynomial trajectory.""" b, n, dof = 2, 5, 7 @@ -282,7 +345,7 @@ def test_neural_planner_finite_diff_helper(): def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -292,7 +355,7 @@ def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeyp cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, ) ) ) @@ -312,7 +375,7 @@ def test_motion_generator_neural_propagates_motion_gen_options(tmp_path, monkeyp def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypatch): - checkpoint_path = _create_fake_checkpoint(tmp_path) + model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() monkeypatch.setattr( SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) @@ -322,7 +385,7 @@ def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypa cfg=MotionGenCfg( planner_cfg=NeuralPlannerCfg( robot_uid="fake_robot", - checkpoint_path=checkpoint_path, + onnx_model_path=model_path, control_part="main_arm", ) ) @@ -344,13 +407,18 @@ def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypa assert options.is_interpolate is True -def test_safe_torch_load_roundtrip(tmp_path): - checkpoint = {"agent": torch.tensor([1.0, 2.0, 3.0])} - path = tmp_path / "checkpoint.pt" - torch.save(checkpoint, path) - - from embodichain.lab.sim.planners.neural_planner import _safe_torch_load +def test_neural_planner_rejects_pytorch_checkpoint(tmp_path, monkeypatch): + pytorch_checkpoint_path = tmp_path / "checkpoint.pt" + pytorch_checkpoint_path.write_bytes(b"not-an-onnx-model") + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) - loaded = _safe_torch_load(path, map_location=torch.device("cpu")) - assert "agent" in loaded - assert torch.equal(loaded["agent"], checkpoint["agent"]) + with pytest.raises(ValueError, match="only accepts standalone .onnx"): + NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=str(pytorch_checkpoint_path), + ) + ) From 1a7658bed45d0774b23ba8c7fda12d98d91e14ba Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Sat, 22 Aug 2026 22:37:53 +0800 Subject: [PATCH 04/11] Enable batched NMG ONNX evaluation --- .../topics/motion-planning/motion-planning.md | 2 +- .../overview/sim/planners/neural_planner.md | 8 ++--- .../lab/sim/planners/neural_planner.py | 1 + scripts/benchmark/motion_generation/README.md | 3 +- .../motion_generation/planners/nmg_onnx.py | 4 ++- tests/sim/planners/test_neural_planner.py | 35 +++++++++++++++++++ 6 files changed, 45 insertions(+), 8 deletions(-) diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index 877fdb9b4..62d9af7aa 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -92,7 +92,7 @@ Learning-based EEF waypoint planner. Franka Panda only. - Use via `MotionGenerator` with `planner_type="neural"` and `plan_opts=NeuralPlanOptions(...)` - Input: `EEF_MOVE` `PlanState` list with batched `xpos:(B, 4, 4)` - Key cfg: `onnx_model_path`, `control_part`, `num_waypoints`, `policy_frame_from_world`, `runtime_tcp_from_policy_tcp` -- The current exported policy has fixed batch size one; dynamic-batch ONNX is not yet supported. +- The NMG exporter produces a dynamic-batch ONNX policy; NeuralPlanner rolls out all environments together. ### CuroboPlanner collision worlds diff --git a/docs/source/overview/sim/planners/neural_planner.md b/docs/source/overview/sim/planners/neural_planner.md index 7df57e4a9..03041102f 100644 --- a/docs/source/overview/sim/planners/neural_planner.md +++ b/docs/source/overview/sim/planners/neural_planner.md @@ -65,10 +65,10 @@ python examples/sim/planners/neural_planner.py \ --onnx-model-path /path/to/best_mean.onnx ``` -The current exported policy has a fixed batch size of one. If the runtime robot -base frame or TCP differs from training, configure `policy_frame_from_world` and -`runtime_tcp_from_policy_tcp` as explicit homogeneous transforms. The conversion -is +The NMG exporter produces a dynamic-batch ONNX policy, so one graph can serve +single-env and env-batched rollout. If the runtime robot base frame or TCP differs +from training, configure `policy_frame_from_world` and +`runtime_tcp_from_policy_tcp` as explicit homogeneous transforms. The conversion is ```text policy_T_policy_tcp = policy_T_world diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index 506767790..482cf79ae 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -278,6 +278,7 @@ def _to_policy_frame(self, xpos: torch.Tensor) -> torch.Tensor: return torch.bmm(torch.bmm(left, xpos), right) @validate_plan_options(options_cls=NeuralPlanOptions) + @torch.no_grad() def plan( self, target_states: list[PlanState], diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index 77b87224a..3f17bf320 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -32,7 +32,7 @@ in that run's `videos/` directory. - 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) +- Direct, batched NMG ONNX adapter (`candidate`, enabled when a model path is supplied) - Lifecycle timing: construct / prepare / cold / warm - Distinct planning, motion-valid, execution, and physical task-success stages - Planning latency, execution wall time, end-to-end time, nominal trajectory @@ -55,7 +55,6 @@ in that run's `videos/` directory. ## Current limits -- Real NMG checkpoint adapter - `collision-deployment` and obstacle-aware common-input tracks - Atomic Task execution is currently `B=1`; the supplied suite covers only Franka + PGI and cuRobo diff --git a/scripts/benchmark/motion_generation/planners/nmg_onnx.py b/scripts/benchmark/motion_generation/planners/nmg_onnx.py index c4d263a4d..69e45a5a1 100644 --- a/scripts/benchmark/motion_generation/planners/nmg_onnx.py +++ b/scripts/benchmark/motion_generation/planners/nmg_onnx.py @@ -42,7 +42,9 @@ class NmgOnnxAdapter(PlannerAdapter): """Run a standalone NMG ONNX policy through :class:`NeuralPlanner`.""" - capabilities = frozenset({"eef_waypoint", "empty_world", "atomic_action"}) + capabilities = frozenset( + {"eef_waypoint", "batched", "empty_world", "atomic_action"} + ) model_revision = "nmg-onnx-v1" def __init__(self, spec: PlannerSpecCfg, context: PlannerContext) -> None: diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index 998c5f46e..4d4553e44 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -238,6 +238,41 @@ def test_neural_planner_returns_velocities_and_accelerations(tmp_path, monkeypat assert torch.isfinite(result.accelerations).all() +def test_neural_planner_disables_grad_for_all_fk_calls(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + grad_states = [] + original_compute_fk = fake_sim.robot.compute_fk + + def checked_compute_fk(*args, **kwargs): + grad_states.append(torch.is_grad_enabled()) + return original_compute_fk(*args, **kwargs) + + monkeypatch.setattr(fake_sim.robot, "compute_fk", checked_compute_fk) + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + + planner.plan( + [PlanState.single(move_type=MoveType.EEF_MOVE, xpos=torch.eye(4))], + NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + max_steps=1, + ), + ) + + assert grad_states + assert not any(grad_states) + + def test_neural_planner_builds_unified_300d_cartesian_observation( tmp_path, monkeypatch ): From a1422c1bf010fc9eb4a1904dbac00b24d90ebed1 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Thu, 27 Aug 2026 10:48:53 +0800 Subject: [PATCH 05/11] fix(planning): preserve failed rollouts for diagnostics --- embodichain/lab/sim/atomic_actions/core.py | 11 ++- embodichain/lab/sim/atomic_actions/engine.py | 14 +++- .../lab/sim/planners/motion_generator.py | 9 ++- .../motion_generation/metrics/trajectory.py | 80 ++++++++++++++++++- scripts/benchmark/motion_generation/models.py | 9 +++ tests/sim/atomic_actions/test_engine.py | 15 ++++ 6 files changed, 133 insertions(+), 5 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index d4b7906bb..2992d4747 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -549,7 +549,16 @@ def build_plan( raise ValueError("Trajectory and planning context batch sizes must match.") if timed.robot_dof != context.robot.robot_dof: raise ValueError("Trajectory robot_dof must match the planning context.") - timed = timed.hold_rows(success_mask, context.robot.qpos) + preserve_failed_positions = ( + getattr( + self.motion_generator.planner, + "preserve_failed_plan_positions", + False, + ) + is True + ) + if not preserve_failed_positions: + timed = timed.hold_rows(success_mask, context.robot.qpos) segments: list[TrajectorySegment] = [] if segment_lengths is not None: diff --git a/embodichain/lab/sim/atomic_actions/engine.py b/embodichain/lab/sim/atomic_actions/engine.py index c52d724ea..5eb7f93c9 100644 --- a/embodichain/lab/sim/atomic_actions/engine.py +++ b/embodichain/lab/sim/atomic_actions/engine.py @@ -372,7 +372,19 @@ def compile( previous_qpos = projected.robot.qpos plan = self.plan(invocation, projected) step_success = alive & plan.plan_success.to(self.device) - trajectory = plan.trajectory.hold_rows(step_success, previous_qpos) + preserve_failed_positions = ( + getattr( + self.motion_generator.planner, + "preserve_failed_plan_positions", + False, + ) + is True + ) + trajectory = ( + plan.trajectory + if preserve_failed_positions + else plan.trajectory.hold_rows(step_success, previous_qpos) + ) plans.append(plan) trajectories.append(trajectory) diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index f26d411da..0ebb289a9 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -745,7 +745,14 @@ def normalize_derivative( else: duration = result.duration - if start_qpos is not None and not success.all(): + preserve_failed_positions = ( + getattr(self.planner, "preserve_failed_plan_positions", False) is True + ) + if ( + start_qpos is not None + and not success.all() + and not preserve_failed_positions + ): held = ( start_qpos.to(dtype=positions.dtype).unsqueeze(1).expand_as(positions) ) diff --git a/scripts/benchmark/motion_generation/metrics/trajectory.py b/scripts/benchmark/motion_generation/metrics/trajectory.py index fff97a6f1..cfad677d5 100644 --- a/scripts/benchmark/motion_generation/metrics/trajectory.py +++ b/scripts/benchmark/motion_generation/metrics/trajectory.py @@ -99,6 +99,10 @@ def match_ordered_waypoints( "arrival_indices": [], "position_errors_m": [], "rotation_errors_rad": [], + "min_position_errors_m": [], + "min_rotation_errors_rad": [], + "min_rotation_errors_at_position_rad": [], + "min_position_errors_at_orientation_m": [], } pos_error, rot_error = _pose_error_matrices(waypoints, trajectory_poses) @@ -126,12 +130,35 @@ def match_ordered_waypoints( ] completed = len(arrival_indices) total = int(waypoints.shape[0]) + min_rotation_at_position: list[float | None] = [] + min_position_at_orientation: list[float | None] = [] + for waypoint_index in range(total): + position_hits = pos_error[waypoint_index] <= position_threshold_m + rotation_hits = rot_error[waypoint_index] <= rotation_threshold_rad + min_rotation_at_position.append( + float(rot_error[waypoint_index, position_hits].min().item()) + if bool(position_hits.any().item()) + else None + ) + min_position_at_orientation.append( + float(pos_error[waypoint_index, rotation_hits].min().item()) + if bool(rotation_hits.any().item()) + else None + ) return { "ordered_waypoints_reached": completed == total, "completed_waypoint_ratio": completed / max(total, 1), "arrival_indices": arrival_indices, "position_errors_m": position_errors, "rotation_errors_rad": rotation_errors, + "min_position_errors_m": [ + float(value) for value in pos_error.min(dim=1).values.tolist() + ], + "min_rotation_errors_rad": [ + float(value) for value in rot_error.min(dim=1).values.tolist() + ], + "min_rotation_errors_at_position_rad": min_rotation_at_position, + "min_position_errors_at_orientation_m": min_position_at_orientation, } @@ -351,18 +378,34 @@ def compute_case_outcomes( name=control_part, to_matrix=True, ) + native_pose_batch = robot.compute_batch_fk( + qpos=positions, + name=control_part, + to_matrix=True, + ) outcomes: list[CaseOutcome] = [] for env_index in range(case.batch_size): native_qpos = positions[env_index] finite = finite_paths[env_index] if finite: validation_qpos = validation_qpos_batch[env_index] - poses = validation_pose_batch[env_index] + validation_poses = validation_pose_batch[env_index] + # Keep the benchmark's established, sample-count-normalized + # trajectory for all success and path metrics. Native planner + # samples are retained below only for rollout diagnostics. + poses = validation_poses + native_poses = native_pose_batch[env_index] else: validation_qpos = torch.empty( (0, positions.shape[-1]), device=robot.device, dtype=positions.dtype ) + validation_poses = torch.empty( + (0, 4, 4), device=robot.device, dtype=positions.dtype + ) poses = torch.empty((0, 4, 4), device=robot.device, dtype=positions.dtype) + native_poses = torch.empty( + (0, 4, 4), device=robot.device, dtype=positions.dtype + ) waypoints = case.target_waypoints[env_index] matching = match_ordered_waypoints( @@ -394,12 +437,23 @@ def compute_case_outcomes( if poses.shape[0] > 0: final_pos_m, final_rot_rad = get_pose_err(poses[-1], waypoints[-1]) + all_pos_error, all_rot_error = _pose_error_matrices( + waypoints, native_poses + ) + min_pos_m = float(all_pos_error[-1].min().item()) + min_rot_rad = float(all_rot_error[-1].min().item()) joint_length, cartesian_length, efficiency = _path_metrics( - validation_qpos, poses, waypoints + validation_qpos, validation_poses, waypoints + ) + trajectory_moved = bool( + torch.linalg.vector_norm(native_qpos - native_qpos[:1], dim=-1).max() + > 1.0e-6 ) else: final_pos_m = final_rot_rad = None + min_pos_m = min_rot_rad = None joint_length = cartesian_length = efficiency = None + trajectory_moved = False if not finite: failure_code = "non_finite_trajectory" @@ -454,6 +508,28 @@ def compute_case_outcomes( path_efficiency=efficiency, failure_code=failure_code, planner_failure_code=planner_failure_code, + min_translation_err_mm=( + min_pos_m * 1000.0 if min_pos_m is not None else None + ), + min_rotation_err_deg=( + min_rot_rad * 180.0 / math.pi if min_rot_rad is not None else None + ), + trajectory_moved=trajectory_moved, + waypoint_min_translation_err_mm=tuple( + float(value) * 1000.0 for value in matching["min_position_errors_m"] + ), + waypoint_min_rotation_err_deg=tuple( + float(value) * 180.0 / math.pi + for value in matching["min_rotation_errors_rad"] + ), + waypoint_min_rotation_err_deg_at_position=tuple( + None if value is None else float(value) * 180.0 / math.pi + for value in matching["min_rotation_errors_at_position_rad"] + ), + waypoint_min_translation_err_mm_at_orientation=tuple( + None if value is None else float(value) * 1000.0 + for value in matching["min_position_errors_at_orientation_m"] + ), ) ) return tuple(outcomes) diff --git a/scripts/benchmark/motion_generation/models.py b/scripts/benchmark/motion_generation/models.py index e670d564d..24b73756c 100644 --- a/scripts/benchmark/motion_generation/models.py +++ b/scripts/benchmark/motion_generation/models.py @@ -128,6 +128,15 @@ class CaseOutcome: joint_tracking_rmse_rad: float | None = None object_lift_delta_m: float | None = None replan_count: int | None = None + min_translation_err_mm: float | None = None + min_rotation_err_deg: float | None = None + trajectory_moved: bool | None = None + waypoint_min_translation_err_mm: tuple[float, ...] = () + waypoint_min_rotation_err_deg: tuple[float, ...] = () + waypoint_min_rotation_err_deg_at_position: tuple[float | None, ...] = () + waypoint_min_translation_err_mm_at_orientation: tuple[float | None, ...] = () + executed_final_translation_err_mm: float | None = None + executed_final_rotation_err_deg: float | None = None @dataclass(frozen=True) diff --git a/tests/sim/atomic_actions/test_engine.py b/tests/sim/atomic_actions/test_engine.py index 8704c357a..95402439c 100644 --- a/tests/sim/atomic_actions/test_engine.py +++ b/tests/sim/atomic_actions/test_engine.py @@ -218,6 +218,21 @@ def test_engine_compile_holds_failed_rows_for_remaining_actions() -> None: assert torch.all(compiled.trajectory.positions[1] == 0.0) +def test_engine_compile_preserves_opted_in_failed_trajectory_for_diagnostics() -> None: + engine = _engine() + engine.motion_generator.planner.preserve_failed_plan_positions = True + engine.register(StubAction()) + target = torch.tensor([[1.0, 1.0, 1.0], [float("nan"), 2.0, 2.0]]) + + compiled = engine.compile((_invocation(target),)) + + assert compiled.plan_success.tolist() == [True, False] + assert torch.equal( + compiled.trajectory.positions[1, -1], torch.tensor([0.0, 2.0, 2.0]) + ) + assert torch.all(compiled.projected_context.robot.qpos[1] == 0.0) + + def test_engine_compile_empty_sequence_is_successful_noop() -> None: engine = _engine() context = engine.initial_context() From 365c25911a2f1265b3ed4aed003225fd2f45a698 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Thu, 27 Aug 2026 10:49:00 +0800 Subject: [PATCH 06/11] feat(nmg): support unified joint waypoint planning --- .../lab/sim/planners/neural_planner.py | 215 +++++++++++++----- .../motion_generation/planners/curobo.py | 8 +- .../motion_generation/planners/nmg_onnx.py | 10 +- tests/sim/planners/test_neural_planner.py | 133 +++++++++++ 4 files changed, 312 insertions(+), 54 deletions(-) diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index 482cf79ae..6d8603ff8 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -119,6 +119,11 @@ def _quat_mul_xyzw(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: ) +def _canonicalize_quat_xyzw(q: torch.Tensor) -> torch.Tensor: + """Map equivalent quaternion signs to the training-time ``w >= 0`` half.""" + return torch.where(q[..., 3:4] < 0.0, -q, q) + + @configclass class NeuralPlannerCfg(BasePlannerCfg): planner_type: str = "neural" @@ -144,6 +149,9 @@ class NeuralPlannerCfg(BasePlannerCfg): use_relative_obs: bool = True """Whether the exported policy uses the unified relative-observation blocks.""" + canonicalize_quat_obs: bool = True + """Whether quaternion observations use the training-time ``w >= 0`` convention.""" + intermediate_orientation: bool = True """Whether every Cartesian waypoint requires its orientation constraint.""" @@ -153,6 +161,9 @@ class NeuralPlannerCfg(BasePlannerCfg): rot_eps: float = 0.1 """Waypoint rotation threshold in radians.""" + joint_eps: float = 0.02 + """Waypoint joint-position threshold in radians.""" + onnx_providers: list[str] | None = None """Optional ONNX Runtime execution-provider priority list.""" @@ -182,8 +193,8 @@ class NeuralPlanner(BasePlanner): r"""Neural motion planner based on an APG waypoint transformer policy. The planner loads a standalone ONNX waypoint policy and rolls it out in - closed loop to drive the arm toward a sequence of - end-effector waypoints. Velocities and accelerations in the returned + closed loop to drive the arm toward a sequence of end-effector and/or + joint-position waypoints. Velocities and accelerations in the returned :class:`PlanResult` are estimated via finite differences over the generated position trajectory. @@ -196,7 +207,11 @@ class NeuralPlanner(BasePlanner): ImportError: If ONNX Runtime is not installed. """ - supported_move_types = frozenset({MoveType.EEF_MOVE}) + supported_move_types = frozenset({MoveType.EEF_MOVE, MoveType.JOINT_MOVE}) + preserve_plan_samples = True + """Keep native closed-loop states instead of distance-resampling them.""" + preserve_failed_plan_positions = True + """Keep closed-loop rollout samples even when not all waypoints converge.""" def __init__(self, cfg: NeuralPlannerCfg): super().__init__(cfg) @@ -236,6 +251,7 @@ def _load_onnx_model(self, model_path: Path) -> None: self._num_waypoints = int(self.cfg.num_waypoints) self._use_relative_obs = bool(self.cfg.use_relative_obs) + self._canonicalize_quat_obs = bool(self.cfg.canonicalize_quat_obs) self._action_dim = int(self.cfg.num_arm_joints) if self._action_dim != 7: raise ValueError( @@ -244,6 +260,7 @@ def _load_onnx_model(self, model_path: Path) -> None: self._max_steps = int(self.cfg.max_steps) self._pos_eps = float(self.cfg.pos_eps) self._rot_eps = float(self.cfg.rot_eps) + self._joint_eps = float(self.cfg.joint_eps) self._intermediate_orientation = bool(self.cfg.intermediate_orientation) self._policy = _OnnxPolicy(model_path, self.cfg.onnx_providers) self._obs_dim = self._policy.obs_dim @@ -290,9 +307,9 @@ def plan( until all waypoints are reached or ``max_steps`` is exhausted. Args: - target_states: List of :class:`PlanState` waypoints. Each entry must - use :attr:`MoveType.EEF_MOVE` and carry an ``xpos`` tensor of - shape ``(B, 4, 4)``. + target_states: List of :class:`PlanState` waypoints. Each entry uses + :attr:`MoveType.EEF_MOVE` with ``xpos`` shape ``(B, 4, 4)`` or + :attr:`MoveType.JOINT_MOVE` with ``qpos`` shape ``(B, 7)``. options: :class:`NeuralPlanOptions` with ``control_part``, ``start_qpos``, and ``max_steps`` overrides. @@ -305,8 +322,8 @@ def plan( via finite differences and are therefore approximate. Raises: - ValueError: If ``control_part`` is not provided, if a target state - is not ``EEF_MOVE``, or if ``start_qpos`` has too few joints. + ValueError: If ``control_part`` is not provided, a target state is + unsupported, or ``start_qpos`` has too few joints. """ if not target_states: return PlanResult( @@ -321,9 +338,16 @@ def plan( ValueError, ) - waypoints_pos, waypoints_quat, valid_mask, episode_k = self._parse_waypoints( - target_states - ) + ( + waypoints_pos, + waypoints_quat, + waypoints_joint, + valid_mask, + pos_mask, + rot_mask, + joint_mask, + episode_k, + ) = self._parse_waypoints(target_states) qpos = self._initial_qpos(control_part, options.start_qpos) b = qpos.shape[0] limits = self.robot.get_qpos_limits(name=control_part)[0].to(self.device) @@ -345,7 +369,11 @@ def plan( ee_pose, waypoints_pos, waypoints_quat, + waypoints_joint, valid_mask, + pos_mask, + rot_mask, + joint_mask, active_idx, last_action, ) @@ -370,7 +398,15 @@ def plan( ee_pose = self._fk_pose_xyzw(qpos, control_part) reached = self._is_active_reached( - ee_pose, waypoints_pos, waypoints_quat, active_idx, episode_k + qpos[:, : self._action_dim], + ee_pose, + waypoints_pos, + waypoints_quat, + waypoints_joint, + pos_mask, + rot_mask, + joint_mask, + active_idx, ) active_idx = torch.where(reached, active_idx + 1, active_idx) converged = converged | (active_idx >= episode_k) @@ -406,9 +442,16 @@ def plan( ), ) - def _parse_waypoints( - self, target_states: list[PlanState] - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + def _parse_waypoints(self, target_states: list[PlanState]) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + int, + ]: if len(target_states) > self._num_waypoints: logger.log_error( f"Received {len(target_states)} waypoints, but the ONNX policy supports " @@ -418,23 +461,68 @@ def _parse_waypoints( b = _infer_batch_size(target_states) or 1 waypoint_pos = torch.zeros(b, self._num_waypoints, 3, device=self.device) waypoint_quat = torch.zeros(b, self._num_waypoints, 4, device=self.device) + waypoint_quat[..., 3] = 1.0 + waypoint_joint = torch.zeros( + b, self._num_waypoints, self._action_dim, device=self.device + ) valid_mask = torch.zeros(b, self._num_waypoints, device=self.device) + pos_mask = torch.zeros_like(valid_mask) + rot_mask = torch.zeros_like(valid_mask) + joint_mask = torch.zeros_like(valid_mask) for idx, target in enumerate(target_states): - if target.move_type != MoveType.EEF_MOVE or target.xpos is None: + if target.move_type == MoveType.EEF_MOVE and target.xpos is not None: + xpos = torch.as_tensor( + target.xpos, dtype=torch.float32, device=self.device + ) + if xpos.dim() == 2: + xpos = xpos.unsqueeze(0) + policy_xpos = self._to_policy_frame(xpos) + waypoint_pos[:, idx] = policy_xpos[:, :3, 3] + quat_xyzw = convert_quat( + quat_from_matrix(policy_xpos[:, :3, :3]), to="xyzw" + ) + waypoint_quat[:, idx] = ( + _canonicalize_quat_xyzw(quat_xyzw) + if getattr(self, "_canonicalize_quat_obs", False) + else quat_xyzw + ) + pos_mask[:, idx] = 1.0 + rot_mask[:, idx] = 1.0 + elif target.move_type == MoveType.JOINT_MOVE and target.qpos is not None: + qpos = torch.as_tensor( + target.qpos, dtype=torch.float32, device=self.device + ) + if qpos.dim() == 1: + qpos = qpos.unsqueeze(0) + if qpos.shape != (b, self._action_dim): + logger.log_error( + "NeuralPlanner JOINT_MOVE qpos must have shape " + f"({b}, {self._action_dim}), got {tuple(qpos.shape)}.", + ValueError, + ) + waypoint_joint[:, idx] = qpos + joint_mask[:, idx] = 1.0 + else: logger.log_error( - "NeuralPlanner expects EEF_MOVE PlanState entries with xpos.", + "NeuralPlanner expects EEF_MOVE entries with xpos or " + "JOINT_MOVE entries with qpos.", ValueError, ) - xpos = torch.as_tensor(target.xpos, dtype=torch.float32, device=self.device) - if xpos.dim() == 2: - xpos = xpos.unsqueeze(0) - policy_xpos = self._to_policy_frame(xpos) - waypoint_pos[:, idx] = policy_xpos[:, :3, 3] - waypoint_quat[:, idx] = convert_quat( - quat_from_matrix(policy_xpos[:, :3, :3]), to="xyzw" - ) valid_mask[:, idx] = 1.0 - return waypoint_pos, waypoint_quat, valid_mask, len(target_states) + if not self._intermediate_orientation: + final_mask = torch.zeros_like(rot_mask) + final_mask[:, len(target_states) - 1] = 1.0 + rot_mask *= final_mask + return ( + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid_mask, + pos_mask, + rot_mask, + joint_mask, + len(target_states), + ) def _initial_qpos( self, control_part: str, start_qpos: torch.Tensor | None @@ -460,6 +548,8 @@ def _fk_pose_xyzw(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: fk = self._to_policy_frame(self._fk_matrix(qpos, control_part)) pos = fk[:, :3, 3] quat_xyzw = convert_quat(quat_from_matrix(fk[:, :3, :3]), to="xyzw") + if getattr(self, "_canonicalize_quat_obs", False): + quat_xyzw = _canonicalize_quat_xyzw(quat_xyzw) return torch.cat([pos, quat_xyzw], dim=-1) def _build_obs( @@ -468,7 +558,11 @@ def _build_obs( ee_pose: torch.Tensor, waypoint_pos: torch.Tensor, waypoint_quat: torch.Tensor, + waypoint_joint: torch.Tensor, valid_mask: torch.Tensor, + pos_mask: torch.Tensor, + rot_mask: torch.Tensor, + joint_mask: torch.Tensor, active_idx: torch.Tensor, last_action: torch.Tensor, ) -> torch.Tensor: @@ -476,23 +570,8 @@ def _build_obs( active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) active_onehot = torch.zeros(b, self._num_waypoints, device=self.device) active_onehot.scatter_(1, active_idx_clamped.unsqueeze(1), 1.0) - pos_mask = valid_mask.clone() - rot_mask = valid_mask.clone() - if not self._intermediate_orientation: - waypoint_ids = torch.arange(self._num_waypoints, device=self.device) - episode_k = valid_mask.sum(dim=-1).long() - rot_mask = (waypoint_ids.unsqueeze(0) == (episode_k - 1).unsqueeze(1)).to( - valid_mask.dtype - ) - joint_mask = torch.zeros_like(valid_mask) - waypoint_joint = torch.zeros( - b, - self._num_waypoints, - self._action_dim, - dtype=joint_pos.dtype, - device=self.device, - ) pos_block = waypoint_pos * pos_mask.unsqueeze(-1) + joint_block = waypoint_joint * joint_mask.unsqueeze(-1) identity = torch.tensor( [0.0, 0.0, 0.0, 1.0], dtype=waypoint_quat.dtype, @@ -508,7 +587,7 @@ def _build_obs( ee_pose, pos_block.reshape(b, self._num_waypoints * 3), quat_block.reshape(b, self._num_waypoints * 4), - waypoint_joint.reshape(b, self._num_waypoints * self._action_dim), + joint_block.reshape(b, self._num_waypoints * self._action_dim), active_onehot, valid_mask, pos_mask, @@ -520,22 +599,33 @@ def _build_obs( idx = torch.arange(b, device=self.device) active_pos = pos_block[idx, active_idx_clamped] active_quat = quat_block[idx, active_idx_clamped] + active_joint = joint_block[idx, active_idx_clamped] inv_eef = _quat_inverse_xyzw(ee_pose[:, 3:7]) active_rel_quat = _quat_mul_xyzw(active_quat, inv_eef) - obs_parts.append( - torch.cat([active_pos - ee_pose[:, :3], active_rel_quat], dim=-1) + if getattr(self, "_canonicalize_quat_obs", False): + active_rel_quat = _canonicalize_quat_xyzw(active_rel_quat) + active_cart_rel = torch.cat( + [active_pos - ee_pose[:, :3], active_rel_quat], dim=-1 ) + active_rel = torch.where( + (joint_mask[idx, active_idx_clamped] > 0.5).unsqueeze(-1), + active_joint - joint_pos, + active_cart_rel, + ) + obs_parts.append(active_rel) rel_pos = (pos_block - ee_pose[:, None, :3]) * pos_mask.unsqueeze(-1) rel_quat = _quat_mul_xyzw( quat_block, inv_eef[:, None, :].expand_as(quat_block), ) + if getattr(self, "_canonicalize_quat_obs", False): + rel_quat = _canonicalize_quat_xyzw(rel_quat) rel_quat = torch.where( rot_mask.unsqueeze(-1) > 0.5, rel_quat, identity.view(1, 1, 4), ) - joint_err = torch.zeros_like(waypoint_joint) + joint_err = (joint_block - joint_pos[:, None]) * joint_mask.unsqueeze(-1) obs_parts.extend( [ rel_pos.reshape(b, self._num_waypoints * 3), @@ -554,6 +644,11 @@ def _build_obs( torch.ones_like(waypoint_type), waypoint_type, ) + waypoint_type = torch.where( + joint_mask > 0.5, + torch.full_like(waypoint_type, 2.0), + waypoint_type, + ) obs_parts.append(waypoint_type) obs = torch.cat(obs_parts, dim=-1) if obs.shape[-1] != self._obs_dim: @@ -564,30 +659,46 @@ def _build_obs( def _is_active_reached( self, + joint_pos: torch.Tensor, ee_pose: torch.Tensor, waypoint_pos: torch.Tensor, waypoint_quat: torch.Tensor, + waypoint_joint: torch.Tensor, + pos_mask: torch.Tensor, + rot_mask: torch.Tensor, + joint_mask: torch.Tensor, active_idx: torch.Tensor, - episode_k: int, ) -> torch.Tensor: b = ee_pose.shape[0] idx = torch.arange(b, device=self.device) active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) active_pos = waypoint_pos[idx, active_idx_clamped] active_quat_xyzw = waypoint_quat[idx, active_idx_clamped] + active_joint = waypoint_joint[idx, active_idx_clamped] + active_pos_mask = pos_mask[idx, active_idx_clamped] > 0.5 + active_rot_mask = rot_mask[idx, active_idx_clamped] > 0.5 + active_joint_mask = joint_mask[idx, active_idx_clamped] > 0.5 pos_dist = (ee_pose[:, :3] - active_pos).norm(dim=-1) ee_quat_wxyz = convert_quat(ee_pose[:, 3:7], to="wxyz") active_quat_wxyz = convert_quat(active_quat_xyzw, to="wxyz") rot_dist = quat_error_magnitude(ee_quat_wxyz, active_quat_wxyz) - orientation_required = self._intermediate_orientation | ( - active_idx >= episode_k - 1 - ) rot_ok = torch.where( - orientation_required, + active_rot_mask, rot_dist < self._rot_eps, torch.ones_like(rot_dist, dtype=torch.bool), ) - reached = (pos_dist < self._pos_eps) & rot_ok + pos_ok = torch.where( + active_pos_mask, + pos_dist < self._pos_eps, + torch.ones_like(pos_dist, dtype=torch.bool), + ) + joint_dist = torch.amax(torch.abs(joint_pos - active_joint), dim=-1) + joint_ok = torch.where( + active_joint_mask, + joint_dist < self._joint_eps, + torch.ones_like(joint_dist, dtype=torch.bool), + ) + reached = pos_ok & rot_ok & joint_ok return reached @staticmethod diff --git a/scripts/benchmark/motion_generation/planners/curobo.py b/scripts/benchmark/motion_generation/planners/curobo.py index 153d903a1..30ff1c1dd 100644 --- a/scripts/benchmark/motion_generation/planners/curobo.py +++ b/scripts/benchmark/motion_generation/planners/curobo.py @@ -47,7 +47,13 @@ class CuroboAdapter(PlannerAdapter): """Run cuRobo with a frozen, empty-world operational configuration.""" capabilities = frozenset( - {"eef_waypoint", "batched", "empty_world", "atomic_action"} + { + "eef_waypoint", + "joint_waypoint", + "batched", + "empty_world", + "atomic_action", + } ) model_revision = "curobo-v2" separate_prepare = True diff --git a/scripts/benchmark/motion_generation/planners/nmg_onnx.py b/scripts/benchmark/motion_generation/planners/nmg_onnx.py index 69e45a5a1..6775488c1 100644 --- a/scripts/benchmark/motion_generation/planners/nmg_onnx.py +++ b/scripts/benchmark/motion_generation/planners/nmg_onnx.py @@ -43,7 +43,13 @@ class NmgOnnxAdapter(PlannerAdapter): """Run a standalone NMG ONNX policy through :class:`NeuralPlanner`.""" capabilities = frozenset( - {"eef_waypoint", "batched", "empty_world", "atomic_action"} + { + "eef_waypoint", + "joint_waypoint", + "batched", + "empty_world", + "atomic_action", + } ) model_revision = "nmg-onnx-v1" @@ -84,9 +90,11 @@ def build(self) -> None: num_arm_joints=int(values.get("num_arm_joints", 7)), num_waypoints=int(values.get("num_waypoints", 8)), use_relative_obs=bool(values.get("use_relative_obs", True)), + canonicalize_quat_obs=bool(values.get("canonicalize_quat_obs", True)), intermediate_orientation=bool(values.get("intermediate_orientation", True)), pos_eps=float(values.get("pos_eps", 0.01)), rot_eps=float(values.get("rot_eps", 0.1)), + joint_eps=float(values.get("joint_eps", 0.02)), onnx_providers=list(providers) if providers is not None else None, policy_frame_from_world=values.get("policy_frame_from_world"), runtime_tcp_from_policy_tcp=values.get("runtime_tcp_from_policy_tcp"), diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index 4d4553e44..d13447ed4 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -99,6 +99,11 @@ def get_robot(self, uid: str) -> FakeRobot: def test_neural_planner_is_registered(): assert MotionGenerator._support_planner_dict["neural"][0] is NeuralPlanner assert MotionGenerator._support_planner_dict["neural"][1] is NeuralPlannerCfg + assert NeuralPlanner.preserve_plan_samples is True + assert NeuralPlanner.preserve_failed_plan_positions is True + assert NeuralPlanner.supported_move_types == frozenset( + {MoveType.EEF_MOVE, MoveType.JOINT_MOVE} + ) def test_neural_planner_generate_with_fake_onnx_model(tmp_path, monkeypatch): @@ -293,14 +298,22 @@ def test_neural_planner_builds_unified_300d_cartesian_observation( waypoint_pos = torch.zeros(1, 8, 3) waypoint_quat = torch.zeros(1, 8, 4) waypoint_quat[..., 3] = 1.0 + waypoint_joint = torch.zeros(1, 8, 7) valid = torch.zeros(1, 8) valid[:, :2] = 1.0 + pos_mask = valid.clone() + rot_mask = valid.clone() + joint_mask = torch.zeros_like(valid) obs = planner._build_obs( joint, eef, waypoint_pos, waypoint_quat, + waypoint_joint, valid, + pos_mask, + rot_mask, + joint_mask, torch.zeros(1, dtype=torch.long), torch.zeros(1, 7), ) @@ -317,6 +330,83 @@ def test_neural_planner_builds_unified_300d_cartesian_observation( assert torch.count_nonzero(obs[:, semantic_start + 32 : semantic_start + 40]) == 0 +def test_neural_planner_builds_joint_constraint_observation(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + target = torch.tensor([[0.1, -0.2, 0.3, 0.0, 0.2, -0.1, 0.4]]) + parsed = planner._parse_waypoints( + [PlanState.from_qpos(target, move_type=MoveType.JOINT_MOVE)] + ) + ( + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + _, + ) = parsed + obs = planner._build_obs( + torch.zeros(1, 7), + torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]), + waypoint_pos, + waypoint_quat, + waypoint_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + torch.zeros(1, dtype=torch.long), + torch.zeros(1, 7), + ) + + semantic_start = 7 + 7 + 8 * (3 + 4 + 7) + relative_start = semantic_start + 5 * 8 + 7 + waypoint_type_start = relative_start + 7 + 8 * (3 + 4 + 7) + assert torch.equal(joint_mask[:, 0], torch.ones(1)) + assert torch.count_nonzero(pos_mask) == 0 + assert torch.count_nonzero(rot_mask) == 0 + assert torch.allclose(obs[:, relative_start : relative_start + 7], target) + assert obs[0, waypoint_type_start].item() == pytest.approx(2.0) + + +def test_neural_planner_accepts_joint_move_goal(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + + result = planner.plan( + [PlanState.from_qpos(torch.zeros(1, 7), move_type=MoveType.JOINT_MOVE)], + NeuralPlanOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + max_steps=1, + ), + ) + + assert result.success.all().item() + + def test_neural_planner_applies_policy_frame_and_tcp_transforms(tmp_path, monkeypatch): model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager() @@ -442,6 +532,49 @@ def test_motion_generator_neural_preserves_native_eef_targets(tmp_path, monkeypa assert options.is_interpolate is True +def test_motion_generator_neural_preserves_failed_rollout_positions( + tmp_path, monkeypatch +): + class MovingFakeOnnxPolicy(FakeOnnxPolicy): + def __call__(self, obs: torch.Tensor) -> torch.Tensor: + self.last_obs = obs.clone() + action = torch.zeros(obs.shape[0], NUM_ARM_JOINTS, device=obs.device) + action[:, 0] = 1.0 + return action + + monkeypatch.setattr(neural_planner_module, "_OnnxPolicy", MovingFakeOnnxPolicy) + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + + motion_generator = MotionGenerator( + cfg=MotionGenCfg( + planner_cfg=NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + max_steps=2, + ) + ) + ) + target = torch.eye(4) + target[0, 3] = 10.0 + + result = motion_generator.generate( + target_states=[PlanState.single(move_type=MoveType.EEF_MOVE, xpos=target)], + options=MotionGenOptions( + control_part="main_arm", + start_qpos=torch.zeros(NUM_ARM_JOINTS), + ), + ) + + assert not result.success.all().item() + assert result.positions is not None + assert result.positions[0, 1, 0] > result.positions[0, 0, 0] + + def test_neural_planner_rejects_pytorch_checkpoint(tmp_path, monkeypatch): pytorch_checkpoint_path = tmp_path / "checkpoint.pt" pytorch_checkpoint_path.write_bytes(b"not-an-onnx-model") From a367b5f49350feaea30af5e4d202bf976b7fb129 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Thu, 27 Aug 2026 10:49:05 +0800 Subject: [PATCH 07/11] feat(benchmark): add six single-arm atomic action cases --- scripts/benchmark/motion_generation/README.md | 6 +- .../benchmark/motion_generation/artifacts.py | 6 +- .../scenarios/atomic_task.py | 718 +++++++++++++++++- .../suites/atomic_franka_pgi_curobo.yaml | 52 +- .../test_atomic_task_benchmark.py | 18 +- 5 files changed, 774 insertions(+), 26 deletions(-) diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index 3f17bf320..b5a0b990f 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -29,8 +29,8 @@ in that run's `videos/` directory. - 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 +- Single-arm Atomic Task slice: Franka + PGI with `MoveEndEffector`, + `MoveJoints`, `PickUp`, `MoveHeldObject`, `Place`, and `Press` - Default matrix: cuRobo (`primary_baseline`); IK / TOPPRA optional diagnostics - Direct, batched NMG ONNX adapter (`candidate`, enabled when a model path is supplied) - Lifecycle timing: construct / prepare / cold / warm @@ -58,5 +58,5 @@ in that run's `videos/` directory. - `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 +- Dual-arm actions and multi-action chains - Latency-budget Pareto sweeps, confidence intervals, subprocess isolation diff --git a/scripts/benchmark/motion_generation/artifacts.py b/scripts/benchmark/motion_generation/artifacts.py index af8928b9d..ecda92707 100644 --- a/scripts/benchmark/motion_generation/artifacts.py +++ b/scripts/benchmark/motion_generation/artifacts.py @@ -141,7 +141,11 @@ def _case_to_dict(case: BenchmarkCase) -> dict[str, Any]: "method": ( "reference_qpos_fk" if case.skill_id == "N/A" - else "independent_sequential_ik" + else ( + "joint_limits_and_reference_fk" + if case.skill_id == "move_joints" + else "independent_sequential_ik" + ) ), "reference_qpos": case.reference_qpos.detach().cpu().tolist(), }, diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py index 2bc1e66c4..19487a7ad 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_task.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -36,9 +36,19 @@ ControlPartCommandProfile, EndEffectorPoseGoal, GraspGoal, + HeldObjectPoseGoal, + HeldObjectState, + JointPositionGoal, MotionPolicy, + MoveHeldObjectOptions, + MoveJointsOptions, ObjectSemantics, PickUpOptions, + PlaceGoal, + PlaceOptions, + PressGoal, + PressOptions, + TaskState, ) from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory from embodichain.lab.sim.planners.utils import PlanResult @@ -71,6 +81,7 @@ (0.0401, 0.0000, -0.9992), ) _TASK_DIFFICULTIES = {"simple", "medium", "hard"} +_GRIPPER_SKILLS = {"pick_up", "move_held_object", "place", "press"} @dataclass(frozen=True) @@ -83,6 +94,8 @@ class _ExecutionObservation: execution_time_ms: float task_completion_time_s: float object_lift_delta_m: torch.Tensor | None = None + final_arm_qpos: torch.Tensor | None = None + final_object_pose: torch.Tensor | None = None @dataclass(frozen=True) @@ -99,6 +112,7 @@ class AtomicSkillCaseProvider(ABC): """Generate and ground one Atomic Action without planner-specific logic.""" skill_id: str + requires_gripper = False @abstractmethod def generate_case( @@ -130,6 +144,13 @@ def lift_segment_start(self, compiled: CompiledTrajectory) -> int | None: """Return the first lift waypoint that should release object dynamics.""" return None + def initial_task_state( + self, scenario: "AtomicTaskScenario", case: BenchmarkCase + ) -> TaskState | None: + """Return an optional symbolic precondition for isolated action testing.""" + del scenario, case + return None + @abstractmethod def task_result( self, @@ -210,6 +231,41 @@ def _difficulty(config: Mapping[str, object]) -> str: return difficulty +def _motion_valid_mask( + motion_outcomes: tuple[CaseOutcome, ...], *, device: torch.device +) -> torch.Tensor: + """Return the common motion-valid result as one device-local mask.""" + return torch.tensor( + [item.motion_valid for item in motion_outcomes], + dtype=torch.bool, + device=device, + ) + + +def _case_pose(case: BenchmarkCase, name: str, *, device: torch.device) -> torch.Tensor: + """Restore one frozen pose tensor from JSON-compatible case parameters.""" + return torch.tensor(case.case_parameters[name], dtype=torch.float32, device=device) + + +def _held_object_state( + scenario: "AtomicTaskScenario", case: BenchmarkCase +) -> HeldObjectState: + """Reconstruct the frozen held-object relation for an isolated case.""" + handle = scenario.object_handle(case.object_id) + semantics = ObjectSemantics( + affordance=Affordance(), + geometry={}, + properties={"benchmark_object_id": handle.object_id}, + label=handle.object_id, + entity=handle.entity, + ) + return HeldObjectState( + semantics=semantics, + object_to_eef=_case_pose(case, "object_to_eef", device=scenario.robot.device), + grasp_xpos=_case_pose(case, "grasp_pose", device=scenario.robot.device), + ) + + class _MoveEndEffectorCases(AtomicSkillCaseProvider): """Deterministic robot-relative MoveEndEffector cases.""" @@ -326,6 +382,7 @@ class _PickUpCases(AtomicSkillCaseProvider): """Explicit-grasp PickUp cases that isolate motion-planner performance.""" skill_id = "pick_up" + requires_gripper = True def generate_case( self, @@ -544,10 +601,565 @@ def task_result( return success, "object_not_grasped" +class _MoveJointsCases(AtomicSkillCaseProvider): + """Deterministic relative joint-space waypoint cases.""" + + skill_id = "move_joints" + + 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_rad") + if not isinstance(raw_offsets, Sequence) or isinstance( + raw_offsets, (str, bytes) + ): + raise TypeError("target_offsets_rad must be a non-empty list.") + offsets = [ + _float_vector(value, name="target_offsets_rad", length=7) + for value in raw_offsets + ] + if not offsets: + raise ValueError("target_offsets_rad must not be empty.") + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + offset_tensor = torch.tensor( + offsets, dtype=start_qpos.dtype, device=start_qpos.device + ) + targets = start_qpos[:, None] + torch.cumsum(offset_tensor, dim=0)[None] + limits = scenario.robot.get_qpos_limits(name=scenario.control_part)[0] + margin = float(config.get("joint_limit_margin_rad", 0.05)) + if bool( + ( + (targets < limits[:, 0][None, None] + margin) + | (targets > limits[:, 1][None, None] - margin) + ) + .any() + .item() + ): + raise RuntimeError( + f"Joint limits rejected MoveJoints case {_case_name(config)!r}." + ) + target_poses = torch.stack( + [ + scenario.robot.compute_fk( + targets[:, index], name=scenario.control_part, to_matrix=True + ) + for index in range(targets.shape[1]) + ], + dim=1, + ) + 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="relative_joint_waypoints", + start_state_bin="pre_action", + start_qpos=start_qpos, + target_waypoints=target_poses, + reference_qpos=targets, + 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_rad": offsets, + "target_qpos": targets.detach().cpu().tolist(), + "joint_threshold_rad": float(config.get("joint_threshold_rad", 0.02)), + "difficulty_factors": dict(config.get("difficulty_factors", {})), + }, + ) + + def build_invocation( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + adapter: "PlannerAdapter", + ) -> ActionInvocation: + target = torch.tensor( + case.case_parameters["target_qpos"], + dtype=torch.float32, + device=scenario.robot.device, + ) + return ActionInvocation( + skill_id=self.skill_id, + goal=JointPositionGoal(target), + 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"]), + ), + skill_options=MoveJointsOptions(), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + del scenario, compiled + if observation.final_arm_qpos is None: + return torch.zeros_like(observation.execution_success), "joint_goal_miss" + target = torch.tensor( + case.case_parameters["target_qpos"], + dtype=observation.final_arm_qpos.dtype, + device=observation.final_arm_qpos.device, + )[:, -1] + error = torch.amax(torch.abs(observation.final_arm_qpos - target), dim=-1) + success = ( + observation.execution_success + & _motion_valid_mask(motion_outcomes, device=error.device) + & (error <= float(case.case_parameters["joint_threshold_rad"])) + ) + return success, "joint_goal_miss" + + +class _HeldObjectCases(AtomicSkillCaseProvider): + """Shared isolated held-object precondition for transport and placement.""" + + requires_gripper = True + + def _prepare_start( + self, + scenario: "AtomicTaskScenario", + config: Mapping[str, object], + ) -> tuple[AtomicObjectHandle, torch.Tensor, torch.Tensor, torch.Tensor]: + object_id = config.get("object") + if not isinstance(object_id, str) or not object_id: + raise ValueError(f"{self.skill_id} cases require an object id.") + handle = scenario.activate_object(object_id) + scenario.restore_base_robot() + object_pose = handle.initial_pose.clone() + held_offset = torch.tensor( + _float_vector( + config.get("held_object_offset_m"), + name="held_object_offset_m", + length=3, + default=(0.0, 0.0, 0.18), + ), + dtype=object_pose.dtype, + device=object_pose.device, + ) + object_pose[:, :3, 3] += held_offset + grasp_pose = object_pose.clone() + grasp_pose[:, :3, :3] = torch.tensor( + config.get("grasp_rotation", _TOP_DOWN_ROTATION), + dtype=grasp_pose.dtype, + device=grasp_pose.device, + ) + grasp_pose[:, :3, 3] += 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, + ) + arm_seed = scenario.robot.get_qpos(name=scenario.control_part) + success, arm_start = scenario.robot.compute_ik( + pose=grasp_pose, + joint_seed=arm_seed, + name=scenario.control_part, + ) + if not bool(torch.as_tensor(success).all().item()): + raise RuntimeError( + f"Independent IK rejected held-object start {_case_name(config)!r}." + ) + scenario.set_robot_start(arm_start, open_gripper=False, grasp_gripper=True) + handle.entity.set_local_pose(object_pose) + handle.entity.clear_dynamics() + if scenario.simulation is not None: + scenario.simulation.update(step=2) + handle.entity.set_local_pose(object_pose) + handle.entity.clear_dynamics() + object_to_eef = torch.bmm(torch.linalg.inv(object_pose), grasp_pose) + return handle, object_pose, grasp_pose, object_to_eef + + def initial_task_state( + self, scenario: "AtomicTaskScenario", case: BenchmarkCase + ) -> TaskState: + return TaskState( + batch_size=case.batch_size, + device=scenario.robot.device, + held_objects={scenario.control_part: _held_object_state(scenario, case)}, + ) + + @staticmethod + def _held_parameters( + handle: AtomicObjectHandle, + object_pose: torch.Tensor, + grasp_pose: torch.Tensor, + object_to_eef: torch.Tensor, + ) -> dict[str, object]: + return { + "object_initial_pose": object_pose.detach().cpu().tolist(), + "grasp_pose": grasp_pose.detach().cpu().tolist(), + "object_to_eef": object_to_eef.detach().cpu().tolist(), + "object_config": dict(handle.config), + } + + +class _MoveHeldObjectCases(_HeldObjectCases): + """Move an already held rigid object to one frozen target pose.""" + + skill_id = "move_held_object" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + handle, object_pose, grasp_pose, object_to_eef = self._prepare_start( + scenario, config + ) + target_object = object_pose.clone() + target_object[:, :3, 3] += torch.tensor( + _float_vector( + config.get("target_object_offset_m"), + name="target_object_offset_m", + length=3, + default=(0.08, 0.08, 0.04), + ), + dtype=object_pose.dtype, + device=object_pose.device, + ) + target_eef = torch.bmm(target_object, object_to_eef) + start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() + references = scenario.solve_reference_qpos(start_qpos, target_eef[:, None]) + 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=1, + path_shape="held_object_transport", + start_state_bin="object_held", + start_qpos=start_qpos, + target_waypoints=target_eef[:, None], + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + object_id=handle.object_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=scenario.robot.get_qpos().clone(), + case_parameters={ + **self._held_parameters(handle, object_pose, grasp_pose, object_to_eef), + "sample_count": int(config.get("sample_count", 80)), + "target_object_pose": target_object.detach().cpu().tolist(), + "object_position_threshold_m": float( + config.get("object_position_threshold_m", 0.04) + ), + "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=HeldObjectPoseGoal( + _case_pose(case, "target_object_pose", device=scenario.robot.device) + ), + 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=MoveHeldObjectOptions(pick_rotate_upright=0.0), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + final_object = observation.final_object_pose + if final_object is None: + return torch.zeros_like(observation.execution_success), "object_goal_miss" + target = _case_pose(case, "target_object_pose", device=final_object.device) + error = torch.linalg.vector_norm( + final_object[:, :3, 3] - target[:, :3, 3], dim=-1 + ) + held = compiled.projected_context.get_held_object(scenario.control_part) + success = ( + observation.execution_success + & _motion_valid_mask(motion_outcomes, device=error.device) + & (held is not None) + & (error <= float(case.case_parameters["object_position_threshold_m"])) + ) + return success, "object_goal_miss" + + +class _PlaceCases(_HeldObjectCases): + """Place an already held rigid object and release it under physics.""" + + skill_id = "place" + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + handle, object_pose, grasp_pose, object_to_eef = self._prepare_start( + scenario, config + ) + target_object = handle.initial_pose.clone() + target_object[:, :3, 3] += torch.tensor( + _float_vector( + config.get("target_object_offset_m"), + name="target_object_offset_m", + length=3, + default=(0.10, 0.10, 0.0), + ), + dtype=object_pose.dtype, + device=object_pose.device, + ) + release = torch.bmm(target_object, object_to_eef) + lift_height = float(config.get("retract_height_m", 0.10)) + approach = release.clone() + approach[:, 2, 3] += lift_height + retract = approach.clone() + targets = torch.stack([approach, release, retract], 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_release_retract", + start_state_bin="object_held", + start_qpos=start_qpos, + target_waypoints=targets, + reference_qpos=references, + robot_id=suite.robot.id, + skill_id=self.skill_id, + object_id=handle.object_id, + task_difficulty=_difficulty(config), + primary_success="task_success", + full_start_qpos=scenario.robot.get_qpos().clone(), + case_parameters={ + **self._held_parameters(handle, object_pose, grasp_pose, object_to_eef), + "sample_count": int(config.get("sample_count", 120)), + "release_pose": release.detach().cpu().tolist(), + "target_object_pose": target_object.detach().cpu().tolist(), + "hand_interp_steps": int(config.get("hand_interp_steps", 12)), + "retract_height_m": lift_height, + "object_position_threshold_m": float( + config.get("object_position_threshold_m", 0.05) + ), + "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=PlaceGoal( + _case_pose(case, "release_pose", device=scenario.robot.device) + ), + 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=PlaceOptions( + hand_interp_steps=int(case.case_parameters["hand_interp_steps"]), + lift_height=float(case.case_parameters["retract_height_m"]), + ), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + final_object = observation.final_object_pose + if final_object is None: + return torch.zeros_like(observation.execution_success), "object_not_placed" + target = _case_pose(case, "target_object_pose", device=final_object.device) + error = torch.linalg.vector_norm( + final_object[:, :3, 3] - target[:, :3, 3], dim=-1 + ) + released = ( + compiled.projected_context.get_held_object(scenario.control_part) is None + ) + success = ( + observation.execution_success + & _motion_valid_mask(motion_outcomes, device=error.device) + & released + & (error <= float(case.case_parameters["object_position_threshold_m"])) + ) + return success, "object_not_placed" + + +class _PressCases(AtomicSkillCaseProvider): + """Close the gripper, reach a contact pose, and retract.""" + + skill_id = "press" + requires_gripper = True + + def generate_case( + self, + scenario: "AtomicTaskScenario", + suite: SuiteCfg, + track: TrackCfg, + config: Mapping[str, object], + *, + seed: int, + batch_size: int, + ) -> BenchmarkCase: + scenario.restore_base_robot() + 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 + ) + press_pose = start_pose.clone() + press_pose[:, :3, 3] += torch.tensor( + _float_vector( + config.get("target_offset_m"), + name="target_offset_m", + length=3, + default=(0.0, 0.0, -0.08), + ), + dtype=start_pose.dtype, + device=start_pose.device, + ) + targets = torch.stack([press_pose, start_pose], dim=1) + 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=2, + path_shape="press_retract", + 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)), + "press_pose": press_pose.detach().cpu().tolist(), + "hand_interp_steps": int(config.get("hand_interp_steps", 8)), + "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=PressGoal( + _case_pose(case, "press_pose", device=scenario.robot.device) + ), + 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=PressOptions( + hand_interp_steps=int(case.case_parameters["hand_interp_steps"]) + ), + ) + + def task_result( + self, + scenario: "AtomicTaskScenario", + case: BenchmarkCase, + compiled: CompiledTrajectory, + observation: _ExecutionObservation, + motion_outcomes: tuple[CaseOutcome, ...], + ) -> tuple[torch.Tensor, str]: + del scenario, case, compiled + success = observation.execution_success & _motion_valid_mask( + motion_outcomes, device=observation.execution_success.device + ) + return success, "press_pose_miss" + + class AtomicTaskScenario(ScenarioProvider): """Run fixed Atomic Actions through an adapter-owned MotionGenerator.""" - required_capabilities = frozenset({"eef_waypoint", "atomic_action"}) + required_capabilities = frozenset( + {"eef_waypoint", "joint_waypoint", "atomic_action"} + ) def __init__(self) -> None: self.simulation: SimulationManager | None = None @@ -604,12 +1216,14 @@ def configure_runtime( simulation.update(step=1) if any( - isinstance(value, Mapping) and value.get("id") == "pick_up" + isinstance(value, Mapping) and value.get("id") in _GRIPPER_SKILLS 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.") + raise ValueError( + "Atomic gripper-action 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.") @@ -715,9 +1329,7 @@ def prepare_planner( """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 any(provider.requires_gripper for provider in self._case_providers.values()): if ( self.end_effector_part is None or self._gripper_open is None @@ -757,7 +1369,18 @@ def reset_case( active_id = case.object_id for index, handle in enumerate(self._objects.values()): if handle.object_id == active_id: - handle.reset() + initial_pose = case.case_parameters.get("object_initial_pose") + if initial_pose is None: + handle.reset() + else: + handle.entity.set_local_pose( + torch.tensor( + initial_pose, + dtype=torch.float32, + device=robot.device, + ) + ) + handle.entity.clear_dynamics() else: handle.park(index) simulation.update(step=2) @@ -776,7 +1399,9 @@ def plan_case(self, adapter: PlannerAdapter, case: BenchmarkCase) -> object: raise RuntimeError( "Atomic Task invocations must pin the adapter motion backend." ) - return self._engine.compile((invocation,)) + task = provider.initial_task_state(self, case) + context = None if task is None else self._engine.initial_context(task=task) + return self._engine.compile((invocation,), context=context) def plan_contract_error(self, result: object) -> str | None: """Accept compiled Atomic Action trajectories instead of raw plans.""" @@ -843,6 +1468,30 @@ def evaluate_case( tracking = observation.joint_tracking_rmse_rad object_lift = observation.object_lift_delta_m + executed_translation_mm: torch.Tensor | None = None + executed_rotation_deg: torch.Tensor | None = None + if observation is not None: + final_target = case.target_waypoints[:, -1] + executed_translation_mm = ( + torch.linalg.vector_norm( + observation.final_tcp_pose[:, :3, 3] - final_target[:, :3, 3], + dim=-1, + ) + * 1000.0 + ) + executed_relative = ( + final_target[:, :3, :3].transpose(-1, -2) + @ observation.final_tcp_pose[:, :3, :3] + ) + executed_trace = torch.diagonal(executed_relative, dim1=-2, dim2=-1).sum( + dim=-1 + ) + executed_rotation_deg = ( + torch.arccos(torch.clamp((executed_trace - 1.0) * 0.5, -1.0, 1.0)) + * 180.0 + / math.pi + ) + durations = result.trajectory.duration.detach().to("cpu") outcomes: list[CaseOutcome] = [] for index, outcome in enumerate(motion_outcomes): @@ -877,6 +1526,16 @@ def evaluate_case( ), replan_count=0, failure_code=failure_code, + executed_final_translation_err_mm=( + None + if executed_translation_mm is None + else float(executed_translation_mm[index].item()) + ), + executed_final_rotation_err_deg=( + None + if executed_rotation_deg is None + else float(executed_rotation_deg[index].item()) + ), ) ) trajectory_duration = ( @@ -919,7 +1578,7 @@ def record_replay( return None self.reset_case(self.simulation, self.robot, case, self.control_part) compiled = result if isinstance(result, CompiledTrajectory) else None - replayable = compiled is not None and self._is_replayable(compiled) + replayable = compiled is not None and self._is_recordable(compiled) provider = self._case_providers.get(case.case_id) video_path = build_video_path( output_dir, algorithm_id, case.skill_id, case.case_id @@ -976,10 +1635,12 @@ def _execute( to_matrix=True, ) object_lift = None + final_object_pose = 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 - ] + final_object_pose = object_handle.entity.get_local_pose( + to_matrix=True + ).clone() + final_object_position = final_object_pose[:, :3, 3] object_lift = final_object_position[:, 2] - initial_object_position[:, 2] return _ExecutionObservation( execution_success=torch.isfinite(observed).all(dim=1) @@ -989,6 +1650,8 @@ def _execute( execution_time_ms=elapsed_ms, task_completion_time_s=simulated_execution_time, object_lift_delta_m=object_lift, + final_arm_qpos=self.robot.get_qpos(name=self.control_part).clone(), + final_object_pose=final_object_pose, ) def _physics_settings(self) -> _PhysicsReplaySettings: @@ -1022,6 +1685,12 @@ def _is_replayable(compiled: CompiledTrajectory) -> bool: and bool(torch.isfinite(trajectory).all().item()) ) + @staticmethod + def _is_recordable(compiled: CompiledTrajectory) -> bool: + """Allow finite failed planner rollouts in diagnostic videos.""" + trajectory = compiled.trajectory.positions + return trajectory.shape[1] > 0 and bool(torch.isfinite(trajectory).all().item()) + def _replay_physics( self, compiled: CompiledTrajectory, @@ -1212,17 +1881,26 @@ def restore_base_robot(self) -> None: self.robot.clear_dynamics() def set_robot_start( - self, manipulator_qpos: torch.Tensor, *, open_gripper: bool + self, + manipulator_qpos: torch.Tensor, + *, + open_gripper: bool, + grasp_gripper: bool = False, ) -> None: - """Install a manipulator start and optional open-gripper command.""" + """Install a manipulator start and optional gripper command.""" if self.robot is None: raise RuntimeError("Atomic Task runtime is not configured.") + if open_gripper and grasp_gripper: + raise ValueError("The gripper cannot start both open and grasping.") 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( + if open_gripper or grasp_gripper: + gripper_qpos = ( + self._gripper_open if open_gripper else self._gripper_grasp + ) + if self.end_effector_part is None or gripper_qpos is None: + raise RuntimeError("Requested gripper state is unavailable.") + command = gripper_qpos.unsqueeze(0).expand( manipulator_qpos.shape[0], -1 ) self.robot.set_qpos(command, name=self.end_effector_part, target=target) @@ -1274,5 +1952,9 @@ def _synchronize() -> None: register_atomic_skill_provider("move_end_effector", _MoveEndEffectorCases) +register_atomic_skill_provider("move_held_object", _MoveHeldObjectCases) +register_atomic_skill_provider("move_joints", _MoveJointsCases) register_atomic_skill_provider("pick_up", _PickUpCases) +register_atomic_skill_provider("place", _PlaceCases) +register_atomic_skill_provider("press", _PressCases) register_scenario_provider("atomic_task", AtomicTaskScenario) diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml index 35ce3e94f..a019216e8 100644 --- a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml @@ -14,15 +14,16 @@ planners: role: candidate enabled: false config: - model_revision: nmg-cartesian-300m + model_revision: nmg-cartesian-k1-3-900m onnx_model_path: null - num_waypoints: 8 + num_waypoints: 3 max_steps: 240 action_scale: 0.2 use_relative_obs: true intermediate_orientation: true pos_eps: 0.01 rot_eps: 0.1 + joint_eps: 0.02 policy_frame_from_world: - [-1.0, 0.0, 0.0, 0.0] - [0.0, -1.0, 0.0, 0.0] @@ -103,6 +104,18 @@ tracks: waypoint_count: 2 maximum_translation_m: 0.13 obstacles: 0 + - id: move_joints + sample_count: 80 + joint_limit_margin_rad: 0.05 + joint_threshold_rad: 0.02 + cases: + - name: relative_multi_joint + task_difficulty: simple + target_offsets_rad: + - [0.12, -0.08, 0.10, 0.12, -0.08, 0.10, -0.10] + difficulty_factors: + waypoint_count: 1 + active_joints: 7 - id: pick_up sample_count: 120 grasp_source: antipodal @@ -124,3 +137,38 @@ tracks: approach: top pre_grasp_clearance_m: 0.15 required_lift_m: 0.04 + - id: move_held_object + sample_count: 80 + held_object_offset_m: [0.0, 0.0, 0.18] + target_object_offset_m: [0.08, 0.08, 0.04] + object_position_threshold_m: 0.04 + cases: + - name: cube_transport + object: cube_50mm + task_difficulty: simple + difficulty_factors: + translation_m: 0.12 + object_preheld: true + - id: place + sample_count: 120 + held_object_offset_m: [0.0, 0.0, 0.18] + target_object_offset_m: [0.10, 0.10, 0.0] + retract_height_m: 0.10 + hand_interp_steps: 12 + object_position_threshold_m: 0.05 + cases: + - name: cube_table_place + object: cube_50mm + task_difficulty: simple + difficulty_factors: + translation_m: 0.14 + object_preheld: true + - id: press + sample_count: 80 + hand_interp_steps: 8 + cases: + - name: vertical_press + task_difficulty: simple + target_offset_m: [0.0, 0.0, -0.08] + difficulty_factors: + press_depth_m: 0.08 diff --git a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py index abe0276e0..061f89bba 100644 --- a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py +++ b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py @@ -124,7 +124,14 @@ def test_atomic_suite_is_franka_pgi_and_curobo_only(): ("atomic-task", "atomic_task") ] skills = suite.enabled_tracks()[0].config["skills"] - assert [item["id"] for item in skills] == ["move_end_effector", "pick_up"] + assert [item["id"] for item in skills] == [ + "move_end_effector", + "move_joints", + "pick_up", + "move_held_object", + "place", + "press", + ] gripper = suite.enabled_tracks()[0].config["gripper"] assert gripper == { "control_part": "hand", @@ -160,7 +167,14 @@ def test_atomic_invocation_pins_motion_generator_and_selected_planner(): def test_atomic_skill_and_object_extensions_are_registry_driven(): - assert atomic_skill_provider_names() == ("move_end_effector", "pick_up") + assert atomic_skill_provider_names() == ( + "move_end_effector", + "move_held_object", + "move_joints", + "pick_up", + "place", + "press", + ) 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"}) From 2e03b3159854a164f4f36d74cde1738b51aa8e37 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Fri, 28 Aug 2026 05:08:34 +0800 Subject: [PATCH 08/11] feat(benchmark): add deterministic atomic task randomization --- scripts/benchmark/motion_generation/README.md | 4 + .../motion_generation/run_benchmark.py | 6 +- .../scenarios/atomic_task.py | 354 ++++++++++++++++-- .../atomic_franka_pgi_curobo_randomized.yaml | 189 ++++++++++ .../test_atomic_task_benchmark.py | 81 ++++ .../test_motion_generation_benchmark.py | 9 + 6 files changed, 601 insertions(+), 42 deletions(-) create mode 100644 scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml diff --git a/scripts/benchmark/motion_generation/README.md b/scripts/benchmark/motion_generation/README.md index b5a0b990f..47671cb82 100644 --- a/scripts/benchmark/motion_generation/README.md +++ b/scripts/benchmark/motion_generation/README.md @@ -15,6 +15,8 @@ python -m scripts.benchmark.motion_generation.run_benchmark \ --suite atomic_franka_pgi_curobo --device cuda python -m scripts.benchmark.motion_generation.run_benchmark \ --suite atomic_franka_pgi_curobo --device cuda --record-video +python -m scripts.benchmark.motion_generation.run_benchmark \ + --suite atomic_franka_pgi_curobo_randomized --device cuda python -m scripts.benchmark.motion_generation.run_benchmark \ --extra-baselines ik_interpolate toppra ``` @@ -31,6 +33,8 @@ in that run's `videos/` directory. - `atomic-task` track with frozen robot/object/task manifests and common physics replay - Single-arm Atomic Task slice: Franka + PGI with `MoveEndEffector`, `MoveJoints`, `PickUp`, `MoveHeldObject`, `Place`, and `Press` +- Deterministic 16-seed generalization sweep for all six single-arm skills, + with bounded robot-start, target, held-object, and object-pose randomization - Default matrix: cuRobo (`primary_baseline`); IK / TOPPRA optional diagnostics - Direct, batched NMG ONNX adapter (`candidate`, enabled when a model path is supplied) - Lifecycle timing: construct / prepare / cold / warm diff --git a/scripts/benchmark/motion_generation/run_benchmark.py b/scripts/benchmark/motion_generation/run_benchmark.py index 49003d8a8..541d50eeb 100644 --- a/scripts/benchmark/motion_generation/run_benchmark.py +++ b/scripts/benchmark/motion_generation/run_benchmark.py @@ -50,7 +50,8 @@ def add_parser_arguments(parser: argparse.ArgumentParser) -> None: "--suite", default="smoke", help=( - "Suite short name (smoke/coverage/atomic_franka_pgi_curobo) " + "Suite short name (smoke/coverage/atomic_franka_pgi_curobo/" + "atomic_franka_pgi_curobo_randomized) " "or an explicit YAML path." ), ) @@ -202,6 +203,9 @@ def _apply_overrides( suite.free_space.start_state_bins = start_state_bins if seeds is not None: suite.free_space.seeds = seeds + for track in suite.tracks: + if track.scenario != "free_space" and "seeds" in track.config: + track.config["seeds"] = list(seeds) if num_trials is not None: suite.protocol.measured_trials = num_trials if warmup_trials is not None: diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py index 19487a7ad..9cfef3da0 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_task.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -82,6 +82,7 @@ ) _TASK_DIFFICULTIES = {"simple", "medium", "hard"} _GRIPPER_SKILLS = {"pick_up", "move_held_object", "place", "press"} +_CASE_QPOS_RESOLUTION_RAD = 1.0e-4 @dataclass(frozen=True) @@ -231,6 +232,80 @@ def _difficulty(config: Mapping[str, object]) -> str: return difficulty +def _seeded_jitter( + amplitude: Sequence[float], + *, + seed: int, + stream: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Sample deterministic independent uniform jitter in ``[-amplitude, +amplitude]``.""" + generator = torch.Generator(device="cpu") + generator.manual_seed((int(seed) * 1_000_003 + int(stream) * 97_409) % (2**63 - 1)) + unit = torch.rand(len(amplitude), generator=generator, dtype=torch.float32) + values = (2.0 * unit - 1.0) * torch.tensor(amplitude, dtype=torch.float32) + return values.to(dtype=dtype, device=device) + + +def _case_generation_seed(seed: int, *, skill_index: int, case_index: int) -> int: + """Derive one stable RNG seed for all stochastic IK work in a case.""" + return ( + int(seed) * 1_000_003 + int(skill_index) * 97_409 + int(case_index) * 13_007 + ) % (2**63 - 1) + + +def _canonical_case_qpos(qpos: torch.Tensor) -> torch.Tensor: + """Remove insignificant device-level IK noise from frozen case states.""" + return torch.round(qpos / _CASE_QPOS_RESOLUTION_RAD) * _CASE_QPOS_RESOLUTION_RAD + + +def _randomized_vector( + base: Sequence[float], + config: Mapping[str, object], + *, + jitter_name: str, + seed: int, + stream: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return one configured vector plus optional deterministic uniform jitter.""" + base_values = [float(value) for value in base] + amplitude = _float_vector( + config.get(jitter_name), + name=jitter_name, + length=len(base_values), + default=(0.0,) * len(base_values), + ) + if any(value < 0.0 for value in amplitude): + raise ValueError(f"{jitter_name} values must be non-negative.") + return torch.tensor(base_values, dtype=dtype, device=device) + _seeded_jitter( + amplitude, + seed=seed, + stream=stream, + dtype=dtype, + device=device, + ) + + +def _randomization_parameters( + config: Mapping[str, object], *, seed: int +) -> dict[str, object]: + """Serialize the deterministic randomization contract into a case manifest.""" + ranges = { + str(key): value + for key, value in config.items() + if str(key).endswith("_jitter_m") or str(key).endswith("_jitter_rad") + } + return { + "enabled": bool(ranges), + "seed": int(seed), + "distribution": "independent_uniform", + "ranges": ranges, + } + + def _motion_valid_mask( motion_outcomes: tuple[CaseOutcome, ...], *, device: torch.device ) -> torch.Tensor: @@ -282,26 +357,40 @@ def generate_case( batch_size: int, ) -> BenchmarkCase: scenario.restore_base_robot() + scenario.randomize_robot_start(config, seed=seed, stream=11) 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 = [ + base_offsets = [ _float_vector(value, name="target_offsets_m", length=3) for value in raw_offsets ] - if not offsets: + if not base_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 ) + offsets_tensor = torch.stack( + [ + _randomized_vector( + offset, + config, + jitter_name="target_offset_jitter_m", + seed=seed, + stream=101 + index, + dtype=start_pose.dtype, + device=start_pose.device, + ) + for index, offset in enumerate(base_offsets) + ] + ) + offsets = offsets_tensor.detach().cpu().tolist() targets = start_pose[:, None].repeat(1, len(offsets), 1, 1) - targets[:, :, :3, 3] += torch.tensor( - offsets, dtype=targets.dtype, device=targets.device - )[None] + targets[:, :, :3, 3] += offsets_tensor[None] references = scenario.solve_reference_qpos(start_qpos, targets) name = _case_name(config) return BenchmarkCase( @@ -321,10 +410,11 @@ def generate_case( skill_id=self.skill_id, task_difficulty=_difficulty(config), primary_success="task_success", - full_start_qpos=scenario.robot.get_qpos().clone(), + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), case_parameters={ "sample_count": int(config.get("sample_count", 80)), "target_offsets_m": offsets, + "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, ) @@ -400,7 +490,9 @@ def generate_case( handle = scenario.activate_object(object_id) scenario.restore_base_robot() - object_pose = handle.entity.get_local_pose(to_matrix=True).clone() + object_pose = scenario.randomize_object_pose( + handle, config, seed=seed, stream=201 + ) 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 @@ -416,6 +508,7 @@ def generate_case( raise RuntimeError( f"Independent IK rejected PickUp case {_case_name(config)!r}." ) + pre_pick_qpos = _canonical_case_qpos(pre_pick_qpos) scenario.set_robot_start(pre_pick_qpos, open_gripper=True) approach = torch.tensor( @@ -464,13 +557,17 @@ def generate_case( grasp_pose[:, :3, :3] = rotation else: raise ValueError("grasp_source must be 'fixed' or 'antipodal'.") - grasp_offset = torch.tensor( + grasp_offset = _randomized_vector( _float_vector( config.get("grasp_offset_m"), name="grasp_offset_m", length=3, default=(0.0, 0.0, 0.0), ), + config, + jitter_name="grasp_offset_jitter_m", + seed=seed, + stream=202, dtype=grasp_pose.dtype, device=grasp_pose.device, ) @@ -501,7 +598,7 @@ def generate_case( object_id=object_id, task_difficulty=_difficulty(config), primary_success="task_success", - full_start_qpos=scenario.robot.get_qpos().clone(), + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), case_parameters={ "sample_count": int(config.get("sample_count", 120)), "grasp_source": grasp_source, @@ -515,6 +612,7 @@ def generate_case( "grasp_pose": grasp_pose.detach().cpu().tolist(), "object_initial_pose": object_pose.detach().cpu().tolist(), "object_config": dict(handle.config), + "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, ) @@ -617,21 +715,34 @@ def generate_case( batch_size: int, ) -> BenchmarkCase: scenario.restore_base_robot() + scenario.randomize_robot_start(config, seed=seed, stream=21) raw_offsets = config.get("target_offsets_rad") if not isinstance(raw_offsets, Sequence) or isinstance( raw_offsets, (str, bytes) ): raise TypeError("target_offsets_rad must be a non-empty list.") - offsets = [ + base_offsets = [ _float_vector(value, name="target_offsets_rad", length=7) for value in raw_offsets ] - if not offsets: + if not base_offsets: raise ValueError("target_offsets_rad must not be empty.") start_qpos = scenario.robot.get_qpos(name=scenario.control_part).clone() - offset_tensor = torch.tensor( - offsets, dtype=start_qpos.dtype, device=start_qpos.device + offset_tensor = torch.stack( + [ + _randomized_vector( + offset, + config, + jitter_name="target_offset_jitter_rad", + seed=seed, + stream=211 + index, + dtype=start_qpos.dtype, + device=start_qpos.device, + ) + for index, offset in enumerate(base_offsets) + ] ) + offsets = offset_tensor.detach().cpu().tolist() targets = start_qpos[:, None] + torch.cumsum(offset_tensor, dim=0)[None] limits = scenario.robot.get_qpos_limits(name=scenario.control_part)[0] margin = float(config.get("joint_limit_margin_rad", 0.05)) @@ -673,12 +784,13 @@ def generate_case( skill_id=self.skill_id, task_difficulty=_difficulty(config), primary_success="task_success", - full_start_qpos=scenario.robot.get_qpos().clone(), + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), case_parameters={ "sample_count": int(config.get("sample_count", 80)), "target_offsets_rad": offsets, "target_qpos": targets.detach().cpu().tolist(), "joint_threshold_rad": float(config.get("joint_threshold_rad", 0.02)), + "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, ) @@ -740,20 +852,35 @@ def _prepare_start( self, scenario: "AtomicTaskScenario", config: Mapping[str, object], - ) -> tuple[AtomicObjectHandle, torch.Tensor, torch.Tensor, torch.Tensor]: + *, + seed: int, + ) -> tuple[ + AtomicObjectHandle, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: object_id = config.get("object") if not isinstance(object_id, str) or not object_id: raise ValueError(f"{self.skill_id} cases require an object id.") handle = scenario.activate_object(object_id) scenario.restore_base_robot() - object_pose = handle.initial_pose.clone() - held_offset = torch.tensor( + table_object_pose = scenario.randomize_object_pose( + handle, config, seed=seed, stream=301 + ) + object_pose = table_object_pose.clone() + held_offset = _randomized_vector( _float_vector( config.get("held_object_offset_m"), name="held_object_offset_m", length=3, default=(0.0, 0.0, 0.18), ), + config, + jitter_name="held_object_offset_jitter_m", + seed=seed, + stream=302, dtype=object_pose.dtype, device=object_pose.device, ) @@ -764,13 +891,17 @@ def _prepare_start( dtype=grasp_pose.dtype, device=grasp_pose.device, ) - grasp_pose[:, :3, 3] += torch.tensor( + grasp_pose[:, :3, 3] += _randomized_vector( _float_vector( config.get("grasp_offset_m"), name="grasp_offset_m", length=3, default=(0.0, 0.0, 0.0), ), + config, + jitter_name="grasp_offset_jitter_m", + seed=seed, + stream=303, dtype=grasp_pose.dtype, device=grasp_pose.device, ) @@ -784,6 +915,7 @@ def _prepare_start( raise RuntimeError( f"Independent IK rejected held-object start {_case_name(config)!r}." ) + arm_start = _canonical_case_qpos(arm_start) scenario.set_robot_start(arm_start, open_gripper=False, grasp_gripper=True) handle.entity.set_local_pose(object_pose) handle.entity.clear_dynamics() @@ -791,8 +923,13 @@ def _prepare_start( scenario.simulation.update(step=2) handle.entity.set_local_pose(object_pose) handle.entity.clear_dynamics() + scenario.set_robot_start( + arm_start, + open_gripper=False, + grasp_gripper=True, + ) object_to_eef = torch.bmm(torch.linalg.inv(object_pose), grasp_pose) - return handle, object_pose, grasp_pose, object_to_eef + return handle, object_pose, grasp_pose, object_to_eef, table_object_pose def initial_task_state( self, scenario: "AtomicTaskScenario", case: BenchmarkCase @@ -833,17 +970,25 @@ def generate_case( seed: int, batch_size: int, ) -> BenchmarkCase: - handle, object_pose, grasp_pose, object_to_eef = self._prepare_start( - scenario, config - ) + ( + handle, + object_pose, + grasp_pose, + object_to_eef, + _table_object_pose, + ) = self._prepare_start(scenario, config, seed=seed) target_object = object_pose.clone() - target_object[:, :3, 3] += torch.tensor( + target_object[:, :3, 3] += _randomized_vector( _float_vector( config.get("target_object_offset_m"), name="target_object_offset_m", length=3, default=(0.08, 0.08, 0.04), ), + config, + jitter_name="target_object_offset_jitter_m", + seed=seed, + stream=311, dtype=object_pose.dtype, device=object_pose.device, ) @@ -869,7 +1014,7 @@ def generate_case( object_id=handle.object_id, task_difficulty=_difficulty(config), primary_success="task_success", - full_start_qpos=scenario.robot.get_qpos().clone(), + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), case_parameters={ **self._held_parameters(handle, object_pose, grasp_pose, object_to_eef), "sample_count": int(config.get("sample_count", 80)), @@ -877,6 +1022,7 @@ def generate_case( "object_position_threshold_m": float( config.get("object_position_threshold_m", 0.04) ), + "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, ) @@ -944,17 +1090,25 @@ def generate_case( seed: int, batch_size: int, ) -> BenchmarkCase: - handle, object_pose, grasp_pose, object_to_eef = self._prepare_start( - scenario, config - ) - target_object = handle.initial_pose.clone() - target_object[:, :3, 3] += torch.tensor( + ( + handle, + object_pose, + grasp_pose, + object_to_eef, + table_object_pose, + ) = self._prepare_start(scenario, config, seed=seed) + target_object = table_object_pose.clone() + target_object[:, :3, 3] += _randomized_vector( _float_vector( config.get("target_object_offset_m"), name="target_object_offset_m", length=3, default=(0.10, 0.10, 0.0), ), + config, + jitter_name="target_object_offset_jitter_m", + seed=seed, + stream=321, dtype=object_pose.dtype, device=object_pose.device, ) @@ -985,7 +1139,7 @@ def generate_case( object_id=handle.object_id, task_difficulty=_difficulty(config), primary_success="task_success", - full_start_qpos=scenario.robot.get_qpos().clone(), + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), case_parameters={ **self._held_parameters(handle, object_pose, grasp_pose, object_to_eef), "sample_count": int(config.get("sample_count", 120)), @@ -996,6 +1150,7 @@ def generate_case( "object_position_threshold_m": float( config.get("object_position_threshold_m", 0.05) ), + "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, ) @@ -1070,18 +1225,23 @@ def generate_case( batch_size: int, ) -> BenchmarkCase: scenario.restore_base_robot() + scenario.randomize_robot_start(config, seed=seed, stream=31) 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 ) press_pose = start_pose.clone() - press_pose[:, :3, 3] += torch.tensor( + press_pose[:, :3, 3] += _randomized_vector( _float_vector( config.get("target_offset_m"), name="target_offset_m", length=3, default=(0.0, 0.0, -0.08), ), + config, + jitter_name="target_offset_jitter_m", + seed=seed, + stream=331, dtype=start_pose.dtype, device=start_pose.device, ) @@ -1105,11 +1265,12 @@ def generate_case( skill_id=self.skill_id, task_difficulty=_difficulty(config), primary_success="task_success", - full_start_qpos=scenario.robot.get_qpos().clone(), + full_start_qpos=_canonical_case_qpos(scenario.robot.get_qpos().clone()), case_parameters={ "sample_count": int(config.get("sample_count", 80)), "press_pose": press_pose.detach().cpu().tolist(), "hand_interp_steps": int(config.get("hand_interp_steps", 8)), + "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, ) @@ -1284,7 +1445,7 @@ def generate_cases( raise ValueError("atomic-task seeds must not be empty.") cases: list[BenchmarkCase] = [] - for skill_value in skill_values: + for skill_index, skill_value in enumerate(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") @@ -1298,20 +1459,33 @@ def generate_cases( for key, value in skill_value.items() if key not in {"id", "cases"} } - for raw_case in raw_cases: + for case_index, raw_case in enumerate(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, + rng_devices = ( + [robot.device] + if torch.device(robot.device).type == "cuda" + else [] ) + with torch.random.fork_rng(devices=rng_devices): + torch.manual_seed( + _case_generation_seed( + seed, + skill_index=skill_index, + case_index=case_index, + ) + ) + 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}." @@ -1771,6 +1945,7 @@ def solve_reference_qpos( raise RuntimeError( f"Independent IK rejected atomic target waypoint {index}." ) + seed = _canonical_case_qpos(seed) references.append(seed.clone()) return torch.stack(references, dim=1) @@ -1865,6 +2040,7 @@ def resolve_antipodal_grasp( if not bool(torch.as_tensor(success).all().item()): feasible = False break + seed = _canonical_case_qpos(seed) if feasible: return grasp raise RuntimeError( @@ -1880,6 +2056,102 @@ def restore_base_robot(self) -> None: self.robot.set_qpos(self._base_full_qpos, target=target) self.robot.clear_dynamics() + def randomize_robot_start( + self, + config: Mapping[str, object], + *, + seed: int, + stream: int, + ) -> torch.Tensor: + """Apply a deterministic, bounded perturbation to the arm start state.""" + if self.robot is None: + raise RuntimeError("Atomic Task runtime is not configured.") + start = self.robot.get_qpos(name=self.control_part).clone() + amplitude = _float_vector( + config.get("start_qpos_jitter_rad"), + name="start_qpos_jitter_rad", + length=start.shape[-1], + default=(0.0,) * start.shape[-1], + ) + if any(value < 0.0 for value in amplitude): + raise ValueError("start_qpos_jitter_rad values must be non-negative.") + randomized = start + _seeded_jitter( + amplitude, + seed=seed, + stream=stream, + dtype=start.dtype, + device=start.device, + ) + limits = self.robot.get_qpos_limits(name=self.control_part)[0] + margin = float(config.get("start_joint_limit_margin_rad", 0.05)) + if not math.isfinite(margin) or margin < 0.0: + raise ValueError("start_joint_limit_margin_rad must be non-negative.") + lower = limits[:, 0] + margin + upper = limits[:, 1] - margin + if bool(((randomized < lower) | (randomized > upper)).any().item()): + raise RuntimeError( + "Randomized arm start violates the configured joint-limit margin." + ) + self.set_robot_start( + randomized, + open_gripper=False, + grasp_gripper=False, + ) + return randomized + + def randomize_object_pose( + self, + handle: AtomicObjectHandle, + config: Mapping[str, object], + *, + seed: int, + stream: int, + ) -> torch.Tensor: + """Install one deterministic position/yaw perturbation for an object.""" + pose = handle.initial_pose.clone() + position_amplitude = _float_vector( + config.get("object_position_jitter_m"), + name="object_position_jitter_m", + length=3, + default=(0.0, 0.0, 0.0), + ) + if any(value < 0.0 for value in position_amplitude): + raise ValueError("object_position_jitter_m values must be non-negative.") + pose[:, :3, 3] += _seeded_jitter( + position_amplitude, + seed=seed, + stream=stream, + dtype=pose.dtype, + device=pose.device, + ) + + yaw_amplitude = float(config.get("object_yaw_jitter_rad", 0.0)) + if not math.isfinite(yaw_amplitude) or yaw_amplitude < 0.0: + raise ValueError("object_yaw_jitter_rad must be finite and non-negative.") + yaw = _seeded_jitter( + (yaw_amplitude,), + seed=seed, + stream=stream + 1, + dtype=pose.dtype, + device=pose.device, + )[0] + cosine = torch.cos(yaw) + sine = torch.sin(yaw) + yaw_rotation = torch.eye(3, dtype=pose.dtype, device=pose.device) + yaw_rotation[0, 0] = cosine + yaw_rotation[0, 1] = -sine + yaw_rotation[1, 0] = sine + yaw_rotation[1, 1] = cosine + pose[:, :3, :3] = yaw_rotation.unsqueeze(0) @ pose[:, :3, :3] + + handle.entity.set_local_pose(pose) + handle.entity.clear_dynamics() + if self.simulation is not None: + self.simulation.update(step=2) + handle.entity.set_local_pose(pose) + handle.entity.clear_dynamics() + return pose + def set_robot_start( self, manipulator_qpos: torch.Tensor, diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml new file mode 100644 index 000000000..59be0c2ae --- /dev/null +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml @@ -0,0 +1,189 @@ +schema_version: 1 +name: atomic_skill_franka_pgi_curobo_randomized +suite_version: atomic_franka_pgi_curobo_randomized_v1 +profile: coverage + +robot: + id: franka_pgi + provider: franka_pgi + config: {} + +planners: + - id: nmg + adapter: nmg_onnx + role: candidate + enabled: false + config: + model_revision: nmg-cartesian-k1-3-900m + onnx_model_path: null + num_waypoints: 3 + max_steps: 240 + action_scale: 0.2 + use_relative_obs: true + intermediate_orientation: true + pos_eps: 0.01 + rot_eps: 0.1 + joint_eps: 0.02 + policy_frame_from_world: + - [-1.0, 0.0, 0.0, 0.0] + - [0.0, -1.0, 0.0, 0.0] + - [0.0, 0.0, 1.0, 0.0] + - [0.0, 0.0, 0.0, 1.0] + runtime_tcp_from_policy_tcp: + - [0.70710678, 0.70710678, 0.0, 0.0] + - [-0.70710678, 0.70710678, 0.0, 0.0] + - [0.0, 0.0, 1.0, -0.0466] + - [0.0, 0.0, 0.0, 1.0] + - 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: 0 + 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-randomized + scenario: atomic_task + enabled: true + config: + batch_sizes: [1] + seeds: [101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116] + 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 + start_qpos_jitter_rad: [0.15, 0.15, 0.15, 0.12, 0.15, 0.12, 0.15] + target_offset_jitter_m: [0.03, 0.03, 0.02] + cases: + - name: randomized_two_waypoint + task_difficulty: medium + target_offsets_m: + - [-0.08, -0.08, 0.08] + - [0.04, 0.12, 0.04] + difficulty_factors: + waypoint_count: 2 + randomized_start: true + randomized_targets: true + - id: move_joints + sample_count: 80 + joint_limit_margin_rad: 0.05 + joint_threshold_rad: 0.02 + start_qpos_jitter_rad: [0.15, 0.15, 0.15, 0.12, 0.15, 0.12, 0.15] + target_offset_jitter_rad: [0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05] + cases: + - name: randomized_multi_joint + task_difficulty: medium + target_offsets_rad: + - [0.12, -0.08, 0.10, 0.12, -0.08, 0.10, -0.10] + difficulty_factors: + waypoint_count: 1 + active_joints: 7 + randomized_start: true + randomized_target: true + - id: pick_up + sample_count: 120 + grasp_source: fixed + 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 + object_position_jitter_m: [0.05, 0.05, 0.0] + object_yaw_jitter_rad: 0.40 + cases: + - name: randomized_cube_pick + object: cube_50mm + task_difficulty: medium + grasp_offset_m: [0.0, 0.0, 0.0] + difficulty_factors: + approach: top + randomized_object_pose: true + - id: move_held_object + sample_count: 80 + held_object_offset_m: [0.0, 0.0, 0.18] + held_object_offset_jitter_m: [0.02, 0.02, 0.015] + target_object_offset_m: [0.08, 0.08, 0.04] + target_object_offset_jitter_m: [0.04, 0.04, 0.02] + object_position_jitter_m: [0.04, 0.04, 0.0] + object_yaw_jitter_rad: 0.40 + object_position_threshold_m: 0.04 + cases: + - name: randomized_cube_transport + object: cube_50mm + task_difficulty: medium + difficulty_factors: + object_preheld: true + randomized_start_and_target: true + - id: place + sample_count: 120 + held_object_offset_m: [0.0, 0.0, 0.18] + held_object_offset_jitter_m: [0.02, 0.02, 0.015] + target_object_offset_m: [0.10, 0.10, 0.0] + target_object_offset_jitter_m: [0.04, 0.04, 0.0] + object_position_jitter_m: [0.04, 0.04, 0.0] + object_yaw_jitter_rad: 0.40 + retract_height_m: 0.10 + hand_interp_steps: 12 + object_position_threshold_m: 0.05 + cases: + - name: randomized_cube_place + object: cube_50mm + task_difficulty: medium + difficulty_factors: + object_preheld: true + randomized_start_and_target: true + - id: press + sample_count: 80 + hand_interp_steps: 8 + start_qpos_jitter_rad: [0.15, 0.15, 0.15, 0.12, 0.15, 0.12, 0.15] + target_offset_jitter_m: [0.03, 0.03, 0.015] + cases: + - name: randomized_vertical_press + task_difficulty: medium + target_offset_m: [0.0, 0.0, -0.08] + difficulty_factors: + randomized_start: true + randomized_target: true diff --git a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py index 061f89bba..c43309b07 100644 --- a/tests/benchmark/motion_generation/test_atomic_task_benchmark.py +++ b/tests/benchmark/motion_generation/test_atomic_task_benchmark.py @@ -46,6 +46,10 @@ ) from scripts.benchmark.motion_generation.scenarios.atomic_task import ( AtomicTaskScenario, + _canonical_case_qpos, + _case_generation_seed, + _randomization_parameters, + _seeded_jitter, atomic_skill_provider_names, create_atomic_skill_provider, ) @@ -140,6 +144,83 @@ def test_atomic_suite_is_franka_pgi_and_curobo_only(): } +def test_randomized_atomic_suite_covers_six_skills_with_fixed_seed_sweep(): + suite = load_suite("atomic_franka_pgi_curobo_randomized") + track = suite.enabled_tracks()[0] + + assert suite.suite_version == "atomic_franka_pgi_curobo_randomized_v1" + assert track.id == "atomic-task-randomized" + assert track.config["seeds"] == list(range(101, 117)) + assert [item["id"] for item in track.config["skills"]] == [ + "move_end_effector", + "move_joints", + "pick_up", + "move_held_object", + "place", + "press", + ] + assert all( + any(str(key).endswith(("_jitter_m", "_jitter_rad")) for key in item) + for item in track.config["skills"] + ) + + +def test_seeded_jitter_is_reproducible_bounded_and_stream_separated(): + kwargs = { + "amplitude": (0.1, 0.2, 0.0), + "seed": 107, + "dtype": torch.float64, + "device": torch.device("cpu"), + } + first = _seeded_jitter(stream=5, **kwargs) + repeated = _seeded_jitter(stream=5, **kwargs) + different_seed = _seeded_jitter(stream=5, **{**kwargs, "seed": 108}) + different_stream = _seeded_jitter(stream=6, **kwargs) + + assert torch.equal(first, repeated) + assert not torch.equal(first, different_seed) + assert not torch.equal(first, different_stream) + assert torch.all(torch.abs(first) <= torch.tensor((0.1, 0.2, 0.0))) + assert first.dtype == torch.float64 + + +def test_case_generation_seed_is_stable_and_stratified(): + value = _case_generation_seed(107, skill_index=2, case_index=3) + + assert value == _case_generation_seed(107, skill_index=2, case_index=3) + assert value != _case_generation_seed(108, skill_index=2, case_index=3) + assert value != _case_generation_seed(107, skill_index=3, case_index=3) + assert value != _case_generation_seed(107, skill_index=2, case_index=4) + + +def test_case_qpos_canonicalization_removes_sub_resolution_noise(): + first = torch.tensor([[0.123421, -0.456779]], dtype=torch.float32) + second = torch.tensor([[0.123419, -0.456781]], dtype=torch.float32) + + assert torch.equal(_canonical_case_qpos(first), _canonical_case_qpos(second)) + + +def test_randomization_manifest_records_contract_and_seed(): + parameters = _randomization_parameters( + { + "target_offset_jitter_m": [0.03, 0.02, 0.01], + "object_yaw_jitter_rad": 0.4, + "sample_count": 80, + }, + seed=113, + ) + + assert parameters == { + "enabled": True, + "seed": 113, + "distribution": "independent_uniform", + "ranges": { + "target_offset_jitter_m": [0.03, 0.02, 0.01], + "object_yaw_jitter_rad": 0.4, + }, + } + + 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() diff --git a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py index d98beaa66..94ed5f67e 100644 --- a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py +++ b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py @@ -299,6 +299,15 @@ def test_nmg_precision_and_external_accuracy_are_independently_configurable(): assert nmg.config["rot_eps"] == pytest.approx(0.20) +def test_seed_override_applies_to_atomic_tracks(): + suite = load_suite("atomic_franka_pgi_curobo_randomized") + + _apply_overrides(suite, seeds=[104]) + + assert suite.free_space.seeds == [104] + assert suite.enabled_tracks()[0].config["seeds"] == [104] + + @pytest.mark.parametrize("override", [{"nmg_pos_eps": 0.0}, {"nmg_rot_eps": -0.1}]) def test_nmg_precision_rejects_non_positive_values(override): suite = load_suite("smoke") From 1f284a4d854662c287782f3376aaf982f4727191 Mon Sep 17 00:00:00 2001 From: matafela Date: Mon, 31 Aug 2026 17:54:33 +0800 Subject: [PATCH 09/11] fix pick_up move_held_object --- .../scenarios/atomic_task.py | 20 +++-- .../test_motion_generation_benchmark.py | 85 +++++++++++++++++++ 2 files changed, 100 insertions(+), 5 deletions(-) diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py index 6b3a3aeb2..2dd320da3 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_task.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -49,6 +49,7 @@ PressGoal, PressOptions, TaskState, + create_simulation_atomic_action_engine, ) from embodichain.lab.sim.atomic_actions.plans import CompiledTrajectory from embodichain.lab.sim.planners.utils import PlanResult @@ -1047,7 +1048,7 @@ def build_invocation( strategy="motion_gen", sample_count=int(case.case_parameters["sample_count"]), ), - skill_options=MoveHeldObjectOptions(pick_rotate_upright=0.0), + skill_options=MoveHeldObjectOptions(), ) def task_result( @@ -1580,10 +1581,19 @@ def prepare_planner( grasp=self._gripper_grasp, ) } - self._engine = AtomicActionEngine( - motion_generator=adapter.require_motion_generator(), - control_profiles=control_profiles, - ) + motion_generator = adapter.require_motion_generator() + scene_entities = tuple(handle.entity for handle in self._objects.values()) + if scene_entities: + self._engine = create_simulation_atomic_action_engine( + motion_generator=motion_generator, + scene_entities=scene_entities, + control_profiles=control_profiles, + ) + else: + self._engine = AtomicActionEngine( + motion_generator=motion_generator, + control_profiles=control_profiles, + ) def close_planner(self, adapter: PlannerAdapter) -> None: """Drop the engine before its adapter closes the shared generator.""" diff --git a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py index 87667fcb2..ab06bdcf1 100644 --- a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py +++ b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py @@ -28,6 +28,8 @@ AtomicActionEngine, ControlPartCommandProfile, GraspGoal, + HeldObjectPoseGoal, + MoveHeldObjectOptions, ) from embodichain.lab.sim.planners.curobo.curobo_planner import CuroboPlanner from embodichain.lab.sim.planners.utils import MoveType, PlanResult @@ -454,6 +456,89 @@ def test_atomic_task_pickup_builds_engine_owned_endpoint_binding(): } +def test_atomic_task_prepare_planner_registers_benchmark_objects_in_scene(): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 7 + robot.control_parts = {"arm": object()} + robot.get_qpos.return_value = torch.zeros(1, 7) + robot.get_qvel.return_value = torch.zeros(1, 7) + motion_generator = Mock(robot=robot, device=torch.device("cpu")) + motion_generator.planner.cfg.planner_type = "stub" + adapter = Mock() + adapter.require_motion_generator.return_value = motion_generator + + object_pose = torch.eye(4).unsqueeze(0) + entity = Mock(uid="atomic_cube") + entity.get_local_pose.return_value = object_pose + scenario = AtomicTaskScenario() + scenario._objects["cube"] = Mock(object_id="cube", entity=entity) + + scenario.prepare_planner(adapter, Mock()) + + scene = scenario.require_engine().initial_context().scene + assert tuple(scene.entities) == ("atomic_cube",) + assert torch.equal(scene.entities["atomic_cube"].pose, object_pose) + + +def test_atomic_task_move_held_object_uses_current_options_contract(): + robot = Mock() + robot.device = torch.device("cpu") + robot.dof = 8 + robot.control_parts = {"arm": object(), "hand": object()} + robot.get_qpos.return_value = torch.zeros(1, 8) + robot.get_qvel.return_value = torch.zeros(1, 8) + robot.get_joint_ids.side_effect = lambda name: ( + list(range(7)) if name == "arm" else [7] + ) + motion_generator = Mock(robot=robot, device=torch.device("cpu")) + motion_generator.planner.cfg.planner_type = "stub" + engine = AtomicActionEngine( + motion_generator, + control_profiles={ + "hand": ControlPartCommandProfile.joint_positions( + open=torch.zeros(1), + grasp=torch.ones(1), + ) + }, + ) + scenario = AtomicTaskScenario() + scenario.robot = robot + scenario.control_part = "arm" + scenario.end_effector_part = "hand" + scenario._engine = engine + target_object_pose = torch.eye(4).unsqueeze(0) + case = BenchmarkCase( + suite_version="test_v1", + track="atomic-task", + scenario_id="move_held_object", + case_id="atomic-task:move_held_object:cube:s11", + seed=11, + batch_size=1, + num_waypoints=1, + path_shape="held_object_transport", + start_state_bin="object_held", + start_qpos=torch.zeros(1, 7), + target_waypoints=target_object_pose.unsqueeze(1), + reference_qpos=torch.zeros(1, 1, 7), + skill_id="move_held_object", + object_id="cube", + case_parameters={ + "sample_count": 12, + "target_object_pose": target_object_pose.tolist(), + }, + ) + + invocation = create_atomic_skill_provider("move_held_object").build_invocation( + scenario, + case, + Mock(), + ) + + assert isinstance(invocation.goal, HeldObjectPoseGoal) + assert type(invocation.skill_options) is MoveHeldObjectOptions + + @pytest.mark.parametrize("override", [{"nmg_pos_eps": 0.0}, {"nmg_rot_eps": -0.1}]) def test_nmg_precision_rejects_non_positive_values(override): suite = load_suite("smoke") From 04ef3e69f5dd1ccc7791de73c79dc0a24da8b46c Mon Sep 17 00:00:00 2001 From: matafela Date: Mon, 31 Aug 2026 18:33:50 +0800 Subject: [PATCH 10/11] fix press benchmark --- .../motion_generation/metrics/trajectory.py | 71 +++++++++++++++++-- .../scenarios/atomic_task.py | 1 + .../test_motion_generation_benchmark.py | 42 +++++++++++ 3 files changed, 109 insertions(+), 5 deletions(-) diff --git a/scripts/benchmark/motion_generation/metrics/trajectory.py b/scripts/benchmark/motion_generation/metrics/trajectory.py index cff281257..068323955 100644 --- a/scripts/benchmark/motion_generation/metrics/trajectory.py +++ b/scripts/benchmark/motion_generation/metrics/trajectory.py @@ -42,7 +42,10 @@ def _pose_error_matrices( - waypoints: torch.Tensor, trajectory_poses: torch.Tensor + waypoints: torch.Tensor, + trajectory_poses: torch.Tensor, + *, + rotation_symmetry: str | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Return waypoint-by-sample translation and geodesic rotation errors.""" waypoints = torch.as_tensor(waypoints, dtype=torch.float64) @@ -57,10 +60,31 @@ def _pose_error_matrices( relative = waypoint_rot.transpose(-1, -2) @ trajectory_rot trace = torch.diagonal(relative, dim1=-2, dim2=-1).sum(dim=-1) rotation_error = torch.arccos(torch.clamp((trace - 1.0) * 0.5, -1.0, 1.0)) + if rotation_symmetry == "half_turn_about_z": + symmetric_rot = waypoint_rot.clone() + symmetric_rot[..., :2] = -symmetric_rot[..., :2] + symmetric_relative = symmetric_rot.transpose(-1, -2) @ trajectory_rot + symmetric_trace = torch.diagonal(symmetric_relative, dim1=-2, dim2=-1).sum( + dim=-1 + ) + symmetric_error = torch.arccos( + torch.clamp((symmetric_trace - 1.0) * 0.5, -1.0, 1.0) + ) + rotation_error = torch.minimum(rotation_error, symmetric_error) + elif rotation_symmetry is not None: + raise ValueError( + "rotation_symmetry must be None or 'half_turn_about_z', " + f"got {rotation_symmetry!r}." + ) return pos_error, rotation_error -def get_pose_err(matrix_a: torch.Tensor, matrix_b: torch.Tensor) -> tuple[float, float]: +def get_pose_err( + matrix_a: torch.Tensor, + matrix_b: torch.Tensor, + *, + rotation_symmetry: str | None = None, +) -> tuple[float, float]: """Return translation (m) and geodesic rotation (rad) pose errors.""" tensor_a = torch.as_tensor(matrix_a, dtype=torch.float64) tensor_b = torch.as_tensor(matrix_b, dtype=torch.float64, device=tensor_a.device) @@ -72,6 +96,24 @@ def get_pose_err(matrix_a: torch.Tensor, matrix_b: torch.Tensor) -> tuple[float, relative = tensor_a[:, :3, :3].transpose(-1, -2) @ tensor_b[:, :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)) + if rotation_symmetry == "half_turn_about_z": + symmetric_b = tensor_b.clone() + symmetric_b[:, :3, :2] = -symmetric_b[:, :3, :2] + symmetric_relative = ( + tensor_a[:, :3, :3].transpose(-1, -2) @ symmetric_b[:, :3, :3] + ) + symmetric_trace = torch.diagonal(symmetric_relative, dim1=-2, dim2=-1).sum( + dim=-1 + ) + symmetric_rotation = torch.arccos( + torch.clamp((symmetric_trace - 1.0) * 0.5, -1.0, 1.0) + ) + rotation = torch.minimum(rotation, symmetric_rotation) + elif rotation_symmetry is not None: + raise ValueError( + "rotation_symmetry must be None or 'half_turn_about_z', " + f"got {rotation_symmetry!r}." + ) return float(translation.mean().item()), float(rotation.mean().item()) @@ -81,6 +123,7 @@ def match_ordered_waypoints( *, position_threshold_m: float, rotation_threshold_rad: float, + rotation_symmetry: str | None = None, ) -> dict[str, object]: """Evaluate ordered arrival and threshold-constrained waypoint errors. @@ -105,7 +148,11 @@ def match_ordered_waypoints( "min_position_errors_at_orientation_m": [], } - pos_error, rot_error = _pose_error_matrices(waypoints, trajectory_poses) + pos_error, rot_error = _pose_error_matrices( + waypoints, + trajectory_poses, + rotation_symmetry=rotation_symmetry, + ) arrival_indices: list[int] = [] next_sample = 0 for waypoint_index in range(waypoints.shape[0]): @@ -168,6 +215,7 @@ def compute_waypoint_errors( *, position_threshold_m: float = 0.01, rotation_threshold_rad: float = 0.1, + rotation_symmetry: str | None = None, ) -> dict[str, float]: """Return ordered, same-sample waypoint errors for one trajectory.""" if isinstance(trajectory_poses, list): @@ -183,6 +231,7 @@ def compute_waypoint_errors( waypoints, position_threshold_m=position_threshold_m, rotation_threshold_rad=rotation_threshold_rad, + rotation_symmetry=rotation_symmetry, ) pos_mm = [float(value) * 1000.0 for value in matched["position_errors_m"]] rot_deg = [ @@ -408,11 +457,15 @@ def compute_case_outcomes( ) waypoints = case.target_waypoints[env_index] + rotation_symmetry = case.case_parameters.get("waypoint_rotation_symmetry") + if rotation_symmetry is not None and not isinstance(rotation_symmetry, str): + raise TypeError("waypoint_rotation_symmetry must be a string or None.") matching = match_ordered_waypoints( poses, waypoints, position_threshold_m=position_threshold_m, rotation_threshold_rad=rotation_threshold_rad, + rotation_symmetry=rotation_symmetry, ) pos_errors_mm = [ float(value) * 1000.0 for value in matching["position_errors_m"] @@ -436,8 +489,16 @@ def compute_case_outcomes( motion_valid = finite and ordered and not joint_violation if poses.shape[0] > 0: - final_pos_m, final_rot_rad = get_pose_err(poses[-1], waypoints[-1]) - all_pos_error, all_rot_error = _pose_error_matrices(waypoints, native_poses) + final_pos_m, final_rot_rad = get_pose_err( + poses[-1], + waypoints[-1], + rotation_symmetry=rotation_symmetry, + ) + all_pos_error, all_rot_error = _pose_error_matrices( + waypoints, + native_poses, + rotation_symmetry=rotation_symmetry, + ) min_pos_m = float(all_pos_error[-1].min().item()) min_rot_rad = float(all_rot_error[-1].min().item()) joint_length, cartesian_length, efficiency = _path_metrics( diff --git a/scripts/benchmark/motion_generation/scenarios/atomic_task.py b/scripts/benchmark/motion_generation/scenarios/atomic_task.py index 2dd320da3..6ce1d9bce 100644 --- a/scripts/benchmark/motion_generation/scenarios/atomic_task.py +++ b/scripts/benchmark/motion_generation/scenarios/atomic_task.py @@ -1311,6 +1311,7 @@ def generate_case( "hand_interp_steps": options.hand_interp_steps, "approach_distance_m": options.approach_distance, "press_distance_m": options.press_distance, + "waypoint_rotation_symmetry": "half_turn_about_z", "randomization": _randomization_parameters(config, seed=seed), "difficulty_factors": dict(config.get("difficulty_factors", {})), }, diff --git a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py index ab06bdcf1..8ff3a2466 100644 --- a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py +++ b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py @@ -125,6 +125,48 @@ def test_ordered_waypoint_requires_position_and_rotation_at_same_sample(): assert result["arrival_indices"] == [] +def test_press_waypoint_validation_accepts_half_turn_about_z_symmetry(): + target = torch.eye(4) + target[:3, 0] = -target[:3, 0] + target[:3, 1] = -target[:3, 1] + positions = torch.zeros(1, 2, 7) + strict_case = replace( + _case(), + scenario_id="press", + skill_id="press", + target_waypoints=target.reshape(1, 1, 4, 4), + ) + symmetric_case = replace( + strict_case, + case_parameters={"waypoint_rotation_symmetry": "half_turn_about_z"}, + ) + + strict = compute_case_outcomes( + _timed_plan_result(positions, success=True), + strict_case, + _MetricRobot(), + "arm", + validation_samples=8, + position_threshold_m=1.0e-4, + rotation_threshold_rad=1.0e-4, + joint_limit_tolerance_rad=1.0e-5, + ) + symmetric = compute_case_outcomes( + _timed_plan_result(positions, success=True), + symmetric_case, + _MetricRobot(), + "arm", + validation_samples=8, + position_threshold_m=1.0e-4, + rotation_threshold_rad=1.0e-4, + joint_limit_tolerance_rad=1.0e-5, + ) + + assert strict[0].failure_code == "waypoint_miss" + assert symmetric[0].motion_valid is True + assert symmetric[0].final_rotation_err_deg == pytest.approx(0.0) + + def test_waypoint_errors_use_threshold_greedy_arrivals(): """Continuous errors must come from the same matching as motion_valid.""" waypoints = torch.stack([_translated_pose(0.0), _translated_pose(0.10)]) From 30f441036058b137a42567bf0fac95614fe3ddc1 Mon Sep 17 00:00:00 2001 From: yangchen73 <122090643@link.cuhk.edu.cn> Date: Tue, 1 Sep 2026 02:35:41 +0800 Subject: [PATCH 11/11] feat(motion-planning): support heterogeneous neural waypoints Allow capable planners to consume ordered pose and joint targets in one plan, update NeuralPlanner constraint masking, and derive NMG model revisions from the runtime ONNX filename instead of suite-specific hardcoding. --- embodichain/lab/sim/atomic_actions/core.py | 12 ++--- embodichain/lab/sim/planners/base_planner.py | 3 ++ .../lab/sim/planners/motion_generator.py | 31 +++++++++--- .../lab/sim/planners/neural_planner.py | 14 +++++- .../motion_generation/planners/base.py | 8 +-- .../motion_generation/planners/nmg_onnx.py | 11 +++- .../suites/atomic_franka_pgi_curobo.yaml | 1 - .../atomic_franka_pgi_curobo_randomized.yaml | 1 - .../motion_generation/suites/coverage.yaml | 1 - .../motion_generation/suites/smoke.yaml | 1 - .../test_motion_generation_benchmark.py | 23 +++++++++ .../planners/test_motion_generator_batched.py | 50 +++++++++++++++++++ tests/sim/planners/test_neural_batched.py | 24 ++++++++- tests/sim/planners/test_neural_planner.py | 40 +++++++++++++++ 14 files changed, 196 insertions(+), 24 deletions(-) diff --git a/embodichain/lab/sim/atomic_actions/core.py b/embodichain/lab/sim/atomic_actions/core.py index 231c4a271..d487e206d 100644 --- a/embodichain/lab/sim/atomic_actions/core.py +++ b/embodichain/lab/sim/atomic_actions/core.py @@ -545,13 +545,13 @@ def build_plan( raise ValueError("Trajectory and planning context batch sizes must match.") if timed.robot_dof != context.robot.robot_dof: raise ValueError("Trajectory robot_dof must match the planning context.") + planner = ( + None + if self._planning_services is None + else self._planning_services.motion_generator.planner + ) preserve_failed_positions = ( - getattr( - self.motion_generator.planner, - "preserve_failed_plan_positions", - False, - ) - is True + getattr(planner, "preserve_failed_plan_positions", False) is True ) if not preserve_failed_positions: timed = timed.hold_rows(success_mask, context.robot.qpos) diff --git a/embodichain/lab/sim/planners/base_planner.py b/embodichain/lab/sim/planners/base_planner.py index 5dabfc930..ab09c57ef 100644 --- a/embodichain/lab/sim/planners/base_planner.py +++ b/embodichain/lab/sim/planners/base_planner.py @@ -224,6 +224,9 @@ def __init__(self, cfg: BasePlannerCfg): waypoints for a joint-only backend. """ + supports_heterogeneous_waypoints: bool = False + """Whether one plan may contain an ordered mixture of movement types.""" + preserve_plan_samples: bool = False """Whether callers must retain this planner's returned sample points exactly. diff --git a/embodichain/lab/sim/planners/motion_generator.py b/embodichain/lab/sim/planners/motion_generator.py index 4139329d3..4f3ab626e 100644 --- a/embodichain/lab/sim/planners/motion_generator.py +++ b/embodichain/lab/sim/planners/motion_generator.py @@ -212,6 +212,11 @@ def supports_joint_trajectory_validation(self) -> bool: is True ) + @property + def supports_heterogeneous_waypoints(self) -> bool: + """Whether the selected backend accepts mixed waypoint movement types.""" + return getattr(self.planner, "supports_heterogeneous_waypoints", False) is True + @property def dynamic_collision_entity_ids(self) -> tuple[str, ...]: """Return canonical dynamic-obstacle IDs declared by the planner.""" @@ -488,11 +493,19 @@ def generate( ) move_types = {state.move_type for state in target_states} - if len(move_types) != 1: + heterogeneous = len(move_types) > 1 + if heterogeneous and not self.supports_heterogeneous_waypoints: names = sorted(move_type.name for move_type in move_types) - raise ValueError(f"All target states must share move_type; got {names}.") + raise ValueError( + f"{type(self.planner).__name__} does not support heterogeneous " + f"waypoints; got {names}." + ) + if heterogeneous and options.strategy == "ik_interp": + raise ValueError( + "strategy='ik_interp' does not support heterogeneous waypoints." + ) move_type = target_states[0].move_type - use_interpolation = ( + use_interpolation = not heterogeneous and ( options.preserve_cartesian_samples or options.strategy == "ik_interp" or ( @@ -512,9 +525,11 @@ def _generate_with_planner( options: MotionGenOptions, ) -> PlanResult: """Dispatch batched targets through the configured planner backend.""" + move_types = {state.move_type for state in target_states} move_type = target_states[0].move_type should_preinterpolate = ( - options.is_interpolate + len(move_types) == 1 + and options.is_interpolate and not self.planner.supports_move_type(MoveType.EEF_MOVE) and self.planner.supports_move_type(MoveType.JOINT_MOVE) ) @@ -566,9 +581,11 @@ def _generate_with_planner( else: target_plan_states = target_states - unsupported_move_types = ( - set() if self.planner.supports_move_type(move_type) else {move_type} - ) + unsupported_move_types = { + candidate + for candidate in move_types + if not self.planner.supports_move_type(candidate) + } if not should_preinterpolate and unsupported_move_types: unsupported_names = sorted( move_type.name for move_type in unsupported_move_types diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index ea0e9beb0..ddae08740 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -208,6 +208,7 @@ class NeuralPlanner(BasePlanner): """ supported_move_types = frozenset({MoveType.EEF_MOVE, MoveType.JOINT_MOVE}) + supports_heterogeneous_waypoints = True preserve_plan_samples = True """Keep native closed-loop states instead of distance-resampling them.""" preserve_failed_plan_positions = True @@ -600,8 +601,19 @@ def _build_obs( active_rel_quat = _quat_mul_xyzw(active_quat, inv_eef) if getattr(self, "_canonicalize_quat_obs", False): active_rel_quat = _canonicalize_quat_xyzw(active_rel_quat) + active_pos_mask = pos_mask[idx, active_idx_clamped].unsqueeze(-1) + active_rot_mask = rot_mask[idx, active_idx_clamped].unsqueeze(-1) + active_rel_quat = torch.where( + active_rot_mask > 0.5, + active_rel_quat, + identity.view(1, 4), + ) active_cart_rel = torch.cat( - [active_pos - ee_pose[:, :3], active_rel_quat], dim=-1 + [ + (active_pos - ee_pose[:, :3]) * active_pos_mask, + active_rel_quat, + ], + dim=-1, ) active_rel = torch.where( (joint_mask[idx, active_idx_clamped] > 0.5).unsqueeze(-1), diff --git a/scripts/benchmark/motion_generation/planners/base.py b/scripts/benchmark/motion_generation/planners/base.py index a8e85481b..534796ec0 100644 --- a/scripts/benchmark/motion_generation/planners/base.py +++ b/scripts/benchmark/motion_generation/planners/base.py @@ -68,13 +68,15 @@ def metadata(self) -> PlannerMetadata: adapter=self.spec.adapter, config_hash=stable_hash(self.spec.config), capabilities=self.capabilities, - model_revision=str( - self.spec.config.get("model_revision", self.model_revision) - ), + model_revision=self._resolved_model_revision(), supported_robots=(self.context.robot_id,), parameters=dict(self.spec.config), ) + def _resolved_model_revision(self) -> str: + """Resolve the recorded model identity for this planner instance.""" + return str(self.spec.config.get("model_revision", self.model_revision)) + def availability(self) -> tuple[bool, str | None]: """Return whether this adapter can run in the current process.""" return True, None diff --git a/scripts/benchmark/motion_generation/planners/nmg_onnx.py b/scripts/benchmark/motion_generation/planners/nmg_onnx.py index 6775488c1..6210a8f2a 100644 --- a/scripts/benchmark/motion_generation/planners/nmg_onnx.py +++ b/scripts/benchmark/motion_generation/planners/nmg_onnx.py @@ -51,12 +51,21 @@ class NmgOnnxAdapter(PlannerAdapter): "atomic_action", } ) - model_revision = "nmg-onnx-v1" def __init__(self, spec: PlannerSpecCfg, context: PlannerContext) -> None: super().__init__(spec, context) self.motion_generator: MotionGenerator | None = None + def _resolved_model_revision(self) -> str: + """Use an explicit revision or derive one from the runtime ONNX path.""" + configured = self.spec.config.get("model_revision") + if configured: + return str(configured) + model_path = self.spec.config.get("onnx_model_path") + if model_path: + return Path(str(model_path)).expanduser().stem + return super()._resolved_model_revision() + def availability(self) -> tuple[bool, str | None]: """Require ONNX Runtime and a resolved standalone policy path.""" if importlib.util.find_spec("onnxruntime") is None: diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml index a019216e8..0d3a53396 100644 --- a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo.yaml @@ -14,7 +14,6 @@ planners: role: candidate enabled: false config: - model_revision: nmg-cartesian-k1-3-900m onnx_model_path: null num_waypoints: 3 max_steps: 240 diff --git a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml index 59be0c2ae..f440671d7 100644 --- a/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml +++ b/scripts/benchmark/motion_generation/suites/atomic_franka_pgi_curobo_randomized.yaml @@ -14,7 +14,6 @@ planners: role: candidate enabled: false config: - model_revision: nmg-cartesian-k1-3-900m onnx_model_path: null num_waypoints: 3 max_steps: 240 diff --git a/scripts/benchmark/motion_generation/suites/coverage.yaml b/scripts/benchmark/motion_generation/suites/coverage.yaml index f34850330..fb8b2cf0d 100644 --- a/scripts/benchmark/motion_generation/suites/coverage.yaml +++ b/scripts/benchmark/motion_generation/suites/coverage.yaml @@ -42,7 +42,6 @@ planners: role: candidate enabled: false config: - model_revision: nmg-onnx-v1 onnx_model_path: null num_waypoints: 8 max_steps: 240 diff --git a/scripts/benchmark/motion_generation/suites/smoke.yaml b/scripts/benchmark/motion_generation/suites/smoke.yaml index cf37174d2..60a98c0e2 100644 --- a/scripts/benchmark/motion_generation/suites/smoke.yaml +++ b/scripts/benchmark/motion_generation/suites/smoke.yaml @@ -42,7 +42,6 @@ planners: role: candidate enabled: false config: - model_revision: nmg-onnx-v1 onnx_model_path: null num_waypoints: 8 max_steps: 240 diff --git a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py index 8ff3a2466..4425f4a08 100644 --- a/tests/benchmark/motion_generation/test_motion_generation_benchmark.py +++ b/tests/benchmark/motion_generation/test_motion_generation_benchmark.py @@ -368,6 +368,29 @@ def test_nmg_precision_and_external_accuracy_are_independently_configurable(): assert nmg.config["rot_eps"] == pytest.approx(0.20) +def test_nmg_model_revision_is_derived_from_runtime_model_path(): + from scripts.benchmark.motion_generation.config import PlannerSpecCfg + from scripts.benchmark.motion_generation.planners.base import PlannerContext + from scripts.benchmark.motion_generation.planners.nmg_onnx import NmgOnnxAdapter + + adapter = NmgOnnxAdapter( + PlannerSpecCfg( + id="nmg", + adapter="nmg_onnx", + role="candidate", + config={"onnx_model_path": "/models/unified-k3.onnx"}, + ), + PlannerContext( + robot=Mock(), + control_part="arm", + device=torch.device("cpu"), + sample_interval=1, + ), + ) + + assert adapter.metadata.model_revision == "unified-k3" + + def test_seed_override_applies_to_atomic_tracks(): suite = load_suite("atomic_franka_pgi_curobo_randomized") diff --git a/tests/sim/planners/test_motion_generator_batched.py b/tests/sim/planners/test_motion_generator_batched.py index 4b84350d9..9887c4c43 100644 --- a/tests/sim/planners/test_motion_generator_batched.py +++ b/tests/sim/planners/test_motion_generator_batched.py @@ -162,6 +162,56 @@ def test_direct_cartesian_planner_requires_joint_fallback_inputs(): ) +def test_motion_generator_dispatches_heterogeneous_waypoints_by_capability(): + planner = Mock() + planner.supports_heterogeneous_waypoints = True + planner.supports_move_type.side_effect = lambda move_type: move_type in { + MoveType.EEF_MOVE, + MoveType.JOINT_MOVE, + } + planner.preserve_plan_samples = True + planner.default_plan_options.return_value = PlanOptions() + planner.with_motion_context.side_effect = ( + lambda options, *, start_qpos, control_part: options + ) + planner.plan.return_value = PlanResult( + success=torch.ones(1, dtype=torch.bool), + positions=torch.zeros(1, 5, 2), + dt=torch.full((1, 5), 0.01), + ) + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = torch.device("cpu") + targets = [ + PlanState.from_xpos(torch.eye(4).unsqueeze(0)), + PlanState.from_qpos(torch.zeros(1, 2)), + ] + + result = generator.generate( + targets, + MotionGenOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + assert result.success.all().item() + assert planner.plan.call_args.kwargs["target_states"] is targets + + +def test_motion_generator_rejects_heterogeneous_waypoints_without_capability(): + planner = _DirectCartesianPlanner() + generator = object.__new__(MotionGenerator) + generator.planner = planner + generator.device = torch.device("cpu") + + with pytest.raises(ValueError, match="does not support heterogeneous"): + generator.generate( + [ + PlanState.from_xpos(torch.eye(4).unsqueeze(0)), + PlanState.from_qpos(torch.zeros(1, 2)), + ], + MotionGenOptions(start_qpos=torch.zeros(1, 2), control_part="arm"), + ) + + def test_bind_collision_world_copies_caller_options() -> None: planner = Mock() planner.collision_world_info = _collision_world_info(("obstacle",)) diff --git a/tests/sim/planners/test_neural_batched.py b/tests/sim/planners/test_neural_batched.py index c1e2acc73..6525ab9ee 100644 --- a/tests/sim/planners/test_neural_batched.py +++ b/tests/sim/planners/test_neural_batched.py @@ -32,6 +32,8 @@ def test_parse_waypoints_batched(self): planner = NeuralPlanner.__new__(NeuralPlanner) planner.device = torch.device("cpu") planner._num_waypoints = 4 + planner._action_dim = 7 + planner._intermediate_orientation = True _set_identity_policy_frames(planner) B = 3 states = [ @@ -41,10 +43,16 @@ def test_parse_waypoints_batched(self): ) for _ in range(2) ] - pos, quat, mask, k = planner._parse_waypoints(states) + pos, quat, joint, mask, pos_mask, rot_mask, joint_mask, k = ( + planner._parse_waypoints(states) + ) assert pos.shape == (B, 4, 3) assert quat.shape == (B, 4, 4) + assert joint.shape == (B, 4, 7) assert mask.shape == (B, 4) + assert torch.equal(pos_mask, mask) + assert torch.equal(rot_mask, mask) + assert torch.count_nonzero(joint_mask) == 0 assert k == 2 @@ -63,6 +71,7 @@ def test_plan_returns_batched_success(self, monkeypatch): planner._max_steps = 5 planner._pos_eps = 1e9 # always reached planner._rot_eps = 1e9 + planner._joint_eps = 1e9 planner._intermediate_orientation = True planner._use_relative_obs = False planner._obs_dim = 101 @@ -130,6 +139,7 @@ def test_converged_env_holds_qpos(self): planner._max_steps = 5 planner._pos_eps = 1e9 planner._rot_eps = 1e9 + planner._joint_eps = 1e9 planner._intermediate_orientation = True planner._use_relative_obs = False planner._obs_dim = 101 @@ -172,7 +182,17 @@ def compute_fk(self, qpos=None, name=None, to_matrix=True): # episode_k = 2 (2 waypoints). Env 0: reached=True every step -> # active_idx 0->1 (step0), 1->2 (step1) => converges after step 1. # Envs 1,2: reached=False always -> never converge within max_steps. - def fake_reach(ee_pose, wp_pos, wp_quat, active_idx, episode_k): + def fake_reach( + joint_pos, + ee_pose, + wp_pos, + wp_quat, + wp_joint, + pos_mask, + rot_mask, + joint_mask, + active_idx, + ): r = torch.zeros(ee_pose.shape[0], dtype=torch.bool, device=ee_pose.device) r[0] = True return r diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index d13447ed4..e7e878ee8 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -381,6 +381,46 @@ def test_neural_planner_builds_joint_constraint_observation(tmp_path, monkeypatc assert obs[0, waypoint_type_start].item() == pytest.approx(2.0) +def test_neural_planner_parses_ordered_pose_then_joint_sequence(tmp_path, monkeypatch): + model_path = _create_fake_onnx_model(tmp_path) + fake_sim = FakeSimulationManager() + monkeypatch.setattr( + SimulationManager, "get_instance", classmethod(lambda cls: fake_sim) + ) + planner = NeuralPlanner( + NeuralPlannerCfg( + robot_uid="fake_robot", + onnx_model_path=model_path, + control_part="main_arm", + ) + ) + pose = torch.eye(4).unsqueeze(0) + joint = torch.tensor([[0.1, -0.2, 0.3, 0.0, 0.2, -0.1, 0.4]]) + + ( + _, + _, + parsed_joint, + valid, + pos_mask, + rot_mask, + joint_mask, + episode_k, + ) = planner._parse_waypoints( + [ + PlanState.from_xpos(pose, move_type=MoveType.EEF_MOVE), + PlanState.from_qpos(joint, move_type=MoveType.JOINT_MOVE), + ] + ) + + assert episode_k == 2 + assert torch.equal(valid[0, :2], torch.ones(2)) + assert torch.equal(pos_mask[0, :2], torch.tensor([1.0, 0.0])) + assert torch.equal(rot_mask[0, :2], torch.tensor([1.0, 0.0])) + assert torch.equal(joint_mask[0, :2], torch.tensor([0.0, 1.0])) + assert torch.allclose(parsed_joint[0, 1], joint[0]) + + def test_neural_planner_accepts_joint_move_goal(tmp_path, monkeypatch): model_path = _create_fake_onnx_model(tmp_path) fake_sim = FakeSimulationManager()