diff --git a/benchmark/benchmark_fast_sac_backends.py b/benchmark/benchmark_fast_sac_backends.py index 826bf1500..b37973f6c 100755 --- a/benchmark/benchmark_fast_sac_backends.py +++ b/benchmark/benchmark_fast_sac_backends.py @@ -12,6 +12,11 @@ ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) +try: + from benchmark.core.task_names import locomotion_env_name, normalize_locomotion_task_id +except ModuleNotFoundError: + from core.task_names import locomotion_env_name, normalize_locomotion_task_id + @dataclass class BackendResult: @@ -41,8 +46,11 @@ def run_backend(task, max_iterations, backend, num_envs): timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") log_dir = os.path.join(ROOT_DIR, "logs", "benchmark", f"{backend}_{timestamp}") + task_id = normalize_locomotion_task_id(task) + env_task_name = locomotion_env_name(task_id) + runner = FastSACRunner( - env_name=task, + env_name=env_task_name, device=device, num_envs=num_envs, replay_buffer_n=cfg.replay_buffer_n, @@ -105,18 +113,22 @@ def run_backend(task, max_iterations, backend, num_envs): def main(): parser = argparse.ArgumentParser() - parser.add_argument("--task", type=str, default="Go1JoystickFlatTerrain") + parser.add_argument("--task", type=str, default="go1_joystick") parser.add_argument("--max_iterations", type=int, default=100) parser.add_argument("--num_envs", type=int, default=4096) parser.add_argument("--backends", type=str, default="torch_mps,torch_cpu,ane") parser.add_argument("--output", type=str, default="benchmark/outputs/fast_sac_backends.json") args = parser.parse_args() + task_id = normalize_locomotion_task_id(args.task) + env_task_name = locomotion_env_name(task_id) backends = [b.strip() for b in args.backends.split(",")] print(f"\n{'=' * 70}") print("Fast SAC Backend Benchmark") - print(f"Task: {args.task}, Iterations: {args.max_iterations}, Envs: {args.num_envs}") + print( + f"Task: {task_id} (env: {env_task_name}), Iterations: {args.max_iterations}, Envs: {args.num_envs}" + ) print(f"Backends: {backends}") print(f"{'=' * 70}\n") @@ -125,7 +137,7 @@ def main(): for backend in backends: print(f"\n>>> Running {backend}...") try: - result = run_backend(args.task, args.max_iterations, backend, args.num_envs) + result = run_backend(task_id, args.max_iterations, backend, args.num_envs) results.append(result) print( f"✓ {backend}: {result.total_time_sec:.1f}s, reward={result.final_reward:.3f}, {result.steps_per_sec:.0f} steps/s" @@ -145,7 +157,8 @@ def main(): output_path.write_text( json.dumps( { - "task": args.task, + "task": task_id, + "env_task_name": env_task_name, "max_iterations": args.max_iterations, "num_envs": args.num_envs, "results": [asdict(r) for r in results], diff --git a/benchmark/benchmark_forward_reset_methods.py b/benchmark/benchmark_forward_reset_methods.py index 628c3ff3f..ba79a5845 100644 --- a/benchmark/benchmark_forward_reset_methods.py +++ b/benchmark/benchmark_forward_reset_methods.py @@ -9,7 +9,7 @@ import mujoco import numpy as np -from mujoco import batch_forward +from mujoco.batch_env import BatchEnvPool try: from benchmark.core.device_info import get_device_info_dict, get_device_info_line @@ -25,17 +25,21 @@ if str(ROOT_DIR) not in sys.path: sys.path.append(str(ROOT_DIR)) -from unilab.envs.locomotion.g1.joystick import G1JoystickCfg -from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg -from unilab.envs.locomotion.go2.joystick import Go2JoystickCfg +try: + from benchmark.core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) +except ModuleNotFoundError: + from core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) -CONSISTENCY_TASKS = ["Go1JoystickFlatTerrain", "Go2JoystickFlatTerrain", "G1JoystickFlatTerrain"] -RESET_TASK = "Go2JoystickFlatTerrain" -TASK_CFG_MAP = { - "Go1JoystickFlatTerrain": Go1JoystickCfg, - "Go2JoystickFlatTerrain": Go2JoystickCfg, - "G1JoystickFlatTerrain": G1JoystickCfg, -} +CONSISTENCY_TASKS = canonical_locomotion_task_ids() +DEFAULT_RESET_TASK = "go2_joystick" @dataclass @@ -57,7 +61,8 @@ class SpeedRecord: def load_model_for_task(task: str) -> mujoco.MjModel: - return mujoco.MjModel.from_xml_path(TASK_CFG_MAP[task]().model_file) + task_key = normalize_locomotion_task_id(task) + return mujoco.MjModel.from_xml_path(locomotion_task_spec(task_key).config_cls().model_file) def _set_state_for_forward(model, data, state): @@ -90,18 +95,8 @@ def python_forward_chunk_loop(model, states, chunk_size): return out -def cpp_batch_forward(model, states, nthread, chunk_size): - workers = [mujoco.MjData(model) for _ in range(nthread)] - with batch_forward.BatchForwardRunner(nthread=nthread) as runner: - sensordata = runner.forward( - model=[model] * states.shape[0], - data=workers, - initial_state=states, - chunk_size=chunk_size, - skipsensor=False, - out_dtype=np.float64, - return_state=False, - ) +def batch_env_forward(pool: BatchEnvPool, states: np.ndarray) -> np.ndarray: + sensordata = pool.forward(states) return np.asarray(sensordata, dtype=np.float64) @@ -148,21 +143,24 @@ def collect_random_reset_states(task, batch_size, random_rounds, seed): def run_consistency(tasks, batch_size, random_rounds, seed, atol, rtol, chunk_size): + _ = chunk_size # BatchEnvPool.forward does not expose chunk-size tuning. records = [] for task in tasks: + task_key = normalize_locomotion_task_id(task) model, states = collect_random_reset_states( - task, batch_size, random_rounds, seed + abs(hash(task)) % 10000 + task_key, batch_size, random_rounds, seed + abs(hash(task_key)) % 10000 ) nthread = min(batch_size, cpu_count()) py_sensor = python_forward_for_loop(model, states) - cpp_sensor = cpp_batch_forward(model, states, nthread, chunk_size) - diff = np.abs(py_sensor - cpp_sensor) + with BatchEnvPool(model, nbatch=states.shape[0], nthread=nthread) as pool: + batch_sensor = batch_env_forward(pool, states) + diff = np.abs(py_sensor - batch_sensor) max_abs = float(np.max(diff)) mean_abs = float(np.mean(diff)) - ok = bool(np.allclose(py_sensor, cpp_sensor, atol=atol, rtol=rtol)) - records.append(ConsistencyRecord(task, batch_size, max_abs, mean_abs, ok)) + ok = bool(np.allclose(py_sensor, batch_sensor, atol=atol, rtol=rtol)) + records.append(ConsistencyRecord(task_key, batch_size, max_abs, mean_abs, ok)) print( - f"[Consistency] {task}: allclose={ok}, max_abs={max_abs:.3e}, mean_abs={mean_abs:.3e}" + f"[Consistency] {task_key}: allclose={ok}, max_abs={max_abs:.3e}, mean_abs={mean_abs:.3e}" ) return records @@ -180,39 +178,42 @@ def _bench_method(func, repeats): def run_reset_speed(task, env_nums, random_rounds, chunk_size, repeats, seed): + _ = chunk_size # BatchEnvPool.forward does not expose chunk-size tuning. + task_key = normalize_locomotion_task_id(task) records = [] for env_num in env_nums: - model, states = collect_random_reset_states(task, env_num, random_rounds, seed + env_num) + model, states = collect_random_reset_states( + task_key, env_num, random_rounds, seed + env_num + ) nthread = min(env_num, cpu_count()) t_for = _bench_method(lambda: python_forward_for_loop(model, states), repeats) t_chunk = _bench_method( lambda: python_forward_chunk_loop(model, states, chunk_size), repeats ) - t_cpp = _bench_method( - lambda: cpp_batch_forward(model, states, nthread, chunk_size), repeats - ) + with BatchEnvPool(model, nbatch=states.shape[0], nthread=nthread) as pool: + t_batch = _bench_method(lambda: batch_env_forward(pool, states), repeats) for method, elapsed in [ ("for_loop", t_for), ("chunk_for_loop", t_chunk), - ("cpp_batch_forward", t_cpp), + ("batch_env_forward", t_batch), ]: records.append( - SpeedRecord(task, method, env_num, elapsed, elapsed * 1e6 / max(env_num, 1)) + SpeedRecord(task_key, method, env_num, elapsed, elapsed * 1e6 / max(env_num, 1)) ) print( - f"[Speed] {task} env={env_num}: for={t_for * 1e3:.3f}ms, chunk={t_chunk * 1e3:.3f}ms, cpp={t_cpp * 1e3:.3f}ms" + f"[Speed] {task_key} env={env_num}: for={t_for * 1e3:.3f}ms, chunk={t_chunk * 1e3:.3f}ms, batch_env={t_batch * 1e3:.3f}ms" ) return records -def plot_speed(records, out_png): +def plot_speed(records, out_png, reset_task): out_png.parent.mkdir(parents=True, exist_ok=True) env_nums = sorted({r.env_num for r in records}) - methods = ["for_loop", "chunk_for_loop", "cpp_batch_forward"] - colors = {"for_loop": "#3B82F6", "chunk_for_loop": "#10B981", "cpp_batch_forward": "#F59E0B"} + methods = ["for_loop", "chunk_for_loop", "batch_env_forward"] + colors = {"for_loop": "#3B82F6", "chunk_for_loop": "#10B981", "batch_env_forward": "#F59E0B"} fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5), sharex=True) x = np.arange(len(env_nums)) @@ -246,7 +247,9 @@ def plot_speed(records, out_png): ax2.set_yscale("log") ax2.set_ylabel("throughput (env/s)") ax2.set_title("Throughput") - fig.suptitle(f"Reset-forward method speed on {RESET_TASK}\n{get_device_info_line()}") + fig.suptitle( + f"Reset-forward method speed on {normalize_locomotion_task_id(reset_task)}\n{get_device_info_line()}" + ) ax2.legend() fig.tight_layout() fig.savefig(out_png, dpi=180) @@ -256,10 +259,16 @@ def plot_speed(records, out_png): def main(): parser = argparse.ArgumentParser(description="Benchmark C++ batch forward vs Python forward") parser.add_argument("--consistency-tasks", type=str, default=",".join(CONSISTENCY_TASKS)) + parser.add_argument("--reset-task", type=str, default=DEFAULT_RESET_TASK) parser.add_argument("--consistency-batch", type=int, default=1024) parser.add_argument("--env-nums", type=str, default="256,512,1024,2048,4096,8192,16384") parser.add_argument("--random-rounds", type=int, default=8) - parser.add_argument("--chunk-size", type=int, default=64) + parser.add_argument( + "--chunk-size", + type=int, + default=64, + help="Legacy no-op retained for CLI compatibility; BatchEnvPool.forward does not use chunk_size.", + ) parser.add_argument("--repeats", type=int, default=8) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--atol", type=float, default=1e-9) @@ -272,7 +281,10 @@ def main(): ) args = parser.parse_args() - consistency_tasks = [t.strip() for t in args.consistency_tasks.split(",") if t.strip()] + consistency_tasks = [ + normalize_locomotion_task_id(t) for t in args.consistency_tasks.split(",") if t.strip() + ] + reset_task = normalize_locomotion_task_id(args.reset_task) env_nums = [int(x.strip()) for x in args.env_nums.split(",") if x.strip()] consistency = run_consistency( @@ -285,7 +297,7 @@ def main(): args.chunk_size, ) speed = run_reset_speed( - RESET_TASK, env_nums, args.random_rounds, args.chunk_size, args.repeats, args.seed + reset_task, env_nums, args.random_rounds, args.chunk_size, args.repeats, args.seed ) out_json = Path(args.out_json) @@ -301,7 +313,7 @@ def main(): "timestamp": datetime.now(timezone.utc).isoformat(), "device_info": get_device_info_dict(), "consistency_tasks": consistency_tasks, - "reset_task": RESET_TASK, + "reset_task": reset_task, "env_nums": env_nums, "chunk_size": args.chunk_size, }, @@ -312,7 +324,7 @@ def main(): indent=2, ) - plot_speed(speed, Path(args.out_png)) + plot_speed(speed, Path(args.out_png), reset_task=reset_task) print(f"Consistency: {all(r.allclose for r in consistency)}") print(f"Saved: {out_json}, {args.out_png}") diff --git a/benchmark/benchmark_physics_step_genesis.py b/benchmark/benchmark_physics_step_genesis.py new file mode 100644 index 000000000..37fb6d53e --- /dev/null +++ b/benchmark/benchmark_physics_step_genesis.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python3 +""" +Benchmark Genesis physics execution. + +Benchmarks Genesis across current locomotion owner task ids +(go1_joystick/go2_joystick/g1_joystick) and outputs JSON + plots +aligned with benchmark/benchmark_physics_step_mj_step.py. +Legacy env names remain accepted as aliases. + +Run without changing repo dependencies: + uv run --with genesis-world \ + python benchmark/benchmark_physics_step_genesis.py +""" + +from __future__ import annotations + +import argparse +import json +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +try: + import genesis as gs +except ImportError: + gs = None + +try: + from benchmark.core.device_info import get_device_info_dict, get_device_info_line + from benchmark.core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) +except ModuleNotFoundError: + from core.device_info import get_device_info_dict, get_device_info_line + from core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) + + +@dataclass +class BenchRecord: + task: str + backend: str + batch_size: int + nstep: int + nthread: int + avg_time_sec: float + sps: float + + +DEFAULT_TASK_IDS = canonical_locomotion_task_ids() +DEFAULT_BATCH_SIZES = [2**k for k in range(8, 15)] # 256 .. 16384 +_GENESIS_INITIALIZED = False + + +def _display_backend(backend: str) -> str: + return backend + + +def _uv_run_hint() -> str: + return "uv run --with genesis-world python benchmark/benchmark_physics_step_genesis.py" + + +def _require_genesis() -> None: + if gs is None: + raise RuntimeError( + "genesis is unavailable in the current environment. " + f"Run with temporary deps instead of editing pyproject.toml:\n {_uv_run_hint()}" + ) + + +def _init_genesis() -> None: + global _GENESIS_INITIALIZED + if _GENESIS_INITIALIZED: + return + gs.init(backend=gs.gpu) + _GENESIS_INITIALIZED = True + + +def _load_task_xml(task_name: str) -> str: + cfg = locomotion_task_spec(task_name).config_cls() + return str(cfg.model_file) + + +def _build_scene(xml_path: str, batch_size: int): + scene = gs.Scene( + show_viewer=False, + rigid_options=gs.options.RigidOptions( + dt=0.01, + constraint_solver=gs.constraint_solver.Newton, + ), + ) + scene.add_entity(gs.morphs.MJCF(file=xml_path)) + scene.build(n_envs=batch_size) + return scene + + +def _run_scene(scene, nstep: int, niter: int) -> float: + t0 = time.perf_counter() + for _ in range(niter): + for _ in range(nstep): + scene.step() + return (time.perf_counter() - t0) / niter + + +def _bench_one_task( + task_name: str, + batch_sizes: List[int], + nstep: int, + warmup: int, + iters: int, +) -> List[BenchRecord]: + task_key = normalize_locomotion_task_id(task_name) + xml_path = _load_task_xml(task_key) + + records: List[BenchRecord] = [] + for batch_size in batch_sizes: + scene = _build_scene(xml_path, batch_size) + try: + _run_scene(scene, nstep=nstep, niter=warmup) + avg_t = _run_scene(scene, nstep=nstep, niter=iters) + finally: + del scene + + records.append( + BenchRecord( + task=task_key, + backend="genesis", + batch_size=batch_size, + nstep=nstep, + nthread=0, + avg_time_sec=avg_t, + sps=batch_size * nstep / avg_t, + ) + ) + print( + f"[{task_key}] batch={batch_size:5d} " + f"genesis={avg_t * 1000:.3f}ms ({batch_size * nstep / avg_t / 1e4:.2f}万fps)" + ) + + return records + + +def _plot_fps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): + out_png.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Total FPS\n{get_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / 1e4 for r in recs], + marker="o", + linestyle="-", + label=f"{locomotion_task_spec(task_name).display_name} ({_display_backend(backend)})", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Total FPS (x1e4)") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() + fig.savefig(out_png, dpi=150) + plt.close(fig) + print(f"Saved FPS plot to {out_png}") + + +def _plot_per_env_sps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): + out_png.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Per-Env SPS\n{get_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / r.batch_size for r in recs], + marker="o", + linestyle="-", + label=f"{locomotion_task_spec(task_name).display_name} ({_display_backend(backend)})", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Per-Env Steps/s") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() + fig.savefig(out_png, dpi=150) + plt.close(fig) + print(f"Saved per-env SPS plot to {out_png}") + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark Genesis physics execution") + parser.add_argument("--nstep", type=int, default=20) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=3) + parser.add_argument( + "--tasks", + type=str, + default=",".join(DEFAULT_TASK_IDS), + ) + parser.add_argument( + "--batch-sizes", type=str, default=",".join(str(x) for x in DEFAULT_BATCH_SIZES) + ) + parser.add_argument( + "--out-json", type=str, default="benchmark/outputs/physics_step/genesis/results.json" + ) + parser.add_argument( + "--out-dir", + type=str, + default="benchmark/outputs/physics_step/genesis", + help="Directory for output plots", + ) + args = parser.parse_args() + + _require_genesis() + _init_genesis() + + task_names = [normalize_locomotion_task_id(x) for x in args.tasks.split(",") if x.strip()] + batch_sizes = [int(x.strip()) for x in args.batch_sizes.split(",") if x.strip()] + + print("Genesis backend available") + print(f"Tasks: {task_names}") + print(f"Batch sizes: {batch_sizes}") + + records: List[BenchRecord] = [] + for task_name in task_names: + records.extend( + _bench_one_task( + task_name=task_name, + batch_sizes=batch_sizes, + nstep=args.nstep, + warmup=args.warmup, + iters=args.iters, + ) + ) + + out_json = Path(args.out_json) + out_json.parent.mkdir(parents=True, exist_ok=True) + payload: Dict[str, object] = { + "meta": { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "device_info": get_device_info_dict(), + "tasks": task_names, + "batch_sizes": batch_sizes, + "nstep": args.nstep, + "nthread": 0, + "warmup": args.warmup, + "iters": args.iters, + "backends": sorted({r.backend for r in records}), + "genesis_available": gs is not None, + }, + "results": [asdict(r) for r in records], + } + out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"Saved results to {out_json}") + + out_dir = Path(args.out_dir) + _plot_fps(records, out_dir / "fps.png", batch_sizes=batch_sizes, task_names=task_names) + _plot_per_env_sps( + records, out_dir / "per_env_sps.png", batch_sizes=batch_sizes, task_names=task_names + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_physics_step_isaacgym.py b/benchmark/benchmark_physics_step_isaacgym.py new file mode 100644 index 000000000..cfa785bfb --- /dev/null +++ b/benchmark/benchmark_physics_step_isaacgym.py @@ -0,0 +1,469 @@ +#!/usr/bin/env python3 +""" +Benchmark Isaac Gym physics execution using URDF assets. + +Benchmarks Isaac Gym across go1/go2/g1 locomotion robots and outputs JSON + plots +aligned with benchmark/benchmark_physics_step_mj_step.py. + +Run without creating any new environment: + PYTHONPATH=/home/admin1/.holosoma_deps/isaacgym/python \ + LD_LIBRARY_PATH=/home/admin1/.holosoma_deps/miniconda3/envs/hsgym/lib \ + uv run --no-project /home/admin1/.holosoma_deps/miniconda3/envs/hsgym/bin/python3.8 \ + benchmark/benchmark_physics_step_isaacgym.py +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import platform +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + +REPO_ROOT = Path(__file__).resolve().parent.parent +DEFAULT_DEPS_ROOT = Path("/home/admin1/.holosoma_deps") +DEFAULT_ISAACGYM_PYTHON = DEFAULT_DEPS_ROOT / "isaacgym" / "python" +DEFAULT_HSGYM_PYTHON = DEFAULT_DEPS_ROOT / "miniconda3" / "envs" / "hsgym" / "bin" / "python3.8" +DEFAULT_HSGYM_LIB = DEFAULT_DEPS_ROOT / "miniconda3" / "envs" / "hsgym" / "lib" +DEFAULT_MODELS_ROOT = Path("/home/admin1/ws/models") + +if str(DEFAULT_ISAACGYM_PYTHON) not in sys.path: + sys.path.insert(0, str(DEFAULT_ISAACGYM_PYTHON)) + +try: + from isaacgym import gymapi +except Exception as _isaacgym_error: + gymapi = None + _ISAACGYM_IMPORT_ERROR = _isaacgym_error +else: + _ISAACGYM_IMPORT_ERROR = None + + +@dataclass(frozen=True) +class TaskSpec: + owner_task_id: str + display_name: str + asset_root: Path + asset_file: str + initial_height: float + + +@dataclass +class BenchRecord: + task: str + backend: str + batch_size: int + nstep: int + nthread: int + avg_time_sec: float + sps: float + + +TASK_SPECS = { + "go1_joystick": TaskSpec( + owner_task_id="go1_joystick", + display_name="go1_joystick", + asset_root=DEFAULT_MODELS_ROOT, + asset_file="go1_description/urdf/go1.urdf", + initial_height=0.40, + ), + "go2_joystick": TaskSpec( + owner_task_id="go2_joystick", + display_name="go2_joystick", + asset_root=DEFAULT_MODELS_ROOT, + asset_file="go2_description/urdf/go2_description.urdf", + initial_height=0.40, + ), + "g1_joystick": TaskSpec( + owner_task_id="g1_joystick", + display_name="g1_joystick", + asset_root=DEFAULT_MODELS_ROOT, + asset_file="g1_description/g1_29dof_rev_1_0.urdf", + initial_height=0.78, + ), +} +TASK_ALIASES = { + "Go1JoystickFlatTerrain": "go1_joystick", + "Go2JoystickFlatTerrain": "go2_joystick", + "G1JoystickFlatTerrain": "g1_joystick", + "task=go1_joystick/isaacgym": "go1_joystick", + "task=go2_joystick/isaacgym": "go2_joystick", + "task=g1_joystick/isaacgym": "g1_joystick", + "go1_joystick/isaacgym": "go1_joystick", + "go2_joystick/isaacgym": "go2_joystick", + "g1_joystick/isaacgym": "g1_joystick", + "go1": "go1_joystick", + "go2": "go2_joystick", + "g1": "g1_joystick", +} +DEFAULT_TASK_IDS = list(TASK_SPECS.keys()) +DEFAULT_BATCH_SIZES = [2**k for k in range(8, 15)] # 256 .. 16384 + + +def _device_info_dict() -> Dict[str, str]: + info = { + "hostname": platform.node(), + "platform": platform.platform(), + "python": platform.python_version(), + } + try: + gpu = subprocess.check_output( + ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + text=True, + ).strip() + if gpu: + info["gpu"] = gpu.splitlines()[0] + except Exception: + info["gpu"] = "unknown" + return info + + +def _device_info_line() -> str: + info = _device_info_dict() + return f"{info.get('gpu', 'unknown')} | Python {info.get('python', 'unknown')}" + + +def _display_backend(backend: str) -> str: + return backend + + +def _runtime_hint() -> str: + return ( + f"PYTHONPATH={DEFAULT_ISAACGYM_PYTHON} " + f"LD_LIBRARY_PATH={DEFAULT_HSGYM_LIB} " + f"uv run --no-project {DEFAULT_HSGYM_PYTHON} " + "benchmark/benchmark_physics_step_isaacgym.py" + ) + + +def _require_isaacgym() -> None: + if gymapi is not None: + return + detail = repr(_ISAACGYM_IMPORT_ERROR) if _ISAACGYM_IMPORT_ERROR is not None else "unknown error" + raise RuntimeError( + "Isaac Gym is unavailable in the current runtime.\n" + f"Import detail: {detail}\n" + "Use the existing /home/admin1/.holosoma_deps runtime directly:\n" + f" {_runtime_hint()}" + ) + + +def _normalize_task_name(task_name: str) -> str: + normalized = task_name.strip() + if normalized.startswith("task="): + normalized = normalized[len("task=") :] + if ( + normalized.endswith("/mujoco") + or normalized.endswith("/motrix") + or normalized.endswith("/genesis") + ): + normalized = normalized.rsplit("/", 1)[0] + if normalized in TASK_SPECS: + return normalized + alias_target = TASK_ALIASES.get(normalized) + if alias_target is not None: + return alias_target + raise ValueError( + f"Unknown task '{task_name}'. Available task ids: {list(TASK_SPECS.keys())}. " + "Accepted aliases include legacy env names and task=/isaacgym forms." + ) + + +def _task_spec(task_name: str) -> TaskSpec: + return TASK_SPECS[_normalize_task_name(task_name)] + + +def _default_dof_targets(dof_props) -> np.ndarray: + lower = np.asarray(dof_props["lower"], dtype=np.float32) + upper = np.asarray(dof_props["upper"], dtype=np.float32) + targets = np.zeros_like(lower, dtype=np.float32) + finite = np.isfinite(lower) & np.isfinite(upper) & (lower <= upper) + targets[finite] = np.clip(targets[finite], lower[finite], upper[finite]) + return targets + + +def _create_sim(compute_device_id: int, graphics_device_id: int, nthread: int): + gym = gymapi.acquire_gym() + sim_params = gymapi.SimParams() + sim_params.dt = 0.01 + sim_params.substeps = 1 + sim_params.up_axis = gymapi.UpAxis.UP_AXIS_Z + sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.81) + sim_params.physx.solver_type = 1 + sim_params.physx.num_position_iterations = 4 + sim_params.physx.num_velocity_iterations = 1 + sim_params.physx.num_threads = nthread + sim_params.physx.use_gpu = True + sim_params.use_gpu_pipeline = True + sim = gym.create_sim(compute_device_id, graphics_device_id, gymapi.SIM_PHYSX, sim_params) + if sim is None: + raise RuntimeError("Failed to create Isaac Gym simulation.") + return gym, sim + + +def _build_task_sim( + task_name: str, + batch_size: int, + nthread: int, + compute_device_id: int, + graphics_device_id: int, +): + spec = _task_spec(task_name) + gym, sim = _create_sim(compute_device_id, graphics_device_id, nthread) + + plane_params = gymapi.PlaneParams() + plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0) + gym.add_ground(sim, plane_params) + + asset_options = gymapi.AssetOptions() + asset_options.flip_visual_attachments = True + asset_options.armature = 0.01 + asset_options.default_dof_drive_mode = int(gymapi.DOF_MODE_POS) + + asset = gym.load_asset(sim, str(spec.asset_root), spec.asset_file, asset_options) + dof_props = gym.get_asset_dof_properties(asset) + dof_props["driveMode"][:].fill(gymapi.DOF_MODE_POS) + dof_props["stiffness"][:].fill(1000.0) + dof_props["damping"][:].fill(10.0) + dof_targets = _default_dof_targets(dof_props) + + spacing = 1.0 + env_lower = gymapi.Vec3(-spacing, 0.0, -spacing) + env_upper = gymapi.Vec3(spacing, spacing, spacing) + num_per_row = max(1, int(math.sqrt(batch_size))) + + pose = gymapi.Transform() + pose.r = gymapi.Quat(0.0, 0.0, 0.0, 1.0) + pose.p = gymapi.Vec3(0.0, 0.0, spec.initial_height) + + for env_idx in range(batch_size): + env = gym.create_env(sim, env_lower, env_upper, num_per_row) + actor = gym.create_actor(env, asset, pose, spec.display_name, env_idx, 2) + gym.set_actor_dof_properties(env, actor, dof_props) + if dof_targets.size > 0: + gym.set_actor_dof_position_targets(env, actor, dof_targets) + + gym.prepare_sim(sim) + return gym, sim + + +def _run_sim(gym, sim, nstep: int, niter: int) -> float: + t0 = time.perf_counter() + for _ in range(niter): + for _ in range(nstep): + gym.simulate(sim) + gym.fetch_results(sim, True) + return (time.perf_counter() - t0) / niter + + +def _bench_one_task( + task_name: str, + batch_sizes: List[int], + nstep: int, + nthread: int, + warmup: int, + iters: int, + compute_device_id: int, + graphics_device_id: int, +) -> List[BenchRecord]: + task_key = _normalize_task_name(task_name) + records: List[BenchRecord] = [] + + for batch_size in batch_sizes: + gym, sim = _build_task_sim( + task_name=task_key, + batch_size=batch_size, + nthread=nthread, + compute_device_id=compute_device_id, + graphics_device_id=graphics_device_id, + ) + try: + _run_sim(gym, sim, nstep=nstep, niter=warmup) + avg_t = _run_sim(gym, sim, nstep=nstep, niter=iters) + finally: + gym.destroy_sim(sim) + + records.append( + BenchRecord( + task=task_key, + backend="isaacgym", + batch_size=batch_size, + nstep=nstep, + nthread=nthread, + avg_time_sec=avg_t, + sps=batch_size * nstep / avg_t, + ) + ) + print( + f"[{task_key}] batch={batch_size:5d} " + f"isaacgym={avg_t * 1000:.3f}ms ({batch_size * nstep / avg_t / 1e4:.2f}万fps)" + ) + + return records + + +def _plot_fps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): + out_png.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Total FPS\n{_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / 1e4 for r in recs], + marker="o", + linestyle="-", + label=f"{_task_spec(task_name).display_name} ({_display_backend(backend)})", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Total FPS (x1e4)") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() + fig.savefig(out_png, dpi=150) + plt.close(fig) + print(f"Saved FPS plot to {out_png}") + + +def _plot_per_env_sps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): + out_png.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Per-Env SPS\n{_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / r.batch_size for r in recs], + marker="o", + linestyle="-", + label=f"{_task_spec(task_name).display_name} ({_display_backend(backend)})", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Per-Env Steps/s") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() + fig.savefig(out_png, dpi=150) + plt.close(fig) + print(f"Saved per-env SPS plot to {out_png}") + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark Isaac Gym physics execution") + parser.add_argument("--nstep", type=int, default=20) + parser.add_argument("--nthread", type=int, default=8) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=3) + parser.add_argument("--compute-device-id", type=int, default=0) + parser.add_argument("--graphics-device-id", type=int, default=-1) + parser.add_argument("--tasks", type=str, default=",".join(DEFAULT_TASK_IDS)) + parser.add_argument( + "--batch-sizes", type=str, default=",".join(str(x) for x in DEFAULT_BATCH_SIZES) + ) + parser.add_argument( + "--out-json", + type=str, + default="benchmark/outputs/physics_step/isaacgym/results.json", + ) + parser.add_argument( + "--out-dir", + type=str, + default="benchmark/outputs/physics_step/isaacgym", + help="Directory for output plots", + ) + args = parser.parse_args() + + _require_isaacgym() + + task_names = [_normalize_task_name(x) for x in args.tasks.split(",") if x.strip()] + batch_sizes = [int(x.strip()) for x in args.batch_sizes.split(",") if x.strip()] + + print("Isaac Gym backend available") + print(f"Tasks: {task_names}") + print(f"Batch sizes: {batch_sizes}") + print(f"Model root: {DEFAULT_MODELS_ROOT}") + + records: List[BenchRecord] = [] + for task_name in task_names: + records.extend( + _bench_one_task( + task_name=task_name, + batch_sizes=batch_sizes, + nstep=args.nstep, + nthread=args.nthread, + warmup=args.warmup, + iters=args.iters, + compute_device_id=args.compute_device_id, + graphics_device_id=args.graphics_device_id, + ) + ) + + out_json = Path(args.out_json) + out_json.parent.mkdir(parents=True, exist_ok=True) + payload: Dict[str, object] = { + "meta": { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "device_info": _device_info_dict(), + "tasks": task_names, + "batch_sizes": batch_sizes, + "nstep": args.nstep, + "nthread": args.nthread, + "warmup": args.warmup, + "iters": args.iters, + "compute_device_id": args.compute_device_id, + "graphics_device_id": args.graphics_device_id, + "backends": sorted({r.backend for r in records}), + "models_root": str(DEFAULT_MODELS_ROOT), + "asset_files": {task: _task_spec(task).asset_file for task in task_names}, + "isaacgym_python_path": str(DEFAULT_ISAACGYM_PYTHON), + "hsgym_python": str(DEFAULT_HSGYM_PYTHON), + }, + "results": [asdict(r) for r in records], + } + out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"Saved results to {out_json}") + + out_dir = Path(args.out_dir) + _plot_fps(records, out_dir / "fps.png", batch_sizes=batch_sizes, task_names=task_names) + _plot_per_env_sps( + records, out_dir / "per_env_sps.png", batch_sizes=batch_sizes, task_names=task_names + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_mj_step.py b/benchmark/benchmark_physics_step_mj_step.py similarity index 54% rename from benchmark/benchmark_mj_step.py rename to benchmark/benchmark_physics_step_mj_step.py index 381e10457..9f0e1ce31 100644 --- a/benchmark/benchmark_mj_step.py +++ b/benchmark/benchmark_physics_step_mj_step.py @@ -1,18 +1,19 @@ #!/usr/bin/env python3 """ -Benchmark parallel physics backends: mujoco.rollout (numpy) and mujoco.mlx_step (macOS). +Benchmark MuJoCo parallel physics execution. -macOS: compares numpy rollout vs native mlx_step, measuring latency and FPS. -Other: compares mujoco.rollout with 1 thread vs N threads. +macOS: compares numpy rollout vs native mlx_step. +Other: benchmarks mujoco.rollout with the configured thread count only. -Sweeps batch sizes across go1/go2/g1 locomotion tasks and outputs JSON + plots. +Sweeps batch sizes across current locomotion owner task ids +(go1_joystick/go2_joystick/g1_joystick). +Legacy env names remain accepted as aliases. """ from __future__ import annotations import argparse import json -import math import platform import time from dataclasses import asdict, dataclass @@ -43,34 +44,37 @@ mx = None mj_mlx_step = None -from unilab.envs.locomotion.g1.joystick import G1JoystickCfg -from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg -from unilab.envs.locomotion.go2.joystick import Go2JoystickCfg - try: from benchmark.core.device_info import get_device_info_dict, get_device_info_line + from benchmark.core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) except ModuleNotFoundError: from core.device_info import get_device_info_dict, get_device_info_line + from core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) @dataclass class BenchRecord: task: str - backend: str # "numpy" | "mlx_native" (macOS) | "rollout_1t" | "rollout_Nt" + backend: str # "numpy" | "mlx_native" (macOS) | "rollout_Nt" batch_size: int nstep: int nthread: int avg_time_sec: float sps: float # steps per second = batch_size * nstep / avg_time_sec - output_shape_mode: str = "n/a" -TASK_CONFIGS = { - "Go1JoystickFlatTerrain": Go1JoystickCfg, - "Go2JoystickFlatTerrain": Go2JoystickCfg, - "G1JoystickFlatTerrain": G1JoystickCfg, -} +DEFAULT_TASK_IDS = canonical_locomotion_task_ids() DEFAULT_BATCH_SIZES = [2**k for k in range(8, 15)] # 256 .. 16384 +TASK_ALPHA = {"go1_joystick": 0.75, "go2_joystick": 0.9, "g1_joystick": 1.0} +TASK_HATCH = {"go1_joystick": "//", "go2_joystick": "\\\\", "g1_joystick": "xx"} def _keyframe0_state_and_ctrl(model: mujoco.MjModel) -> tuple[np.ndarray, np.ndarray]: @@ -142,25 +146,19 @@ def _run_mlx( return (time.perf_counter() - t0) / niter -def _infer_output_shape_mode(out) -> str: - state_mx = out[0] if isinstance(out, tuple) else out.state_mx - return "last_only" if state_mx.ndim == 2 else "full_traj" - - def _has_native_mujoco_mlx_step() -> bool: return mj_mlx_step is not None and hasattr(mj_mlx_step, "MlxStepRunner") def _load_task_model(task_name: str) -> mujoco.MjModel: - cfg = TASK_CONFIGS[task_name]() + cfg = locomotion_task_spec(task_name).config_cls() return mujoco.MjModel.from_xml_path(cfg.model_file) -def _geomean(values: List[float]) -> float: - vals = [v for v in values if v > 0.0] - if not vals: - return 0.0 - return float(math.exp(sum(math.log(v) for v in vals) / len(vals))) +def _display_backend(backend: str) -> str: + if backend.startswith("rollout_") and backend.endswith("t"): + return f"rollout ({backend[len('rollout_') : -1]} threads)" + return backend def _bench_one_task( @@ -172,8 +170,9 @@ def _bench_one_task( iters: int, chunk_size: int, ) -> List[BenchRecord]: + task_key = normalize_locomotion_task_id(task_name) np.random.seed(42) - model = _load_task_model(task_name) + model = _load_task_model(task_key) nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) state0, ctrl0 = _keyframe0_state_and_ctrl(model) @@ -219,16 +218,6 @@ def _bench_one_task( chunk_size, ) - probe = mlx_runner.step( - model=model_list, - data=data_list, - initial_state=initial_state_mx, - control=control_mx, - nstep=nstep, - out_dtype=mx.float32, - ) - output_shape_mode = _infer_output_shape_mode(probe) - numpy_t = _run_numpy( numpy_runner, model_list, @@ -253,51 +242,34 @@ def _bench_one_task( records.append( BenchRecord( - task=task_name, + task=task_key, backend="numpy", batch_size=batch_size, nstep=nstep, nthread=actual_nthread, avg_time_sec=numpy_t, sps=batch_size * nstep / numpy_t, - output_shape_mode="n/a", ) ) records.append( BenchRecord( - task=task_name, + task=task_key, backend="mlx_native", batch_size=batch_size, nstep=nstep, nthread=actual_nthread, avg_time_sec=mlx_t, sps=batch_size * nstep / mlx_t, - output_shape_mode=output_shape_mode, ) ) print( - f"[{task_name}] batch={batch_size:5d} " + f"[{task_key}] batch={batch_size:5d} " f"numpy={numpy_t * 1000:.3f}ms ({batch_size * nstep / numpy_t / 1e4:.2f}万fps) " f"mlx={mlx_t * 1000:.3f}ms ({batch_size * nstep / mlx_t / 1e4:.2f}万fps)" ) else: - data_list_1 = [mujoco.MjData(model) for _ in range(1)] data_list_n = [mujoco.MjData(model) for _ in range(nthread)] - with ( - mj_rollout.Rollout(nthread=1) as runner1, - mj_rollout.Rollout(nthread=nthread) as runner_n, - ): - _run_numpy( - runner1, - model_list, - data_list_1, - initial_state, - control, - state_buf, - sensordata_buf, - nstep, - warmup, - ) + with mj_rollout.Rollout(nthread=nthread) as runner_n: _run_numpy( runner_n, model_list, @@ -309,17 +281,6 @@ def _bench_one_task( nstep, warmup, ) - t1 = _run_numpy( - runner1, - model_list, - data_list_1, - initial_state, - control, - state_buf, - sensordata_buf, - nstep, - iters, - ) tn = _run_numpy( runner_n, model_list, @@ -334,211 +295,129 @@ def _bench_one_task( records.append( BenchRecord( - task=task_name, - backend="rollout_1t", - batch_size=batch_size, - nstep=nstep, - nthread=1, - avg_time_sec=t1, - sps=batch_size * nstep / t1, - output_shape_mode="n/a", - ) - ) - records.append( - BenchRecord( - task=task_name, + task=task_key, backend=f"rollout_{nthread}t", batch_size=batch_size, nstep=nstep, nthread=nthread, avg_time_sec=tn, sps=batch_size * nstep / tn, - output_shape_mode="n/a", ) ) print( - f"[{task_name}] batch={batch_size:5d} " - f"rollout(1t)={t1 * 1000:.3f}ms rollout({nthread}t)={tn * 1000:.3f}ms" + f"[{task_key}] batch={batch_size:5d} " + f"rollout({nthread}t)={tn * 1000:.3f}ms ({batch_size * nstep / tn / 1e4:.2f}万fps)" ) return records -def _plot_latency(records: List[BenchRecord], out_png: Path, batch_sizes: List[int], nthread: int): +def _plot_fps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): out_png.parent.mkdir(parents=True, exist_ok=True) - if _IS_MACOS: - backend_colors = {"numpy": "#4C78A8", "mlx_native": "#F58518"} - backend_display = {"numpy": "Backend:NumPy", "mlx_native": "Backend:MLX"} - else: - backend_colors = {"rollout_1t": "#4C78A8", f"rollout_{nthread}t": "#F58518"} - backend_display = { - "rollout_1t": "rollout (1 thread)", - f"rollout_{nthread}t": f"rollout ({nthread} threads)", - } - - all_backends = list(backend_colors.keys()) - task_alpha = { - "Go1JoystickFlatTerrain": 0.75, - "Go2JoystickFlatTerrain": 0.9, - "G1JoystickFlatTerrain": 1.0, - } - task_hatch = { - "Go1JoystickFlatTerrain": "//", - "Go2JoystickFlatTerrain": "\\\\", - "G1JoystickFlatTerrain": "xx", - } - task_names = list(TASK_CONFIGS.keys()) - x = np.arange(len(batch_sizes), dtype=np.float64) - bar_width = 0.10 - pair_inner_gap = 0.02 - task_group_gap = 0.09 - pair_span = len(all_backends) * bar_width + (len(all_backends) - 1) * pair_inner_gap - total_span = len(task_names) * pair_span + (len(task_names) - 1) * task_group_gap - left_edge = -0.5 * total_span - value_map = {(r.task, r.backend, r.batch_size): r.avg_time_sec * 1000.0 for r in records} - - fig, ax = plt.subplots(figsize=(13, 7)) - for task_idx, task_name in enumerate(task_names): - pair_start = left_edge + task_idx * (pair_span + task_group_gap) - for b_idx, backend in enumerate(all_backends): - offset = pair_start + b_idx * (bar_width + pair_inner_gap) + 0.5 * bar_width - y = [value_map.get((task_name, backend, b), np.nan) for b in batch_sizes] - ax.bar( - x + offset, - y, - width=bar_width, - color=backend_colors[backend], - alpha=task_alpha[task_name], - hatch=task_hatch[task_name], - edgecolor="black", - linewidth=0.2, + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Total FPS\n{get_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / 1e4 for r in recs], + marker="o", + linestyle="--" if backend == "numpy" else "-", + label=f"{locomotion_task_spec(task_name).display_name} ({_display_backend(backend)})", ) - ax.set_xticks(x) - ax.set_xticklabels([str(v) for v in batch_sizes]) - ax.set_xlabel("Batch Size") - ax.set_ylabel("Average Time per Rollout Call (ms)") - ax.set_yscale("log") - ax.grid(True, which="major", axis="y", alpha=0.3) - backend_handles = [ - Patch(facecolor=backend_colors[b], edgecolor="black", label=backend_display[b]) - for b in all_backends - ] - task_handles = [ - Patch( - facecolor="white", - edgecolor="black", - hatch=task_hatch["Go1JoystickFlatTerrain"], - label="Task:Go1", - ), - Patch( - facecolor="white", - edgecolor="black", - hatch=task_hatch["Go2JoystickFlatTerrain"], - label="Task:Go2", - ), - Patch( - facecolor="white", - edgecolor="black", - hatch=task_hatch["G1JoystickFlatTerrain"], - label="Task:G1", - ), - ] - fig.suptitle( - f"Rollout Time by Backend and Task\n{get_device_info_line()}", - y=0.965, - fontsize=13, - ) - fig.legend( - handles=backend_handles + task_handles, - fontsize=9, - ncol=len(backend_handles + task_handles), - loc="upper center", - bbox_to_anchor=(0.5, 0.895), - frameon=True, - handlelength=1.8, - columnspacing=1.2, - ) - fig.tight_layout(rect=[0.0, 0.0, 1.0, 0.89]) + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Total FPS (x1e4)") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() fig.savefig(out_png, dpi=150) plt.close(fig) - print(f"Saved latency plot to {out_png}") + print(f"Saved FPS plot to {out_png}") -def _plot_fps(records: List[BenchRecord], out_png: Path, batch_sizes: List[int]): +def _plot_per_env_sps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): out_png.parent.mkdir(parents=True, exist_ok=True) fig, ax = plt.subplots(figsize=(9, 5)) - ax.set_title(f"Parallel Physics — Total FPS\n{get_device_info_line()}", fontsize=9) + ax.set_title(f"Parallel Physics — Per-Env SPS\n{get_device_info_line()}", fontsize=9) - backend_linestyle = {"numpy": "--", "rollout_1t": "--"} - - for task_name in TASK_CONFIGS: + for task_name in task_names: for backend in sorted({r.backend for r in records if r.task == task_name}): recs = sorted( [r for r in records if r.task == task_name and r.backend == backend], key=lambda r: r.batch_size, ) - short = task_name.split("Joystick")[0] ax.plot( [r.batch_size for r in recs], - [r.sps / 1e4 for r in recs], + [r.sps / r.batch_size for r in recs], marker="o", - linestyle=backend_linestyle.get(backend, "-"), - label=f"{short} ({backend})", + linestyle="--" if backend == "numpy" else "-", + label=f"{locomotion_task_spec(task_name).display_name} ({_display_backend(backend)})", ) ax.set_xscale("log", base=2) ax.set_xticks(batch_sizes) ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") ax.set_xlabel("Batch Size (Num Envs)") - ax.set_ylabel("Total FPS (x1e4)") + ax.set_ylabel("Per-Env Steps/s") ax.grid(True, alpha=0.3) ax.legend(fontsize=7) fig.tight_layout() fig.savefig(out_png, dpi=150) plt.close(fig) - print(f"Saved FPS plot to {out_png}") + print(f"Saved per-env SPS plot to {out_png}") def main(): parser = argparse.ArgumentParser( description="Benchmark parallel physics backends (rollout / mlx_step)" ) - parser.add_argument("--nstep", type=int, default=1) - parser.add_argument("--nthread", type=int, default=8) - parser.add_argument("--warmup", type=int, default=5) - parser.add_argument("--iters", type=int, default=10) + parser.add_argument("--nstep", type=int, default=20) + parser.add_argument("--nthread", type=int, default=cpu_count() * 2) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=3) parser.add_argument( "--chunk-size", type=int, default=16, help="MLX step chunk size (macOS only)" ) parser.add_argument( "--tasks", type=str, - default="Go1JoystickFlatTerrain,Go2JoystickFlatTerrain,G1JoystickFlatTerrain", + default=",".join(DEFAULT_TASK_IDS), ) parser.add_argument( "--batch-sizes", type=str, default=",".join(str(x) for x in DEFAULT_BATCH_SIZES) ) - parser.add_argument("--out-json", type=str, default="benchmark/outputs/mlx_step/results.json") + parser.add_argument( + "--out-json", type=str, default="benchmark/outputs/physics_step/mj_step/results.json" + ) parser.add_argument( "--out-dir", type=str, - default="benchmark/outputs/mlx_step", + default="benchmark/outputs/physics_step/mj_step", help="Directory for output plots", ) args = parser.parse_args() - task_names = [x.strip() for x in args.tasks.split(",") if x.strip()] + task_names = [normalize_locomotion_task_id(x) for x in args.tasks.split(",") if x.strip()] batch_sizes = [int(x.strip()) for x in args.batch_sizes.split(",") if x.strip()] - for name in task_names: - if name not in TASK_CONFIGS: - raise ValueError(f"Unknown task '{name}'. Available: {list(TASK_CONFIGS.keys())}") if _IS_MACOS: if not _has_native_mujoco_mlx_step(): raise RuntimeError("Native MLX step backend unavailable. Requires mujoco.mlx_step.") - print("macOS: native mlx_step available") + print(f"macOS: native mlx_step available, nthread={args.nthread}") else: - print(f"Non-macOS: using mujoco.rollout (1t vs {args.nthread}t)") + print(f"Non-macOS: using mujoco.rollout ({args.nthread} threads)") print(f"Tasks: {task_names}") print(f"Batch sizes: {batch_sizes}") @@ -557,11 +436,6 @@ def main(): ) ) - if _IS_MACOS: - backend_a, backend_b = "numpy", "mlx_native" - else: - backend_a, backend_b = "rollout_1t", f"rollout_{args.nthread}t" - out_json = Path(args.out_json) out_json.parent.mkdir(parents=True, exist_ok=True) payload: Dict[str, object] = { @@ -576,40 +450,7 @@ def main(): "iters": args.iters, "chunk_size": args.chunk_size, "native_mlx_step_available": _has_native_mujoco_mlx_step(), - "mlx_output_shape_modes": sorted( - {r.output_shape_mode for r in records if r.backend == "mlx_native"} - ) - if _IS_MACOS - else [], - }, - "summary": { - "speedup_a_over_b": { - task: { - "backend_a": backend_a, - "backend_b": backend_b, - "geomean": _geomean( - [ - next( - r.avg_time_sec - for r in records - if r.task == task and r.backend == backend_a and r.batch_size == b - ) - / max( - next( - r.avg_time_sec - for r in records - if r.task == task - and r.backend == backend_b - and r.batch_size == b - ), - 1e-12, - ) - for b in batch_sizes - ] - ), - } - for task in task_names - } + "backends": sorted({r.backend for r in records}), }, "results": [asdict(r) for r in records], } @@ -617,8 +458,10 @@ def main(): print(f"Saved results to {out_json}") out_dir = Path(args.out_dir) - _plot_latency(records, out_dir / "latency.png", batch_sizes=batch_sizes, nthread=args.nthread) - _plot_fps(records, out_dir / "fps.png", batch_sizes=batch_sizes) + _plot_fps(records, out_dir / "fps.png", batch_sizes=batch_sizes, task_names=task_names) + _plot_per_env_sps( + records, out_dir / "per_env_sps.png", batch_sizes=batch_sizes, task_names=task_names + ) if __name__ == "__main__": diff --git a/benchmark/benchmark_physics_step_mujoco_warp.py b/benchmark/benchmark_physics_step_mujoco_warp.py new file mode 100644 index 000000000..2f5b47348 --- /dev/null +++ b/benchmark/benchmark_physics_step_mujoco_warp.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python3 +""" +Benchmark MuJoCo Warp physics execution. + +Benchmarks mujoco_warp across current locomotion owner task ids +(go1_joystick/go2_joystick/g1_joystick) and outputs JSON + plots +aligned with benchmark/benchmark_physics_step_mj_step.py. +Legacy env names remain accepted as aliases. + +Run without changing repo dependencies: + uv run --with mujoco-warp --with warp-lang \ + python benchmark/benchmark_physics_step_mujoco_warp.py +""" + +from __future__ import annotations + +import argparse +import json +import time +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Dict, List + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + +try: + import mujoco +except ImportError: + mujoco = None + +try: + import mujoco_warp as mj_warp +except ImportError: + mj_warp = None + +try: + import warp +except ImportError: + warp = None + +try: + from benchmark.core.device_info import get_device_info_dict, get_device_info_line + from benchmark.core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) +except ModuleNotFoundError: + from core.device_info import get_device_info_dict, get_device_info_line + from core.task_names import ( + canonical_locomotion_task_ids, + locomotion_task_spec, + normalize_locomotion_task_id, + ) + + +@dataclass +class BenchRecord: + task: str + backend: str + batch_size: int + nstep: int + nthread: int + avg_time_sec: float + sps: float + + +DEFAULT_TASK_IDS = canonical_locomotion_task_ids() +DEFAULT_BATCH_SIZES = [2**k for k in range(8, 15)] # 256 .. 16384 +DEFAULT_NJMAX_BY_TASK = { + "go1_joystick": 100, + "go2_joystick": 100, + "g1_joystick": 150, +} + + +def _display_backend(backend: str) -> str: + return backend + + +def _uv_run_hint() -> str: + return ( + "uv run --with mujoco-warp --with warp-lang " + "python benchmark/benchmark_physics_step_mujoco_warp.py" + ) + + +def _require_mujoco_warp() -> None: + if mujoco is None: + raise RuntimeError( + "mujoco is unavailable in the current environment. " + f"Run with temporary deps instead of editing pyproject.toml:\n {_uv_run_hint()}" + ) + if mj_warp is None: + raise RuntimeError( + "mujoco_warp is unavailable in the current environment. " + f"Run with temporary deps instead of editing pyproject.toml:\n {_uv_run_hint()}" + ) + if warp is None: + raise RuntimeError( + "warp is unavailable in the current environment. " + f"Run with temporary deps instead of editing pyproject.toml:\n {_uv_run_hint()}" + ) + + +def _load_task_model(task_name: str) -> "mujoco.MjModel": + cfg = locomotion_task_spec(task_name).config_cls() + return mujoco.MjModel.from_xml_path(cfg.model_file) + + +def _task_njmax(task_name: str) -> int: + return DEFAULT_NJMAX_BY_TASK.get(task_name, -1) + + +def _make_warp_data(model: "mujoco.MjModel", batch_size: int, njmax: int): + try: + if njmax > 0: + return mj_warp.make_data(model, nworld=batch_size, njmax=njmax, nconmax=njmax) + return mj_warp.make_data(model, nworld=batch_size) + except TypeError: + return mj_warp.make_data(model, nworld=batch_size) + + +def _run_warp(warp_model, warp_data, nstep: int, niter: int) -> float: + t0 = time.perf_counter() + for _ in range(niter): + for _ in range(nstep): + mj_warp.step(warp_model, warp_data) + warp.synchronize() + return (time.perf_counter() - t0) / niter + + +def _bench_one_task( + task_name: str, + batch_sizes: List[int], + nstep: int, + warmup: int, + iters: int, +) -> List[BenchRecord]: + task_key = normalize_locomotion_task_id(task_name) + model = _load_task_model(task_key) + warp_model = mj_warp.put_model(model) + njmax = _task_njmax(task_key) + + records: List[BenchRecord] = [] + for batch_size in batch_sizes: + warp_data = _make_warp_data(model, batch_size, njmax) + + _run_warp(warp_model, warp_data, nstep=nstep, niter=warmup) + avg_t = _run_warp(warp_model, warp_data, nstep=nstep, niter=iters) + + records.append( + BenchRecord( + task=task_key, + backend="mujoco_warp", + batch_size=batch_size, + nstep=nstep, + nthread=0, + avg_time_sec=avg_t, + sps=batch_size * nstep / avg_t, + ) + ) + print( + f"[{task_key}] batch={batch_size:5d} " + f"mujoco_warp={avg_t * 1000:.3f}ms ({batch_size * nstep / avg_t / 1e4:.2f}万fps)" + ) + + return records + + +def _plot_fps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): + out_png.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Total FPS\n{get_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / 1e4 for r in recs], + marker="o", + linestyle="-", + label=f"{locomotion_task_spec(task_name).display_name} ({_display_backend(backend)})", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Total FPS (x1e4)") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() + fig.savefig(out_png, dpi=150) + plt.close(fig) + print(f"Saved FPS plot to {out_png}") + + +def _plot_per_env_sps( + records: List[BenchRecord], out_png: Path, batch_sizes: List[int], task_names: List[str] +): + out_png.parent.mkdir(parents=True, exist_ok=True) + fig, ax = plt.subplots(figsize=(9, 5)) + ax.set_title(f"Parallel Physics — Per-Env SPS\n{get_device_info_line()}", fontsize=9) + + for task_name in task_names: + for backend in sorted({r.backend for r in records if r.task == task_name}): + recs = sorted( + [r for r in records if r.task == task_name and r.backend == backend], + key=lambda r: r.batch_size, + ) + ax.plot( + [r.batch_size for r in recs], + [r.sps / r.batch_size for r in recs], + marker="o", + linestyle="-", + label=f"{locomotion_task_spec(task_name).display_name} ({_display_backend(backend)})", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(batch_sizes) + ax.set_xticklabels([str(b) for b in batch_sizes], rotation=30, ha="right") + ax.set_xlabel("Batch Size (Num Envs)") + ax.set_ylabel("Per-Env Steps/s") + ax.grid(True, alpha=0.3) + ax.legend(fontsize=7) + fig.tight_layout() + fig.savefig(out_png, dpi=150) + plt.close(fig) + print(f"Saved per-env SPS plot to {out_png}") + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark MuJoCo Warp physics execution") + parser.add_argument("--nstep", type=int, default=20) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=3) + parser.add_argument( + "--tasks", + type=str, + default=",".join(DEFAULT_TASK_IDS), + ) + parser.add_argument( + "--batch-sizes", type=str, default=",".join(str(x) for x in DEFAULT_BATCH_SIZES) + ) + parser.add_argument( + "--out-json", + type=str, + default="benchmark/outputs/physics_step/mujoco_warp/results.json", + ) + parser.add_argument( + "--out-dir", + type=str, + default="benchmark/outputs/physics_step/mujoco_warp", + help="Directory for output plots", + ) + args = parser.parse_args() + + _require_mujoco_warp() + + task_names = [normalize_locomotion_task_id(x) for x in args.tasks.split(",") if x.strip()] + batch_sizes = [int(x.strip()) for x in args.batch_sizes.split(",") if x.strip()] + + print("MuJoCo Warp backend available") + print(f"Tasks: {task_names}") + print(f"Batch sizes: {batch_sizes}") + + records: List[BenchRecord] = [] + for task_name in task_names: + records.extend( + _bench_one_task( + task_name=task_name, + batch_sizes=batch_sizes, + nstep=args.nstep, + warmup=args.warmup, + iters=args.iters, + ) + ) + + out_json = Path(args.out_json) + out_json.parent.mkdir(parents=True, exist_ok=True) + payload: Dict[str, object] = { + "meta": { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "device_info": get_device_info_dict(), + "tasks": task_names, + "batch_sizes": batch_sizes, + "nstep": args.nstep, + "nthread": 0, + "warmup": args.warmup, + "iters": args.iters, + "backends": sorted({r.backend for r in records}), + "task_njmax": {task: _task_njmax(task) for task in task_names}, + "warp_available": mj_warp is not None and warp is not None, + }, + "results": [asdict(r) for r in records], + } + out_json.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"Saved results to {out_json}") + + out_dir = Path(args.out_dir) + _plot_fps(records, out_dir / "fps.png", batch_sizes=batch_sizes, task_names=task_names) + _plot_per_env_sps( + records, out_dir / "per_env_sps.png", batch_sizes=batch_sizes, task_names=task_names + ) + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_postprocess.py b/benchmark/benchmark_postprocess.py index 30fedcbd6..e57831f76 100644 --- a/benchmark/benchmark_postprocess.py +++ b/benchmark/benchmark_postprocess.py @@ -29,8 +29,10 @@ def ensure_registries() -> None: try: from benchmark.core.device_info import get_device_info_dict, get_device_info_line + from benchmark.core.task_names import locomotion_env_name except ModuleNotFoundError: from core.device_info import get_device_info_dict, get_device_info_line + from core.task_names import locomotion_env_name try: import unilab.envs.locomotion @@ -51,6 +53,8 @@ def ensure_registries() -> None: OUTPUT_JSON = OUTPUT_DIR / "latest_postprocess_benchmark.json" OUTPUT_PNG = OUTPUT_DIR / "latest_postprocess_latency.png" TORCH_DEVICE = "mps" +OWNER_TASK_ID = "go1_joystick" +ENV_TASK_NAME = locomotion_env_name(OWNER_TASK_ID) def sync_torch_mps(): @@ -72,7 +76,7 @@ def geomean(values: list[float]) -> float: def build_go1_layout() -> dict: - env = registry.make("Go1JoystickFlatTerrain", num_envs=1, sim_backend="mujoco") + env = registry.make(ENV_TASK_NAME, num_envs=1, sim_backend="mujoco") layout = { "sensor_dim": int(env.model.nsensordata), "physics_dim": int(env.physics_state_dim), @@ -99,7 +103,7 @@ def build_go1_layout() -> dict: def measure_physics_step_ms(num_envs: int, iters: int) -> float: - env = registry.make("Go1JoystickFlatTerrain", num_envs=num_envs, sim_backend="mujoco") + env = registry.make(ENV_TASK_NAME, num_envs=num_envs, sim_backend="mujoco") try: initial_state, _, _ = env.reset(np.arange(env.num_envs)) action_low = env.action_space.low.astype(np.float32) @@ -131,7 +135,7 @@ def measure_physics_step_mlx_native_ms(num_envs: int, iters: int) -> float: """Measure pure MuJoCo physics stepping via native mujoco.mlx_step.""" if mj_mlx_step is None: raise RuntimeError("Native mujoco.mlx_step is unavailable.") - env = registry.make("Go1JoystickFlatTerrain", num_envs=num_envs, sim_backend="mujoco") + env = registry.make(ENV_TASK_NAME, num_envs=num_envs, sim_backend="mujoco") try: initial_state, _, _ = env.reset(np.arange(env.num_envs)) action_low = env.action_space.low.astype(np.float32) @@ -192,7 +196,7 @@ def measure_rollout_bridge_mlx_pipeline_ms( """End-to-end: MLX action -> MuJoCo rollout -> MLX sensordata/reward.""" if mj_mlx_step is None: raise RuntimeError("Native mujoco.mlx_step is unavailable.") - env = registry.make("Go1JoystickFlatTerrain", num_envs=num_envs, sim_backend="mujoco") + env = registry.make(ENV_TASK_NAME, num_envs=num_envs, sim_backend="mujoco") try: initial_state, _, reset_info = env.reset(np.arange(env.num_envs)) action_low = env.action_space.low.astype(np.float32) @@ -1116,8 +1120,12 @@ def main(): "mps_available": bool(torch.backends.mps.is_available()), "env_list": env_list, "iters": args.iters, - "task": "Go1JoystickFlatTerrain", - "note": "Postprocess logic synchronized with latest Go1JoystickFlatTerrain; physics step uses native mujoco.mlx_step.", + "task": OWNER_TASK_ID, + "env_task_name": ENV_TASK_NAME, + "note": ( + "Postprocess logic synchronized with latest " + f"{OWNER_TASK_ID} (env: {ENV_TASK_NAME}); physics step uses native mujoco.mlx_step." + ), }, "summary": { "geomean_speedup_cpu_over_mlx_postprocess": geomean( diff --git a/benchmark/core/task_names.py b/benchmark/core/task_names.py new file mode 100644 index 000000000..10af6a148 --- /dev/null +++ b/benchmark/core/task_names.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from unilab.envs.locomotion.g1.joystick import G1JoystickPPOCfg +from unilab.envs.locomotion.go1.joystick import Go1JoystickCfg +from unilab.envs.locomotion.go2.joystick import Go2JoystickCfg + + +@dataclass(frozen=True) +class LocomotionTaskSpec: + owner_task_id: str + env_task_name: str + display_name: str + config_cls: type + + +_TASK_SPECS = { + "go1_joystick": LocomotionTaskSpec( + owner_task_id="go1_joystick", + env_task_name="Go1JoystickFlatTerrain", + display_name="go1_joystick", + config_cls=Go1JoystickCfg, + ), + "go2_joystick": LocomotionTaskSpec( + owner_task_id="go2_joystick", + env_task_name="Go2JoystickFlatTerrain", + display_name="go2_joystick", + config_cls=Go2JoystickCfg, + ), + "g1_joystick": LocomotionTaskSpec( + owner_task_id="g1_joystick", + env_task_name="G1JoystickFlatTerrain", + display_name="g1_joystick", + config_cls=G1JoystickPPOCfg, + ), +} +_TASK_ALIASES = {spec.env_task_name: spec.owner_task_id for spec in _TASK_SPECS.values()} +_TASK_ALIASES.update({f"task={task_id}/mujoco": task_id for task_id in _TASK_SPECS}) +_TASK_ALIASES.update({f"{task_id}/mujoco": task_id for task_id in _TASK_SPECS}) + + +def canonical_locomotion_task_ids() -> list[str]: + return list(_TASK_SPECS.keys()) + + +def normalize_locomotion_task_id(task_name: str) -> str: + normalized = task_name.strip() + if normalized.startswith("task="): + normalized = normalized[len("task=") :] + if normalized.endswith("/motrix"): + raise ValueError( + f"Task '{task_name}' targets motrix, but this benchmark only measures MuJoCo paths." + ) + if normalized in _TASK_SPECS: + return normalized + alias_target = _TASK_ALIASES.get(normalized) + if alias_target is not None: + return alias_target + raise ValueError( + f"Unknown task '{task_name}'. Available task ids: {canonical_locomotion_task_ids()}. " + "Accepted aliases also include the legacy env names and task=/mujoco forms." + ) + + +def locomotion_task_spec(task_name: str) -> LocomotionTaskSpec: + return _TASK_SPECS[normalize_locomotion_task_id(task_name)] + + +def locomotion_env_name(task_name: str) -> str: + return locomotion_task_spec(task_name).env_task_name diff --git a/benchmark/smoke_test.py b/benchmark/smoke_test.py index d3e952dfa..ff5cb8def 100644 --- a/benchmark/smoke_test.py +++ b/benchmark/smoke_test.py @@ -12,7 +12,10 @@ "benchmark_ane_inference", "benchmark_mlp_inference", "benchmark_forward_reset_methods", - "benchmark_mj_step", + "benchmark_physics_step_genesis", + "benchmark_physics_step_isaacgym", + "benchmark_physics_step_mj_step", + "benchmark_physics_step_mujoco_warp", "benchmark_mlx_compile", "benchmark_postprocess", "benchmark_reset_batch_vs_loop",