diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a4053108c..4c26b0078 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -9,10 +9,18 @@ on: - ".github/CODEOWNERS" - ".github/ISSUE_TEMPLATE/**" - ".github/pull_request_template.md" + pull_request_target: + branches: [main] + types: [opened, synchronize, reopened, ready_for_review] + paths-ignore: + - "LICENSE" + - ".github/CODEOWNERS" + - ".github/ISSUE_TEMPLATE/**" + - ".github/pull_request_template.md" workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.ref_name }}-${{ github.event.pull_request.number }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }} cancel-in-progress: true permissions: @@ -23,11 +31,26 @@ env: jobs: ruff-lint: + if: | + github.event_name == 'workflow_dispatch' || + ( + github.event_name == 'pull_request' && + ( + github.event.pull_request.head.repo.full_name == github.repository || + github.event.pull_request.head.ref != 'main' + ) + ) || + ( + github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name != github.repository && + github.event.pull_request.head.ref == 'main' + ) timeout-minutes: 5 runs-on: ubuntu-slim steps: - uses: actions/checkout@v6.0.2 with: + ref: ${{ github.event_name == 'workflow_dispatch' && github.ref || format('refs/pull/{0}/merge', github.event.pull_request.number) }} persist-credentials: false - name: Install uv uses: astral-sh/setup-uv@v8.0.0 @@ -37,11 +60,26 @@ jobs: run: uv run --no-sync ruff check --output-format=github . ruff-format: + if: | + github.event_name == 'workflow_dispatch' || + ( + github.event_name == 'pull_request' && + ( + github.event.pull_request.head.repo.full_name == github.repository || + github.event.pull_request.head.ref != 'main' + ) + ) || + ( + github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name != github.repository && + github.event.pull_request.head.ref == 'main' + ) timeout-minutes: 5 runs-on: ubuntu-slim steps: - uses: actions/checkout@v6.0.2 with: + ref: ${{ github.event_name == 'workflow_dispatch' && github.ref || format('refs/pull/{0}/merge', github.event.pull_request.number) }} persist-credentials: false - name: Install uv uses: astral-sh/setup-uv@v8.0.0 @@ -51,11 +89,26 @@ jobs: run: uv run --no-sync ruff format --check . mypy: + if: | + github.event_name == 'workflow_dispatch' || + ( + github.event_name == 'pull_request' && + ( + github.event.pull_request.head.repo.full_name == github.repository || + github.event.pull_request.head.ref != 'main' + ) + ) || + ( + github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name != github.repository && + github.event.pull_request.head.ref == 'main' + ) timeout-minutes: 10 runs-on: ubuntu-slim steps: - uses: actions/checkout@v6.0.2 with: + ref: ${{ github.event_name == 'workflow_dispatch' && github.ref || format('refs/pull/{0}/merge', github.event.pull_request.number) }} persist-credentials: false - name: Install uv uses: astral-sh/setup-uv@v8.0.0 @@ -65,11 +118,26 @@ jobs: run: uv run mypy src/unilab pyright: + if: | + github.event_name == 'workflow_dispatch' || + ( + github.event_name == 'pull_request' && + ( + github.event.pull_request.head.repo.full_name == github.repository || + github.event.pull_request.head.ref != 'main' + ) + ) || + ( + github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name != github.repository && + github.event.pull_request.head.ref == 'main' + ) timeout-minutes: 10 runs-on: ubuntu-slim steps: - uses: actions/checkout@v6.0.2 with: + ref: ${{ github.event_name == 'workflow_dispatch' && github.ref || format('refs/pull/{0}/merge', github.event.pull_request.number) }} persist-credentials: false - name: Install uv uses: astral-sh/setup-uv@v8.0.0 @@ -79,6 +147,20 @@ jobs: run: uv run pyright test: + if: | + github.event_name == 'workflow_dispatch' || + ( + github.event_name == 'pull_request' && + ( + github.event.pull_request.head.repo.full_name == github.repository || + github.event.pull_request.head.ref != 'main' + ) + ) || + ( + github.event_name == 'pull_request_target' && + github.event.pull_request.head.repo.full_name != github.repository && + github.event.pull_request.head.ref == 'main' + ) strategy: matrix: os: [ubuntu-slim] @@ -87,6 +169,7 @@ jobs: steps: - uses: actions/checkout@v6.0.2 with: + ref: ${{ github.event_name == 'workflow_dispatch' && github.ref || format('refs/pull/{0}/merge', github.event.pull_request.number) }} persist-credentials: false - name: Install uv uses: astral-sh/setup-uv@v8.0.0 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cf1790847..998bc9c85 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -97,7 +97,7 @@ uv run pytest -m veryslow -v ## CI Workflow -PRs targeting `main` trigger five jobs automatically: `ruff-lint`, `ruff-format`, `mypy`, `pyright`, and `test`. The workflow also enables manual runs through `workflow_dispatch`, validates docs through the pytest suite on docs changes, and cancels older in-progress runs for the same PR branch. +PRs targeting `main` trigger five jobs automatically: `ruff-lint`, `ruff-format`, `mypy`, `pyright`, and `test`. Fork PRs whose head branch is also named `main` are routed through a restricted `pull_request_target` lane so CI still starts automatically instead of remaining silent. The workflow also enables manual runs through `workflow_dispatch`, validates docs through the pytest suite on docs changes, and cancels older in-progress runs for the same PR branch. Coverage policy: the default CI lane enforces a minimum non-slow coverage floor, and that floor should only ratchet upward when the suite meaningfully expands. diff --git a/benchmark/benchmark_env_step.py b/benchmark/benchmark_env_step.py index b832cea23..13f4bc345 100644 --- a/benchmark/benchmark_env_step.py +++ b/benchmark/benchmark_env_step.py @@ -2,21 +2,61 @@ Usage: # Run all combinations (go1/go2/g1/g1_motion_tracking x mujoco/motrix): - uv run python benchmark/benchmark_env_step.py + uv run benchmark/benchmark_env_step.py # Single task + backend: - uv run python benchmark/benchmark_env_step.py task=g1_joystick/motrix + uv run benchmark/benchmark_env_step.py task=g1_joystick/motrix # Override bench params: - uv run python benchmark/benchmark_env_step.py task=go1_joystick/mujoco num_envs=4096 num_steps=500 + uv run benchmark/benchmark_env_step.py task=go1_joystick/mujoco num_envs=4096 num_steps=500 + + # Save to custom locations: + uv run benchmark/benchmark_env_step.py --out-json tmp/env_step.json --plot-dir tmp/env_step_plots """ +import importlib.util import sys from pathlib import Path +from typing import Any, cast import numpy as np +plt: Any = None +mpatches: Any = None + +try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.patches as _mpatches + import matplotlib.pyplot as _plt + + mpatches = _mpatches + plt = _plt +except Exception: + pass + ROOT_DIR = Path(__file__).parent.parent +CORE_DIR = ROOT_DIR / "benchmark" / "core" +DEFAULT_OUTPUT_JSON = ROOT_DIR / "benchmark" / "outputs" / "env_step" / "results.json" + + +def _load_helper_module(module_name: str, relative_path: str): + spec = importlib.util.spec_from_file_location(module_name, CORE_DIR / relative_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Failed to load helper module {module_name} from {relative_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +_DEVICE_INFO = _load_helper_module("bench_env_step_device_info", "device_info.py") +_OUTPUT = _load_helper_module("bench_env_step_output", "output.py") + +get_device_info_dict = _DEVICE_INFO.get_device_info_dict +get_device_info_line = _DEVICE_INFO.get_device_info_line +save_json = _OUTPUT.save_json TASK_CONFIGS = { "go1": "task=go1_joystick", @@ -31,6 +71,30 @@ DEFAULT_WARMUP_STEPS = 10 BACKENDS = ["mujoco", "motrix"] +TASK_COLORS = { + "go1": "#4C78A8", + "go2": "#54A24B", + "g1": "#F58518", + "g1_mt": "#B279A2", +} +BACKEND_STYLES = { + "mujoco": {"marker": "o", "linestyle": "-", "hatch": "//"}, + "motrix": {"marker": "s", "linestyle": "--", "hatch": "xx"}, +} +BACKEND_TICK_LABELS = { + "mujoco": "mj", + "motrix": "mx", +} +BREAKDOWN_SEGMENTS = [ + ("apply_action_ms", "apply_action", "#4C78A8"), + ("backend_set_ctrl_ms", "set_ctrl", "#F58518"), + ("backend_physics_ms", "physics", "#54A24B"), + ("backend_refresh_cache_ms", "refresh_cache", "#E45756"), + ("step_core_other_ms", "step_core_other", "#9D755D"), + ("update_state_ms", "update_state", "#72B7B2"), + ("reset_done_ms", "reset_done", "#B279A2"), + ("env_step_other_ms", "env_step_other", "#BAB0AC"), +] def _is_matrix_mode(argv: list[str]) -> bool: @@ -41,26 +105,67 @@ def _is_matrix_mode(argv: list[str]) -> bool: return True -def _parse_bench_args(args: list[str]) -> tuple[dict[str, str], list[str]]: - """Parse benchmark-specific args and return (bench_kwargs, hydra_overrides). +def _parse_cli_args(args: list[str]) -> tuple[dict[str, str], dict[str, Any], list[str]]: + """Parse benchmark/output args and return (bench_kwargs, output_kwargs, hydra_overrides). + + Benchmark args: + num_envs=XXX, num_steps=XXX, warmup_steps=XXX + Output args: + --out-json PATH / --out-json=PATH / out_json=PATH + --plot-dir PATH / --plot-dir=PATH / plot_dir=PATH + --skip-plots / skip_plots=true|false - Benchmark args: num_envs=XXX, num_steps=XXX, warmup_steps=XXX - These are not part of Hydra config, so we extract them separately. + Non-Hydra args are extracted before config composition. """ bench_kwargs: dict[str, str] = {} + output_kwargs: dict[str, Any] = { + "out_json": None, + "plot_dir": None, + "skip_plots": False, + } hydra_overrides: list[str] = [] - for arg in args: + i = 0 + while i < len(args): + arg = args[i] if arg.startswith("num_envs="): bench_kwargs["num_envs"] = arg.split("=", 1)[1] elif arg.startswith("num_steps="): bench_kwargs["num_steps"] = arg.split("=", 1)[1] elif arg.startswith("warmup_steps="): bench_kwargs["warmup_steps"] = arg.split("=", 1)[1] + elif arg.startswith("out_json="): + output_kwargs["out_json"] = arg.split("=", 1)[1] + elif arg.startswith("plot_dir="): + output_kwargs["plot_dir"] = arg.split("=", 1)[1] + elif arg.startswith("skip_plots="): + output_kwargs["skip_plots"] = arg.split("=", 1)[1].lower() in { + "1", + "true", + "yes", + "on", + } + elif arg == "--out-json": + i += 1 + if i >= len(args): + raise ValueError("--out-json requires a path") + output_kwargs["out_json"] = args[i] + elif arg.startswith("--out-json="): + output_kwargs["out_json"] = arg.split("=", 1)[1] + elif arg == "--plot-dir": + i += 1 + if i >= len(args): + raise ValueError("--plot-dir requires a path") + output_kwargs["plot_dir"] = args[i] + elif arg.startswith("--plot-dir="): + output_kwargs["plot_dir"] = arg.split("=", 1)[1] + elif arg == "--skip-plots": + output_kwargs["skip_plots"] = True else: hydra_overrides.append(arg) + i += 1 - return bench_kwargs, hydra_overrides + return bench_kwargs, output_kwargs, hydra_overrides def _compose_cfg(extra_args: list[str]): @@ -82,11 +187,11 @@ def _compose_cfg(extra_args: list[str]): return compose(config_name="config", overrides=overrides) -def _run_single(extra_args: list[str]) -> dict: +def _run_single(extra_args: list[str]) -> dict[str, Any]: """Run a single bench in-process via Hydra and return timing records.""" from unilab.training import BackendAdapter, create_env, ensure_registries - bench_kwargs, hydra_overrides = _parse_bench_args(extra_args) + bench_kwargs, _, hydra_overrides = _parse_cli_args(extra_args) cfg = _compose_cfg(hydra_overrides) ensure_registries() @@ -101,34 +206,37 @@ def _run_single(extra_args: list[str]) -> dict: adapter = BackendAdapter(cfg, root_dir=ROOT_DIR, algo_name="ppo") env_cfg_override = adapter.build_task_env_cfg_override() - env = create_env( - cfg, - num_envs=num_envs, - env_cfg_override=env_cfg_override, - sim_backend=sim_backend, - task_name=task_name, - ) - - nu = env._backend.num_actuators # type: ignore[reportAttributeAccessIssue] - env.init_state() - - for _ in range(warmup_steps): - actions = np.random.uniform(-1, 1, size=(num_envs, nu)).astype(np.float32) - env.step(actions) - + env = None timing_records: dict[str, list[float]] = {} - for _ in range(num_steps): - actions = np.random.uniform(-1, 1, size=(num_envs, nu)).astype(np.float32) - state = env.step(actions) - timing = state.info.get("timing", {}) - for k, v in timing.items(): - timing_records.setdefault(k, []).append(v) - - env.close() + try: + env = create_env( + cfg, + num_envs=num_envs, + env_cfg_override=env_cfg_override, + sim_backend=sim_backend, + task_name=task_name, + ) + + nu = env._backend.num_actuators # type: ignore[reportAttributeAccessIssue] + env.init_state() + + for _ in range(warmup_steps): + actions = np.random.uniform(-1, 1, size=(num_envs, nu)).astype(np.float32) + env.step(actions) + + for _ in range(num_steps): + actions = np.random.uniform(-1, 1, size=(num_envs, nu)).astype(np.float32) + state = env.step(actions) + timing = state.info.get("timing", {}) + for k, v in timing.items(): + timing_records.setdefault(k, []).append(float(v)) + finally: + if env is not None: + env.close() return { - "task_name": task_name, - "sim_backend": sim_backend, + "task_name": str(task_name), + "sim_backend": str(sim_backend), "num_envs": num_envs, "num_steps": num_steps, "warmup_steps": warmup_steps, @@ -136,11 +244,56 @@ def _run_single(extra_args: list[str]) -> dict: } -def _print_single_report(result: dict) -> None: +def _result_label(result: dict[str, Any]) -> str: + return ( + f"{result.get('task_key', _short_task_label(result['task_name']))}/{result['sim_backend']}" + ) + + +def _compute_result_summary(result: dict[str, Any]) -> dict[str, Any]: tr = result["timing_records"] - total_arr = np.array(tr.get("env_step_total_ms", [])) - total_s = total_arr.sum() / 1000.0 + total_arr = np.array(tr.get("env_step_total_ms", []), dtype=np.float64) + total_s = float(total_arr.sum() / 1000.0) steps_per_env = result["num_steps"] * result["num_envs"] + throughput = float(steps_per_env / total_s) if total_s > 0 else 0.0 + + summary: dict[str, Any] = { + "label": _result_label(result), + "total_time_s": total_s, + "mean_step_ms": float(total_arr.mean()) if total_arr.size else 0.0, + "median_step_ms": float(np.median(total_arr)) if total_arr.size else 0.0, + "std_step_ms": float(total_arr.std()) if total_arr.size else 0.0, + "min_step_ms": float(total_arr.min()) if total_arr.size else 0.0, + "max_step_ms": float(total_arr.max()) if total_arr.size else 0.0, + "throughput_env_steps_per_s": throughput, + "timing_mean_ms": {}, + "timing_median_ms": {}, + } + for key, values in tr.items(): + arr = np.array(values, dtype=np.float64) + summary["timing_mean_ms"][key] = float(arr.mean()) if arr.size else 0.0 + summary["timing_median_ms"][key] = float(np.median(arr)) if arr.size else 0.0 + return summary + + +def _serialize_result(result: dict[str, Any]) -> dict[str, Any]: + return { + "task_name": result["task_name"], + "task_key": result.get("task_key"), + "sim_backend": result["sim_backend"], + "num_envs": result["num_envs"], + "num_steps": result["num_steps"], + "warmup_steps": result["warmup_steps"], + "summary": _compute_result_summary(result), + "timing_records": { + key: [float(v) for v in values] for key, values in result["timing_records"].items() + }, + } + + +def _print_single_report(result: dict[str, Any]) -> None: + tr = result["timing_records"] + summary = _compute_result_summary(result) print(f"\n{'=' * 60}") print(f" Task: {result['task_name']}") @@ -148,13 +301,13 @@ def _print_single_report(result: dict) -> None: print(f" Num envs: {result['num_envs']}") print(f" Steps: {result['num_steps']} (warmup: {result['warmup_steps']})") print(f"{'=' * 60}") - print(f" Total time: {total_s:.3f}s") - print(f" Mean step time: {total_arr.mean():.3f}ms") - print(f" Median step time: {np.median(total_arr):.3f}ms") - print(f" Std step time: {total_arr.std():.3f}ms") - print(f" Min step time: {total_arr.min():.3f}ms") - print(f" Max step time: {total_arr.max():.3f}ms") - print(f" Throughput: {steps_per_env / total_s:.0f} env-steps/s") + print(f" Total time: {summary['total_time_s']:.3f}s") + print(f" Mean step time: {summary['mean_step_ms']:.3f}ms") + print(f" Median step time: {summary['median_step_ms']:.3f}ms") + print(f" Std step time: {summary['std_step_ms']:.3f}ms") + print(f" Min step time: {summary['min_step_ms']:.3f}ms") + print(f" Max step time: {summary['max_step_ms']:.3f}ms") + print(f" Throughput: {summary['throughput_env_steps_per_s']:.0f} env-steps/s") if tr: print(f"{'- ' * 30}") print(" Breakdown:") @@ -177,17 +330,18 @@ def _short_task_label(task_name: str) -> str: return task_name[:8] -def _print_comparison_table(results: list[dict]) -> None: +def _print_comparison_table(results: list[dict[str, Any]]) -> None: rows_spec: list[tuple[str, str]] = [ # (display_label, timing_key_or_special) ("total", "env_step_total_ms"), (" apply_action", "apply_action_ms"), - (" step_core", "step_core_ms"), - (" set_ctrl", "backend_set_ctrl_ms"), - (" physics", "backend_physics_ms"), - (" refresh", "backend_refresh_cache_ms"), + (" set_ctrl", "backend_set_ctrl_ms"), + (" physics", "backend_physics_ms"), + (" refresh_cache", "backend_refresh_cache_ms"), + (" step_core_other", "step_core_other_ms"), (" update_state", "update_state_ms"), (" reset_done", "reset_done_ms"), + (" env_step_other", "env_step_other_ms"), ("throughput", "__throughput__"), ] @@ -220,14 +374,14 @@ def hline(left: str, mid: str, right: str, fill: str = "─") -> str: cells: list[str] = [] if key == "__throughput__": for r in results: - total_arr = np.array(r["timing_records"].get("env_step_total_ms", [])) + total_arr = _timing_array(r, "env_step_total_ms") total_s = total_arr.sum() / 1000.0 steps_per_env = r["num_steps"] * r["num_envs"] cells.append(f"{steps_per_env / total_s:,.0f}" if total_s > 0 else "-") else: for r in results: - arr = r["timing_records"].get(key) - cells.append(f"{np.median(arr):.3f}" if arr else "-") + arr = _timing_array(r, key) + cells.append(f"{np.median(arr):.3f}" if arr.size else "-") row = "│" + label.ljust(metric_w) + "│" row += "│".join(v.rjust(col_w - 1) + " " for v in cells) + "│" @@ -244,7 +398,390 @@ def hline(left: str, mid: str, right: str, fill: str = "─") -> str: print(f" ({r0['num_envs']} envs, {r0['num_steps']} steps, throughput = env-steps/s)") -def _run_matrix(extra_args: list[str]) -> None: +def _flatten_step_series(results: list[dict[str, Any]]) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [] + for result in results: + label = _result_label(result) + total_steps = result["timing_records"].get("env_step_total_ms", []) + for step_idx, step_ms in enumerate(total_steps, start=1): + records.append( + { + "label": label, + "step_idx": step_idx, + "step_time_ms": float(step_ms), + } + ) + return records + + +def _task_key(result: dict[str, Any]) -> str: + return cast(str, result.get("task_key", _short_task_label(result["task_name"]))) + + +def _task_color(task_key: str) -> str: + return TASK_COLORS.get(task_key, "#4C78A8") + + +def _backend_style(backend: str) -> dict[str, str]: + return BACKEND_STYLES.get(backend, {"marker": "o", "linestyle": "-", "hatch": ""}) + + +def _ordered_results(results: list[dict[str, Any]]) -> list[dict[str, Any]]: + task_order = {task_key: idx for idx, task_key in enumerate(TASK_CONFIGS)} + backend_order = {backend: idx for idx, backend in enumerate(BACKENDS)} + return sorted( + results, + key=lambda result: ( + task_order.get(_task_key(result), len(task_order)), + backend_order.get(result["sim_backend"], len(backend_order)), + result["task_name"], + ), + ) + + +def _align_timing_array(arr: np.ndarray, target_size: int) -> np.ndarray: + if arr.size == target_size: + return arr + if arr.size == 0: + return np.zeros(target_size, dtype=np.float64) + if arr.size > target_size: + return arr[:target_size] + padded = np.zeros(target_size, dtype=np.float64) + padded[: arr.size] = arr + return padded + + +def _timing_array(result: dict[str, Any], key: str) -> np.ndarray: + timing_records = result["timing_records"] + if key in timing_records: + return np.asarray(timing_records[key], dtype=np.float64) + + if key == "step_core_other_ms": + step_core = _timing_array(result, "step_core_ms") + if step_core.size == 0: + return step_core + backend_total = ( + _align_timing_array(_timing_array(result, "backend_set_ctrl_ms"), step_core.size) + + _align_timing_array(_timing_array(result, "backend_physics_ms"), step_core.size) + + _align_timing_array(_timing_array(result, "backend_refresh_cache_ms"), step_core.size) + ) + return cast(np.ndarray, np.clip(step_core - backend_total, a_min=0.0, a_max=None)) + + if key == "env_step_other_ms": + total = _timing_array(result, "env_step_total_ms") + if total.size == 0: + return total + measured = ( + _align_timing_array(_timing_array(result, "apply_action_ms"), total.size) + + _align_timing_array(_timing_array(result, "step_core_ms"), total.size) + + _align_timing_array(_timing_array(result, "update_state_ms"), total.size) + + _align_timing_array(_timing_array(result, "reset_done_ms"), total.size) + ) + return cast(np.ndarray, np.clip(total - measured, a_min=0.0, a_max=None)) + + return np.array([], dtype=np.float64) + + +def _median_timing_ms(result: dict[str, Any], key: str) -> float: + arr = _timing_array(result, key) + return float(np.median(arr)) if arr.size else 0.0 + + +def _grouped_positions( + results: list[dict[str, Any]], + *, + bar_spacing: float = 1.0, + group_gap: float = 1.1, +) -> tuple[list[dict[str, Any]], np.ndarray, list[dict[str, Any]]]: + ordered = _ordered_results(results) + positions: list[float] = [] + groups: list[dict[str, Any]] = [] + cursor = 0.0 + group_start = 0 + + for idx, result in enumerate(ordered): + positions.append(cursor) + task_key = _task_key(result) + next_result = ordered[idx + 1] if idx + 1 < len(ordered) else None + cursor += bar_spacing + if next_result is None or _task_key(next_result) != task_key: + group_positions = positions[group_start : idx + 1] + groups.append( + { + "task_key": task_key, + "label": task_key, + "start": group_positions[0], + "end": group_positions[-1], + "center": float(sum(group_positions) / len(group_positions)), + } + ) + group_start = idx + 1 + if next_result is not None: + cursor += group_gap + + return ordered, np.array(positions, dtype=np.float64), groups + + +def _decorate_grouped_xaxis( + ax, + ordered_results: list[dict[str, Any]], + positions: np.ndarray, + groups: list[dict[str, Any]], +) -> None: + if not ordered_results: + return + + ax.set_xticks(positions) + ax.set_xticklabels( + [ + BACKEND_TICK_LABELS.get(result["sim_backend"], result["sim_backend"]) + for result in ordered_results + ], + rotation=0, + ha="center", + fontsize=9, + ) + + for idx, group in enumerate(groups): + ax.axvspan( + group["start"] - 0.48, + group["end"] + 0.48, + color=_task_color(group["task_key"]), + alpha=0.05, + zorder=0, + ) + ax.text( + group["center"], + -0.18, + group["label"], + transform=ax.get_xaxis_transform(), + ha="center", + va="top", + fontsize=9, + fontweight="bold", + ) + if idx < len(groups) - 1: + boundary = (group["end"] + groups[idx + 1]["start"]) / 2.0 + ax.axvline(boundary, color="#9AA0A6", linewidth=1.0, alpha=0.7) + + ax.set_xlim(float(positions[0] - 0.7), float(positions[-1] + 0.7)) + + +def _backend_legend_handles() -> list[Any]: + if mpatches is None: + return [] + return [ + mpatches.Patch( + facecolor="#FFFFFF", + edgecolor="#444444", + hatch=_backend_style(backend)["hatch"], + label=backend, + ) + for backend in BACKENDS + ] + + +def _breakdown_segments_for_results(results: list[dict[str, Any]]) -> list[tuple[str, str, str]]: + visible_segments: list[tuple[str, str, str]] = [] + for key, label, color in BREAKDOWN_SEGMENTS: + if key in {"step_core_other_ms", "env_step_other_ms"}: + if max((_median_timing_ms(result, key) for result in results), default=0.0) <= 0.05: + continue + visible_segments.append((key, label, color)) + return visible_segments + + +def _save_total_series_plot(results: list[dict[str, Any]], output_path: Path) -> bool: + if plt is None or not results: + return False + + output_path.parent.mkdir(parents=True, exist_ok=True) + ordered = _ordered_results(results) + fig, ax = plt.subplots(figsize=(10.5, 6.0)) + + for result in ordered: + step_times = _timing_array(result, "env_step_total_ms") + if step_times.size == 0: + continue + step_idx = np.arange(1, step_times.size + 1) + style = _backend_style(result["sim_backend"]) + ax.plot( + step_idx, + step_times, + color=_task_color(_task_key(result)), + linestyle=style["linestyle"], + marker=style["marker"], + linewidth=1.5, + markersize=3.4, + markevery=max(1, step_times.size // 12), + alpha=0.92, + label=_result_label(result), + ) + + ax.set_title(f"Env step total time per iteration\n{get_device_info_line()}") + ax.set_xlabel("step iteration") + ax.set_ylabel("env.step total time (ms)") + ax.grid(True, alpha=0.3) + ax.legend(ncol=2, fontsize=8) + fig.tight_layout() + fig.savefig(output_path, dpi=160, bbox_inches="tight") + plt.close(fig) + print(f"Saved: {output_path.resolve()}") + return True + + +def _save_summary_plot(results: list[dict[str, Any]], output_path: Path) -> bool: + if plt is None or not results: + return False + + output_path.parent.mkdir(parents=True, exist_ok=True) + ordered, x, groups = _grouped_positions(results) + summaries = [_compute_result_summary(result) for result in ordered] + labels = [summary["label"] for summary in summaries] + width = 0.78 + + fig, axes = plt.subplots(1, 2, figsize=(max(10, len(labels) * 1.6), 5.5)) + median_ms = [summary["median_step_ms"] for summary in summaries] + throughput = [summary["throughput_env_steps_per_s"] for summary in summaries] + + for idx, result in enumerate(ordered): + style = _backend_style(result["sim_backend"]) + color = _task_color(_task_key(result)) + axes[0].bar( + x[idx], + median_ms[idx], + width=width, + color=color, + edgecolor="#444444", + hatch=style["hatch"], + alpha=0.92, + ) + axes[1].bar( + x[idx], + throughput[idx], + width=width, + color=color, + edgecolor="#444444", + hatch=style["hatch"], + alpha=0.92, + ) + + axes[0].set_ylabel("median env.step time (ms)") + axes[0].set_title("Latency") + axes[0].grid(axis="y", alpha=0.3) + + axes[1].set_ylabel("throughput (env-steps/s)") + axes[1].set_title("Throughput") + axes[1].grid(axis="y", alpha=0.3) + + for ax in axes: + _decorate_grouped_xaxis(ax, ordered, x, groups) + + fig.suptitle(f"Env step benchmark summary\n{get_device_info_line()}") + handles = _backend_legend_handles() + if handles: + axes[1].legend(handles=handles, title="backend", loc="upper right") + fig.subplots_adjust(bottom=0.24, top=0.82, wspace=0.22) + fig.savefig(output_path, dpi=160, bbox_inches="tight") + plt.close(fig) + print(f"Saved: {output_path.resolve()}") + return True + + +def _save_breakdown_plot(results: list[dict[str, Any]], output_path: Path) -> bool: + if plt is None or not results: + return False + + output_path.parent.mkdir(parents=True, exist_ok=True) + ordered, x, groups = _grouped_positions(results) + width = 0.78 + cumulative = np.zeros(len(ordered), dtype=np.float64) + segments = _breakdown_segments_for_results(ordered) + + fig, ax = plt.subplots(figsize=(max(10, len(ordered) * 1.6), 5.8)) + for timing_key, display_name, color in segments: + values_np = np.array( + [_median_timing_ms(result, timing_key) for result in ordered], dtype=np.float64 + ) + for idx, result in enumerate(ordered): + style = _backend_style(result["sim_backend"]) + ax.bar( + x[idx], + values_np[idx], + width, + bottom=cumulative[idx], + label=display_name if idx == 0 else "_nolegend_", + color=color, + edgecolor="#444444", + hatch=style["hatch"], + alpha=0.94, + ) + cumulative += values_np + + _decorate_grouped_xaxis(ax, ordered, x, groups) + ax.set_ylabel("median step time (ms)") + ax.set_title(f"Env step median breakdown\n{get_device_info_line()}") + ax.grid(axis="y", alpha=0.3) + handles, legend_labels = ax.get_legend_handles_labels() + backend_handles = _backend_legend_handles() + if backend_handles: + handles.extend(backend_handles) + legend_labels.extend([handle.get_label() for handle in backend_handles]) + ax.legend(handles, legend_labels, ncol=2, fontsize=8, loc="upper left") + fig.subplots_adjust(bottom=0.24, top=0.86) + fig.savefig(output_path, dpi=160, bbox_inches="tight") + plt.close(fig) + print(f"Saved: {output_path.resolve()}") + return True + + +def _persist_outputs( + results: list[dict[str, Any]], + *, + mode: str, + out_json: Path, + plot_dir: Path | None, + skip_plots: bool, + failures: list[dict[str, str]] | None = None, +) -> None: + plot_files: list[str] = [] + effective_plot_dir = plot_dir or out_json.parent + + if not skip_plots: + series_path = effective_plot_dir / "env_step_total_series.png" + summary_path = effective_plot_dir / "env_step_summary.png" + breakdown_path = effective_plot_dir / "env_step_breakdown.png" + + if _save_total_series_plot(results, series_path): + plot_files.append(str(series_path.resolve())) + elif results: + print("matplotlib unavailable; skipped env step series plot.") + + if _save_summary_plot(results, summary_path): + plot_files.append(str(summary_path.resolve())) + elif plt is None and results: + print("matplotlib unavailable; skipped env step summary plot.") + + if _save_breakdown_plot(results, breakdown_path): + plot_files.append(str(breakdown_path.resolve())) + elif plt is None and results: + print("matplotlib unavailable; skipped env step breakdown plot.") + + meta: dict[str, Any] = { + "mode": mode, + "device_info": get_device_info_dict(), + "matplotlib_available": plt is not None, + "plot_files": plot_files, + } + if failures: + meta["failures"] = failures + + save_json(out_json, [_serialize_result(result) for result in results], meta) + + +def _run_matrix( + extra_args: list[str], *, out_json: Path, plot_dir: Path | None, skip_plots: bool +) -> None: """Run all task x backend combinations and print comparison.""" from unilab.base.backend.motrix_backend import MOTRIX_AVAILABLE @@ -253,6 +790,7 @@ def _run_matrix(extra_args: list[str]) -> None: print("Note: motrixsim not available, running mujoco only\n") results: list[dict] = [] + failures: list[dict[str, str]] = [] for task_key, task_override in TASK_CONFIGS.items(): for backend in backends: label = f"{task_key}/{backend}" @@ -266,19 +804,39 @@ def _run_matrix(extra_args: list[str]) -> None: _print_single_report(result) except Exception as e: print(f" FAILED: {e}\n") + failures.append({"label": label, "error": str(e)}) if len(results) > 1: _print_comparison_table(results) + _persist_outputs( + results, + mode="matrix", + out_json=out_json, + plot_dir=plot_dir, + skip_plots=skip_plots, + failures=failures, + ) def main() -> None: argv = sys.argv[1:] + _, output_kwargs, _ = _parse_cli_args(argv) + out_json = Path(output_kwargs["out_json"]) if output_kwargs["out_json"] else DEFAULT_OUTPUT_JSON + plot_dir = Path(output_kwargs["plot_dir"]) if output_kwargs["plot_dir"] else None + skip_plots = bool(output_kwargs["skip_plots"]) if _is_matrix_mode(argv): - _run_matrix(argv) + _run_matrix(argv, out_json=out_json, plot_dir=plot_dir, skip_plots=skip_plots) else: result = _run_single(argv) _print_single_report(result) + _persist_outputs( + [result], + mode="single", + out_json=out_json, + plot_dir=plot_dir, + skip_plots=skip_plots, + ) if __name__ == "__main__": diff --git a/benchmark/benchmark_mujoco_backend_step_detail.py b/benchmark/benchmark_mujoco_backend_step_detail.py new file mode 100644 index 000000000..7e9cc50a6 --- /dev/null +++ b/benchmark/benchmark_mujoco_backend_step_detail.py @@ -0,0 +1,586 @@ +#!/usr/bin/env python3 +# pyright: reportAttributeAccessIssue=false +""" +Benchmark detailed MuJoCo backend step overhead. + +This benchmark mirrors the hot path inside `MujocoBackend.step()` and splits it +into: + 1. control broadcast (`set_ctrl`) + 2. `BatchEnvPool.step` (`pool_step`) + 3. physics-state cast/copy (`state_copy`) + 4. `BatchEnvPool.forward` (`forward`) + 5. sensor-data cast/copy (`sensor_copy`) + +It sweeps current locomotion owner tasks across MuJoCo only, with environment +counts from 2^8 to 2^13 by default, and supports multiple `nstep` values so the +per-substep cost can also be compared. + +Usage: + uv run benchmark/benchmark_mujoco_backend_step_detail.py + + uv run benchmark/benchmark_mujoco_backend_step_detail.py \ + --tasks go1_joystick,go2_joystick,g1_joystick \ + --env-nums 256,512,1024,2048,4096,8192 \ + --nsteps 1,2,3,4 + + uv run benchmark/benchmark_mujoco_backend_step_detail.py \ + --tasks g1_joystick --env-nums 2048 --nsteps 1,2,3 --iters 20 +""" + +from __future__ import annotations + +import argparse +import importlib.util +import os +import sys +import time +from dataclasses import asdict, dataclass +from multiprocessing import cpu_count +from pathlib import Path +from typing import Sequence + +import matplotlib +import mujoco +import numpy as np +from matplotlib.patches import Rectangle +from mujoco.batch_env import BatchEnvPool + +from unilab.base.dtype_config import get_global_dtype +from unilab.utils.xml_utils import create_discardvisual_xml + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +ROOT_DIR = Path(__file__).resolve().parents[1] +CORE_DIR = ROOT_DIR / "benchmark" / "core" + + +def _load_helper_module(module_name: str, relative_path: str): + spec = importlib.util.spec_from_file_location(module_name, CORE_DIR / relative_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Failed to load helper module {module_name} from {relative_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +_DEVICE_INFO = _load_helper_module("bench_device_info", "device_info.py") +_OUTPUT = _load_helper_module("bench_output", "output.py") +_TASK_NAMES = _load_helper_module("bench_task_names", "task_names.py") + +get_device_info_dict = _DEVICE_INFO.get_device_info_dict +get_device_info_line = _DEVICE_INFO.get_device_info_line +save_json = _OUTPUT.save_json +canonical_locomotion_task_ids = _TASK_NAMES.canonical_locomotion_task_ids +locomotion_task_spec = _TASK_NAMES.locomotion_task_spec +normalize_locomotion_task_id = _TASK_NAMES.normalize_locomotion_task_id + + +DEFAULT_TASK_IDS = canonical_locomotion_task_ids() +DEFAULT_ENV_NUMS = [2**k for k in range(8, 14)] # 256 .. 8192 +DEFAULT_NSTEPS = [1, 2, 3, 4] +DEFAULT_WARMUP = 3 +DEFAULT_ITERS = 10 +DEFAULT_OUTPUT_JSON = ( + ROOT_DIR / "benchmark" / "outputs" / "mujoco_backend_step_detail" / "results.json" +) + +TASK_COLORS = { + "go1_joystick": "#4C78A8", + "go2_joystick": "#54A24B", + "g1_joystick": "#F58518", +} +COMPONENT_COLORS = { + "set_ctrl_ms": "#9AA0A6", + "pool_step_ms": "#4C78A8", + "state_copy_ms": "#9EC5FE", + "forward_ms": "#F58518", + "sensor_copy_ms": "#FFBF79", +} + + +@dataclass +class BenchRecord: + task: str + env_num: int + nstep: int + nthread: int + warmup: int + iters: int + control_dim: int + state_dim: int + sensor_dim: int + set_ctrl_ms: float + pool_step_ms: float + state_copy_ms: float + physics_total_ms: float + forward_ms: float + sensor_copy_ms: float + refresh_total_ms: float + backend_total_ms: float + physics_per_substep_ms: float + refresh_per_substep_ms: float + total_per_substep_ms: float + + +def _parse_csv_ints(text: str) -> list[int]: + values = [int(part.strip()) for part in text.split(",") if part.strip()] + if not values: + raise ValueError("Expected at least one integer value.") + return values + + +def _parse_csv_tasks(text: str) -> list[str]: + values = [ + normalize_locomotion_task_id(part.strip()) for part in text.split(",") if part.strip() + ] + if not values: + raise ValueError("Expected at least one task id.") + return values + + +def _keyframe0_state_and_ctrl(model: mujoco.MjModel) -> tuple[np.ndarray, np.ndarray]: + data = mujoco.MjData(model) + if model.nkey > 0: + mujoco.mj_resetDataKeyframe(model, data, 0) + else: + mujoco.mj_resetData(model, data) + + nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + state0 = np.empty((nstate,), dtype=np.float64) + mujoco.mj_getState(model, data, state0, mujoco.mjtState.mjSTATE_FULLPHYSICS) + + if model.nu == 0: + ctrl0 = np.empty((0,), dtype=np.float64) + elif model.nkey > 0: + ctrl0 = np.asarray(model.key_ctrl[0], dtype=np.float64).copy() + else: + ctrl0 = np.zeros((model.nu,), dtype=np.float64) + return state0, ctrl0 + + +def _load_discardvisual_model(model_file: str) -> mujoco.MjModel: + model_path = create_discardvisual_xml(model_file) + try: + return mujoco.MjModel.from_xml_path(model_path) + finally: + os.remove(model_path) + + +def _control_limits(model: mujoco.MjModel) -> tuple[np.ndarray, np.ndarray]: + if model.nu == 0: + empty = np.empty((0,), dtype=np.float64) + return empty, empty + + low = np.full((model.nu,), -1.0, dtype=np.float64) + high = np.full((model.nu,), 1.0, dtype=np.float64) + if hasattr(model, "actuator_ctrllimited"): + limited = np.asarray(model.actuator_ctrllimited, dtype=bool) + if np.any(limited): + ctrl_range = np.asarray(model.actuator_ctrlrange, dtype=np.float64) + low[limited] = ctrl_range[limited, 0] + high[limited] = ctrl_range[limited, 1] + return low, high + + +def _sample_controls( + env_num: int, + ctrl_low: np.ndarray, + ctrl_high: np.ndarray, + rng: np.random.Generator, +) -> np.ndarray: + if ctrl_low.size == 0: + return np.empty((env_num, 0), dtype=np.float64) + controls = rng.uniform(ctrl_low, ctrl_high, size=(env_num, ctrl_low.shape[0])) + return np.ascontiguousarray(controls, dtype=np.float64) + + +def _median_ms(samples: list[float]) -> float: + return float(np.median(np.asarray(samples, dtype=np.float64))) + + +def _benchmark_one( + task: str, + env_num: int, + nstep: int, + warmup: int, + iters: int, + seed: int, +) -> BenchRecord: + spec = locomotion_task_spec(task) + cfg = spec.config_cls() + model = _load_discardvisual_model(cfg.model_file) + np_dtype = get_global_dtype() + + state0, _ = _keyframe0_state_and_ctrl(model) + ctrl_low, ctrl_high = _control_limits(model) + nthread = min(env_num, cpu_count() * 2) + rng = np.random.default_rng(seed) + + total_iters = warmup + iters + controls = [_sample_controls(env_num, ctrl_low, ctrl_high, rng) for _ in range(total_iters)] + + set_ctrl_samples: list[float] = [] + pool_step_samples: list[float] = [] + state_copy_samples: list[float] = [] + forward_samples: list[float] = [] + sensor_copy_samples: list[float] = [] + + with BatchEnvPool(model, nbatch=env_num, nthread=nthread) as pool: + physics_state = np.broadcast_to(state0.astype(np_dtype), (env_num, state0.shape[0])).copy() + sensor_data = np.zeros((env_num, model.nsensordata), dtype=np_dtype) + sensor_init = pool.forward(physics_state) + sensor_data[:] = sensor_init.astype(np_dtype) + + for iteration_idx, ctrl in enumerate(controls): + t0 = time.perf_counter() + control_traj = np.broadcast_to(ctrl[:, None, :], (env_num, nstep, ctrl.shape[-1])) + set_ctrl_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + state_np = pool.step( + physics_state, + nstep=nstep, + control=control_traj, + control_spec=int(mujoco.mjtState.mjSTATE_CTRL), + ) + pool_step_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + physics_state[:] = state_np.astype(np_dtype) + state_copy_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + sensor_np = pool.forward(physics_state) + forward_ms = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + sensor_data[:] = sensor_np.astype(np_dtype) + sensor_copy_ms = (time.perf_counter() - t0) * 1000.0 + + if iteration_idx >= warmup: + set_ctrl_samples.append(set_ctrl_ms) + pool_step_samples.append(pool_step_ms) + state_copy_samples.append(state_copy_ms) + forward_samples.append(forward_ms) + sensor_copy_samples.append(sensor_copy_ms) + + set_ctrl_ms = _median_ms(set_ctrl_samples) + pool_step_ms = _median_ms(pool_step_samples) + state_copy_ms = _median_ms(state_copy_samples) + forward_ms = _median_ms(forward_samples) + sensor_copy_ms = _median_ms(sensor_copy_samples) + physics_total_ms = pool_step_ms + state_copy_ms + refresh_total_ms = forward_ms + sensor_copy_ms + backend_total_ms = set_ctrl_ms + physics_total_ms + refresh_total_ms + + return BenchRecord( + task=task, + env_num=env_num, + nstep=nstep, + nthread=nthread, + warmup=warmup, + iters=iters, + control_dim=model.nu, + state_dim=state0.shape[0], + sensor_dim=model.nsensordata, + set_ctrl_ms=set_ctrl_ms, + pool_step_ms=pool_step_ms, + state_copy_ms=state_copy_ms, + physics_total_ms=physics_total_ms, + forward_ms=forward_ms, + sensor_copy_ms=sensor_copy_ms, + refresh_total_ms=refresh_total_ms, + backend_total_ms=backend_total_ms, + physics_per_substep_ms=physics_total_ms / nstep, + refresh_per_substep_ms=refresh_total_ms / nstep, + total_per_substep_ms=backend_total_ms / nstep, + ) + + +def _grouped_bar_positions( + env_nums: Sequence[int], + nsteps: Sequence[int], + *, + bar_width: float = 0.22, + intra_gap: float = 0.06, + group_gap: float = 0.48, +) -> tuple[dict[tuple[int, int], float], list[float], float]: + positions: dict[tuple[int, int], float] = {} + group_centers: list[float] = [] + cursor = 0.0 + stride = bar_width + intra_gap + + for env_num in env_nums: + group_positions: list[float] = [] + for nstep in nsteps: + positions[(env_num, nstep)] = cursor + group_positions.append(cursor) + cursor += stride + group_centers.append(float(sum(group_positions) / len(group_positions))) + cursor += group_gap + + return positions, group_centers, bar_width + + +def _decorate_grouped_env_axis( + ax, + env_nums: Sequence[int], + nsteps: Sequence[int], + group_centers: Sequence[float], + positions: dict[tuple[int, int], float], + bar_width: float, +) -> None: + tick_positions = [positions[(env_num, nstep)] for env_num in env_nums for nstep in nsteps] + tick_labels = [f"n={nstep}" for env_num in env_nums for nstep in nsteps] + ax.set_xticks(tick_positions) + ax.set_xticklabels(tick_labels, rotation=0, ha="center", fontsize=8) + + for idx, env_num in enumerate(env_nums): + left = positions[(env_num, nsteps[0])] - bar_width * 0.7 + right = positions[(env_num, nsteps[-1])] + bar_width * 0.7 + ax.axvspan(left, right, color="#000000", alpha=0.025, zorder=0) + ax.text( + group_centers[idx], + -0.16, + str(env_num), + transform=ax.get_xaxis_transform(), + ha="center", + va="top", + fontsize=9, + fontweight="bold", + ) + if idx < len(env_nums) - 1: + boundary = (right + positions[(env_nums[idx + 1], nsteps[0])] - bar_width * 0.7) / 2.0 + ax.axvline(boundary, color="#B0B0B0", linewidth=1.0, alpha=0.8) + + +def _plot_component_bars( + records: Sequence[BenchRecord], + output_path: Path, + *, + amortized_per_substep: bool, +) -> bool: + if not records: + return False + + output_path.parent.mkdir(parents=True, exist_ok=True) + tasks = [task for task in DEFAULT_TASK_IDS if any(record.task == task for record in records)] + nsteps = sorted({record.nstep for record in records}) + env_nums = sorted({record.env_num for record in records}) + positions, group_centers, bar_width = _grouped_bar_positions(env_nums, nsteps) + component_keys = [ + "set_ctrl_ms", + "pool_step_ms", + "state_copy_ms", + "forward_ms", + "sensor_copy_ms", + ] + + fig, axes = plt.subplots( + len(tasks), 1, figsize=(max(11.5, len(env_nums) * 1.9), 4.2 * len(tasks)) + ) + if len(tasks) == 1: + axes = [axes] + + ylabel = ( + "amortized time per substep (ms)" + if amortized_per_substep + else "time per env.step call (ms)" + ) + title_suffix = "amortized-per-substep" if amortized_per_substep else "total" + + for ax, task in zip(axes, tasks, strict=False): + task_records = [record for record in records if record.task == task] + bars = {(record.env_num, record.nstep): record for record in task_records} + x_values = [positions[(env_num, nstep)] for env_num in env_nums for nstep in nsteps] + cumulative = np.zeros(len(x_values), dtype=np.float64) + for component_key in component_keys: + values = [] + for env_num in env_nums: + for nstep in nsteps: + record = bars.get((env_num, nstep)) + if record is None: + values.append(0.0) + continue + value = getattr(record, component_key) + if amortized_per_substep: + value /= nstep + values.append(value) + values_np = np.asarray(values, dtype=np.float64) + ax.bar( + x_values, + values_np, + width=bar_width, + bottom=cumulative, + label=component_key.removesuffix("_ms"), + color=COMPONENT_COLORS[component_key], + edgecolor="#444444", + linewidth=0.5, + alpha=0.95, + ) + cumulative += values_np + + ax.set_title(f"{task} — {title_suffix}", fontsize=10) + ax.set_ylabel(ylabel) + ax.grid(axis="y", alpha=0.25) + _decorate_grouped_env_axis(ax, env_nums, nsteps, group_centers, positions, bar_width) + + component_handles = [ + Rectangle( + (0, 0), + 1, + 1, + facecolor=COMPONENT_COLORS[key], + edgecolor="#444444", + label=key.removesuffix("_ms"), + ) + for key in component_keys + ] + axes[0].legend( + handles=component_handles, ncol=min(3, len(component_handles)), fontsize=8, loc="upper left" + ) + axes[-1].set_xlabel("num_envs") + + fig.suptitle( + ( + "MuJoCo backend step detail — stacked component bars " + "(amortized over nstep)\n" + f"{get_device_info_line()}" + ) + if amortized_per_substep + else f"MuJoCo backend step detail — stacked component bars\n{get_device_info_line()}", + fontsize=11, + ) + fig.subplots_adjust(bottom=0.14, top=0.9, hspace=0.32) + fig.savefig(output_path, dpi=160, bbox_inches="tight") + plt.close(fig) + print(f"Saved: {output_path.resolve()}") + return True + + +def _print_table(records: Sequence[BenchRecord]) -> None: + if not records: + print("No records.") + return + + headers = ( + "task", + "envs", + "nstep", + "backend_total", + "physics_total", + "refresh_total", + "pool_step", + "state_copy", + "forward", + "sensor_copy", + ) + rows = [] + for record in records: + rows.append( + ( + record.task, + str(record.env_num), + str(record.nstep), + f"{record.backend_total_ms:.3f}", + f"{record.physics_total_ms:.3f}", + f"{record.refresh_total_ms:.3f}", + f"{record.pool_step_ms:.3f}", + f"{record.state_copy_ms:.3f}", + f"{record.forward_ms:.3f}", + f"{record.sensor_copy_ms:.3f}", + ) + ) + + widths = [len(header) for header in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def _fmt(values: Sequence[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + print() + print(_fmt(headers)) + print("-+-".join("-" * width for width in widths)) + for row in rows: + print(_fmt(row)) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Benchmark detailed MuJoCo backend step overhead.") + parser.add_argument("--tasks", type=str, default=",".join(DEFAULT_TASK_IDS)) + parser.add_argument("--env-nums", type=str, default=",".join(str(v) for v in DEFAULT_ENV_NUMS)) + parser.add_argument("--nsteps", type=str, default=",".join(str(v) for v in DEFAULT_NSTEPS)) + parser.add_argument("--warmup", type=int, default=DEFAULT_WARMUP) + parser.add_argument("--iters", type=int, default=DEFAULT_ITERS) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out-json", type=Path, default=DEFAULT_OUTPUT_JSON) + parser.add_argument("--plot-dir", type=Path, default=None) + parser.add_argument("--skip-plots", action="store_true") + args = parser.parse_args() + + tasks = _parse_csv_tasks(args.tasks) + env_nums = _parse_csv_ints(args.env_nums) + nsteps = _parse_csv_ints(args.nsteps) + plot_dir = args.plot_dir if args.plot_dir is not None else args.out_json.parent + + print(f"tasks: {tasks}") + print(f"env_nums: {env_nums}") + print(f"nsteps: {nsteps}") + print(f"warmup={args.warmup}, iters={args.iters}") + + records: list[BenchRecord] = [] + for task_idx, task in enumerate(tasks): + for env_num in env_nums: + for nstep in nsteps: + seed = args.seed + task_idx * 100_000 + env_num * 17 + nstep * 997 + print(f"[run] task={task} envs={env_num} nstep={nstep}", flush=True) + record = _benchmark_one( + task=task, + env_num=env_num, + nstep=nstep, + warmup=args.warmup, + iters=args.iters, + seed=seed, + ) + records.append(record) + print( + f" total={record.backend_total_ms:.3f}ms " + f"physics={record.physics_total_ms:.3f}ms " + f"refresh={record.refresh_total_ms:.3f}ms " + f"per_substep={record.total_per_substep_ms:.3f}ms" + ) + + _print_table(records) + + plot_files: list[str] = [] + if not args.skip_plots: + total_plot = plot_dir / "component_breakdown_total_ms.png" + amortized_plot = plot_dir / "component_breakdown_amortized_per_substep_ms.png" + if _plot_component_bars(records, total_plot, amortized_per_substep=False): + plot_files.append(str(total_plot.resolve())) + if _plot_component_bars(records, amortized_plot, amortized_per_substep=True): + plot_files.append(str(amortized_plot.resolve())) + + save_json( + args.out_json, + [asdict(record) for record in records], + { + "device_info": get_device_info_dict(), + "xml_preprocess": "create_discardvisual_xml", + "tasks": tasks, + "env_nums": env_nums, + "nsteps": nsteps, + "warmup": args.warmup, + "iters": args.iters, + "dtype": str(get_global_dtype()), + "plot_files": plot_files, + }, + ) + + +if __name__ == "__main__": + main() diff --git a/docs/zh_CN/CONTRIBUTING.md b/docs/zh_CN/CONTRIBUTING.md index 5c96902bf..d924db04b 100644 --- a/docs/zh_CN/CONTRIBUTING.md +++ b/docs/zh_CN/CONTRIBUTING.md @@ -97,7 +97,7 @@ uv run pytest -m veryslow -v ## CI Workflow -指向 `main` 的 PR 会自动触发五个 job: `ruff-lint`、`ruff-format`、`mypy`、`pyright` 和 `test`。workflow 也支持通过 `workflow_dispatch` 手动触发；文档改动会通过 pytest 套件参与校验，并且会自动取消同一 PR 分支上较早的进行中运行。 +指向 `main` 的 PR 会自动触发五个 job: `ruff-lint`、`ruff-format`、`mypy`、`pyright` 和 `test`。如果 PR 来自 fork，且 head branch 也叫 `main`，仓库会通过受限的 `pull_request_target` 通道来启动 CI，避免 GitHub Actions 继续静默不跑。workflow 也支持通过 `workflow_dispatch` 手动触发；文档改动会通过 pytest 套件参与校验，并且会自动取消同一 PR 分支上较早的进行中运行。 覆盖率策略: 默认 CI lane 会对非 slow 测试施加最低覆盖率门槛，这个门槛只应随着测试护栏增强而逐步上调，不应回退。