# ----------------------------------------------------------------------------
# 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.
# ----------------------------------------------------------------------------
"""Run Motion Planning and Atomic Skills suites through the shared benchmark.
cuRobo is the default primary baseline. IK interpolation and TOPPRA are
optional diagnostic baselines. NMG is loaded from a standalone ONNX policy.
Run: ``embodichain benchmark motion-generation --suite
motion_panda_free_space_smoke`` or select the Franka + PGI Atomic Skills suite
with ``--suite atomic_franka_pgi_smoke``.
"""
from __future__ import annotations
import argparse
from pathlib import Path
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
__all__ = [
"add_parser_arguments",
"run_all_benchmarks",
"run_from_args",
"main",
]
[docs]
def add_parser_arguments(parser: argparse.ArgumentParser) -> None:
"""Add planner benchmark options to an existing argument parser."""
parser.add_argument(
"--suite",
default="motion_panda_free_space_smoke",
help=(
"Suite file stem or explicit YAML path. Motion Planning examples: "
"motion_panda_free_space_smoke, motion_panda_free_space_coverage, "
"motion_ur5_free_space_coverage. Atomic Skills examples: "
"atomic_franka_pgi_smoke, atomic_franka_pgi_randomized, "
"atomic_franka_pgi_pose_batch, atomic_ur5_pgi_randomized. "
"Default: motion_panda_free_space_smoke."
),
)
parser.add_argument(
"--algorithms",
nargs="+",
default=None,
help="Override enabled suite algorithms by id.",
)
parser.add_argument(
"--extra-baselines",
nargs="+",
choices=("ik_interpolate", "toppra"),
default=[],
help="Enable optional diagnostic baselines.",
)
parser.add_argument(
"--device",
choices=("auto", "cpu", "cuda"),
default="auto",
help="Simulation device; cuRobo itself always requires CUDA.",
)
parser.add_argument("--batch-sizes", nargs="+", type=int, default=None)
parser.add_argument("--num-waypoints", nargs="+", type=int, default=None)
parser.add_argument("--path-shapes", nargs="+", default=None)
parser.add_argument("--start-state-bins", nargs="+", default=None)
parser.add_argument("--seeds", nargs="+", type=int, default=None)
parser.add_argument("--num-trials", type=int, default=None)
parser.add_argument("--warmup-trials", type=int, default=None)
parser.add_argument("--sample-interval", type=int, default=None)
parser.add_argument("--validation-samples", type=int, default=None)
parser.add_argument("--position-threshold-m", type=float, default=None)
parser.add_argument("--rotation-threshold-rad", type=float, default=None)
parser.add_argument(
"--nmg-pos-eps",
type=float,
default=None,
help="NMG internal waypoint position threshold in metres.",
)
parser.add_argument(
"--nmg-rot-eps",
type=float,
default=None,
help="NMG internal waypoint rotation threshold in radians.",
)
parser.add_argument(
"--nmg-onnx-path",
default=None,
help="Path to the standalone NMG ONNX policy.",
)
parser.add_argument(
"--output-root", default="outputs/benchmarks", help="Artifact root directory."
)
parser.add_argument(
"--headless", action="store_true", default=True, help="Run headlessly."
)
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 Skills 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 <run_dir>/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:
"""Resolve one planner id from the suite with an actionable error."""
for spec in suite.planners:
if spec.id == planner_id:
return spec
raise ValueError(
f"Suite {suite.name!r} does not declare planner {planner_id!r}; "
f"available ids: {[spec.id for spec in suite.planners]}."
)
def _resolve_planners(
suite: SuiteCfg,
algorithms: list[str] | None,
extra_baselines: list[str],
) -> list[PlannerSpecCfg]:
"""Resolve enabled algorithms while retaining suite ordering and roles."""
selected_ids = (
list(algorithms)
if algorithms is not None
else [spec.id for spec in suite.planners if spec.enabled]
)
selected_ids = list(dict.fromkeys(selected_ids))
for planner_id in extra_baselines:
if planner_id not in selected_ids:
selected_ids.append(planner_id)
if not selected_ids:
raise ValueError("No algorithms were selected for the benchmark.")
return [_planner_by_id(suite, planner_id) for planner_id in selected_ids]
def _apply_overrides(
suite: SuiteCfg,
*,
batch_sizes: list[int] | None = None,
num_waypoints: list[int] | None = None,
path_shapes: list[str] | None = None,
start_state_bins: list[str] | None = None,
seeds: list[int] | None = None,
num_trials: int | None = None,
warmup_trials: int | None = None,
sample_interval: int | None = None,
validation_samples: int | None = None,
position_threshold_m: float | None = None,
rotation_threshold_rad: float | None = None,
nmg_pos_eps: float | None = None,
nmg_rot_eps: float | None = None,
nmg_onnx_path: str | None = None,
) -> None:
"""Apply explicit CLI/programmatic overrides to a loaded suite."""
free_space_overrides = {
key: value
for key, value in {
"batch_sizes": batch_sizes,
"waypoint_counts": num_waypoints,
"path_shapes": path_shapes,
"start_state_bins": start_state_bins,
"seeds": seeds,
}.items()
if value is not None
}
if free_space_overrides:
for track in suite.tracks:
if track.scenario == "free_space":
for key, value in free_space_overrides.items():
track.config[key] = list(value)
elif track.scenario == "atomic_skill" and seeds is not None:
track.config["seeds"] = list(seeds)
if num_trials is not None:
suite.protocol.measured_trials = num_trials
if warmup_trials is not None:
suite.protocol.warmup_trials = warmup_trials
if sample_interval is not None:
suite.protocol.sample_interval = sample_interval
if validation_samples is not None:
suite.protocol.validation_samples = validation_samples
if position_threshold_m is not None:
suite.protocol.position_threshold_m = position_threshold_m
if rotation_threshold_rad is not None:
suite.protocol.rotation_threshold_rad = rotation_threshold_rad
# Every NMG row, not just the one named "nmg": a suite can pair a nominal
# and a retimed arm, and leaving the paired row without a checkpoint would
# silently drop it from availability and run no comparison at all. The
# paired rows must also keep identical tolerances, or the comparison would
# measure a tolerance difference alongside the timing one.
for nmg in (spec for spec in suite.planners if spec.adapter == "nmg_onnx"):
if nmg_pos_eps is not None:
nmg.config["pos_eps"] = nmg_pos_eps
if nmg_rot_eps is not None:
nmg.config["rot_eps"] = nmg_rot_eps
if nmg_onnx_path is not None:
nmg.config["onnx_model_path"] = str(Path(nmg_onnx_path))
suite.validate_benchmark()
[docs]
def run_all_benchmarks(
num_waypoints_list: list[int] | None = None,
sim_device: str = "auto",
headless: bool = True,
nmg_onnx_path: str | None = None,
*,
suite_name: str = "motion_panda_free_space_smoke",
algorithms: list[str] | None = None,
extra_baselines: list[str] | None = None,
batch_sizes: list[int] | None = None,
path_shapes: list[str] | None = None,
start_state_bins: list[str] | None = None,
seeds: list[int] | None = None,
num_trials: int | None = None,
warmup_trials: int | None = None,
sample_interval: int | None = None,
validation_samples: int | None = None,
position_threshold_m: float | None = None,
rotation_threshold_rad: float | None = None,
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 .api import run_benchmark
suite = load_suite(suite_name)
_apply_overrides(
suite,
batch_sizes=batch_sizes,
num_waypoints=num_waypoints_list,
path_shapes=path_shapes,
start_state_bins=start_state_bins,
seeds=seeds,
num_trials=num_trials,
warmup_trials=warmup_trials,
sample_interval=sample_interval,
validation_samples=validation_samples,
position_threshold_m=position_threshold_m,
rotation_threshold_rad=rotation_threshold_rad,
nmg_pos_eps=nmg_pos_eps,
nmg_rot_eps=nmg_rot_eps,
nmg_onnx_path=nmg_onnx_path,
)
specs = _resolve_planners(suite, algorithms, list(extra_baselines or []))
return run_benchmark(
suite=suite,
planners=specs,
device=sim_device,
headless=headless,
output_root=output_root,
video=video,
)
[docs]
def run_from_args(args: argparse.Namespace) -> BenchmarkRunResult:
"""Run the benchmark from parsed unified-CLI arguments."""
return run_all_benchmarks(
num_waypoints_list=args.num_waypoints,
sim_device=args.device,
headless=args.headless,
nmg_onnx_path=args.nmg_onnx_path,
suite_name=args.suite,
algorithms=args.algorithms,
extra_baselines=args.extra_baselines,
batch_sizes=args.batch_sizes,
path_shapes=args.path_shapes,
start_state_bins=args.start_state_bins,
seeds=args.seeds,
num_trials=args.num_trials,
warmup_trials=args.warmup_trials,
sample_interval=args.sample_interval,
validation_samples=args.validation_samples,
position_threshold_m=args.position_threshold_m,
rotation_threshold_rad=args.rotation_threshold_rad,
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),
)
def _parse_args() -> argparse.Namespace:
"""Parse standalone module arguments using the unified option schema."""
parser = argparse.ArgumentParser(
description="Compare planners with Motion Planning or Atomic Skills suites."
)
add_parser_arguments(parser)
return parser.parse_args()
[docs]
def main() -> None:
"""Run the motion-generation CLI using process arguments."""
run_from_args(_parse_args())
if __name__ == "__main__":
main()