From c9693d463e38a4ed63d38d67f6bb1b81bcbad0e6 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 15:56:35 +0800 Subject: [PATCH 01/26] feat: add g1 joystick numba acceleration --- benchmark/benchmark_g1_joystick_numba.py | 812 ++++++++++++++++++ pyproject.rocm.toml | 1 + pyproject.toml | 1 + src/unilab/envs/locomotion/g1/joystick.py | 39 +- .../envs/locomotion/g1/joystick_numba.py | 656 ++++++++++++++ .../test_g1_joystick_numba_benchmark.py | 21 + .../envs/locomotion/g1/test_joystick_numba.py | 296 +++++++ uv.lock | 315 ++++--- uv.rocm.lock | 55 ++ 9 files changed, 2059 insertions(+), 137 deletions(-) create mode 100644 benchmark/benchmark_g1_joystick_numba.py create mode 100644 src/unilab/envs/locomotion/g1/joystick_numba.py create mode 100644 tests/benchmark/test_g1_joystick_numba_benchmark.py create mode 100644 tests/envs/locomotion/g1/test_joystick_numba.py diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py new file mode 100644 index 000000000..c2eabb9ea --- /dev/null +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -0,0 +1,812 @@ +"""Benchmark the task-specific Numba backend for G1 joystick rewards. + +This benchmark intentionally avoids MuJoCo/Motrix construction. It measures the +hot slice owned by ``src/unilab/envs/locomotion/g1/joystick.py``: + +* baseline: real ``G1WalkEnv._compute_reward`` reward dispatch + termination; +* accelerated: ``G1WalkNumbaAccelerator.compute``. + +Synthetic backend arrays keep the benchmark deterministic while still using the +same reward functions, reward config fields, sensor names, and accelerator entry +point as the task. + +Run: + uv run python -m benchmark.benchmark_g1_joystick_numba + uv run python benchmark/benchmark_g1_joystick_numba.py + uv run python -m benchmark.benchmark_g1_joystick_numba --quick +""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import time +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from statistics import mean, stdev +from typing import Any + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import numpy as np +from numba import get_num_threads, set_num_threads + +try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt +except Exception: # pragma: no cover - plotting is optional in benchmark scripts + plt = None + +from benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.dtype_config import get_global_dtype +from unilab.envs.locomotion.g1.joystick import ( + G1RewardConfig, + G1WalkEnv, + build_upper_body_pose_weights, +) +from unilab.envs.locomotion.g1.joystick_numba import G1WalkNumbaAccelerator + +NUM_ACTION = 29 +DEFAULT_POSE_WEIGHTS = [ + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 0.01, + 1.0, + 5.0, + 0.01, + 5.0, + 5.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, + 50.0, +] + +PPO_SCALES = { + "tracking_lin_vel": 2.0, + "tracking_ang_vel": 0.2, + "feet_phase": 1.0, + "lin_vel_z": -1.0, + "ang_vel_xy": -0.25, + "base_height": -500.0, + "orientation": -5.0, + "action_rate": -0.01, + "pose": -0.1, +} + +SAC_SCALES = { + "tracking_lin_vel": 2.0, + "tracking_ang_vel": 1.5, + "penalty_ang_vel_xy": -1.0, + "penalty_orientation": -10.0, + "penalty_action_rate": -4.0, + "pose": -0.5, + "penalty_feet_ori": -20.0, + "feet_phase": 5.0, + "alive": 10.0, +} + +FULL_SUPPORTED_SCALES = { + "tracking_lin_vel": 2.0, + "tracking_ang_vel": 1.5, + "forward_progress": 0.5, + "under_speed": -0.5, + "lin_vel_z": -1.0, + "orientation": -5.0, + "penalty_orientation": -10.0, + "ang_vel_xy": -0.25, + "penalty_ang_vel_xy": -1.0, + "action_rate": -0.01, + "penalty_action_rate": -4.0, + "base_height": -500.0, + "pose": -0.5, + "upper_body_pose": -0.1, + "penalty_close_feet_xy": -2.0, + "penalty_feet_ori": -20.0, + "feet_phase": 5.0, + "feet_phase_contrast": 1.0, + "feet_phase_contact": 1.0, + "feet_double_stance": -0.2, + "feet_air_time": 0.5, + "alive": 10.0, +} + + +@dataclass(frozen=True) +class ProfileSpec: + name: str + scales: dict[str, float] + reward_cfg: G1RewardConfig + + +@dataclass +class BenchCase: + profile: str + num_envs: int + path: str + threads: int | None + mean_ms: float + min_ms: float + std_ms: float + env_per_s: float + speedup_vs_numpy: float + compile_ms: float | None = None + + +@dataclass +class SyntheticBackend: + base_pos: np.ndarray + sensors: dict[str, np.ndarray] = field(default_factory=dict) + + def get_base_pos(self) -> np.ndarray: + return self.base_pos + + def get_sensor_data(self, name: str) -> np.ndarray: + return self.sensors[name] + + +@dataclass +class SyntheticBatch: + env: G1WalkEnv + info: dict[str, Any] + linvel: np.ndarray + gyro: np.ndarray + gravity: np.ndarray + dof_pos: np.ndarray + dof_vel: np.ndarray + + +class _Cfg: + ctrl_dt = 0.02 + + +def make_profile_specs() -> dict[str, ProfileSpec]: + return { + "ppo_default": ProfileSpec( + name="ppo_default", + scales=PPO_SCALES, + reward_cfg=G1RewardConfig( + scales=PPO_SCALES, + tracking_sigma=0.25, + gait_frequency=1.5, + feet_phase_swing_height=0.09, + feet_phase_tracking_sigma=0.008, + base_height_target=0.754, + min_base_height=0.55, + max_tilt_deg=25.0, + pose_weights=DEFAULT_POSE_WEIGHTS, + ), + ), + "sac_default": ProfileSpec( + name="sac_default", + scales=SAC_SCALES, + reward_cfg=G1RewardConfig( + scales=SAC_SCALES, + tracking_sigma=0.25, + gait_frequency=1.5, + feet_phase_swing_height=0.09, + feet_phase_tracking_sigma=0.04, + base_height_target=0.754, + min_base_height=0.3, + max_tilt_deg=65.0, + close_feet_threshold=0.15, + pose_weights=DEFAULT_POSE_WEIGHTS, + ), + ), + "full_supported": ProfileSpec( + name="full_supported", + scales=FULL_SUPPORTED_SCALES, + reward_cfg=G1RewardConfig( + scales=FULL_SUPPORTED_SCALES, + tracking_sigma=0.25, + gait_frequency=1.5, + feet_phase_swing_height=0.09, + feet_phase_tracking_sigma=0.04, + base_height_target=0.754, + min_base_height=0.3, + max_tilt_deg=65.0, + min_forward_speed_for_gait_reward=0.0, + close_feet_threshold=0.15, + pose_weights=DEFAULT_POSE_WEIGHTS, + ), + ), + } + + +def _make_fake_env( + num_envs: int, reward_cfg: G1RewardConfig, backend: SyntheticBackend +) -> G1WalkEnv: + env = object.__new__(G1WalkEnv) + env._num_envs = num_envs + env._num_action = NUM_ACTION + env._cfg = _Cfg() + env._reward_cfg = reward_cfg + env._backend = backend + env._enable_reward_log = True + env.default_angles = np.zeros((NUM_ACTION,), dtype=get_global_dtype()) + env._pose_weights = np.asarray(reward_cfg.pose_weights, dtype=get_global_dtype()) + env._upper_body_pose_weights = build_upper_body_pose_weights(reward_cfg.pose_weights) + env._init_reward_functions() + return env + + +def make_batch(num_envs: int, spec: ProfileSpec, seed: int) -> SyntheticBatch: + rng = np.random.default_rng(seed) + dtype = get_global_dtype() + + linvel = rng.normal(loc=(0.8, 0.0, 0.0), scale=(0.25, 0.08, 0.08), size=(num_envs, 3)).astype( + dtype + ) + gyro = rng.normal(loc=(0.0, 0.0, 0.1), scale=(0.08, 0.08, 0.15), size=(num_envs, 3)).astype( + dtype + ) + gravity = rng.normal(loc=(0.0, 0.0, 0.98), scale=(0.02, 0.02, 0.01), size=(num_envs, 3)) + gravity[:, 2] = np.clip(gravity[:, 2], -1.0, 1.0) + gravity = gravity.astype(dtype) + dof_pos = rng.normal(loc=0.0, scale=0.03, size=(num_envs, NUM_ACTION)).astype(dtype) + dof_vel = rng.normal(loc=0.0, scale=0.2, size=(num_envs, NUM_ACTION)).astype(dtype) + current_actions = rng.normal(loc=0.0, scale=0.2, size=(num_envs, NUM_ACTION)).astype(dtype) + last_actions = (current_actions + rng.normal(0.0, 0.05, size=current_actions.shape)).astype( + dtype + ) + commands = rng.normal(loc=(0.8, 0.0, 0.1), scale=(0.2, 0.05, 0.1), size=(num_envs, 3)).astype( + dtype + ) + commands[:, 0] = np.maximum(commands[:, 0], 0.05) + gait_phase = rng.uniform(0.0, 2.0 * np.pi, size=(num_envs, 2)).astype(dtype) + feet_air_time = rng.uniform(0.0, 0.8, size=(num_envs, 2)).astype(dtype) + + base_pos = np.zeros((num_envs, 3), dtype=dtype) + base_pos[:, 2] = spec.reward_cfg.base_height_target + rng.normal(0.0, 0.015, size=num_envs) + left_foot_pos = np.column_stack( + [ + np.full(num_envs, -0.1, dtype=dtype), + rng.normal(-0.06, 0.01, size=num_envs), + rng.uniform(0.0, 0.1, size=num_envs), + ] + ).astype(dtype) + right_foot_pos = np.column_stack( + [ + np.full(num_envs, 0.1, dtype=dtype), + rng.normal(0.06, 0.01, size=num_envs), + rng.uniform(0.0, 0.1, size=num_envs), + ] + ).astype(dtype) + left_foot_quat = np.tile(np.array([1.0, 0.01, 0.02, 0.0], dtype=dtype), (num_envs, 1)) + right_foot_quat = np.tile(np.array([1.0, 0.02, 0.01, 0.0], dtype=dtype), (num_envs, 1)) + + sensors = { + "upvector": gravity, + "left_foot_pos": left_foot_pos, + "right_foot_pos": right_foot_pos, + "left_foot_quat": left_foot_quat, + "right_foot_quat": right_foot_quat, + } + for side in ("left", "right"): + for idx in range(4): + sensors[f"{side}_foot_contact_{idx}"] = (rng.random(num_envs) > 0.45).astype(dtype) + backend = SyntheticBackend(base_pos=base_pos, sensors=sensors) + env = _make_fake_env(num_envs, spec.reward_cfg, backend) + info = { + "steps": np.zeros((num_envs,), dtype=np.uint32), + "commands": commands, + "current_actions": current_actions, + "last_actions": last_actions, + "gait_phase": gait_phase, + "feet_air_time": feet_air_time, + "log": {}, + } + return SyntheticBatch( + env=env, + info=info, + linvel=linvel, + gyro=gyro, + gravity=gravity, + dof_pos=dof_pos, + dof_vel=dof_vel, + ) + + +def compute_numpy(batch: SyntheticBatch) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: + env = batch.env + info = {**batch.info, "log": {}} + max_tilt_rad = np.deg2rad(env._reward_cfg.max_tilt_deg) + tilt = np.arccos(np.clip(batch.gravity[:, 2], -1.0, 1.0)) + terminated = np.logical_or( + tilt > max_tilt_rad, + env._backend.get_base_pos()[:, 2] < env._reward_cfg.min_base_height, + ) + reward = env._compute_reward( + info, + batch.linvel, + batch.gyro, + batch.gravity, + batch.dof_pos, + batch.dof_vel, + ) + return reward, terminated, info.get("log", {}) + + +def compute_numba( + batch: SyntheticBatch, accelerator: G1WalkNumbaAccelerator +) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: + info = {**batch.info, "log": {}} + out = accelerator.compute( + env=batch.env, + info=info, + linvel=batch.linvel, + gyro=batch.gyro, + gravity=batch.gravity, + dof_pos=batch.dof_pos, + dof_vel=batch.dof_vel, + scales=batch.env._reward_cfg.scales, + enable_log=True, + ) + return out.reward, out.terminated, out.log + + +def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: + for _ in range(warmup): + fn() + samples = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + samples.append((time.perf_counter() - t0) * 1e3) + return mean(samples), min(samples), stdev(samples) if len(samples) > 1 else 0.0 + + +def check_parity(batch: SyntheticBatch, accelerator: G1WalkNumbaAccelerator) -> dict[str, float]: + reward_np, terminated_np, log_np = compute_numpy(batch) + reward_nb, terminated_nb, log_nb = compute_numba(batch, accelerator) + np.testing.assert_allclose(reward_nb, reward_np, rtol=5e-5, atol=5e-5) + np.testing.assert_array_equal(terminated_nb, terminated_np) + for key, value in log_np.items(): + if key in log_nb: + np.testing.assert_allclose(log_nb[key], value, rtol=5e-5, atol=5e-5) + return { + "max_abs_reward_diff": float(np.max(np.abs(reward_nb - reward_np))), + "termination_mismatch": float(np.count_nonzero(terminated_nb != terminated_np)), + } + + +def bench_one( + *, + profile: ProfileSpec, + num_envs: int, + thread_counts: list[int], + iters: int, + warmup: int, + seed: int, +) -> tuple[list[BenchCase], dict[str, float]]: + batch = make_batch(num_envs, profile, seed) + max_threads = get_num_threads() + + numpy_mean, numpy_min, numpy_std = time_call( + lambda: compute_numpy(batch), iters=iters, warmup=warmup + ) + records = [ + BenchCase( + profile=profile.name, + num_envs=num_envs, + path="numpy_dispatch", + threads=None, + mean_ms=numpy_mean, + min_ms=numpy_min, + std_ms=numpy_std, + env_per_s=num_envs / (numpy_mean * 1e-3), + speedup_vs_numpy=1.0, + ) + ] + + compile_driver = G1WalkNumbaAccelerator.from_env(batch.env, num_threads=1) + t0 = time.perf_counter() + compute_numba(batch, compile_driver) + compile_ms = (time.perf_counter() - t0) * 1e3 + parity = check_parity(batch, compile_driver) + + for threads in [1, *thread_counts]: + if threads > max_threads: + continue + accelerator = G1WalkNumbaAccelerator.from_env(batch.env, num_threads=threads) + numba_mean, numba_min, numba_std = time_call( + lambda: compute_numba(batch, accelerator), iters=iters, warmup=warmup + ) + records.append( + BenchCase( + profile=profile.name, + num_envs=num_envs, + path="numba_accelerator", + threads=threads, + mean_ms=numba_mean, + min_ms=numba_min, + std_ms=numba_std, + env_per_s=num_envs / (numba_mean * 1e-3), + speedup_vs_numpy=numpy_mean / numba_mean, + compile_ms=compile_ms if threads == 1 else None, + ) + ) + set_num_threads(max_threads) + return records, parity + + +def _case_to_dict(case: BenchCase) -> dict[str, Any]: + return { + "profile": case.profile, + "num_envs": case.num_envs, + "path": case.path, + "threads": case.threads, + "mean_ms": case.mean_ms, + "min_ms": case.min_ms, + "std_ms": case.std_ms, + "env_per_s": case.env_per_s, + "speedup_vs_numpy": case.speedup_vs_numpy, + "compile_ms": case.compile_ms, + } + + +def _format_table(records: list[BenchCase]) -> str: + headers = [ + "profile", + "envs", + "path", + "threads", + "mean_ms", + "min_ms", + "speedup", + "M env/s", + ] + rows = [] + for r in records: + rows.append( + [ + r.profile, + str(r.num_envs), + r.path, + "-" if r.threads is None else str(r.threads), + f"{r.mean_ms:.3f}", + f"{r.min_ms:.3f}", + f"{r.speedup_vs_numpy:.2f}x", + f"{r.env_per_s / 1e6:.2f}", + ] + ) + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + lines = [fmt(headers), "-+-".join("-" * width for width in widths)] + lines.extend(fmt(row) for row in rows) + return "\n".join(lines) + + +def _best_numba_by_case(records: list[BenchCase]) -> dict[tuple[str, int], BenchCase]: + best_by_case: dict[tuple[str, int], BenchCase] = {} + for record in records: + if record.path != "numba_accelerator": + continue + key = (record.profile, record.num_envs) + if key not in best_by_case or record.mean_ms < best_by_case[key].mean_ms: + best_by_case[key] = record + return best_by_case + + +def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) -> list[str]: + if plt is None or not records: + print("Plotting skipped: matplotlib is not available.") + return [] + + output_dir.mkdir(parents=True, exist_ok=True) + profiles = sorted({record.profile for record in records}) + num_envs = sorted({record.num_envs for record in records}) + best_by_case = _best_numba_by_case(records) + + fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(21, 6)) + fig.suptitle(f"G1 joystick Numba reward+termination benchmark\n{device_info}", fontsize=13) + + # 1) Best speedup vs env count. + ax1 = axes[0] + for profile in profiles: + x = [] + y = [] + labels = [] + for env_count in num_envs: + best = best_by_case.get((profile, env_count)) + if best is None: + continue + x.append(env_count) + y.append(best.speedup_vs_numpy) + labels.append(best.threads) + if not x: + continue + ax1.plot(x, y, marker="o", label=profile) + for x_val, y_val, threads in zip(x, y, labels): + ax1.annotate( + f"{threads}T", + xy=(x_val, y_val), + xytext=(0, 6), + textcoords="offset points", + ha="center", + fontsize=8, + ) + ax1.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") + ax1.set_title("Best speedup vs numpy dispatch") + ax1.set_xlabel("num_envs") + ax1.set_ylabel("Speedup") + ax1.set_xscale("log", base=2) + ax1.set_xticks(num_envs) + ax1.set_xticklabels([str(value) for value in num_envs]) + ax1.grid(True, alpha=0.3) + ax1.legend(fontsize=8) + + # 2) Latency: numpy dispatch vs best numba. + ax2 = axes[1] + for profile in profiles: + numpy_subset = sorted( + [ + record + for record in records + if record.profile == profile and record.path == "numpy_dispatch" + ], + key=lambda record: record.num_envs, + ) + if numpy_subset: + ax2.plot( + [record.num_envs for record in numpy_subset], + [record.mean_ms for record in numpy_subset], + marker="o", + linestyle="--", + label=f"{profile} numpy", + ) + best_subset = [ + best_by_case[(profile, env_count)] + for env_count in num_envs + if (profile, env_count) in best_by_case + ] + if best_subset: + ax2.plot( + [record.num_envs for record in best_subset], + [record.mean_ms for record in best_subset], + marker="s", + linestyle="-", + label=f"{profile} numba best", + ) + ax2.set_title("Latency: numpy dispatch vs best numba") + ax2.set_xlabel("num_envs") + ax2.set_ylabel("Mean latency (ms)") + ax2.set_xscale("log", base=2) + ax2.set_yscale("log") + ax2.set_xticks(num_envs) + ax2.set_xticklabels([str(value) for value in num_envs]) + ax2.grid(True, alpha=0.3) + ax2.legend(fontsize=7) + + # 3) Throughput: numpy dispatch vs best numba. + ax3 = axes[2] + for profile in profiles: + numpy_subset = sorted( + [ + record + for record in records + if record.profile == profile and record.path == "numpy_dispatch" + ], + key=lambda record: record.num_envs, + ) + if numpy_subset: + ax3.plot( + [record.num_envs for record in numpy_subset], + [record.env_per_s / 1e6 for record in numpy_subset], + marker="o", + linestyle="--", + label=f"{profile} numpy", + ) + best_subset = [ + best_by_case[(profile, env_count)] + for env_count in num_envs + if (profile, env_count) in best_by_case + ] + if best_subset: + ax3.plot( + [record.num_envs for record in best_subset], + [record.env_per_s / 1e6 for record in best_subset], + marker="s", + linestyle="-", + label=f"{profile} numba best", + ) + ax3.set_title("Throughput: numpy dispatch vs best numba") + ax3.set_xlabel("num_envs") + ax3.set_ylabel("Throughput (M env/s)") + ax3.set_xscale("log", base=2) + ax3.set_xticks(num_envs) + ax3.set_xticklabels([str(value) for value in num_envs]) + ax3.grid(True, alpha=0.3) + ax3.legend(fontsize=7) + + fig.tight_layout(rect=(0, 0, 1, 0.9)) + summary_path = output_dir / "g1_joystick_numba_summary.png" + fig.savefig(summary_path, dpi=160) + plt.close(fig) + print(f"Saved summary plot: {summary_path.resolve()}") + return [str(summary_path.resolve())] + + +def write_report( + *, + output_dir: Path, + records: list[BenchCase], + parity: dict[str, dict[str, float]], + args: argparse.Namespace, +) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + device_info_line = get_device_info_line() + plot_paths = save_plots(records, output_dir, device_info=device_info_line) + meta = { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "python": platform.python_version(), + "platform": platform.platform(), + "device_info": get_device_info_dict(), + "iters": args.iters, + "warmup": args.warmup, + "profiles": args.profiles, + "num_envs": args.num_envs, + "threads": args.threads, + "scope": "G1 joystick reward+termination hot slice; synthetic backend arrays", + } + payload = { + "meta": meta, + "results": [_case_to_dict(record) for record in records], + "parity": parity, + "plots": plot_paths, + } + json_path = output_dir / "results.json" + json_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + + best_by_case = _best_numba_by_case(records) + + summary_lines = [ + "# G1 joystick Numba benchmark", + "", + "Scope: reward dispatch plus termination for `G1WalkEnv.update_state`, using", + "deterministic synthetic backend arrays. Physics stepping, obs assembly, reset,", + "and policy inference are intentionally out of scope.", + "", + "## Summary", + "", + ] + for key in sorted(best_by_case): + best = best_by_case[key] + summary_lines.append( + f"- `{best.profile}` / {best.num_envs} envs: best {best.mean_ms:.3f} ms " + f"at {best.threads} threads, {best.speedup_vs_numpy:.2f}x vs numpy dispatch " + f"({best.env_per_s / 1e6:.2f}M env/s)." + ) + if plot_paths: + summary_lines.extend(["", "## Plots", ""]) + for path in plot_paths: + rel_path = Path(path).name + title = rel_path.removesuffix(".png").replace("_", " ") + summary_lines.append(f"![{title}]({rel_path})") + summary_lines.append("") + summary_lines.extend( + [ + "", + "## Detailed Results", + "", + "```text", + _format_table(records), + "```", + "", + "## Parity", + "", + "Reward is checked with `rtol=5e-5, atol=5e-5`; termination is exact.", + "", + "```json", + json.dumps(parity, indent=2), + "```", + "", + "## Interpretation", + "", + "- `numba 1 thread` isolates fusion/codegen benefit over Python reward dispatch.", + "- Higher thread counts add row-parallel speedup over the same fused kernel.", + "- End-to-end training speedup will be lower because this benchmark excludes physics,", + " obs assembly, reset/RNG, policy inference, and learner work.", + ] + ) + md_path = output_dir / "report.md" + md_path.write_text("\n".join(summary_lines) + "\n", encoding="utf-8") + print(f"Saved JSON: {json_path.resolve()}") + print(f"Saved report: {md_path.resolve()}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--profiles", + nargs="+", + default=["ppo_default", "sac_default", "full_supported"], + choices=sorted(make_profile_specs()), + ) + parser.add_argument("--num-envs", nargs="+", type=int, default=[512, 2048, 8192]) + parser.add_argument("--threads", nargs="+", type=int, default=[2, 4, 8, 16]) + parser.add_argument("--iters", type=int, default=80) + parser.add_argument("--warmup", type=int, default=8) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("benchmark/outputs/g1_joystick_numba"), + ) + parser.add_argument( + "--quick", + action="store_true", + help="Short smoke run: sac_default at 512 and 2048 envs.", + ) + args = parser.parse_args() + if args.quick: + args.profiles = ["sac_default"] + args.num_envs = [512, 2048] + args.threads = [2, 4] + args.iters = 10 + args.warmup = 2 + return args + + +def main() -> None: + args = parse_args() + specs = make_profile_specs() + all_records: list[BenchCase] = [] + parity: dict[str, dict[str, float]] = {} + + print("=" * 80) + print("G1 joystick Numba benchmark: reward dispatch + termination") + print("=" * 80) + print(f"host numba threads: {get_num_threads()}") + print(f"profiles={args.profiles} num_envs={args.num_envs} threads={args.threads}") + for profile_name in args.profiles: + spec = specs[profile_name] + for num_envs in args.num_envs: + records, parity_result = bench_one( + profile=spec, + num_envs=num_envs, + thread_counts=args.threads, + iters=args.iters, + warmup=args.warmup, + seed=args.seed, + ) + all_records.extend(records) + parity[f"{profile_name}:{num_envs}"] = parity_result + print() + print(_format_table(records)) + + write_report(output_dir=args.output_dir, records=all_records, parity=parity, args=args) + + +if __name__ == "__main__": + main() diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index b66a673c0..3617ec873 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -12,6 +12,7 @@ license-files = ["LICENSE"] requires-python = ">=3.10,<3.14" dependencies = [ "numpy", + "numba", "torch==2.11.0", "triton-rocm==3.6.0 ; sys_platform == 'linux' and platform_machine == 'x86_64'", "gymnasium", diff --git a/pyproject.toml b/pyproject.toml index 10256d1b1..5940c60fc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ license-files = ["LICENSE"] requires-python = ">=3.10,<3.14" dependencies = [ "numpy", + "numba", "torch==2.9.0 ; sys_platform == 'linux' and platform_machine == 'aarch64'", "torch==2.7.0 ; sys_platform != 'linux' or platform_machine != 'aarch64'", "gymnasium", diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 229391da1..90201dbd2 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -207,6 +207,8 @@ class G1WalkEnvCfg(G1BaseCfg): gait_phase_init_mode: str = "offset_phase" reset_base_qvel_limit: float = 0.5 curriculum: CurriculumConfig = field(default_factory=CurriculumConfig) + numba_acceleration: bool = False + numba_num_threads: int | None = None class G1WalkDomainRandomizationProvider(LocomotionDRProvider): @@ -304,6 +306,13 @@ def __init__(self, cfg: G1WalkEnvCfg, num_envs=1, backend_type="mujoco"): ) self._init_reward_functions() + self._numba_accelerator = None + if cfg.numba_acceleration: + from unilab.envs.locomotion.g1.joystick_numba import G1WalkNumbaAccelerator + + self._numba_accelerator = G1WalkNumbaAccelerator.from_env( + self, num_threads=cfg.numba_num_threads + ) if cfg.domain_rand.randomize_kp or cfg.domain_rand.randomize_kd: base_kp, base_kd = backend.get_actuator_gains() dr_provider = G1WalkDomainRandomizationProvider(base_kp=base_kp, base_kd=base_kd) @@ -352,14 +361,30 @@ def update_state(self, state: NpEnvState) -> NpEnvState: dof_pos = self.get_dof_pos() dof_vel = self.get_dof_vel() - max_tilt_rad = np.deg2rad(self._reward_cfg.max_tilt_deg) - tilt = np.arccos(np.clip(gravity[:, 2], -1, 1)) - terminated = np.logical_or( - tilt > max_tilt_rad, - self._terrain_relative_base_height() < self._reward_cfg.min_base_height, - ) + if self._numba_accelerator is not None: + accel_result = self._numba_accelerator.compute( + env=self, + info=state.info, + linvel=linvel, + gyro=gyro, + gravity=gravity, + dof_pos=dof_pos, + dof_vel=dof_vel, + scales=self._reward_cfg.scales, + enable_log=self._enable_reward_log, + ) + terminated = accel_result.terminated + reward = accel_result.reward + state.info["log"] = accel_result.log + else: + max_tilt_rad = np.deg2rad(self._reward_cfg.max_tilt_deg) + tilt = np.arccos(np.clip(gravity[:, 2], -1, 1)) + terminated = np.logical_or( + tilt > max_tilt_rad, + self._terrain_relative_base_height() < self._reward_cfg.min_base_height, + ) + reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel) state = state.replace(obs=obs, reward=reward, terminated=terminated) diff --git a/src/unilab/envs/locomotion/g1/joystick_numba.py b/src/unilab/envs/locomotion/g1/joystick_numba.py new file mode 100644 index 000000000..e9eec25e5 --- /dev/null +++ b/src/unilab/envs/locomotion/g1/joystick_numba.py @@ -0,0 +1,656 @@ +"""Optional Numba hot path for the G1 joystick locomotion task. + +This module is task-owned on purpose: it mirrors the reward/termination math in +``joystick.py`` without changing the base env contract or the shared locomotion +reward dispatcher. Importing this file must be cheap and safe when ``numba`` is +not installed. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Mapping + +import numpy as np + +try: # pragma: no cover - exercised in environments with numba installed + from numba import get_num_threads, get_thread_id, njit, prange, set_num_threads + + NUMBA_AVAILABLE = True +except Exception: # pragma: no cover - default test env may not install numba + get_num_threads = get_thread_id = njit = prange = set_num_threads = None # type: ignore[assignment] + NUMBA_AVAILABLE = False + + +TERM_ORDER: tuple[str, ...] = ( + "tracking_lin_vel", + "tracking_ang_vel", + "forward_progress", + "under_speed", + "lin_vel_z", + "orientation", + "penalty_orientation", + "ang_vel_xy", + "penalty_ang_vel_xy", + "action_rate", + "penalty_action_rate", + "base_height", + "pose", + "upper_body_pose", + "penalty_close_feet_xy", + "penalty_feet_ori", + "feet_phase", + "feet_phase_contrast", + "feet_phase_contact", + "feet_double_stance", + "feet_air_time", + "alive", +) +TERM_INDEX = {name: i for i, name in enumerate(TERM_ORDER)} +SUPPORTED_TERMS = frozenset(TERM_ORDER) + +FOOT_POSITION_TERMS = frozenset(("penalty_close_feet_xy", "feet_phase", "feet_phase_contrast")) +FOOT_QUAT_TERMS = frozenset(("penalty_feet_ori",)) +FOOT_CONTACT_TERMS = frozenset(("feet_phase_contact", "feet_double_stance")) +FEET_AIR_TIME_TERMS = frozenset(("feet_air_time",)) + + +@dataclass(frozen=True) +class G1WalkNumbaResult: + reward: np.ndarray + terminated: np.ndarray + log: dict[str, float] + + +def _active_terms(scales: Mapping[str, float]) -> frozenset[str]: + return frozenset(name for name, scale in scales.items() if scale != 0.0) + + +def unsupported_terms(scales: Mapping[str, float]) -> frozenset[str]: + """Return nonzero reward terms this task-specific kernel cannot compute.""" + return _active_terms(scales) - SUPPORTED_TERMS + + +def is_available(scales: Mapping[str, float]) -> bool: + return NUMBA_AVAILABLE and not unsupported_terms(scales) + + +def _scalar_sensor(sensor_values: np.ndarray) -> np.ndarray: + arr = np.asarray(sensor_values) + if arr.ndim == 1: + return arr + if arr.ndim == 2 and arr.shape[1] == 1: + return arr[:, 0] + raise ValueError(f"Expected scalar sensor values, got shape {arr.shape}") + + +def _aggregated_contact(backend: Any, sensor_names: tuple[str, ...]) -> np.ndarray: + contacts = [_scalar_sensor(backend.get_sensor_data(name)) for name in sensor_names] + return np.asarray(np.any(np.stack(contacts, axis=1) > 0.5, axis=1), dtype=np.bool_) + + +if NUMBA_AVAILABLE: + + def _dev(fn): + return njit(inline="always", fastmath=True, cache=True, nogil=True)(fn) + + @_dev + def _positive(x): + return x if x > 0.0 else 0.0 + + @_dev + def tracking_lin_vel_i(linvel, commands, tracking_sigma, i): + dx = commands[i, 0] - linvel[i, 0] + dy = commands[i, 1] - linvel[i, 1] + return math.exp(-(dx * dx + dy * dy) / tracking_sigma) + + @_dev + def tracking_ang_vel_i(gyro, commands, tracking_sigma, i): + dz = commands[i, 2] - gyro[i, 2] + return math.exp(-(dz * dz) / tracking_sigma) + + @_dev + def forward_progress_i(linvel, commands, i): + commanded_speed = commands[i, 0] if commands[i, 0] > 1.0e-6 else 1.0e-6 + progress = _positive(linvel[i, 0]) / commanded_speed + return 1.0 if progress > 1.0 else progress + + @_dev + def under_speed_i(linvel, commands, i): + commanded_speed = commands[i, 0] if commands[i, 0] > 1.0e-6 else 1.0e-6 + gap = commands[i, 0] - _positive(linvel[i, 0]) + return _positive(gap) / commanded_speed + + @_dev + def lin_vel_z_i(linvel, i): + return linvel[i, 2] * linvel[i, 2] + + @_dev + def orientation_i(gravity, i): + return gravity[i, 0] * gravity[i, 0] + gravity[i, 1] * gravity[i, 1] + + @_dev + def ang_vel_xy_i(gyro, i): + return gyro[i, 0] * gyro[i, 0] + gyro[i, 1] * gyro[i, 1] + + @_dev + def action_rate_i(current_actions, last_actions, n_action, i): + acc = 0.0 + for j in range(n_action): + d = current_actions[i, j] - last_actions[i, j] + acc += d * d + return acc + + @_dev + def weighted_pose_i(dof_pos, default_angles, weights, n_action, i): + acc = 0.0 + for j in range(n_action): + d = dof_pos[i, j] - default_angles[j] + acc += weights[j] * d * d + return acc + + @_dev + def base_height_i(base_height, base_height_target, i): + d = base_height[i] - base_height_target + return d * d + + @_dev + def close_feet_xy_i(left_foot_pos, right_foot_pos, close_feet_threshold, i): + dx = left_foot_pos[i, 0] - right_foot_pos[i, 0] + dy = left_foot_pos[i, 1] - right_foot_pos[i, 1] + feet_dist = math.sqrt(dx * dx + dy * dy) + if feet_dist >= close_feet_threshold: + return 0.0 + d = feet_dist - close_feet_threshold + return d * d + + @_dev + def feet_ori_i(left_foot_quat, right_foot_quat, i): + return ( + left_foot_quat[i, 1] * left_foot_quat[i, 1] + + left_foot_quat[i, 2] * left_foot_quat[i, 2] + + right_foot_quat[i, 1] * right_foot_quat[i, 1] + + right_foot_quat[i, 2] * right_foot_quat[i, 2] + ) + + @_dev + def _bezier_height(phi, swing_height): + phi_normalized = (phi + math.pi) % (2.0 * math.pi) - math.pi + x = (phi_normalized + math.pi) / (2.0 * math.pi) + if x <= 0.5: + t = 2.0 * x + bezier = t * t * t + 3.0 * (t * t * (1.0 - t)) + return swing_height * bezier + t = 2.0 * x - 1.0 + bezier = t * t * t + 3.0 * (t * t * (1.0 - t)) + return swing_height * (1.0 - bezier) + + @_dev + def _gait_gate(linvel, min_forward_speed_for_gait_reward, i): + return 1.0 if _positive(linvel[i, 0]) >= min_forward_speed_for_gait_reward else 0.0 + + @_dev + def feet_phase_i( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + swing_height, + tracking_sigma, + min_forward_speed_for_gait_reward, + i, + ): + left_target = _bezier_height(gait_phase[i, 0], swing_height) + right_target = _bezier_height(gait_phase[i, 1], swing_height) + left_err = left_foot_pos[i, 2] - left_target + right_err = right_foot_pos[i, 2] - right_target + reward = math.exp(-((left_err * left_err) + (right_err * right_err)) / tracking_sigma) + return reward * _gait_gate(linvel, min_forward_speed_for_gait_reward, i) + + @_dev + def feet_phase_contrast_i( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + swing_height, + tracking_sigma, + min_forward_speed_for_gait_reward, + i, + ): + left_target = _bezier_height(gait_phase[i, 0], swing_height) + right_target = _bezier_height(gait_phase[i, 1], swing_height) + actual_delta = left_foot_pos[i, 2] - right_foot_pos[i, 2] + target_delta = left_target - right_target + err = actual_delta - target_delta + reward = math.exp(-(err * err) / tracking_sigma) + return reward * _gait_gate(linvel, min_forward_speed_for_gait_reward, i) + + @_dev + def feet_phase_contact_i( + linvel, + gait_phase, + left_contact, + right_contact, + swing_height, + min_forward_speed_for_gait_reward, + i, + ): + contact_height_threshold = swing_height * 0.5 + left_target_contact = ( + _bezier_height(gait_phase[i, 0], swing_height) <= contact_height_threshold + ) + right_target_contact = ( + _bezier_height(gait_phase[i, 1], swing_height) <= contact_height_threshold + ) + left_match = 1.0 if left_contact[i] == left_target_contact else 0.0 + right_match = 1.0 if right_contact[i] == right_target_contact else 0.0 + return ( + 0.5 + * (left_match + right_match) + * _gait_gate(linvel, min_forward_speed_for_gait_reward, i) + ) + + @_dev + def feet_double_stance_i(commands, left_contact, right_contact, i): + forward_command = 1.0 if _positive(commands[i, 0]) > 1.0e-6 else 0.0 + double_stance = 1.0 if left_contact[i] and right_contact[i] else 0.0 + return double_stance * forward_command + + @_dev + def feet_air_time_i(feet_air_time, i): + acc = 0.0 + if feet_air_time[i, 0] > 0.05 and feet_air_time[i, 0] < 0.5: + acc += 1.0 + if feet_air_time[i, 1] > 0.05 and feet_air_time[i, 1] < 0.5: + acc += 1.0 + return acc + + @_dev + def terminated_i(gravity, base_height, max_tilt_rad, min_base_height, i): + gz = gravity[i, 2] + if gz < -1.0: + gz = -1.0 + elif gz > 1.0: + gz = 1.0 + return math.acos(gz) > max_tilt_rad or base_height[i] < min_base_height + + @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _compute_reward_termination_kernel( + linvel, + gyro, + gravity, + dof_pos, + dof_vel, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + default_angles, + pose_weights, + upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + scale, + ctrl_dt, + tracking_sigma, + base_height_target, + min_base_height, + max_tilt_rad, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + close_feet_threshold, + reward, + terminated, + log_scratch, + ): + n = reward.shape[0] + n_action = dof_pos.shape[1] + for i in prange(n): + tid = get_thread_id() + r = 0.0 + + w = tracking_lin_vel_i(linvel, commands, tracking_sigma, i) * scale[0] + r += w + log_scratch[tid, 0] += w + + w = tracking_ang_vel_i(gyro, commands, tracking_sigma, i) * scale[1] + r += w + log_scratch[tid, 1] += w + + w = forward_progress_i(linvel, commands, i) * scale[2] + r += w + log_scratch[tid, 2] += w + + w = under_speed_i(linvel, commands, i) * scale[3] + r += w + log_scratch[tid, 3] += w + + w = lin_vel_z_i(linvel, i) * scale[4] + r += w + log_scratch[tid, 4] += w + + orientation = orientation_i(gravity, i) + w = orientation * scale[5] + r += w + log_scratch[tid, 5] += w + w = orientation * scale[6] + r += w + log_scratch[tid, 6] += w + + ang_vel_xy = ang_vel_xy_i(gyro, i) + w = ang_vel_xy * scale[7] + r += w + log_scratch[tid, 7] += w + w = ang_vel_xy * scale[8] + r += w + log_scratch[tid, 8] += w + + action_rate = action_rate_i(current_actions, last_actions, n_action, i) + w = action_rate * scale[9] + r += w + log_scratch[tid, 9] += w + w = action_rate * scale[10] + r += w + log_scratch[tid, 10] += w + + w = base_height_i(base_height, base_height_target, i) * scale[11] + r += w + log_scratch[tid, 11] += w + + pose = weighted_pose_i(dof_pos, default_angles, pose_weights, n_action, i) + w = pose * scale[12] + r += w + log_scratch[tid, 12] += w + upper_body_pose = weighted_pose_i( + dof_pos, default_angles, upper_body_pose_weights, n_action, i + ) + w = upper_body_pose * scale[13] + r += w + log_scratch[tid, 13] += w + + w = close_feet_xy_i(left_foot_pos, right_foot_pos, close_feet_threshold, i) * scale[14] + r += w + log_scratch[tid, 14] += w + + w = feet_ori_i(left_foot_quat, right_foot_quat, i) * scale[15] + r += w + log_scratch[tid, 15] += w + + w = ( + feet_phase_i( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + i, + ) + * scale[16] + ) + r += w + log_scratch[tid, 16] += w + + w = ( + feet_phase_contrast_i( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + i, + ) + * scale[17] + ) + r += w + log_scratch[tid, 17] += w + + w = ( + feet_phase_contact_i( + linvel, + gait_phase, + left_contact, + right_contact, + feet_phase_swing_height, + min_forward_speed_for_gait_reward, + i, + ) + * scale[18] + ) + r += w + log_scratch[tid, 18] += w + + w = feet_double_stance_i(commands, left_contact, right_contact, i) * scale[19] + r += w + log_scratch[tid, 19] += w + + w = feet_air_time_i(feet_air_time, i) * scale[20] + r += w + log_scratch[tid, 20] += w + + w = scale[21] + r += w + log_scratch[tid, 21] += w + + reward[i] = r * ctrl_dt + terminated[i] = terminated_i(gravity, base_height, max_tilt_rad, min_base_height, i) + + +class G1WalkNumbaAccelerator: + """Driver that keeps config-derived arrays and calls the fused kernel.""" + + def __init__( + self, + *, + num_envs: int, + num_action: int, + ctrl_dt: float, + tracking_sigma: float, + base_height_target: float, + min_base_height: float, + max_tilt_deg: float, + feet_phase_swing_height: float, + feet_phase_tracking_sigma: float, + min_forward_speed_for_gait_reward: float, + close_feet_threshold: float, + default_angles: np.ndarray, + pose_weights: np.ndarray, + upper_body_pose_weights: np.ndarray, + num_threads: int | None = None, + ) -> None: + self.num_envs = int(num_envs) + self.num_action = int(num_action) + self.ctrl_dt = float(ctrl_dt) + self.tracking_sigma = float(tracking_sigma) + self.base_height_target = float(base_height_target) + self.min_base_height = float(min_base_height) + self.max_tilt_rad = float(np.deg2rad(max_tilt_deg)) + self.feet_phase_swing_height = float(feet_phase_swing_height) + self.feet_phase_tracking_sigma = float(feet_phase_tracking_sigma) + self.min_forward_speed_for_gait_reward = float(min_forward_speed_for_gait_reward) + self.close_feet_threshold = float(close_feet_threshold) + self.default_angles = np.asarray(default_angles, dtype=np.float64) + self.pose_weights = np.asarray(pose_weights, dtype=np.float64) + self.upper_body_pose_weights = np.asarray(upper_body_pose_weights, dtype=np.float64) + self.num_threads = num_threads + self.scale = np.zeros((len(TERM_ORDER),), dtype=np.float64) + self._zero_vec2 = np.zeros((self.num_envs, 2), dtype=np.float64) + self._zero_vec3 = np.zeros((self.num_envs, 3), dtype=np.float64) + self._zero_vec4 = np.zeros((self.num_envs, 4), dtype=np.float64) + self._zero_bool = np.zeros((self.num_envs,), dtype=np.bool_) + + @classmethod + def from_env(cls, env: Any, num_threads: int | None = None) -> "G1WalkNumbaAccelerator": + if not NUMBA_AVAILABLE: + raise RuntimeError( + "G1Walk numba_acceleration=True requires numba. Install it or run through " + "`uv run --with numba ...`; disable numba_acceleration to use the numpy path." + ) + return cls( + num_envs=env.num_envs, + num_action=env._num_action, + ctrl_dt=env._cfg.ctrl_dt, + tracking_sigma=env._reward_cfg.tracking_sigma, + base_height_target=env._reward_cfg.base_height_target, + min_base_height=env._reward_cfg.min_base_height, + max_tilt_deg=env._reward_cfg.max_tilt_deg, + feet_phase_swing_height=env._reward_cfg.feet_phase_swing_height, + feet_phase_tracking_sigma=env._reward_cfg.feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward=getattr( + env._reward_cfg, "min_forward_speed_for_gait_reward", 0.0 + ), + close_feet_threshold=getattr(env._reward_cfg, "close_feet_threshold", 0.15), + default_angles=env.default_angles, + pose_weights=env._pose_weights, + upper_body_pose_weights=env._upper_body_pose_weights, + num_threads=num_threads, + ) + + def _sync_scales(self, scales: Mapping[str, float]) -> None: + unsupported = unsupported_terms(scales) + if unsupported: + raise ValueError( + "G1Walk Numba accelerator does not support active reward terms " + f"{sorted(unsupported)}. Disable numba_acceleration or add these terms " + "to src/unilab/envs/locomotion/g1/joystick_numba.py." + ) + self.scale.fill(0.0) + for name, value in scales.items(): + idx = TERM_INDEX.get(name) + if idx is not None: + self.scale[idx] = float(value) + + def compute( + self, + *, + env: Any, + info: dict[str, Any], + linvel: np.ndarray, + gyro: np.ndarray, + gravity: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + scales: Mapping[str, float], + enable_log: bool, + ) -> G1WalkNumbaResult: + if not NUMBA_AVAILABLE: + raise RuntimeError( + "G1Walk Numba accelerator was constructed while numba is unavailable; " + "this indicates an invalid accelerator state." + ) + self._sync_scales(scales) + + active = _active_terms(scales) + backend = env._backend + dtype = linvel.dtype + base_height = np.asarray(backend.get_base_pos()[:, 2], dtype=dtype) + commands = np.asarray(info["commands"], dtype=dtype) + current_actions = np.asarray( + info.get("current_actions", np.zeros((self.num_envs, self.num_action), dtype=dtype)), + dtype=dtype, + ) + last_actions = np.asarray( + info.get("last_actions", np.zeros((self.num_envs, self.num_action), dtype=dtype)), + dtype=dtype, + ) + gait_phase = np.asarray(info.get("gait_phase", self._zero_vec2), dtype=dtype) + feet_air_time = np.asarray(info.get("feet_air_time", self._zero_vec2), dtype=dtype) + + if active & FOOT_POSITION_TERMS: + left_foot_pos = np.asarray(backend.get_sensor_data("left_foot_pos"), dtype=dtype) + right_foot_pos = np.asarray(backend.get_sensor_data("right_foot_pos"), dtype=dtype) + else: + left_foot_pos = right_foot_pos = self._zero_vec3 + + if active & FOOT_QUAT_TERMS: + left_foot_quat = np.asarray(backend.get_sensor_data("left_foot_quat"), dtype=dtype) + right_foot_quat = np.asarray(backend.get_sensor_data("right_foot_quat"), dtype=dtype) + else: + left_foot_quat = right_foot_quat = self._zero_vec4 + + if active & FOOT_CONTACT_TERMS: + left_contact = _aggregated_contact( + backend, + ( + "left_foot_contact_0", + "left_foot_contact_1", + "left_foot_contact_2", + "left_foot_contact_3", + ), + ) + right_contact = _aggregated_contact( + backend, + ( + "right_foot_contact_0", + "right_foot_contact_1", + "right_foot_contact_2", + "right_foot_contact_3", + ), + ) + else: + left_contact = right_contact = self._zero_bool + + if self.num_threads is not None: + set_num_threads(self.num_threads) + nthreads = get_num_threads() + reward = np.empty((linvel.shape[0],), dtype=dtype) + terminated = np.empty((linvel.shape[0],), dtype=np.bool_) + log_scratch = np.zeros((nthreads, len(TERM_ORDER)), dtype=np.float64) + + _compute_reward_termination_kernel( + linvel, + gyro, + gravity, + dof_pos, + dof_vel, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + self.default_angles, + self.pose_weights, + self.upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + self.scale, + self.ctrl_dt, + self.tracking_sigma, + self.base_height_target, + self.min_base_height, + self.max_tilt_rad, + self.feet_phase_swing_height, + self.feet_phase_tracking_sigma, + self.min_forward_speed_for_gait_reward, + self.close_feet_threshold, + reward, + terminated, + log_scratch, + ) + + step_count = info.get("steps", np.zeros((linvel.shape[0],), dtype=np.uint32)) + should_log = enable_log and int(step_count[0]) % 4 == 0 + log = {} if should_log else info.get("log", {}) + if should_log: + term_sums = log_scratch.sum(axis=0) + for idx, name in enumerate(TERM_ORDER): + if self.scale[idx] != 0.0: + log[f"reward/{name}"] = float(term_sums[idx] / linvel.shape[0]) + return G1WalkNumbaResult(reward=reward, terminated=terminated, log=log) diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py new file mode 100644 index 000000000..a6a3bc409 --- /dev/null +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +from benchmark import benchmark_g1_joystick_numba as bench + + +def test_g1_joystick_numba_benchmark_builds_records_and_matches_numpy() -> None: + spec = bench.make_profile_specs()["sac_default"] + + records, parity = bench.bench_one( + profile=spec, + num_envs=64, + thread_counts=[1], + iters=1, + warmup=0, + seed=0, + ) + + assert parity["termination_mismatch"] == 0.0 + assert parity["max_abs_reward_diff"] < 1.0e-5 + assert {record.path for record in records} == {"numpy_dispatch", "numba_accelerator"} + assert any(record.path == "numba_accelerator" and record.threads == 1 for record in records) diff --git a/tests/envs/locomotion/g1/test_joystick_numba.py b/tests/envs/locomotion/g1/test_joystick_numba.py new file mode 100644 index 000000000..dc3402303 --- /dev/null +++ b/tests/envs/locomotion/g1/test_joystick_numba.py @@ -0,0 +1,296 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from unilab.envs.locomotion.g1.joystick_numba import ( + NUMBA_AVAILABLE, + G1WalkNumbaAccelerator, + unsupported_terms, +) + + +def test_unsupported_terms_only_reports_active_unknown_terms(): + assert unsupported_terms({"tracking_lin_vel": 1.0, "custom": 0.0}) == frozenset() + assert unsupported_terms({"tracking_lin_vel": 1.0, "custom": 1.0}) == frozenset({"custom"}) + + +def test_numba_accelerator_requires_numba_when_declared(monkeypatch): + import unilab.envs.locomotion.g1.joystick_numba as joystick_numba + + monkeypatch.setattr(joystick_numba, "NUMBA_AVAILABLE", False) + with pytest.raises(RuntimeError, match="requires numba"): + G1WalkNumbaAccelerator.from_env(_Env(8, 29)) + + +class _Backend: + def __init__(self, n: int): + self.base = np.full((n, 3), [0.0, 0.0, 0.754], dtype=np.float32) + self.sensors = { + "left_foot_pos": np.column_stack( + [np.full(n, -0.1), np.zeros(n), np.full(n, 0.02)] + ).astype(np.float32), + "right_foot_pos": np.column_stack( + [np.full(n, 0.1), np.zeros(n), np.full(n, 0.02)] + ).astype(np.float32), + "left_foot_quat": np.tile(np.array([1.0, 0.01, 0.02, 0.0], dtype=np.float32), (n, 1)), + "right_foot_quat": np.tile(np.array([1.0, 0.02, 0.01, 0.0], dtype=np.float32), (n, 1)), + } + for side in ("left", "right"): + for idx in range(4): + self.sensors[f"{side}_foot_contact_{idx}"] = np.ones((n,), dtype=np.float32) + + def get_base_pos(self): + return self.base + + def get_sensor_data(self, name: str): + return self.sensors[name] + + +class _Cfg: + ctrl_dt = 0.02 + + +class _RewardCfg: + tracking_sigma = 0.25 + base_height_target = 0.754 + min_base_height = 0.3 + max_tilt_deg = 65.0 + feet_phase_swing_height = 0.09 + feet_phase_tracking_sigma = 0.04 + min_forward_speed_for_gait_reward = 0.0 + close_feet_threshold = 0.15 + + +class _Env: + def __init__(self, n: int, n_action: int): + self.num_envs = n + self._num_action = n_action + self._cfg = _Cfg() + self._reward_cfg = _RewardCfg() + self.default_angles = np.zeros((n_action,), dtype=np.float32) + self._pose_weights = np.ones((n_action,), dtype=np.float32) + self._upper_body_pose_weights = np.ones((n_action,), dtype=np.float32) + self._backend = _Backend(n) + + +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_walk_numba_unsupported_active_term_raises(): + n = 8 + n_action = 29 + env = _Env(n, n_action) + accel = G1WalkNumbaAccelerator.from_env(env, num_threads=1) + info = { + "steps": np.zeros((n,), dtype=np.uint32), + "commands": np.zeros((n, 3), dtype=np.float32), + } + + with pytest.raises(ValueError, match="does not support active reward terms"): + accel.compute( + env=env, + info=info, + linvel=np.zeros((n, 3), dtype=np.float32), + gyro=np.zeros((n, 3), dtype=np.float32), + gravity=np.tile(np.array([0.0, 0.0, 1.0], dtype=np.float32), (n, 1)), + dof_pos=np.zeros((n, n_action), dtype=np.float32), + dof_vel=np.zeros((n, n_action), dtype=np.float32), + scales={"tracking_lin_vel": 1.0, "custom": 1.0}, + enable_log=False, + ) + + +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_walk_numba_basic_reward_parity(): + n = 512 + n_action = 29 + rng = np.random.default_rng(0) + env = _Env(n, n_action) + info = { + "steps": np.zeros((n,), dtype=np.uint32), + "commands": rng.normal(size=(n, 3)).astype(np.float32), + "current_actions": rng.normal(size=(n, n_action)).astype(np.float32), + "last_actions": rng.normal(size=(n, n_action)).astype(np.float32), + "gait_phase": rng.uniform(0.0, 2 * np.pi, size=(n, 2)).astype(np.float32), + } + linvel = rng.normal(size=(n, 3)).astype(np.float32) + gyro = rng.normal(size=(n, 3)).astype(np.float32) + gravity = np.tile(np.array([0.01, -0.02, 0.99], dtype=np.float32), (n, 1)) + dof_pos = rng.normal(scale=0.01, size=(n, n_action)).astype(np.float32) + dof_vel = rng.normal(size=(n, n_action)).astype(np.float32) + scales = { + "tracking_lin_vel": 2.0, + "tracking_ang_vel": 1.5, + "penalty_orientation": -10.0, + "penalty_action_rate": -4.0, + "pose": -0.5, + "alive": 10.0, + } + + accel = G1WalkNumbaAccelerator.from_env(env, num_threads=2) + assert accel is not None + out = accel.compute( + env=env, + info=info, + linvel=linvel, + gyro=gyro, + gravity=gravity, + dof_pos=dof_pos, + dof_vel=dof_vel, + scales=scales, + enable_log=True, + ) + assert out is not None + + reward = np.zeros((n,), dtype=np.float32) + err = np.sum((info["commands"][:, :2] - linvel[:, :2]) ** 2, axis=1) + reward += np.exp(-err / _RewardCfg.tracking_sigma) * scales["tracking_lin_vel"] + err = (info["commands"][:, 2] - gyro[:, 2]) ** 2 + reward += np.exp(-err / _RewardCfg.tracking_sigma) * scales["tracking_ang_vel"] + reward += (gravity[:, 0] ** 2 + gravity[:, 1] ** 2) * scales["penalty_orientation"] + reward += ( + np.sum((info["current_actions"] - info["last_actions"]) ** 2, axis=1) + * scales["penalty_action_rate"] + ) + reward += np.sum((dof_pos - env.default_angles) ** 2, axis=1) * scales["pose"] + reward += np.ones((n,), dtype=np.float32) * scales["alive"] + reward *= _Cfg.ctrl_dt + + terminated = ( + np.arccos(np.clip(gravity[:, 2], -1.0, 1.0)) > np.deg2rad(_RewardCfg.max_tilt_deg) + ) | (env._backend.get_base_pos()[:, 2] < _RewardCfg.min_base_height) + + np.testing.assert_allclose(out.reward, reward, rtol=2e-5, atol=2e-5) + np.testing.assert_array_equal(out.terminated, terminated) + assert set(out.log) == {f"reward/{name}" for name in scales} + + +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_walk_numba_term_py_funcs_match_numpy_math(): + from unilab.envs.locomotion.g1 import joystick_numba as T + + n = 3 + n_action = 29 + rng = np.random.default_rng(1) + linvel = rng.normal(size=(n, 3)).astype(np.float32) + gyro = rng.normal(size=(n, 3)).astype(np.float32) + gravity = rng.normal(size=(n, 3)).astype(np.float32) + commands = rng.normal(size=(n, 3)).astype(np.float32) + dof_pos = rng.normal(size=(n, n_action)).astype(np.float32) + current_actions = rng.normal(size=(n, n_action)).astype(np.float32) + last_actions = rng.normal(size=(n, n_action)).astype(np.float32) + default_angles = rng.normal(size=(n_action,)).astype(np.float32) + weights = rng.uniform(0.1, 2.0, size=(n_action,)).astype(np.float32) + base_height = rng.uniform(0.4, 0.9, size=(n,)).astype(np.float32) + gait_phase = rng.uniform(0.0, 2 * np.pi, size=(n, 2)).astype(np.float32) + left_foot_pos = rng.normal(size=(n, 3)).astype(np.float32) + right_foot_pos = rng.normal(size=(n, 3)).astype(np.float32) + left_foot_quat = rng.normal(size=(n, 4)).astype(np.float32) + right_foot_quat = rng.normal(size=(n, 4)).astype(np.float32) + left_contact = np.array([True, False, True]) + right_contact = np.array([True, True, False]) + feet_air_time = rng.uniform(0.0, 0.8, size=(n, 2)).astype(np.float32) + + i = 1 + tracking_sigma = 0.25 + base_height_target = 0.754 + swing_height = 0.09 + feet_sigma = 0.04 + min_forward_speed = 0.0 + close_feet_threshold = 0.15 + + np.testing.assert_allclose( + T.tracking_lin_vel_i.py_func(linvel, commands, tracking_sigma, i), + np.exp(-np.sum((commands[i, :2] - linvel[i, :2]) ** 2) / tracking_sigma), + ) + np.testing.assert_allclose( + T.tracking_ang_vel_i.py_func(gyro, commands, tracking_sigma, i), + np.exp(-((commands[i, 2] - gyro[i, 2]) ** 2) / tracking_sigma), + ) + commanded_speed = max(commands[i, 0], 1.0e-6) + forward_speed = max(linvel[i, 0], 0.0) + assert T.forward_progress_i.py_func(linvel, commands, i) == pytest.approx( + min(forward_speed / commanded_speed, 1.0) + ) + assert T.under_speed_i.py_func(linvel, commands, i) == pytest.approx( + max(commands[i, 0] - forward_speed, 0.0) / commanded_speed + ) + assert T.lin_vel_z_i.py_func(linvel, i) == pytest.approx(linvel[i, 2] ** 2) + assert T.orientation_i.py_func(gravity, i) == pytest.approx( + gravity[i, 0] ** 2 + gravity[i, 1] ** 2 + ) + assert T.ang_vel_xy_i.py_func(gyro, i) == pytest.approx(gyro[i, 0] ** 2 + gyro[i, 1] ** 2) + np.testing.assert_allclose( + T.action_rate_i.py_func(current_actions, last_actions, n_action, i), + np.sum((current_actions[i] - last_actions[i]) ** 2), + ) + np.testing.assert_allclose( + T.weighted_pose_i.py_func(dof_pos, default_angles, weights, n_action, i), + np.sum(weights * (dof_pos[i] - default_angles) ** 2), + ) + assert T.base_height_i.py_func(base_height, base_height_target, i) == pytest.approx( + (base_height[i] - base_height_target) ** 2 + ) + + feet_dist = np.linalg.norm(left_foot_pos[i, :2] - right_foot_pos[i, :2]) + close_feet = ( + (feet_dist - close_feet_threshold) ** 2 if feet_dist < close_feet_threshold else 0.0 + ) + assert T.close_feet_xy_i.py_func( + left_foot_pos, right_foot_pos, close_feet_threshold, i + ) == pytest.approx(close_feet) + np.testing.assert_allclose( + T.feet_ori_i.py_func(left_foot_quat, right_foot_quat, i), + left_foot_quat[i, 1] ** 2 + + left_foot_quat[i, 2] ** 2 + + right_foot_quat[i, 1] ** 2 + + right_foot_quat[i, 2] ** 2, + ) + assert T.feet_air_time_i.py_func(feet_air_time, i) == pytest.approx( + np.sum((feet_air_time[i] > 0.05) & (feet_air_time[i] < 0.5)) + ) + assert T.terminated_i.py_func(gravity, base_height, np.deg2rad(65.0), 0.3, i) == pytest.approx( + np.arccos(np.clip(gravity[i, 2], -1.0, 1.0)) > np.deg2rad(65.0) or base_height[i] < 0.3 + ) + + # Phase terms share the same scalar Bezier helper; checking they return + # finite values catches drift in the less vector-friendly gait math. + assert np.isfinite( + T.feet_phase_i.py_func( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + swing_height, + feet_sigma, + min_forward_speed, + i, + ) + ) + assert np.isfinite( + T.feet_phase_contrast_i.py_func( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + swing_height, + feet_sigma, + min_forward_speed, + i, + ) + ) + assert np.isfinite( + T.feet_phase_contact_i.py_func( + linvel, + gait_phase, + left_contact, + right_contact, + swing_height, + min_forward_speed, + i, + ) + ) + assert T.feet_double_stance_i.py_func( + commands, left_contact, right_contact, i + ) == pytest.approx( + float(left_contact[i] and right_contact[i]) * float(max(commands[i, 0], 0.0) > 1.0e-6) + ) diff --git a/uv.lock b/uv.lock index 9dc86e5a3..946de3f0a 100644 --- a/uv.lock +++ b/uv.lock @@ -389,7 +389,7 @@ name = "coloredlogs" version = "15.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "humanfriendly" }, + { name = "humanfriendly", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/c7/eed8f27100517e8c0e6b923d5f0845d0cb99763da6fdee00478f91db7325/coloredlogs-15.0.1.tar.gz", hash = "sha256:7c991aa71a4577af2f82600d8f8f3a89f936baeaf9b50a9c197da014e5bf16b0", size = 278520, upload-time = "2021-06-11T10:22:45.202Z" } wheels = [ @@ -430,7 +430,7 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" } wheels = [ @@ -517,7 +517,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } wheels = [ @@ -717,8 +717,8 @@ name = "embreex" version = "2.17.7.post7" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x') or (python_full_version < '3.11' and platform_machine != 's390x' and sys_platform != 'linux')" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x') or (python_full_version >= '3.11' and platform_machine != 's390x' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/0d/63/4aa2cb2b1051f4b8cdf2ca8c67493f872717c72ecfbb2988c2d88c0b3ba1/embreex-2.17.7.post7-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d52ee8e69ccac4fbdc21227158363a52dedd4ccdf4baf6d77ac3cf71d7ca3b77", size = 10739984, upload-time = "2025-10-22T20:07:14.81Z" }, @@ -754,10 +754,10 @@ wheels = [ [package.optional-dependencies] epath = [ - { name = "fsspec" }, - { name = "importlib-resources" }, - { name = "typing-extensions" }, - { name = "zipp" }, + { name = "fsspec", marker = "python_full_version < '3.11'" }, + { name = "importlib-resources", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "zipp", marker = "python_full_version < '3.11'" }, ] [[package]] @@ -791,9 +791,9 @@ wheels = [ [package.optional-dependencies] epath = [ - { name = "fsspec" }, - { name = "typing-extensions" }, - { name = "zipp" }, + { name = "fsspec", marker = "python_full_version >= '3.11'" }, + { name = "typing-extensions", marker = "python_full_version >= '3.11'" }, + { name = "zipp", marker = "python_full_version >= '3.11'" }, ] [[package]] @@ -801,7 +801,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -1121,7 +1121,7 @@ name = "humanfriendly" version = "10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pyreadline3", marker = "sys_platform == 'win32'" }, + { name = "pyreadline3", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cc/3f/2c29224acb2e2df4d2046e4c73ee2662023c58ff5b113c4c1adac0886c43/humanfriendly-10.0.tar.gz", hash = "sha256:6b0b831ce8f15f7300721aa49829fc4e83921a9a301cc7f606be6686a2288ddc", size = 360702, upload-time = "2021-09-17T21:40:43.31Z" } wheels = [ @@ -1248,17 +1248,17 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "decorator" }, - { name = "exceptiongroup" }, - { name = "jedi" }, - { name = "matplotlib-inline" }, - { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit" }, - { name = "pygments" }, - { name = "stack-data" }, - { name = "traitlets" }, - { name = "typing-extensions" }, + { name = "colorama", marker = "python_full_version < '3.11' and sys_platform == 'win32'" }, + { name = "decorator", marker = "python_full_version < '3.11'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "jedi", marker = "python_full_version < '3.11'" }, + { name = "matplotlib-inline", marker = "python_full_version < '3.11'" }, + { name = "pexpect", marker = "python_full_version < '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit", marker = "python_full_version < '3.11'" }, + { name = "pygments", marker = "python_full_version < '3.11'" }, + { name = "stack-data", marker = "python_full_version < '3.11'" }, + { name = "traitlets", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/40/18/f8598d287006885e7136451fdea0755af4ebcbfe342836f24deefaed1164/ipython-8.39.0.tar.gz", hash = "sha256:4110ae96012c379b8b6db898a07e186c40a2a1ef5d57a7fa83166047d9da7624", size = 5513971, upload-time = "2026-03-27T10:02:13.94Z" } wheels = [ @@ -1278,17 +1278,17 @@ resolution-markers = [ "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "decorator" }, - { name = "ipython-pygments-lexers" }, - { name = "jedi" }, - { name = "matplotlib-inline" }, - { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit" }, - { name = "pygments" }, - { name = "stack-data" }, - { name = "traitlets" }, - { name = "typing-extensions" }, + { name = "colorama", marker = "python_full_version == '3.11.*' and sys_platform == 'win32'" }, + { name = "decorator", marker = "python_full_version == '3.11.*'" }, + { name = "ipython-pygments-lexers", marker = "python_full_version == '3.11.*'" }, + { name = "jedi", marker = "python_full_version == '3.11.*'" }, + { name = "matplotlib-inline", marker = "python_full_version == '3.11.*'" }, + { name = "pexpect", marker = "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit", marker = "python_full_version == '3.11.*'" }, + { name = "pygments", marker = "python_full_version == '3.11.*'" }, + { name = "stack-data", marker = "python_full_version == '3.11.*'" }, + { name = "traitlets", marker = "python_full_version == '3.11.*'" }, + { name = "typing-extensions", marker = "python_full_version == '3.11.*'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c5/25/daae0e764047b0a2480c7bbb25d48f4f509b5818636562eeac145d06dfee/ipython-9.10.1.tar.gz", hash = "sha256:e170e9b2a44312484415bdb750492699bf329233b03f2557a9692cce6466ada4", size = 4426663, upload-time = "2026-03-27T09:53:26.244Z" } wheels = [ @@ -1314,16 +1314,16 @@ resolution-markers = [ "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "decorator" }, - { name = "ipython-pygments-lexers" }, - { name = "jedi" }, - { name = "matplotlib-inline" }, - { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "prompt-toolkit" }, - { name = "pygments" }, - { name = "stack-data" }, - { name = "traitlets" }, + { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, + { name = "decorator", marker = "python_full_version >= '3.12'" }, + { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.12'" }, + { name = "jedi", marker = "python_full_version >= '3.12'" }, + { name = "matplotlib-inline", marker = "python_full_version >= '3.12'" }, + { name = "pexpect", marker = "python_full_version >= '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit", marker = "python_full_version >= '3.12'" }, + { name = "pygments", marker = "python_full_version >= '3.12'" }, + { name = "stack-data", marker = "python_full_version >= '3.12'" }, + { name = "traitlets", marker = "python_full_version >= '3.12'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3a/73/7114f80a8f9cabdb13c27732dce24af945b2923dcab80723602f7c8bc2d8/ipython-9.12.0.tar.gz", hash = "sha256:01daa83f504b693ba523b5a407246cabde4eb4513285a3c6acaff11a66735ee4", size = 4428879, upload-time = "2026-03-27T09:42:45.312Z" } wheels = [ @@ -1335,7 +1335,7 @@ name = "ipython-pygments-lexers" version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pygments" }, + { name = "pygments", marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } wheels = [ @@ -1753,6 +1753,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d3/97/68f80ca3ac4924f250cdfa6e20142a803e5e50fca96ef5148c52ee8c10ea/librt-0.8.1-cp313-cp313-win_arm64.whl", hash = "sha256:924817ab3141aca17893386ee13261f1d100d1ef410d70afe4389f2359fea4f0", size = 52495, upload-time = "2026-02-17T16:12:11.633Z" }, ] +[[package]] +name = "llvmlite" +version = "0.48.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/dc/a0/acc8ffcd5bdc63df0097e22c719bfcd61b604358343089313a8aebbb24ab/llvmlite-0.48.0.tar.gz", hash = "sha256:543b19f9ef8f3c7c60d1468191e4ee1b1537bf9f8a3d56f64c0ddd98de92edd2", size = 184016, upload-time = "2026-07-02T20:20:05.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/4e/32543c42568fb321b3bdfcf9106e4116ab8f5a7bbcfd9ecf5569b0c07d83/llvmlite-0.48.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:614aad57df707e3172efd5165f2aa7da6a0c6897e40dce590bf756396815ba76", size = 40480650, upload-time = "2026-07-01T18:41:01.945Z" }, + { url = "https://files.pythonhosted.org/packages/a9/0d/6aa48abd423067139a129d1434b77bbcc56080db51d12a88510bb491ca3d/llvmlite-0.48.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:13532f248960ba888ad5ab8150494e2f3a3d20e5f59f264e63741ea5b0ba844c", size = 59890118, upload-time = "2026-07-01T18:41:10.608Z" }, + { url = "https://files.pythonhosted.org/packages/5a/c7/aa917444d871a79608af49149de1b28764e87d2ab41f933c5cd02431d03d/llvmlite-0.48.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ee0c77685a18f5fca994ae21d0763007fca5c5c64b41de37accc78b69079176", size = 58343459, upload-time = "2026-07-01T18:41:06.21Z" }, + { url = "https://files.pythonhosted.org/packages/c5/2b/ceee1cdc263617109d514ac4d1b31f10a282662740ff7d5777baae25b3b5/llvmlite-0.48.0-cp310-cp310-win_amd64.whl", hash = "sha256:02853fe4214acb3780fc920c3fee10564b61d58a35e1b78afcc8a546c2deaba3", size = 41864734, upload-time = "2026-07-01T18:41:14.746Z" }, + { url = "https://files.pythonhosted.org/packages/9a/55/595981f14fbae9ba966feb12af552b1fe69889e44e64ac883a731ed335e0/llvmlite-0.48.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:56a7e24607d3f02d7b1bae8d29c7e1e423d53143d68b072999777f19678fe77b", size = 40480651, upload-time = "2026-07-01T18:41:18.438Z" }, + { url = "https://files.pythonhosted.org/packages/26/08/0109d1b9cb3f4603f3890e30bc66c65332b79185f12a045343b2ae431f67/llvmlite-0.48.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6fa532d6bb3fd3f0803567c736401c54aecfe1a396d3ad25d2440d220e09f0e7", size = 59890118, upload-time = "2026-07-01T18:41:28.184Z" }, + { url = "https://files.pythonhosted.org/packages/02/eb/c5281be180c789cdffbf45b671884c57d7e61345ef3b0f643a4965e108e8/llvmlite-0.48.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:979a66a3f28a02565383ff463527dce78e9b856298872a361283132488e83591", size = 58343458, upload-time = "2026-07-01T18:41:23.397Z" }, + { url = "https://files.pythonhosted.org/packages/aa/f7/b3222b13f2d424dae3c9e63fde476af25ebccf1f3faf0b52d1b79fc15c70/llvmlite-0.48.0-cp311-cp311-win_amd64.whl", hash = "sha256:efaee0276e5e17c2b99b92e0c974bd484ef5977cf5dbc9168e82b71578edb47f", size = 41864734, upload-time = "2026-07-01T18:41:31.932Z" }, + { url = "https://files.pythonhosted.org/packages/92/a2/28696a9e61e245d1a79816d29d106692a90a2b6e7d78c98b326db70827af/llvmlite-0.48.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:d66c3beb4209087ddd4cf4ed2a0856b6887e6a913bdcf1aacfec9851cf2cba4e", size = 40480651, upload-time = "2026-07-01T18:41:35.694Z" }, + { url = "https://files.pythonhosted.org/packages/80/f2/72409351db66d0a317ec5087e076f31fb7b773a640db8a90ce6b5cac9edd/llvmlite-0.48.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:416fa4c2c66c2c6dc6d0a402648c19206e548efa0aa1eff01ad5cdad0af8217d", size = 59890118, upload-time = "2026-07-01T18:41:44.886Z" }, + { url = "https://files.pythonhosted.org/packages/3a/27/5ae2f3722606360480707adb47f001ad89df8251d06b14ee80336e660b66/llvmlite-0.48.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f5e5a5131045b72345c71062ea1a91910dde913792b6c9b28ebb2c1c0a712e98", size = 58343459, upload-time = "2026-07-01T18:41:40.306Z" }, + { url = "https://files.pythonhosted.org/packages/16/78/d824ffff7521cd140dc2006e44ce2bc82e64b48d1b32e90e956308c85a74/llvmlite-0.48.0-cp312-cp312-win_amd64.whl", hash = "sha256:d45c7541a80934ec6d8ab0defe67439494ecd2193cbf852a44ba827808976ac1", size = 41865022, upload-time = "2026-07-01T18:41:48.663Z" }, + { url = "https://files.pythonhosted.org/packages/9c/23/fe9316d14626b42c73ef0b502e724705a6ee9450afe53759c0a99c37c2d7/llvmlite-0.48.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:a83a99ef0c05b4ccddf9b6218ed9fe84b653a0caf7c1d9dbe148d6d16c67f518", size = 40480652, upload-time = "2026-07-01T18:41:52.216Z" }, + { url = "https://files.pythonhosted.org/packages/1b/4a/90715fa12006d681270b08d881195b6fab3ec39572e048764a1f7f59fed7/llvmlite-0.48.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8761b9e522f55207e24424fcd98370289eec2710bf8e915c82d1053f642450dc", size = 59890120, upload-time = "2026-07-01T18:42:00.748Z" }, + { url = "https://files.pythonhosted.org/packages/70/5e/7b3e20d64650ca3c80af0cdb664ec4b575ec83d9d4dd05bea8bd31f9bbb6/llvmlite-0.48.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2fe5cb59b2063bfa039dcb8ca6481c0181bf552f340d10dcf61d7996a665556e", size = 58343457, upload-time = "2026-07-01T18:41:56.41Z" }, + { url = "https://files.pythonhosted.org/packages/17/97/5a430055d1838cf1fb7a01cfa943300f5e4c026fc6333a522c5e4a03b0c1/llvmlite-0.48.0-cp313-cp313-win_amd64.whl", hash = "sha256:91c7e24e74cde3f02b88aa5acca678373f9e069f3b98531b3dbb3a142d9d10bb", size = 41865022, upload-time = "2026-07-01T18:42:04.57Z" }, +] + [[package]] name = "lxml" version = "6.0.4" @@ -2155,7 +2179,7 @@ name = "mlx" version = "0.31.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mlx-metal" }, + { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/9b/f9/f1663dafd45af02467f4f41777c13ec34b9104b2b0450d870c3f906285cd/mlx-0.31.1-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:bc46c911cc060d2eaf21b9e24a1712dc56763b660b53631b9057a32ab1c0271a", size = 574137, upload-time = "2026-03-12T02:15:54.996Z" }, @@ -2514,6 +2538,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f9/33/bd5b9137445ea4b680023eb0469b2bb969d61303dedb2aac6560ff3d14a1/notebook_shim-0.2.4-py3-none-any.whl", hash = "sha256:411a5be4e9dc882a074ccbcae671eda64cceb068767e9a3419096986560e1cef", size = 13307, upload-time = "2024-02-14T23:35:16.286Z" }, ] +[[package]] +name = "numba" +version = "0.66.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/a0/570e3dc53e5602b49108f62a13e529f1eec8bfc7ef37d49c825924dcf546/numba-0.66.0.tar.gz", hash = "sha256:b900e63a0e26c05ea9a6d5a3a5a0a177cb64c5011887bf43edb8c3ed2c38d363", size = 2806181, upload-time = "2026-07-01T23:12:46.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/48/d139bde40f2359351bfe26ee1b261937f458ac177ab810d4f045ae1c9d92/numba-0.66.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:27951c47e0def9bf8afe580eb961102902e2fd23cb77924b7d9d7cc0f8b444cb", size = 2727368, upload-time = "2026-07-01T23:12:04.282Z" }, + { url = "https://files.pythonhosted.org/packages/36/e4/b780bfa9191410da50ba249cb3248a75014e17f611e72709cbddcb21f42d/numba-0.66.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc408c54b450f41582f4be1608f8981c1dcc44c7f40355cc150dd93015753407", size = 3803554, upload-time = "2026-07-01T23:12:06.379Z" }, + { url = "https://files.pythonhosted.org/packages/1c/b2/a051b96626bdf5c4d8fa6b8d450605c09638d85dc872ab63ef9a67096dca/numba-0.66.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c14c044c06b453ec3fa7715dfe75425e2ba72c73377a7ffde6d9ec511dfd94c", size = 3510065, upload-time = "2026-07-01T23:12:08.051Z" }, + { url = "https://files.pythonhosted.org/packages/34/01/24dcdc3e919522e2efbd92969c281ff40deb1d5f8a994bcd0057081c158c/numba-0.66.0-cp310-cp310-win_amd64.whl", hash = "sha256:2338cc0d43609fe448930848fd35a5bc688761b986f81b597a6f45cc0f8c9577", size = 2780379, upload-time = "2026-07-01T23:12:09.772Z" }, + { url = "https://files.pythonhosted.org/packages/9e/02/970796b4daa709604cde22e87a7cda9bde473c278ea4a75f59fe38cee47f/numba-0.66.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:bbd531c327557a9004507fa6bff06c53ab51a7a5776b75261bb9cef1efe2b2ea", size = 2727049, upload-time = "2026-07-01T23:12:11.296Z" }, + { url = "https://files.pythonhosted.org/packages/8c/99/33a6ed9c1a0b5e42efa98eb0edf617d61dca576c82625947377b1d4540c9/numba-0.66.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc6629becb21a867d85401ec89f426dd24c484a4193ade8a38309debfd1529ca", size = 3808870, upload-time = "2026-07-01T23:12:12.944Z" }, + { url = "https://files.pythonhosted.org/packages/04/20/8c51126025211659235b8de2866dfa226984ae0c8273461a3cf374716741/numba-0.66.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aac69f3ccb8af100f5913c1241edc9692bad1cdd2508721713f426eb06c9a659", size = 3514498, upload-time = "2026-07-01T23:12:15.307Z" }, + { url = "https://files.pythonhosted.org/packages/5e/c9/9476940bc6d5caf5c0cf2e4c5feecbf01244bbe6f914614082dd7a3e520e/numba-0.66.0-cp311-cp311-win_amd64.whl", hash = "sha256:fb601841d9e02e6237bb6522e36d0741614be3cfe2b482a6f00a41b5ba209443", size = 2780225, upload-time = "2026-07-01T23:12:16.924Z" }, + { url = "https://files.pythonhosted.org/packages/62/a3/70deb7f88461c1cd5d16aa990c2380604102661a427667b8950dcdccc27f/numba-0.66.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:53ca5900b7cab15109796030113a6b28576bae5ad7bb507ad6dd1360ddd81ba4", size = 2727264, upload-time = "2026-07-01T23:12:18.669Z" }, + { url = "https://files.pythonhosted.org/packages/2d/55/25c319845e9a4e08f16611ddbda56a192eb7b6ed13e1a2bff2da272ffb97/numba-0.66.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0999e3ee1b18c48e1fb51d11af35ef59852c7f4f50569c9550c25faef0616ad1", size = 3866252, upload-time = "2026-07-01T23:12:20.429Z" }, + { url = "https://files.pythonhosted.org/packages/71/ef/a82d6fd6bf1b0fe461651e924d3647eeec9ac17f8eee4896264bf7480930/numba-0.66.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:efe0d2d5099790df945e0cb6e1b3104bd965d7bbfac50d62f1d5d1d6ade0825d", size = 3566974, upload-time = "2026-07-01T23:12:22.116Z" }, + { url = "https://files.pythonhosted.org/packages/fc/eb/9e6171e378822ab191c7abcfd3d8cfc8644516f6c7834c22e210e4acc070/numba-0.66.0-cp312-cp312-win_amd64.whl", hash = "sha256:b075a4e7ebc43dc6294f223e2821659656209fd5e0ce53245877c23d66d6e1a9", size = 2797403, upload-time = "2026-07-01T23:12:23.724Z" }, + { url = "https://files.pythonhosted.org/packages/03/52/176c02d005c5c5143cde10a85bbcdcb6236d9e34c3aac089380e0506cd1d/numba-0.66.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:380b2556a2019ccd1e956ae77dd257eaa39403f7520768b626d44b755112785e", size = 2727084, upload-time = "2026-07-01T23:12:25.434Z" }, + { url = "https://files.pythonhosted.org/packages/44/b5/e930010965568fe7f2c6c962fd2849d458cb9f62c3ab7584af8a19a2b40a/numba-0.66.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:939316d5d8619751207b8972a67852b5a7646665298cb4de693cd6bf135152f4", size = 3873663, upload-time = "2026-07-01T23:12:27.308Z" }, + { url = "https://files.pythonhosted.org/packages/d0/ec/5b51457cbe96e4831141d83e892e65191b23a1b78728456c62909d231ace/numba-0.66.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdf506775d9f02eb92a87bf5c5b1e0d25506fd18cafd769f4ed914a8feac73e7", size = 3573529, upload-time = "2026-07-01T23:12:28.944Z" }, + { url = "https://files.pythonhosted.org/packages/83/7e/cea7710e96913d3c7f2999f16db1b28e6c5be5171cbf40f77f98333a7243/numba-0.66.0-cp313-cp313-win_amd64.whl", hash = "sha256:c5bfe5350284509ab0474390321454c3a8627a188af5b68c910e83df3e2db4a7", size = 2797247, upload-time = "2026-07-01T23:12:30.774Z" }, +] + [[package]] name = "numpy" version = "2.2.6" @@ -2731,7 +2784,7 @@ name = "nvidia-cudnn-cu12" version = "9.7.1.26" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/25/dc/dc825c4b1c83b538e207e34f48f86063c88deaa35d46c651c7c181364ba2/nvidia_cudnn_cu12-9.7.1.26-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:6d011159a158f3cfc47bf851aea79e31bcff60d530b70ef70474c84cac484d07", size = 726851421, upload-time = "2025-02-06T22:18:29.812Z" }, @@ -2742,7 +2795,7 @@ name = "nvidia-cudnn-cu13" version = "9.13.0.50" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8a/9c/9e99c00dc23db324244ec257d1e84d79539202ee2f185dee2c1fa97c9549/nvidia_cudnn_cu13-9.13.0.50-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:33f0aa0b64230101b348648fd0693342188071d3f8a137c0cf50051c24b3584b", size = 412337597, upload-time = "2025-09-04T20:22:31.535Z" }, @@ -2753,7 +2806,7 @@ name = "nvidia-cufft" version = "12.0.0.15" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/9b/e9/4e49b1baf6899e42eeec324a49d7aa2219fec42076327c4e468000dd375a/nvidia_cufft-12.0.0.15-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:1885731254835797572ff075f3daf43a2a0a2801210dea26971940dae7e1a367", size = 214053580, upload-time = "2025-08-04T10:20:45.781Z" }, @@ -2764,7 +2817,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.41" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ac/26/b53c493c38dccb1f1a42e1a21dc12cba2a77fbe36c652f7726d9ec4aba28/nvidia_cufft_cu12-11.3.3.41-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:da650080ab79fcdf7a4b06aa1b460e99860646b176a43f6208099bdc17836b6a", size = 193118795, upload-time = "2025-01-23T17:56:30.536Z" }, @@ -2807,9 +2860,9 @@ name = "nvidia-cusolver" version = "12.0.3.29" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas" }, - { name = "nvidia-cusparse" }, - { name = "nvidia-nvjitlink" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/a7/bb/2e60de9bb1f0c3395eabd91ccad00f4ba3ef736dc9190a158a9d268419f5/nvidia_cusolver-12.0.3.29-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:3bb6e65ce0beaeafdd069b320246e8f17c1cd30ddb27a0539143a3706733a4d8", size = 193104180, upload-time = "2025-08-04T10:22:19.821Z" }, @@ -2820,9 +2873,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.2.55" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12" }, - { name = "nvidia-cusparse-cu12" }, - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c2/08/953675873a136d96bb12f93b49ba045d1107bc94d2551c52b12fa6c7dec3/nvidia_cusolver_cu12-11.7.2.55-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4d1354102f1e922cee9db51920dba9e2559877cf6ff5ad03a00d853adafb191b", size = 260373342, upload-time = "2025-01-23T17:58:56.406Z" }, @@ -2833,7 +2886,7 @@ name = "nvidia-cusparse" version = "12.6.2.49" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/fc/30/f32023427f2ef4ec27e8293dfddb5068de566912cd0a45eccfd400017a62/nvidia_cusparse-12.6.2.49-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:5d3269c19283a0057fb5ebfb003ae2a10c97a28a6958f4238354826b055827c7", size = 155888587, upload-time = "2025-08-04T10:23:04.091Z" }, @@ -2844,7 +2897,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.7.53" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c2/ab/31e8149c66213b846c082a3b41b1365b831f41191f9f40c6ddbc8a7d550e/nvidia_cusparse_cu12-12.5.7.53-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3c1b61eb8c85257ea07e9354606b26397612627fdcd327bfd91ccf6155e7c86d", size = 292064180, upload-time = "2025-01-23T18:00:23.233Z" }, @@ -3002,12 +3055,12 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "coloredlogs" }, - { name = "flatbuffers" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "sympy" }, + { name = "coloredlogs", marker = "python_full_version < '3.11'" }, + { name = "flatbuffers", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "packaging", marker = "python_full_version < '3.11'" }, + { name = "protobuf", marker = "python_full_version < '3.11'" }, + { name = "sympy", marker = "python_full_version < '3.11'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/39/18/272d3d7406909141d3c9943796e3e97cafa53f4342d9231c0cfd8cb05702/onnxruntime-1.19.2-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:84fa57369c06cadd3c2a538ae2a26d76d583e7c34bdecd5769d71ca5c0fc750e", size = 16776408, upload-time = "2024-09-04T06:37:02.431Z" }, @@ -3052,11 +3105,11 @@ resolution-markers = [ "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "flatbuffers" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "protobuf" }, - { name = "sympy" }, + { name = "flatbuffers", marker = "python_full_version >= '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "packaging", marker = "python_full_version >= '3.11'" }, + { name = "protobuf", marker = "python_full_version >= '3.11'" }, + { name = "sympy", marker = "python_full_version >= '3.11'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/60/69/6c40720201012c6af9aa7d4ecdd620e521bd806dc6269d636fdd5c5aeebe/onnxruntime-1.24.4-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:0bdfce8e9a6497cec584aab407b71bf697dac5e1b7b7974adc50bf7533bdb3a2", size = 17332131, upload-time = "2026-03-17T22:05:49.005Z" }, @@ -3212,7 +3265,7 @@ name = "pexpect" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "ptyprocess" }, + { name = "ptyprocess", marker = "sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } wheels = [ @@ -3994,7 +4047,7 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } wheels = [ @@ -4070,7 +4123,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" } }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } wheels = [ @@ -4417,14 +4470,14 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'win32'", ] dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "setuptools", marker = "python_full_version >= '3.12'" }, - { name = "sympy" }, - { name = "typing-extensions" }, + { name = "filelock", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "fsspec", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "jinja2", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.11' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.12' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "sympy", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "typing-extensions", marker = "(platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'linux' and sys_platform != 'win32')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/46/c2/3fb87940fa160d956ee94d644d37b99a24b9c05a4222bf34f94c71880e28/torch-2.7.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:c9afea41b11e1a1ab1b258a5c31afbd646d6319042bfe4f231b408034b51128b", size = 99158447, upload-time = "2025-04-23T14:35:10.557Z" }, @@ -4463,29 +4516,29 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'win32'", ] dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "nvidia-cublas-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cuda-cupti-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cudnn-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cufft-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cufile-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-curand-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cusolver-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-cusparselt-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-nccl-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'win32'" }, - { name = "nvidia-nvtx-cu12", marker = "sys_platform != 'win32'" }, - { name = "setuptools", marker = "python_full_version >= '3.12'" }, - { name = "sympy" }, - { name = "triton", version = "3.3.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'win32'" }, - { name = "typing-extensions" }, + { name = "filelock", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, + { name = "fsspec", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, + { name = "jinja2", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.12' and sys_platform == 'win32')" }, + { name = "sympy", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, + { name = "triton", version = "3.3.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:ac1849553ee673dfafb44c610c60cb60a2890f0e117f43599a526cf777eb8b8c", upload-time = "2025-04-22T18:19:25Z" }, @@ -4511,30 +4564,30 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", ] dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "nvidia-cublas" }, - { name = "nvidia-cuda-cupti" }, - { name = "nvidia-cuda-nvrtc" }, - { name = "nvidia-cuda-runtime" }, - { name = "nvidia-cudnn-cu13" }, - { name = "nvidia-cufft" }, - { name = "nvidia-cufile" }, - { name = "nvidia-curand" }, - { name = "nvidia-cusolver" }, - { name = "nvidia-cusparse" }, - { name = "nvidia-cusparselt-cu13" }, - { name = "nvidia-nccl-cu13" }, - { name = "nvidia-nvjitlink" }, - { name = "nvidia-nvshmem-cu13" }, - { name = "nvidia-nvtx" }, - { name = "setuptools", marker = "python_full_version >= '3.12'" }, - { name = "sympy" }, - { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" } }, - { name = "typing-extensions" }, + { name = "filelock", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "fsspec", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "jinja2", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cublas", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu13", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-curand", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu13", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu13", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu13", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "python_full_version >= '3.12' and platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "sympy", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "triton", version = "3.5.0", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.9.0%2Bcu130-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:46004a346db6bfd69ecd2e42dce48e0fce2ad0e5a910f8203db5206f5515387e", upload-time = "2025-10-14T17:30:04Z" }, @@ -4629,7 +4682,7 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 'x86_64' and sys_platform == 'linux'", ] dependencies = [ - { name = "setuptools" }, + { name = "setuptools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/76/04/d54d3a6d077c646624dc9461b0059e23fd5d30e0dbe67471e3654aec81f9/triton-3.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fad99beafc860501d7fcc1fb7045d9496cbe2c882b1674640304949165a916e7", size = 156441993, upload-time = "2025-04-09T20:27:25.107Z" }, @@ -4719,6 +4772,7 @@ dependencies = [ { name = "mlx", marker = "sys_platform == 'darwin'" }, { name = "ninja", marker = "sys_platform == 'linux'" }, { name = "notebook" }, + { name = "numba" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "onnxruntime", version = "1.19.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -4773,6 +4827,7 @@ requires-dist = [ { name = "mujoco-uni", marker = "extra == 'mujoco'", specifier = "==3.8.0" }, { name = "ninja", marker = "sys_platform == 'linux'" }, { name = "notebook", specifier = ">=7.5.5" }, + { name = "numba" }, { name = "numpy" }, { name = "onnxruntime", marker = "python_full_version < '3.11'", specifier = "<1.20" }, { name = "onnxruntime", marker = "python_full_version >= '3.11'", specifier = ">=1.20" }, diff --git a/uv.rocm.lock b/uv.rocm.lock index 4541dec4d..08eff2869 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -1858,6 +1858,30 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d3/97/68f80ca3ac4924f250cdfa6e20142a803e5e50fca96ef5148c52ee8c10ea/librt-0.8.1-cp313-cp313-win_arm64.whl", hash = "sha256:924817ab3141aca17893386ee13261f1d100d1ef410d70afe4389f2359fea4f0", size = 52495, upload-time = "2026-02-17T16:12:11.633Z" }, ] +[[package]] +name = "llvmlite" +version = "0.48.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/dc/a0/acc8ffcd5bdc63df0097e22c719bfcd61b604358343089313a8aebbb24ab/llvmlite-0.48.0.tar.gz", hash = "sha256:543b19f9ef8f3c7c60d1468191e4ee1b1537bf9f8a3d56f64c0ddd98de92edd2", size = 184016, upload-time = "2026-07-02T20:20:05.308Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/4e/32543c42568fb321b3bdfcf9106e4116ab8f5a7bbcfd9ecf5569b0c07d83/llvmlite-0.48.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:614aad57df707e3172efd5165f2aa7da6a0c6897e40dce590bf756396815ba76", size = 40480650, upload-time = "2026-07-01T18:41:01.945Z" }, + { url = "https://files.pythonhosted.org/packages/a9/0d/6aa48abd423067139a129d1434b77bbcc56080db51d12a88510bb491ca3d/llvmlite-0.48.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:13532f248960ba888ad5ab8150494e2f3a3d20e5f59f264e63741ea5b0ba844c", size = 59890118, upload-time = "2026-07-01T18:41:10.608Z" }, + { url = "https://files.pythonhosted.org/packages/5a/c7/aa917444d871a79608af49149de1b28764e87d2ab41f933c5cd02431d03d/llvmlite-0.48.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3ee0c77685a18f5fca994ae21d0763007fca5c5c64b41de37accc78b69079176", size = 58343459, upload-time = "2026-07-01T18:41:06.21Z" }, + { url = "https://files.pythonhosted.org/packages/c5/2b/ceee1cdc263617109d514ac4d1b31f10a282662740ff7d5777baae25b3b5/llvmlite-0.48.0-cp310-cp310-win_amd64.whl", hash = "sha256:02853fe4214acb3780fc920c3fee10564b61d58a35e1b78afcc8a546c2deaba3", size = 41864734, upload-time = "2026-07-01T18:41:14.746Z" }, + { url = "https://files.pythonhosted.org/packages/9a/55/595981f14fbae9ba966feb12af552b1fe69889e44e64ac883a731ed335e0/llvmlite-0.48.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:56a7e24607d3f02d7b1bae8d29c7e1e423d53143d68b072999777f19678fe77b", size = 40480651, upload-time = "2026-07-01T18:41:18.438Z" }, + { url = "https://files.pythonhosted.org/packages/26/08/0109d1b9cb3f4603f3890e30bc66c65332b79185f12a045343b2ae431f67/llvmlite-0.48.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6fa532d6bb3fd3f0803567c736401c54aecfe1a396d3ad25d2440d220e09f0e7", size = 59890118, upload-time = "2026-07-01T18:41:28.184Z" }, + { url = "https://files.pythonhosted.org/packages/02/eb/c5281be180c789cdffbf45b671884c57d7e61345ef3b0f643a4965e108e8/llvmlite-0.48.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:979a66a3f28a02565383ff463527dce78e9b856298872a361283132488e83591", size = 58343458, upload-time = "2026-07-01T18:41:23.397Z" }, + { url = "https://files.pythonhosted.org/packages/aa/f7/b3222b13f2d424dae3c9e63fde476af25ebccf1f3faf0b52d1b79fc15c70/llvmlite-0.48.0-cp311-cp311-win_amd64.whl", hash = "sha256:efaee0276e5e17c2b99b92e0c974bd484ef5977cf5dbc9168e82b71578edb47f", size = 41864734, upload-time = "2026-07-01T18:41:31.932Z" }, + { url = "https://files.pythonhosted.org/packages/92/a2/28696a9e61e245d1a79816d29d106692a90a2b6e7d78c98b326db70827af/llvmlite-0.48.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:d66c3beb4209087ddd4cf4ed2a0856b6887e6a913bdcf1aacfec9851cf2cba4e", size = 40480651, upload-time = "2026-07-01T18:41:35.694Z" }, + { url = "https://files.pythonhosted.org/packages/80/f2/72409351db66d0a317ec5087e076f31fb7b773a640db8a90ce6b5cac9edd/llvmlite-0.48.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:416fa4c2c66c2c6dc6d0a402648c19206e548efa0aa1eff01ad5cdad0af8217d", size = 59890118, upload-time = "2026-07-01T18:41:44.886Z" }, + { url = "https://files.pythonhosted.org/packages/3a/27/5ae2f3722606360480707adb47f001ad89df8251d06b14ee80336e660b66/llvmlite-0.48.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f5e5a5131045b72345c71062ea1a91910dde913792b6c9b28ebb2c1c0a712e98", size = 58343459, upload-time = "2026-07-01T18:41:40.306Z" }, + { url = "https://files.pythonhosted.org/packages/16/78/d824ffff7521cd140dc2006e44ce2bc82e64b48d1b32e90e956308c85a74/llvmlite-0.48.0-cp312-cp312-win_amd64.whl", hash = "sha256:d45c7541a80934ec6d8ab0defe67439494ecd2193cbf852a44ba827808976ac1", size = 41865022, upload-time = "2026-07-01T18:41:48.663Z" }, + { url = "https://files.pythonhosted.org/packages/9c/23/fe9316d14626b42c73ef0b502e724705a6ee9450afe53759c0a99c37c2d7/llvmlite-0.48.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:a83a99ef0c05b4ccddf9b6218ed9fe84b653a0caf7c1d9dbe148d6d16c67f518", size = 40480652, upload-time = "2026-07-01T18:41:52.216Z" }, + { url = "https://files.pythonhosted.org/packages/1b/4a/90715fa12006d681270b08d881195b6fab3ec39572e048764a1f7f59fed7/llvmlite-0.48.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8761b9e522f55207e24424fcd98370289eec2710bf8e915c82d1053f642450dc", size = 59890120, upload-time = "2026-07-01T18:42:00.748Z" }, + { url = "https://files.pythonhosted.org/packages/70/5e/7b3e20d64650ca3c80af0cdb664ec4b575ec83d9d4dd05bea8bd31f9bbb6/llvmlite-0.48.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2fe5cb59b2063bfa039dcb8ca6481c0181bf552f340d10dcf61d7996a665556e", size = 58343457, upload-time = "2026-07-01T18:41:56.41Z" }, + { url = "https://files.pythonhosted.org/packages/17/97/5a430055d1838cf1fb7a01cfa943300f5e4c026fc6333a522c5e4a03b0c1/llvmlite-0.48.0-cp313-cp313-win_amd64.whl", hash = "sha256:91c7e24e74cde3f02b88aa5acca678373f9e069f3b98531b3dbb3a142d9d10bb", size = 41865022, upload-time = "2026-07-01T18:42:04.57Z" }, +] + [[package]] name = "lxml" version = "6.0.4" @@ -2627,6 +2651,35 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f9/33/bd5b9137445ea4b680023eb0469b2bb969d61303dedb2aac6560ff3d14a1/notebook_shim-0.2.4-py3-none-any.whl", hash = "sha256:411a5be4e9dc882a074ccbcae671eda64cceb068767e9a3419096986560e1cef", size = 13307, upload-time = "2024-02-14T23:35:16.286Z" }, ] +[[package]] +name = "numba" +version = "0.66.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ae/a0/570e3dc53e5602b49108f62a13e529f1eec8bfc7ef37d49c825924dcf546/numba-0.66.0.tar.gz", hash = "sha256:b900e63a0e26c05ea9a6d5a3a5a0a177cb64c5011887bf43edb8c3ed2c38d363", size = 2806181, upload-time = "2026-07-01T23:12:46.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2b/48/d139bde40f2359351bfe26ee1b261937f458ac177ab810d4f045ae1c9d92/numba-0.66.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:27951c47e0def9bf8afe580eb961102902e2fd23cb77924b7d9d7cc0f8b444cb", size = 2727368, upload-time = "2026-07-01T23:12:04.282Z" }, + { url = "https://files.pythonhosted.org/packages/36/e4/b780bfa9191410da50ba249cb3248a75014e17f611e72709cbddcb21f42d/numba-0.66.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc408c54b450f41582f4be1608f8981c1dcc44c7f40355cc150dd93015753407", size = 3803554, upload-time = "2026-07-01T23:12:06.379Z" }, + { url = "https://files.pythonhosted.org/packages/1c/b2/a051b96626bdf5c4d8fa6b8d450605c09638d85dc872ab63ef9a67096dca/numba-0.66.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7c14c044c06b453ec3fa7715dfe75425e2ba72c73377a7ffde6d9ec511dfd94c", size = 3510065, upload-time = "2026-07-01T23:12:08.051Z" }, + { url = "https://files.pythonhosted.org/packages/34/01/24dcdc3e919522e2efbd92969c281ff40deb1d5f8a994bcd0057081c158c/numba-0.66.0-cp310-cp310-win_amd64.whl", hash = "sha256:2338cc0d43609fe448930848fd35a5bc688761b986f81b597a6f45cc0f8c9577", size = 2780379, upload-time = "2026-07-01T23:12:09.772Z" }, + { url = "https://files.pythonhosted.org/packages/9e/02/970796b4daa709604cde22e87a7cda9bde473c278ea4a75f59fe38cee47f/numba-0.66.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:bbd531c327557a9004507fa6bff06c53ab51a7a5776b75261bb9cef1efe2b2ea", size = 2727049, upload-time = "2026-07-01T23:12:11.296Z" }, + { url = "https://files.pythonhosted.org/packages/8c/99/33a6ed9c1a0b5e42efa98eb0edf617d61dca576c82625947377b1d4540c9/numba-0.66.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc6629becb21a867d85401ec89f426dd24c484a4193ade8a38309debfd1529ca", size = 3808870, upload-time = "2026-07-01T23:12:12.944Z" }, + { url = "https://files.pythonhosted.org/packages/04/20/8c51126025211659235b8de2866dfa226984ae0c8273461a3cf374716741/numba-0.66.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:aac69f3ccb8af100f5913c1241edc9692bad1cdd2508721713f426eb06c9a659", size = 3514498, upload-time = "2026-07-01T23:12:15.307Z" }, + { url = "https://files.pythonhosted.org/packages/5e/c9/9476940bc6d5caf5c0cf2e4c5feecbf01244bbe6f914614082dd7a3e520e/numba-0.66.0-cp311-cp311-win_amd64.whl", hash = "sha256:fb601841d9e02e6237bb6522e36d0741614be3cfe2b482a6f00a41b5ba209443", size = 2780225, upload-time = "2026-07-01T23:12:16.924Z" }, + { url = "https://files.pythonhosted.org/packages/62/a3/70deb7f88461c1cd5d16aa990c2380604102661a427667b8950dcdccc27f/numba-0.66.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:53ca5900b7cab15109796030113a6b28576bae5ad7bb507ad6dd1360ddd81ba4", size = 2727264, upload-time = "2026-07-01T23:12:18.669Z" }, + { url = "https://files.pythonhosted.org/packages/2d/55/25c319845e9a4e08f16611ddbda56a192eb7b6ed13e1a2bff2da272ffb97/numba-0.66.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0999e3ee1b18c48e1fb51d11af35ef59852c7f4f50569c9550c25faef0616ad1", size = 3866252, upload-time = "2026-07-01T23:12:20.429Z" }, + { url = "https://files.pythonhosted.org/packages/71/ef/a82d6fd6bf1b0fe461651e924d3647eeec9ac17f8eee4896264bf7480930/numba-0.66.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:efe0d2d5099790df945e0cb6e1b3104bd965d7bbfac50d62f1d5d1d6ade0825d", size = 3566974, upload-time = "2026-07-01T23:12:22.116Z" }, + { url = "https://files.pythonhosted.org/packages/fc/eb/9e6171e378822ab191c7abcfd3d8cfc8644516f6c7834c22e210e4acc070/numba-0.66.0-cp312-cp312-win_amd64.whl", hash = "sha256:b075a4e7ebc43dc6294f223e2821659656209fd5e0ce53245877c23d66d6e1a9", size = 2797403, upload-time = "2026-07-01T23:12:23.724Z" }, + { url = "https://files.pythonhosted.org/packages/03/52/176c02d005c5c5143cde10a85bbcdcb6236d9e34c3aac089380e0506cd1d/numba-0.66.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:380b2556a2019ccd1e956ae77dd257eaa39403f7520768b626d44b755112785e", size = 2727084, upload-time = "2026-07-01T23:12:25.434Z" }, + { url = "https://files.pythonhosted.org/packages/44/b5/e930010965568fe7f2c6c962fd2849d458cb9f62c3ab7584af8a19a2b40a/numba-0.66.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:939316d5d8619751207b8972a67852b5a7646665298cb4de693cd6bf135152f4", size = 3873663, upload-time = "2026-07-01T23:12:27.308Z" }, + { url = "https://files.pythonhosted.org/packages/d0/ec/5b51457cbe96e4831141d83e892e65191b23a1b78728456c62909d231ace/numba-0.66.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdf506775d9f02eb92a87bf5c5b1e0d25506fd18cafd769f4ed914a8feac73e7", size = 3573529, upload-time = "2026-07-01T23:12:28.944Z" }, + { url = "https://files.pythonhosted.org/packages/83/7e/cea7710e96913d3c7f2999f16db1b28e6c5be5171cbf40f77f98333a7243/numba-0.66.0-cp313-cp313-win_amd64.whl", hash = "sha256:c5bfe5350284509ab0474390321454c3a8627a188af5b68c910e83df3e2db4a7", size = 2797247, upload-time = "2026-07-01T23:12:30.774Z" }, +] + [[package]] name = "numpy" version = "2.2.6" @@ -4694,6 +4747,7 @@ dependencies = [ { name = "mlx", marker = "sys_platform == 'darwin'" }, { name = "ninja", marker = "sys_platform == 'linux'" }, { name = "notebook" }, + { name = "numba" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "onnxruntime", version = "1.19.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -4748,6 +4802,7 @@ requires-dist = [ { name = "mujoco-uni", marker = "extra == 'mujoco'", specifier = "==3.8.0" }, { name = "ninja", marker = "sys_platform == 'linux'" }, { name = "notebook", specifier = ">=7.5.5" }, + { name = "numba" }, { name = "numpy" }, { name = "onnxruntime", marker = "python_full_version < '3.11'", specifier = "<1.20" }, { name = "onnxruntime", marker = "python_full_version >= '3.11'", specifier = ">=1.20" }, From 4b272c15ca622bd0467a1707597b0b797a420741 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 16:14:42 +0800 Subject: [PATCH 02/26] bench: extend g1 joystick numba thread sweep --- benchmark/benchmark_g1_joystick_numba.py | 26 ++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index c2eabb9ea..d42c4509b 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -54,6 +54,8 @@ from unilab.envs.locomotion.g1.joystick_numba import G1WalkNumbaAccelerator NUM_ACTION = 29 +DEFAULT_THREADS = [2, 4, 8, 16, 32, 64] +QUICK_THREADS = [2, 4] DEFAULT_POSE_WEIGHTS = [ 0.01, 1.0, @@ -677,6 +679,10 @@ def write_report( "profiles": args.profiles, "num_envs": args.num_envs, "threads": args.threads, + "requested_threads": args.threads, + "measured_threads": getattr(args, "measured_threads", None), + "skipped_threads": getattr(args, "skipped_threads", None), + "numba_max_threads": getattr(args, "numba_max_threads", None), "scope": "G1 joystick reward+termination hot slice; synthetic backend arrays", } payload = { @@ -754,7 +760,7 @@ def parse_args() -> argparse.Namespace: choices=sorted(make_profile_specs()), ) parser.add_argument("--num-envs", nargs="+", type=int, default=[512, 2048, 8192]) - parser.add_argument("--threads", nargs="+", type=int, default=[2, 4, 8, 16]) + parser.add_argument("--threads", nargs="+", type=int, default=None) parser.add_argument("--iters", type=int, default=80) parser.add_argument("--warmup", type=int, default=8) parser.add_argument("--seed", type=int, default=0) @@ -772,9 +778,12 @@ def parse_args() -> argparse.Namespace: if args.quick: args.profiles = ["sac_default"] args.num_envs = [512, 2048] - args.threads = [2, 4] + if args.threads is None: + args.threads = QUICK_THREADS args.iters = 10 args.warmup = 2 + elif args.threads is None: + args.threads = DEFAULT_THREADS return args @@ -783,12 +792,21 @@ def main() -> None: specs = make_profile_specs() all_records: list[BenchCase] = [] parity: dict[str, dict[str, float]] = {} + max_threads = get_num_threads() + args.numba_max_threads = max_threads + args.measured_threads = sorted({1, *(threads for threads in args.threads if threads <= max_threads)}) + args.skipped_threads = sorted({threads for threads in args.threads if threads > max_threads}) print("=" * 80) print("G1 joystick Numba benchmark: reward dispatch + termination") print("=" * 80) - print(f"host numba threads: {get_num_threads()}") - print(f"profiles={args.profiles} num_envs={args.num_envs} threads={args.threads}") + print(f"host numba threads: {max_threads}") + print( + f"profiles={args.profiles} num_envs={args.num_envs} " + f"requested_threads={args.threads} measured_threads={args.measured_threads}" + ) + if args.skipped_threads: + print(f"skipped threads above numba max: {args.skipped_threads}") for profile_name in args.profiles: spec = specs[profile_name] for num_envs in args.num_envs: From 72ce7dad540697435fd56b752a806f212f889bf7 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 16:23:53 +0800 Subject: [PATCH 03/26] bench: extend g1 joystick num env sweep --- benchmark/benchmark_g1_joystick_numba.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index d42c4509b..a1f4cee3f 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -759,7 +759,12 @@ def parse_args() -> argparse.Namespace: default=["ppo_default", "sac_default", "full_supported"], choices=sorted(make_profile_specs()), ) - parser.add_argument("--num-envs", nargs="+", type=int, default=[512, 2048, 8192]) + parser.add_argument( + "--num-envs", + nargs="+", + type=int, + default=[512, 1024, 2048, 4096, 8192, 16384, 32768], + ) parser.add_argument("--threads", nargs="+", type=int, default=None) parser.add_argument("--iters", type=int, default=80) parser.add_argument("--warmup", type=int, default=8) From e86a3911ee04be323d0b6c295fb7aebdf958ea8c Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 16:42:25 +0800 Subject: [PATCH 04/26] bench: add g1 joystick numba e2e comparison --- benchmark/benchmark_g1_joystick_numba.py | 238 +++++++++++++++++- .../test_g1_joystick_numba_benchmark.py | 27 ++ 2 files changed, 263 insertions(+), 2 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index a1f4cee3f..c5970aee9 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -14,6 +14,7 @@ uv run python -m benchmark.benchmark_g1_joystick_numba uv run python benchmark/benchmark_g1_joystick_numba.py uv run python -m benchmark.benchmark_g1_joystick_numba --quick + uv run python -m benchmark.benchmark_g1_joystick_numba --quick --e2e """ from __future__ import annotations @@ -159,6 +160,24 @@ class BenchCase: compile_ms: float | None = None +@dataclass +class EndToEndCase: + case: str + path: str + num_envs: int + warmup_steps: int + measure_steps: int + numba_acceleration: bool + numba_threads: int | None + collector_active_steps_per_sec: float + total_active_ms: float + env_step_ms: float + update_state_ms: float | None + speedup_vs_numpy: float = 1.0 + env_step_speedup_vs_numpy: float | None = None + update_state_speedup_vs_numpy: float | None = None + + @dataclass class SyntheticBackend: base_pos: np.ndarray @@ -471,6 +490,109 @@ def _case_to_dict(case: BenchCase) -> dict[str, Any]: } +def _e2e_case_to_dict(case: EndToEndCase) -> dict[str, Any]: + return { + "case": case.case, + "path": case.path, + "num_envs": case.num_envs, + "warmup_steps": case.warmup_steps, + "measure_steps": case.measure_steps, + "numba_acceleration": case.numba_acceleration, + "numba_threads": case.numba_threads, + "collector_active_steps_per_sec": case.collector_active_steps_per_sec, + "total_active_ms": case.total_active_ms, + "env_step_ms": case.env_step_ms, + "update_state_ms": case.update_state_ms, + "speedup_vs_numpy": case.speedup_vs_numpy, + "env_step_speedup_vs_numpy": case.env_step_speedup_vs_numpy, + "update_state_speedup_vs_numpy": case.update_state_speedup_vs_numpy, + } + + +def _timing_mean_ms(result: Any, key: str) -> float | None: + stat = result.env_step_timing_ms_per_vector_step.get(key) + return float(stat.mean_ms) if stat is not None else None + + +def _run_e2e_collector_case( + *, + num_envs: int, + warmup_steps: int, + measure_steps: int, + numba_threads: int | None, +) -> list[EndToEndCase]: + """Run a real collector active-window A/B test using training construction paths. + + This mirrors benchmark_offpolicy_collector_active.py: Hydra owner config, + create_env, actor action sampling, env.step, terminal-observation handling, + replay writes, and bookkeeping are all included. It is intentionally optional + because it constructs a real MuJoCo env and is much heavier than the synthetic + reward+termination hot-slice benchmark above. + """ + from benchmark.benchmark_offpolicy_collector_active import _build_and_run_case + + case_name = "sac/g1_walk_flat/mujoco" + common = { + "warmup_steps": warmup_steps, + "measure_steps": measure_steps, + "replay_capacity_steps": max(2, measure_steps + warmup_steps + 1), + "num_envs": num_envs, + } + variants = [ + ( + "training_collector_numpy", + False, + [ + "++env.numba_acceleration=false", + ], + ), + ( + "training_collector_numba", + True, + [ + "++env.numba_acceleration=true", + f"++env.numba_num_threads={numba_threads}" if numba_threads is not None else "", + ], + ), + ] + + records: list[EndToEndCase] = [] + for path, enabled, overrides in variants: + result = _build_and_run_case( + case_name, + extra_overrides=[override for override in overrides if override], + **common, + ) + env_step_ms = float(result.phase_ms_per_vector_step["env_step_ms"].mean_ms) + records.append( + EndToEndCase( + case=case_name, + path=path, + num_envs=num_envs, + warmup_steps=warmup_steps, + measure_steps=measure_steps, + numba_acceleration=enabled, + numba_threads=numba_threads if enabled else None, + collector_active_steps_per_sec=float(result.collector_active_steps_per_sec), + total_active_ms=float(result.total_active_ms), + env_step_ms=env_step_ms, + update_state_ms=_timing_mean_ms(result, "update_state_ms"), + ) + ) + + baseline = next(record for record in records if not record.numba_acceleration) + for record in records: + record.speedup_vs_numpy = ( + record.collector_active_steps_per_sec / baseline.collector_active_steps_per_sec + ) + record.env_step_speedup_vs_numpy = ( + baseline.env_step_ms / record.env_step_ms if record.env_step_ms > 0.0 else None + ) + if baseline.update_state_ms is not None and record.update_state_ms: + record.update_state_speedup_vs_numpy = baseline.update_state_ms / record.update_state_ms + return records + + def _format_table(records: list[BenchCase]) -> str: headers = [ "profile", @@ -509,6 +631,56 @@ def fmt(values: list[str]) -> str: return "\n".join(lines) +def _format_e2e_table(records: list[EndToEndCase]) -> str: + headers = [ + "case", + "envs", + "path", + "threads", + "collector M steps/s", + "speedup", + "env_step ms", + "env_step speedup", + "update_state ms", + "update_state speedup", + ] + rows = [] + for record in records: + rows.append( + [ + record.case, + str(record.num_envs), + record.path, + "-" if record.numba_threads is None else str(record.numba_threads), + f"{record.collector_active_steps_per_sec / 1e6:.3f}", + f"{record.speedup_vs_numpy:.2f}x", + f"{record.env_step_ms:.3f}", + ( + "-" + if record.env_step_speedup_vs_numpy is None + else f"{record.env_step_speedup_vs_numpy:.2f}x" + ), + "-" if record.update_state_ms is None else f"{record.update_state_ms:.3f}", + ( + "-" + if record.update_state_speedup_vs_numpy is None + else f"{record.update_state_speedup_vs_numpy:.2f}x" + ), + ] + ) + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + lines = [fmt(headers), "-+-".join("-" * width for width in widths)] + lines.extend(fmt(row) for row in rows) + return "\n".join(lines) + + def _best_numba_by_case(records: list[BenchCase]) -> dict[tuple[str, int], BenchCase]: best_by_case: dict[tuple[str, int], BenchCase] = {} for record in records: @@ -664,6 +836,7 @@ def write_report( output_dir: Path, records: list[BenchCase], parity: dict[str, dict[str, float]], + e2e_records: list[EndToEndCase], args: argparse.Namespace, ) -> None: output_dir.mkdir(parents=True, exist_ok=True) @@ -684,10 +857,12 @@ def write_report( "skipped_threads": getattr(args, "skipped_threads", None), "numba_max_threads": getattr(args, "numba_max_threads", None), "scope": "G1 joystick reward+termination hot slice; synthetic backend arrays", + "e2e_enabled": args.e2e, } payload = { "meta": meta, "results": [_case_to_dict(record) for record in records], + "end_to_end_results": [_e2e_case_to_dict(record) for record in e2e_records], "parity": parity, "plots": plot_paths, } @@ -703,7 +878,11 @@ def write_report( "deterministic synthetic backend arrays. Physics stepping, obs assembly, reset,", "and policy inference are intentionally out of scope.", "", - "## Summary", + "## Numba-specific hot slice", + "", + "This section measures only the part this task-specific Numba backend accelerates:", + "`G1WalkEnv` reward dispatch plus termination inside `update_state`. It excludes", + "physics, observation assembly, reset/RNG, policy inference, learner work, and replay.", "", ] for key in sorted(best_by_case): @@ -720,6 +899,32 @@ def write_report( title = rel_path.removesuffix(".png").replace("_", " ") summary_lines.append(f"![{title}]({rel_path})") summary_lines.append("") + if e2e_records: + summary_lines.extend( + [ + "", + "## End-to-end training-flow comparison", + "", + "This section mirrors `benchmark_offpolicy_collector_active.py`: Hydra owner config,", + "`create_env`, actor action sampling, `env.step`, terminal-observation handling,", + "replay writes, and collector-side bookkeeping are included. It is a collector", + "active-window comparison, not a full learner convergence benchmark.", + "", + "```text", + _format_e2e_table(e2e_records), + "```", + ] + ) + else: + summary_lines.extend( + [ + "", + "## End-to-end training-flow comparison", + "", + "Not run. Pass `--e2e` to add a real off-policy collector active-window A/B", + "comparison for `sac/g1_walk_flat/mujoco` with `numba_acceleration=false/true`.", + ] + ) summary_lines.extend( [ "", @@ -769,6 +974,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--iters", type=int, default=80) parser.add_argument("--warmup", type=int, default=8) parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--e2e", + action="store_true", + help="Also run a real off-policy collector active-window baseline vs numba comparison.", + ) + parser.add_argument("--e2e-num-envs", type=int, default=512) + parser.add_argument("--e2e-warmup-steps", type=int, default=2) + parser.add_argument("--e2e-measure-steps", type=int, default=8) + parser.add_argument("--e2e-numba-threads", type=int, default=None) parser.add_argument( "--output-dir", type=Path, @@ -796,6 +1010,7 @@ def main() -> None: args = parse_args() specs = make_profile_specs() all_records: list[BenchCase] = [] + e2e_records: list[EndToEndCase] = [] parity: dict[str, dict[str, float]] = {} max_threads = get_num_threads() args.numba_max_threads = max_threads @@ -828,7 +1043,26 @@ def main() -> None: print() print(_format_table(records)) - write_report(output_dir=args.output_dir, records=all_records, parity=parity, args=args) + if args.e2e: + print() + print("=" * 80) + print("End-to-end training-flow comparison: off-policy collector active window") + print("=" * 80) + e2e_records = _run_e2e_collector_case( + num_envs=args.e2e_num_envs, + warmup_steps=args.e2e_warmup_steps, + measure_steps=args.e2e_measure_steps, + numba_threads=args.e2e_numba_threads, + ) + print(_format_e2e_table(e2e_records)) + + write_report( + output_dir=args.output_dir, + records=all_records, + parity=parity, + e2e_records=e2e_records, + args=args, + ) if __name__ == "__main__": diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index a6a3bc409..b7aaa62fc 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -19,3 +19,30 @@ def test_g1_joystick_numba_benchmark_builds_records_and_matches_numpy() -> None: assert parity["max_abs_reward_diff"] < 1.0e-5 assert {record.path for record in records} == {"numpy_dispatch", "numba_accelerator"} assert any(record.path == "numba_accelerator" and record.threads == 1 for record in records) + + +def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: + record = bench.EndToEndCase( + case="sac/g1_walk_flat/mujoco", + path="training_collector_numba", + num_envs=64, + warmup_steps=1, + measure_steps=2, + numba_acceleration=True, + numba_threads=4, + collector_active_steps_per_sec=30_000.0, + total_active_ms=4.2, + env_step_ms=1.5, + update_state_ms=0.2, + speedup_vs_numpy=1.25, + env_step_speedup_vs_numpy=1.5, + update_state_speedup_vs_numpy=2.0, + ) + + payload = bench._e2e_case_to_dict(record) + table = bench._format_e2e_table([record]) + + assert payload["numba_acceleration"] is True + assert payload["numba_threads"] == 4 + assert "training_collector_numba" in table + assert "1.25x" in table From 076367bdb9d58edbdb2f1dac6bf1ef09746e9bf3 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 16:53:22 +0800 Subject: [PATCH 05/26] bench: sweep g1 joystick collector e2e gains --- benchmark/benchmark_g1_joystick_numba.py | 123 +++++++++++++++--- .../test_g1_joystick_numba_benchmark.py | 17 +++ 2 files changed, 124 insertions(+), 16 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index c5970aee9..1037afd8b 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -57,6 +57,8 @@ NUM_ACTION = 29 DEFAULT_THREADS = [2, 4, 8, 16, 32, 64] QUICK_THREADS = [2, 4] +DEFAULT_NUM_ENVS = [512, 1024, 2048, 4096, 8192, 16384, 32768] +DEFAULT_E2E_NUM_ENVS = [1024, 2048, 4096, 8192, 16384, 32768] DEFAULT_POSE_WEIGHTS = [ 0.01, 1.0, @@ -514,7 +516,7 @@ def _timing_mean_ms(result: Any, key: str) -> float | None: return float(stat.mean_ms) if stat is not None else None -def _run_e2e_collector_case( +def _run_e2e_collector_pair( *, num_envs: int, warmup_steps: int, @@ -525,9 +527,9 @@ def _run_e2e_collector_case( This mirrors benchmark_offpolicy_collector_active.py: Hydra owner config, create_env, actor action sampling, env.step, terminal-observation handling, - replay writes, and bookkeeping are all included. It is intentionally optional - because it constructs a real MuJoCo env and is much heavier than the synthetic - reward+termination hot-slice benchmark above. + replay writes, and bookkeeping are all included. Learner updates are excluded. + It is intentionally optional because it constructs a real MuJoCo env and is + much heavier than the synthetic reward+termination hot-slice benchmark above. """ from benchmark.benchmark_offpolicy_collector_active import _build_and_run_case @@ -593,6 +595,41 @@ def _run_e2e_collector_case( return records +def _best_threads_for_profile( + records: list[BenchCase], *, profile: str, num_envs: list[int] +) -> dict[int, int]: + best_by_case = _best_numba_by_case(records) + selected: dict[int, int] = {} + for env_count in num_envs: + best = best_by_case.get((profile, env_count)) + if best is not None and best.threads is not None: + selected[env_count] = int(best.threads) + return selected + + +def _run_e2e_collector_sweep( + *, + num_envs: list[int], + warmup_steps: int, + measure_steps: int, + selected_threads: dict[int, int], + fallback_numba_threads: int | None, +) -> list[EndToEndCase]: + records: list[EndToEndCase] = [] + for env_count in num_envs: + numba_threads = selected_threads.get(env_count, fallback_numba_threads) + print(f"e2e collector case: num_envs={env_count} numba_threads={numba_threads}") + records.extend( + _run_e2e_collector_pair( + num_envs=env_count, + warmup_steps=warmup_steps, + measure_steps=measure_steps, + numba_threads=numba_threads, + ) + ) + return records + + def _format_table(records: list[BenchCase]) -> str: headers = [ "profile", @@ -858,6 +895,10 @@ def write_report( "numba_max_threads": getattr(args, "numba_max_threads", None), "scope": "G1 joystick reward+termination hot slice; synthetic backend arrays", "e2e_enabled": args.e2e, + "e2e_num_envs": args.e2e_num_envs, + "e2e_warmup_steps": args.e2e_warmup_steps, + "e2e_measure_steps": args.e2e_measure_steps, + "e2e_numba_threads_source": "best sac_default hot-slice thread per num_env", } payload = { "meta": meta, @@ -884,6 +925,15 @@ def write_report( "`G1WalkEnv` reward dispatch plus termination inside `update_state`. It excludes", "physics, observation assembly, reset/RNG, policy inference, learner work, and replay.", "", + "Profile meanings:", + "", + "- `ppo_default`: reward terms and thresholds matching the PPO G1 walk-flat style.", + "- `sac_default`: reward terms and thresholds matching the SAC G1 walk-flat owner config;", + " this is also the source used to choose Numba threads for the collector A/B run.", + "- `full_supported`: synthetic stress profile with every reward term currently supported", + " by `joystick_numba.py` enabled, used to estimate the upper bound when reward work is", + " heavier than the default training configs.", + "", ] for key in sorted(best_by_case): best = best_by_case[key] @@ -903,12 +953,13 @@ def write_report( summary_lines.extend( [ "", - "## End-to-end training-flow comparison", + "## End-to-end collector comparison", "", "This section mirrors `benchmark_offpolicy_collector_active.py`: Hydra owner config,", "`create_env`, actor action sampling, `env.step`, terminal-observation handling,", - "replay writes, and collector-side bookkeeping are included. It is a collector", - "active-window comparison, not a full learner convergence benchmark.", + "replay writes, and collector-side bookkeeping are included. Learner updates are", + "not run. The Numba variant uses the best `sac_default` hot-slice thread count", + "found above for each `num_envs`.", "", "```text", _format_e2e_table(e2e_records), @@ -919,10 +970,11 @@ def write_report( summary_lines.extend( [ "", - "## End-to-end training-flow comparison", + "## End-to-end collector comparison", "", "Not run. Pass `--e2e` to add a real off-policy collector active-window A/B", "comparison for `sac/g1_walk_flat/mujoco` with `numba_acceleration=false/true`.", + "This collector comparison does not run learner updates.", ] ) summary_lines.extend( @@ -946,8 +998,10 @@ def write_report( "", "- `numba 1 thread` isolates fusion/codegen benefit over Python reward dispatch.", "- Higher thread counts add row-parallel speedup over the same fused kernel.", - "- End-to-end training speedup will be lower because this benchmark excludes physics,", - " obs assembly, reset/RNG, policy inference, and learner work.", + "- Hot-slice speedup is an upper bound for collector speedup because collector timing", + " also includes physics, observation assembly, reset/RNG, policy inference, replay,", + " and bookkeeping.", + "- The optional collector comparison still excludes learner updates.", ] ) md_path = output_dir / "report.md" @@ -968,7 +1022,7 @@ def parse_args() -> argparse.Namespace: "--num-envs", nargs="+", type=int, - default=[512, 1024, 2048, 4096, 8192, 16384, 32768], + default=DEFAULT_NUM_ENVS, ) parser.add_argument("--threads", nargs="+", type=int, default=None) parser.add_argument("--iters", type=int, default=80) @@ -979,10 +1033,15 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Also run a real off-policy collector active-window baseline vs numba comparison.", ) - parser.add_argument("--e2e-num-envs", type=int, default=512) + parser.add_argument("--e2e-num-envs", nargs="+", type=int, default=DEFAULT_E2E_NUM_ENVS) parser.add_argument("--e2e-warmup-steps", type=int, default=2) parser.add_argument("--e2e-measure-steps", type=int, default=8) - parser.add_argument("--e2e-numba-threads", type=int, default=None) + parser.add_argument( + "--e2e-numba-threads", + type=int, + default=None, + help="Fallback Numba thread count when a requested e2e num_env lacks hot-slice data.", + ) parser.add_argument( "--output-dir", type=Path, @@ -997,6 +1056,7 @@ def parse_args() -> argparse.Namespace: if args.quick: args.profiles = ["sac_default"] args.num_envs = [512, 2048] + args.e2e_num_envs = [1024, 2048] if args.threads is None: args.threads = QUICK_THREADS args.iters = 10 @@ -1044,15 +1104,46 @@ def main() -> None: print(_format_table(records)) if args.e2e: + missing_e2e_hot_slice = [ + num_envs + for num_envs in args.e2e_num_envs + if ("sac_default", num_envs) not in _best_numba_by_case(all_records) + ] + if missing_e2e_hot_slice: + print() + print("=" * 80) + print("Completing sac_default hot-slice data for e2e thread selection") + print("=" * 80) + spec = specs["sac_default"] + for num_envs in missing_e2e_hot_slice: + records, parity_result = bench_one( + profile=spec, + num_envs=num_envs, + thread_counts=args.threads, + iters=args.iters, + warmup=args.warmup, + seed=args.seed, + ) + all_records.extend(records) + parity[f"sac_default:{num_envs}"] = parity_result + print() + print(_format_table(records)) + + selected_threads = _best_threads_for_profile( + all_records, profile="sac_default", num_envs=args.e2e_num_envs + ) print() print("=" * 80) - print("End-to-end training-flow comparison: off-policy collector active window") + print("End-to-end collector comparison: off-policy collector active window") print("=" * 80) - e2e_records = _run_e2e_collector_case( + print(f"e2e num_envs={args.e2e_num_envs}") + print(f"e2e numba threads from sac_default hot-slice best: {selected_threads}") + e2e_records = _run_e2e_collector_sweep( num_envs=args.e2e_num_envs, warmup_steps=args.e2e_warmup_steps, measure_steps=args.e2e_measure_steps, - numba_threads=args.e2e_numba_threads, + selected_threads=selected_threads, + fallback_numba_threads=args.e2e_numba_threads, ) print(_format_e2e_table(e2e_records)) diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index b7aaa62fc..0ec422558 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -46,3 +46,20 @@ def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: assert payload["numba_threads"] == 4 assert "training_collector_numba" in table assert "1.25x" in table + + +def test_g1_joystick_numba_benchmark_selects_best_hot_slice_threads() -> None: + records = [ + bench.BenchCase("sac_default", 1024, "numpy_dispatch", None, 1.0, 1.0, 0.0, 1000.0, 1.0), + bench.BenchCase( + "sac_default", 1024, "numba_accelerator", 2, 0.8, 0.8, 0.0, 1250.0, 1.25 + ), + bench.BenchCase( + "sac_default", 1024, "numba_accelerator", 4, 0.6, 0.6, 0.0, 1666.0, 1.67 + ), + bench.BenchCase("ppo_default", 1024, "numba_accelerator", 8, 0.5, 0.5, 0.0, 2000.0, 2.0), + ] + + assert bench._best_threads_for_profile(records, profile="sac_default", num_envs=[1024]) == { + 1024: 4 + } From d345fbc30c385770dc52ddde786a3cdc1988d433 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 16:59:57 +0800 Subject: [PATCH 06/26] bench: default g1 joystick e2e to motrixsim --- benchmark/benchmark_g1_joystick_numba.py | 11 +++++++++-- tests/benchmark/test_g1_joystick_numba_benchmark.py | 3 ++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index 1037afd8b..87e70159b 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -59,6 +59,7 @@ QUICK_THREADS = [2, 4] DEFAULT_NUM_ENVS = [512, 1024, 2048, 4096, 8192, 16384, 32768] DEFAULT_E2E_NUM_ENVS = [1024, 2048, 4096, 8192, 16384, 32768] +DEFAULT_E2E_CASE = "sac/g1_walk_flat/motrixsim" DEFAULT_POSE_WEIGHTS = [ 0.01, 1.0, @@ -518,6 +519,7 @@ def _timing_mean_ms(result: Any, key: str) -> float | None: def _run_e2e_collector_pair( *, + case_name: str, num_envs: int, warmup_steps: int, measure_steps: int, @@ -533,7 +535,6 @@ def _run_e2e_collector_pair( """ from benchmark.benchmark_offpolicy_collector_active import _build_and_run_case - case_name = "sac/g1_walk_flat/mujoco" common = { "warmup_steps": warmup_steps, "measure_steps": measure_steps, @@ -609,6 +610,7 @@ def _best_threads_for_profile( def _run_e2e_collector_sweep( *, + case_name: str, num_envs: list[int], warmup_steps: int, measure_steps: int, @@ -621,6 +623,7 @@ def _run_e2e_collector_sweep( print(f"e2e collector case: num_envs={env_count} numba_threads={numba_threads}") records.extend( _run_e2e_collector_pair( + case_name=case_name, num_envs=env_count, warmup_steps=warmup_steps, measure_steps=measure_steps, @@ -896,6 +899,7 @@ def write_report( "scope": "G1 joystick reward+termination hot slice; synthetic backend arrays", "e2e_enabled": args.e2e, "e2e_num_envs": args.e2e_num_envs, + "e2e_case": args.e2e_case, "e2e_warmup_steps": args.e2e_warmup_steps, "e2e_measure_steps": args.e2e_measure_steps, "e2e_numba_threads_source": "best sac_default hot-slice thread per num_env", @@ -973,7 +977,7 @@ def write_report( "## End-to-end collector comparison", "", "Not run. Pass `--e2e` to add a real off-policy collector active-window A/B", - "comparison for `sac/g1_walk_flat/mujoco` with `numba_acceleration=false/true`.", + f"comparison for `{args.e2e_case}` with `numba_acceleration=false/true`.", "This collector comparison does not run learner updates.", ] ) @@ -1034,6 +1038,7 @@ def parse_args() -> argparse.Namespace: help="Also run a real off-policy collector active-window baseline vs numba comparison.", ) parser.add_argument("--e2e-num-envs", nargs="+", type=int, default=DEFAULT_E2E_NUM_ENVS) + parser.add_argument("--e2e-case", default=DEFAULT_E2E_CASE) parser.add_argument("--e2e-warmup-steps", type=int, default=2) parser.add_argument("--e2e-measure-steps", type=int, default=8) parser.add_argument( @@ -1136,9 +1141,11 @@ def main() -> None: print("=" * 80) print("End-to-end collector comparison: off-policy collector active window") print("=" * 80) + print(f"e2e case={args.e2e_case}") print(f"e2e num_envs={args.e2e_num_envs}") print(f"e2e numba threads from sac_default hot-slice best: {selected_threads}") e2e_records = _run_e2e_collector_sweep( + case_name=args.e2e_case, num_envs=args.e2e_num_envs, warmup_steps=args.e2e_warmup_steps, measure_steps=args.e2e_measure_steps, diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index 0ec422558..64928fa76 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -23,7 +23,7 @@ def test_g1_joystick_numba_benchmark_builds_records_and_matches_numpy() -> None: def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: record = bench.EndToEndCase( - case="sac/g1_walk_flat/mujoco", + case=bench.DEFAULT_E2E_CASE, path="training_collector_numba", num_envs=64, warmup_steps=1, @@ -46,6 +46,7 @@ def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: assert payload["numba_threads"] == 4 assert "training_collector_numba" in table assert "1.25x" in table + assert "motrixsim" in table def test_g1_joystick_numba_benchmark_selects_best_hot_slice_threads() -> None: From 9ca3d402a5e6d616fac1fb4e194d7a5363b65485 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 17:09:30 +0800 Subject: [PATCH 07/26] bench: add collector e2e timing breakdown --- benchmark/benchmark_g1_joystick_numba.py | 36 +++++++++++++++++-- .../test_g1_joystick_numba_benchmark.py | 4 +++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index 87e70159b..e0274abc8 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -174,8 +174,11 @@ class EndToEndCase: numba_threads: int | None collector_active_steps_per_sec: float total_active_ms: float + collector_step_ms: float env_step_ms: float + physics_step_ms: float | None update_state_ms: float | None + other_ms: float speedup_vs_numpy: float = 1.0 env_step_speedup_vs_numpy: float | None = None update_state_speedup_vs_numpy: float | None = None @@ -504,8 +507,11 @@ def _e2e_case_to_dict(case: EndToEndCase) -> dict[str, Any]: "numba_threads": case.numba_threads, "collector_active_steps_per_sec": case.collector_active_steps_per_sec, "total_active_ms": case.total_active_ms, + "collector_step_ms": case.collector_step_ms, "env_step_ms": case.env_step_ms, + "physics_step_ms": case.physics_step_ms, "update_state_ms": case.update_state_ms, + "other_ms": case.other_ms, "speedup_vs_numpy": case.speedup_vs_numpy, "env_step_speedup_vs_numpy": case.env_step_speedup_vs_numpy, "update_state_speedup_vs_numpy": case.update_state_speedup_vs_numpy, @@ -567,6 +573,18 @@ def _run_e2e_collector_pair( **common, ) env_step_ms = float(result.phase_ms_per_vector_step["env_step_ms"].mean_ms) + physics_step_ms = ( + float(result.physics_ms_per_vector_step.mean_ms) + if result.physics_ms_per_vector_step is not None + else None + ) + update_state_ms = _timing_mean_ms(result, "update_state_ms") + collector_step_ms = float(result.total_active_ms) / float(result.measure_steps) + other_ms = collector_step_ms + if physics_step_ms is not None: + other_ms -= physics_step_ms + if update_state_ms is not None: + other_ms -= update_state_ms records.append( EndToEndCase( case=case_name, @@ -578,8 +596,11 @@ def _run_e2e_collector_pair( numba_threads=numba_threads if enabled else None, collector_active_steps_per_sec=float(result.collector_active_steps_per_sec), total_active_ms=float(result.total_active_ms), + collector_step_ms=collector_step_ms, env_step_ms=env_step_ms, - update_state_ms=_timing_mean_ms(result, "update_state_ms"), + physics_step_ms=physics_step_ms, + update_state_ms=update_state_ms, + other_ms=other_ms, ) ) @@ -679,9 +700,12 @@ def _format_e2e_table(records: list[EndToEndCase]) -> str: "threads", "collector M steps/s", "speedup", + "collector step ms", "env_step ms", - "env_step speedup", + "physics ms", "update_state ms", + "other ms", + "env_step speedup", "update_state speedup", ] rows = [] @@ -694,13 +718,16 @@ def _format_e2e_table(records: list[EndToEndCase]) -> str: "-" if record.numba_threads is None else str(record.numba_threads), f"{record.collector_active_steps_per_sec / 1e6:.3f}", f"{record.speedup_vs_numpy:.2f}x", + f"{record.collector_step_ms:.3f}", f"{record.env_step_ms:.3f}", + "-" if record.physics_step_ms is None else f"{record.physics_step_ms:.3f}", + "-" if record.update_state_ms is None else f"{record.update_state_ms:.3f}", + f"{record.other_ms:.3f}", ( "-" if record.env_step_speedup_vs_numpy is None else f"{record.env_step_speedup_vs_numpy:.2f}x" ), - "-" if record.update_state_ms is None else f"{record.update_state_ms:.3f}", ( "-" if record.update_state_speedup_vs_numpy is None @@ -964,6 +991,9 @@ def write_report( "replay writes, and collector-side bookkeeping are included. Learner updates are", "not run. The Numba variant uses the best `sac_default` hot-slice thread count", "found above for each `num_envs`.", + "`other_ms` is the collector active step remainder after subtracting reported", + "`physics_ms` and `update_state_ms`; if the backend does not report physics timing,", + "only `update_state_ms` is subtracted.", "", "```text", _format_e2e_table(e2e_records), diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index 64928fa76..8ebef5a7a 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -32,8 +32,11 @@ def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: numba_threads=4, collector_active_steps_per_sec=30_000.0, total_active_ms=4.2, + collector_step_ms=2.1, env_step_ms=1.5, + physics_step_ms=0.9, update_state_ms=0.2, + other_ms=1.0, speedup_vs_numpy=1.25, env_step_speedup_vs_numpy=1.5, update_state_speedup_vs_numpy=2.0, @@ -47,6 +50,7 @@ def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: assert "training_collector_numba" in table assert "1.25x" in table assert "motrixsim" in table + assert "physics ms" in table def test_g1_joystick_numba_benchmark_selects_best_hot_slice_threads() -> None: From 2066ef90eac27cb64648ff3f183ab5504e9bec9c Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 17:35:20 +0800 Subject: [PATCH 08/26] perf: add g1 motion tracking numba backend --- .../benchmark_g1_motion_tracking_numba.py | 1096 +++++++++++++++++ .../g1/motion_tracking_numba.py | 552 +++++++++ .../envs/motion_tracking/g1/tracking.py | 60 +- ...test_g1_motion_tracking_numba_benchmark.py | 70 ++ tests/envs/test_g1_motion_tracking_numba.py | 308 +++++ 5 files changed, 2072 insertions(+), 14 deletions(-) create mode 100644 benchmark/benchmark_g1_motion_tracking_numba.py create mode 100644 src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py create mode 100644 tests/benchmark/test_g1_motion_tracking_numba_benchmark.py create mode 100644 tests/envs/test_g1_motion_tracking_numba.py diff --git a/benchmark/benchmark_g1_motion_tracking_numba.py b/benchmark/benchmark_g1_motion_tracking_numba.py new file mode 100644 index 000000000..cf976d222 --- /dev/null +++ b/benchmark/benchmark_g1_motion_tracking_numba.py @@ -0,0 +1,1096 @@ +#!/usr/bin/env python3 +"""Benchmark the task-specific Numba backend for G1 motion tracking. + +This benchmark has two scopes: + +* default hot slice: reward plus termination, using deterministic synthetic + arrays and the real ``G1MotionTrackingEnv`` reward/termination methods; +* optional ``--e2e``: collector-side A/B through + ``benchmark_offpolicy_collector_active.py`` without learner updates. + +Run: + uv run python -m benchmark.benchmark_g1_motion_tracking_numba + uv run python benchmark/benchmark_g1_motion_tracking_numba.py + uv run python -m benchmark.benchmark_g1_motion_tracking_numba --quick + uv run python -m benchmark.benchmark_g1_motion_tracking_numba --quick --e2e +""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import time +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from statistics import mean, stdev +from typing import Any, cast + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import numpy as np +from numba import get_num_threads, set_num_threads + +try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt +except Exception: # pragma: no cover - plotting is optional in benchmark scripts + plt = None + +from benchmark.core.device_info import get_device_info_dict, get_device_info_line +from unilab.dtype_config import get_global_dtype +from unilab.envs.motion_tracking.g1.motion_tracking_numba import ( + G1MotionTrackingNumbaAccelerator, +) +from unilab.envs.motion_tracking.g1.tracking import ( + G1MotionTrackingCfg, + G1MotionTrackingEnv, + RewardConfig, +) + +NUM_ACTION = 29 +DEFAULT_THREADS = [2, 4, 8, 16, 32, 64] +QUICK_THREADS = [2, 4] +DEFAULT_NUM_ENVS = [512, 1024, 2048, 4096, 8192, 16384, 32768] +DEFAULT_E2E_NUM_ENVS = [1024, 2048, 4096, 8192, 16384, 32768] +DEFAULT_E2E_CASE = "sac/g1_motion_tracking/motrixsim" + +PPO_SCALES = { + "motion_global_root_pos": 1.0, + "motion_global_root_ori": 0.5, + "motion_body_pos": 1.0, + "motion_body_ori": 1.0, + "motion_body_lin_vel": 1.0, + "motion_body_ang_vel": 1.0, + "motion_joint_pos": 0.0, + "motion_joint_vel": 0.0, + "action_rate_l2": -0.05, + "joint_limit": -10.0, + "undesired_contacts": -0.1, +} +SAC_SCALES = { + "motion_global_root_pos": 0.5, + "motion_global_root_ori": 0.5, + "motion_body_pos": 1.0, + "motion_body_ori": 1.0, + "motion_body_lin_vel": 1.0, + "motion_body_ang_vel": 1.0, + "motion_joint_pos": 0.0, + "motion_joint_vel": 0.0, + "action_rate_l2": -0.1, + "joint_limit": -2.0, + "undesired_contacts": -0.1, +} +FULL_SUPPORTED_SCALES = { + "motion_global_root_pos": 1.0, + "motion_global_root_ori": 0.5, + "motion_body_pos": 1.0, + "motion_body_ori": 1.0, + "motion_body_lin_vel": 1.0, + "motion_body_ang_vel": 1.0, + "motion_ee_body_pos_z": 0.3, + "motion_joint_pos": 0.4, + "motion_joint_vel": 0.2, + "action_rate_l2": -0.1, + "joint_limit": -10.0, + "undesired_contacts": -0.1, +} + + +@dataclass(frozen=True) +class ProfileSpec: + name: str + scales: dict[str, float] + reward_cfg: RewardConfig + + +@dataclass +class BenchCase: + profile: str + num_envs: int + path: str + threads: int | None + mean_ms: float + min_ms: float + std_ms: float + env_per_s: float + speedup_vs_numpy: float + compile_ms: float | None = None + + +@dataclass +class EndToEndCase: + case: str + path: str + num_envs: int + warmup_steps: int + measure_steps: int + numba_acceleration: bool + numba_threads: int | None + collector_active_steps_per_sec: float + total_active_ms: float + collector_step_ms: float + env_step_ms: float + physics_step_ms: float | None + update_state_ms: float | None + other_ms: float + speedup_vs_numpy: float = 1.0 + env_step_speedup_vs_numpy: float | None = None + update_state_speedup_vs_numpy: float | None = None + + +@dataclass +class MotionDataBatch: + body_pos_w: np.ndarray + body_quat_w: np.ndarray + body_lin_vel_w: np.ndarray + body_ang_vel_w: np.ndarray + joint_pos: np.ndarray + joint_vel: np.ndarray + + +@dataclass +class SyntheticBatch: + env: G1MotionTrackingEnv + info: dict[str, Any] + motion_data: MotionDataBatch + robot_body_pos_w: np.ndarray + robot_body_quat_w: np.ndarray + robot_body_lin_vel_w: np.ndarray + robot_body_ang_vel_w: np.ndarray + dof_pos: np.ndarray + dof_vel: np.ndarray + + +@dataclass +class SyntheticCfg: + ctrl_dt: float = 0.02 + reward_config: RewardConfig = field(default_factory=RewardConfig) + body_names: tuple[str, ...] = G1MotionTrackingCfg.body_names + anchor_body_name: str = G1MotionTrackingCfg.anchor_body_name + anchor_pos_z_threshold: float = G1MotionTrackingCfg.anchor_pos_z_threshold + anchor_ori_threshold: float = G1MotionTrackingCfg.anchor_ori_threshold + ee_body_pos_z_threshold: float = G1MotionTrackingCfg.ee_body_pos_z_threshold + ee_body_names: tuple[str, ...] = G1MotionTrackingCfg.ee_body_names + undesired_contact_z_threshold: float = G1MotionTrackingCfg.undesired_contact_z_threshold + terminate_on_undesired_contacts: bool = G1MotionTrackingCfg.terminate_on_undesired_contacts + + +def make_profile_specs() -> dict[str, ProfileSpec]: + return { + "ppo_default": ProfileSpec( + name="ppo_default", + scales=PPO_SCALES, + reward_cfg=RewardConfig(scales=PPO_SCALES), + ), + "sac_default": ProfileSpec( + name="sac_default", + scales=SAC_SCALES, + reward_cfg=RewardConfig(scales=SAC_SCALES), + ), + "full_supported": ProfileSpec( + name="full_supported", + scales=FULL_SUPPORTED_SCALES, + reward_cfg=RewardConfig(scales=FULL_SUPPORTED_SCALES), + ), + } + + +def _make_fake_env(num_envs: int, reward_cfg: RewardConfig) -> G1MotionTrackingEnv: + env = cast(G1MotionTrackingEnv, object.__new__(G1MotionTrackingEnv)) + cfg = SyntheticCfg(reward_config=reward_cfg) + env._cfg = cfg + env._num_envs = num_envs + env._num_action = NUM_ACTION + env.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) + env.ee_body_indices = np.array( + [cfg.body_names.index(name) for name in cfg.ee_body_names], dtype=np.int32 + ) + ee_set = set(cfg.ee_body_names) + env.undesired_contact_body_indices = np.array( + [idx for idx, name in enumerate(cfg.body_names) if name not in ee_set], dtype=np.int32 + ) + env._has_ee_body_indices = bool(env.ee_body_indices.size) + env._has_undesired_contact_body_indices = bool(env.undesired_contact_body_indices.size) + env._n_motion_bodies = len(cfg.body_names) + env._joint_range = np.column_stack( + [np.full(NUM_ACTION, -2.5), np.full(NUM_ACTION, 2.5)] + ).astype(get_global_dtype()) + env._joint_lower = env._joint_range[:, 0] + env._joint_upper = env._joint_range[:, 1] + dtype = get_global_dtype() + n_body = env._n_motion_bodies + env.body_pos_relative_w = np.zeros((num_envs, n_body, 3), dtype=dtype) + env.body_quat_relative_w = np.zeros((num_envs, n_body, 4), dtype=dtype) + env.body_quat_relative_w[:, :, 0] = 1.0 + env._delta_pos_w = np.empty((num_envs, 3), dtype=dtype) + env._delta_ori_w = np.empty((num_envs, 4), dtype=dtype) + env._body_vec_error = np.empty((num_envs, n_body, 3), dtype=dtype) + env._env_error = np.empty((num_envs,), dtype=dtype) + env._env_error2 = np.empty((num_envs,), dtype=dtype) + env._reward_term = np.empty((num_envs,), dtype=dtype) + env._weighted_reward = np.empty((num_envs,), dtype=dtype) + env._terminated = np.empty((num_envs,), dtype=bool) + env._env_bool = np.empty((num_envs,), dtype=bool) + env._quat_error_w = np.empty((num_envs, n_body), dtype=dtype) + env._quat_error_x = np.empty((num_envs, n_body), dtype=dtype) + env._joint_error = np.empty((num_envs, NUM_ACTION), dtype=dtype) + env._joint_error_upper = np.empty((num_envs, NUM_ACTION), dtype=dtype) + env._ee_pos_error_z = np.empty((num_envs, env.ee_body_indices.size), dtype=dtype) + env._ee_terminated = np.empty((num_envs, env.ee_body_indices.size), dtype=bool) + env._undesired_contact_mask = np.empty( + (num_envs, env.undesired_contact_body_indices.size), dtype=bool + ) + env._enable_reward_log = True + env._init_reward_functions() + env._active_reward_fns = { + name: reward_fn + for name, reward_fn in env._reward_fns.items() + if env._reward_term_is_active(name) + } + return env + + +def _unit_quats(rng: np.random.Generator, shape: tuple[int, ...]) -> np.ndarray: + q = rng.standard_normal((*shape, 4)).astype(get_global_dtype()) + q /= np.linalg.norm(q, axis=-1, keepdims=True) + return q + + +def _perturb_quats(rng: np.random.Generator, q: np.ndarray, sigma: float) -> np.ndarray: + out = q + sigma * rng.standard_normal(q.shape).astype(q.dtype) + out /= np.linalg.norm(out, axis=-1, keepdims=True) + return out + + +def make_batch(num_envs: int, spec: ProfileSpec, seed: int) -> SyntheticBatch: + rng = np.random.default_rng(seed) + dtype = get_global_dtype() + n_body = len(G1MotionTrackingCfg.body_names) + + def f32(value: np.ndarray) -> np.ndarray: + return np.ascontiguousarray(value, dtype=dtype) + + target_pos = rng.uniform(-1.0, 1.0, (num_envs, n_body, 3)).astype(dtype) + target_quat = _unit_quats(rng, (num_envs, n_body)) + target_lin_vel = rng.uniform(-2.0, 2.0, (num_envs, n_body, 3)).astype(dtype) + target_ang_vel = rng.uniform(-3.0, 3.0, (num_envs, n_body, 3)).astype(dtype) + target_joint_pos = rng.uniform(-1.0, 1.0, (num_envs, NUM_ACTION)).astype(dtype) + target_joint_vel = rng.uniform(-2.0, 2.0, (num_envs, NUM_ACTION)).astype(dtype) + + motion_data = MotionDataBatch( + body_pos_w=f32(target_pos + 0.04 * rng.standard_normal((num_envs, n_body, 3))), + body_quat_w=f32(_perturb_quats(rng, target_quat, 0.02)), + body_lin_vel_w=f32(target_lin_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3))), + body_ang_vel_w=f32(target_ang_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3))), + joint_pos=f32(target_joint_pos + 0.03 * rng.standard_normal((num_envs, NUM_ACTION))), + joint_vel=f32(target_joint_vel + 0.05 * rng.standard_normal((num_envs, NUM_ACTION))), + ) + robot_body_pos_w = f32(target_pos + 0.06 * rng.standard_normal((num_envs, n_body, 3))) + robot_body_quat_w = f32(_perturb_quats(rng, target_quat, 0.03)) + robot_body_lin_vel_w = f32( + target_lin_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3)) + ) + robot_body_ang_vel_w = f32( + target_ang_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3)) + ) + dof_pos = f32(target_joint_pos + 0.05 * rng.standard_normal((num_envs, NUM_ACTION))) + dof_vel = f32(target_joint_vel + 0.1 * rng.standard_normal((num_envs, NUM_ACTION))) + current_actions = f32(rng.uniform(-1.0, 1.0, (num_envs, NUM_ACTION))) + last_actions = f32(current_actions + 0.1 * rng.standard_normal((num_envs, NUM_ACTION))) + + env = _make_fake_env(num_envs, spec.reward_cfg) + env._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + term_count = max(1, int(0.02 * num_envs)) + robot_body_pos_w[:term_count, env.anchor_body_idx, 2] += 0.6 + if spec.scales.get("undesired_contacts", 0.0) != 0.0: + robot_body_pos_w[: max(1, term_count // 2), env.undesired_contact_body_indices[0], 2] = 0.0 + + info = { + "steps": np.zeros((num_envs,), dtype=np.uint32), + "current_actions": current_actions, + "last_actions": last_actions, + "log": {}, + } + return SyntheticBatch( + env=env, + info=info, + motion_data=motion_data, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=robot_body_quat_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, + dof_pos=dof_pos, + dof_vel=dof_vel, + ) + + +def compute_numpy(batch: SyntheticBatch) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: + env = batch.env + info = {**batch.info, "log": {}} + terminated = env._compute_terminations( + batch.motion_data, batch.robot_body_pos_w, batch.robot_body_quat_w + ) + reward = env._compute_reward( + info, + batch.motion_data, + batch.robot_body_pos_w, + batch.robot_body_quat_w, + batch.robot_body_lin_vel_w, + batch.robot_body_ang_vel_w, + batch.dof_pos, + batch.dof_vel, + ) + return reward, terminated, info.get("log", {}) + + +def compute_numba( + batch: SyntheticBatch, accelerator: G1MotionTrackingNumbaAccelerator +) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: + info = {**batch.info, "log": {}} + out = accelerator.compute( + info=info, + motion_data=batch.motion_data, + ref_body_pos_w=batch.env.body_pos_relative_w, + ref_body_quat_w=batch.env.body_quat_relative_w, + robot_body_pos_w=batch.robot_body_pos_w, + robot_body_quat_w=batch.robot_body_quat_w, + robot_body_lin_vel_w=batch.robot_body_lin_vel_w, + robot_body_ang_vel_w=batch.robot_body_ang_vel_w, + dof_pos=batch.dof_pos, + dof_vel=batch.dof_vel, + scales=batch.env._cfg.reward_config.scales, + enable_log=True, + ) + return out.reward, out.terminated, out.log + + +def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: + for _ in range(warmup): + fn() + samples = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + samples.append((time.perf_counter() - t0) * 1e3) + return mean(samples), min(samples), stdev(samples) if len(samples) > 1 else 0.0 + + +def check_parity( + batch: SyntheticBatch, accelerator: G1MotionTrackingNumbaAccelerator +) -> dict[str, float]: + reward_np, terminated_np, log_np = compute_numpy(batch) + reward_nb, terminated_nb, log_nb = compute_numba(batch, accelerator) + np.testing.assert_allclose(reward_nb, reward_np, rtol=1e-4, atol=1e-5) + np.testing.assert_array_equal(terminated_nb, terminated_np) + for key, value in log_np.items(): + if key in log_nb: + np.testing.assert_allclose(log_nb[key], value, rtol=1e-3, atol=1e-6) + return { + "max_abs_reward_diff": float(np.max(np.abs(reward_nb - reward_np))), + "termination_mismatch": float(np.count_nonzero(terminated_nb != terminated_np)), + } + + +def bench_one( + *, + profile: ProfileSpec, + num_envs: int, + thread_counts: list[int], + iters: int, + warmup: int, + seed: int, +) -> tuple[list[BenchCase], dict[str, float]]: + batch = make_batch(num_envs, profile, seed) + max_threads = get_num_threads() + numpy_mean, numpy_min, numpy_std = time_call( + lambda: compute_numpy(batch), iters=iters, warmup=warmup + ) + records = [ + BenchCase( + profile=profile.name, + num_envs=num_envs, + path="numpy_dispatch", + threads=None, + mean_ms=numpy_mean, + min_ms=numpy_min, + std_ms=numpy_std, + env_per_s=num_envs / (numpy_mean * 1e-3), + speedup_vs_numpy=1.0, + ) + ] + + compile_driver = G1MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=1) + t0 = time.perf_counter() + compute_numba(batch, compile_driver) + compile_ms = (time.perf_counter() - t0) * 1e3 + parity = check_parity(batch, compile_driver) + + for threads in [1, *thread_counts]: + if threads > max_threads: + continue + accelerator = G1MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=threads) + numba_mean, numba_min, numba_std = time_call( + lambda: compute_numba(batch, accelerator), iters=iters, warmup=warmup + ) + records.append( + BenchCase( + profile=profile.name, + num_envs=num_envs, + path="numba_accelerator", + threads=threads, + mean_ms=numba_mean, + min_ms=numba_min, + std_ms=numba_std, + env_per_s=num_envs / (numba_mean * 1e-3), + speedup_vs_numpy=numpy_mean / numba_mean, + compile_ms=compile_ms if threads == 1 else None, + ) + ) + set_num_threads(max_threads) + return records, parity + + +def _case_to_dict(case: BenchCase) -> dict[str, Any]: + return { + "profile": case.profile, + "num_envs": case.num_envs, + "path": case.path, + "threads": case.threads, + "mean_ms": case.mean_ms, + "min_ms": case.min_ms, + "std_ms": case.std_ms, + "env_per_s": case.env_per_s, + "speedup_vs_numpy": case.speedup_vs_numpy, + "compile_ms": case.compile_ms, + } + + +def _e2e_case_to_dict(case: EndToEndCase) -> dict[str, Any]: + return { + "case": case.case, + "path": case.path, + "num_envs": case.num_envs, + "warmup_steps": case.warmup_steps, + "measure_steps": case.measure_steps, + "numba_acceleration": case.numba_acceleration, + "numba_threads": case.numba_threads, + "collector_active_steps_per_sec": case.collector_active_steps_per_sec, + "total_active_ms": case.total_active_ms, + "collector_step_ms": case.collector_step_ms, + "env_step_ms": case.env_step_ms, + "physics_step_ms": case.physics_step_ms, + "update_state_ms": case.update_state_ms, + "other_ms": case.other_ms, + "speedup_vs_numpy": case.speedup_vs_numpy, + "env_step_speedup_vs_numpy": case.env_step_speedup_vs_numpy, + "update_state_speedup_vs_numpy": case.update_state_speedup_vs_numpy, + } + + +def _timing_mean_ms(result: Any, key: str) -> float | None: + stat = result.env_step_timing_ms_per_vector_step.get(key) + return float(stat.mean_ms) if stat is not None else None + + +def _run_e2e_collector_pair( + *, + case_name: str, + num_envs: int, + warmup_steps: int, + measure_steps: int, + numba_threads: int | None, +) -> list[EndToEndCase]: + from benchmark.benchmark_offpolicy_collector_active import _build_and_run_case + + common = { + "warmup_steps": warmup_steps, + "measure_steps": measure_steps, + "replay_capacity_steps": max(2, measure_steps + warmup_steps + 1), + "num_envs": num_envs, + } + variants = [ + ("training_collector_numpy", False, ["++env.numba_acceleration=false"]), + ( + "training_collector_numba", + True, + [ + "++env.numba_acceleration=true", + f"++env.numba_num_threads={numba_threads}" if numba_threads is not None else "", + ], + ), + ] + records: list[EndToEndCase] = [] + for path, enabled, overrides in variants: + result = _build_and_run_case( + case_name, + extra_overrides=[override for override in overrides if override], + **common, + ) + env_step_ms = float(result.phase_ms_per_vector_step["env_step_ms"].mean_ms) + physics_step_ms = ( + float(result.physics_ms_per_vector_step.mean_ms) + if result.physics_ms_per_vector_step is not None + else None + ) + update_state_ms = _timing_mean_ms(result, "update_state_ms") + collector_step_ms = float(result.total_active_ms) / float(result.measure_steps) + other_ms = collector_step_ms + if physics_step_ms is not None: + other_ms -= physics_step_ms + if update_state_ms is not None: + other_ms -= update_state_ms + records.append( + EndToEndCase( + case=case_name, + path=path, + num_envs=num_envs, + warmup_steps=warmup_steps, + measure_steps=measure_steps, + numba_acceleration=enabled, + numba_threads=numba_threads if enabled else None, + collector_active_steps_per_sec=float(result.collector_active_steps_per_sec), + total_active_ms=float(result.total_active_ms), + collector_step_ms=collector_step_ms, + env_step_ms=env_step_ms, + physics_step_ms=physics_step_ms, + update_state_ms=update_state_ms, + other_ms=other_ms, + ) + ) + baseline = next(record for record in records if not record.numba_acceleration) + for record in records: + record.speedup_vs_numpy = ( + record.collector_active_steps_per_sec / baseline.collector_active_steps_per_sec + ) + record.env_step_speedup_vs_numpy = ( + baseline.env_step_ms / record.env_step_ms if record.env_step_ms > 0.0 else None + ) + if baseline.update_state_ms is not None and record.update_state_ms: + record.update_state_speedup_vs_numpy = baseline.update_state_ms / record.update_state_ms + return records + + +def _best_numba_by_case(records: list[BenchCase]) -> dict[tuple[str, int], BenchCase]: + best_by_case: dict[tuple[str, int], BenchCase] = {} + for record in records: + if record.path != "numba_accelerator": + continue + key = (record.profile, record.num_envs) + if key not in best_by_case or record.mean_ms < best_by_case[key].mean_ms: + best_by_case[key] = record + return best_by_case + + +def _best_threads_for_profile( + records: list[BenchCase], *, profile: str, num_envs: list[int] +) -> dict[int, int]: + best_by_case = _best_numba_by_case(records) + selected: dict[int, int] = {} + for env_count in num_envs: + best = best_by_case.get((profile, env_count)) + if best is not None and best.threads is not None: + selected[env_count] = int(best.threads) + return selected + + +def _run_e2e_collector_sweep( + *, + case_name: str, + num_envs: list[int], + warmup_steps: int, + measure_steps: int, + selected_threads: dict[int, int], + fallback_numba_threads: int | None, +) -> list[EndToEndCase]: + records: list[EndToEndCase] = [] + for env_count in num_envs: + numba_threads = selected_threads.get(env_count, fallback_numba_threads) + print(f"e2e collector case: num_envs={env_count} numba_threads={numba_threads}") + records.extend( + _run_e2e_collector_pair( + case_name=case_name, + num_envs=env_count, + warmup_steps=warmup_steps, + measure_steps=measure_steps, + numba_threads=numba_threads, + ) + ) + return records + + +def _format_table(records: list[BenchCase]) -> str: + headers = ["profile", "envs", "path", "threads", "mean_ms", "min_ms", "speedup", "M env/s"] + rows = [ + [ + r.profile, + str(r.num_envs), + r.path, + "-" if r.threads is None else str(r.threads), + f"{r.mean_ms:.3f}", + f"{r.min_ms:.3f}", + f"{r.speedup_vs_numpy:.2f}x", + f"{r.env_per_s / 1e6:.2f}", + ] + for r in records + ] + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + +def _format_e2e_table(records: list[EndToEndCase]) -> str: + headers = [ + "case", + "envs", + "path", + "threads", + "collector M steps/s", + "speedup", + "collector step ms", + "env_step ms", + "physics ms", + "update_state ms", + "other ms", + "env_step speedup", + "update_state speedup", + ] + rows = [] + for record in records: + rows.append( + [ + record.case, + str(record.num_envs), + record.path, + "-" if record.numba_threads is None else str(record.numba_threads), + f"{record.collector_active_steps_per_sec / 1e6:.3f}", + f"{record.speedup_vs_numpy:.2f}x", + f"{record.collector_step_ms:.3f}", + f"{record.env_step_ms:.3f}", + "-" if record.physics_step_ms is None else f"{record.physics_step_ms:.3f}", + "-" if record.update_state_ms is None else f"{record.update_state_ms:.3f}", + f"{record.other_ms:.3f}", + ( + "-" + if record.env_step_speedup_vs_numpy is None + else f"{record.env_step_speedup_vs_numpy:.2f}x" + ), + ( + "-" + if record.update_state_speedup_vs_numpy is None + else f"{record.update_state_speedup_vs_numpy:.2f}x" + ), + ] + ) + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + +def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) -> list[str]: + if plt is None or not records: + print("Plotting skipped: matplotlib is not available.") + return [] + output_dir.mkdir(parents=True, exist_ok=True) + profiles = sorted({record.profile for record in records}) + num_envs = sorted({record.num_envs for record in records}) + best_by_case = _best_numba_by_case(records) + + fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(21, 6)) + fig.suptitle(f"G1 motion tracking Numba reward+termination benchmark\n{device_info}", fontsize=13) + + ax1, ax2, ax3 = axes + for profile in profiles: + x, y, labels = [], [], [] + for env_count in num_envs: + best = best_by_case.get((profile, env_count)) + if best is None: + continue + x.append(env_count) + y.append(best.speedup_vs_numpy) + labels.append(best.threads) + if x: + ax1.plot(x, y, marker="o", label=profile) + for x_val, y_val, threads in zip(x, y, labels): + ax1.annotate( + f"{threads}T", + xy=(x_val, y_val), + xytext=(0, 6), + textcoords="offset points", + ha="center", + fontsize=8, + ) + ax1.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") + ax1.set_title("Best speedup vs numpy") + ax1.set_xlabel("num_envs") + ax1.set_ylabel("Speedup") + ax1.set_xscale("log", base=2) + ax1.set_xticks(num_envs) + ax1.set_xticklabels([str(value) for value in num_envs]) + ax1.grid(True, alpha=0.3) + ax1.legend(fontsize=8) + + for profile in profiles: + numpy_subset = sorted( + [r for r in records if r.profile == profile and r.path == "numpy_dispatch"], + key=lambda r: r.num_envs, + ) + if numpy_subset: + ax2.plot( + [r.num_envs for r in numpy_subset], + [r.mean_ms for r in numpy_subset], + marker="o", + linestyle="--", + label=f"{profile} numpy", + ) + best_subset = [ + best_by_case[(profile, env_count)] + for env_count in num_envs + if (profile, env_count) in best_by_case + ] + if best_subset: + ax2.plot( + [r.num_envs for r in best_subset], + [r.mean_ms for r in best_subset], + marker="s", + linestyle="-", + label=f"{profile} numba best", + ) + ax3.plot( + [r.num_envs for r in best_subset], + [r.env_per_s / 1e6 for r in best_subset], + marker="s", + linestyle="-", + label=f"{profile} numba best", + ) + if numpy_subset: + ax3.plot( + [r.num_envs for r in numpy_subset], + [r.env_per_s / 1e6 for r in numpy_subset], + marker="o", + linestyle="--", + label=f"{profile} numpy", + ) + ax2.set_title("Latency: numpy vs best numba") + ax2.set_xlabel("num_envs") + ax2.set_ylabel("Mean latency (ms)") + ax2.set_xscale("log", base=2) + ax2.set_yscale("log") + ax2.set_xticks(num_envs) + ax2.set_xticklabels([str(value) for value in num_envs]) + ax2.grid(True, alpha=0.3) + ax2.legend(fontsize=7) + + ax3.set_title("Throughput: numpy vs best numba") + ax3.set_xlabel("num_envs") + ax3.set_ylabel("Throughput (M env/s)") + ax3.set_xscale("log", base=2) + ax3.set_xticks(num_envs) + ax3.set_xticklabels([str(value) for value in num_envs]) + ax3.grid(True, alpha=0.3) + ax3.legend(fontsize=7) + + fig.tight_layout(rect=(0, 0, 1, 0.9)) + summary_path = output_dir / "g1_motion_tracking_numba_summary.png" + fig.savefig(summary_path, dpi=160) + plt.close(fig) + print(f"Saved summary plot: {summary_path.resolve()}") + return [str(summary_path.resolve())] + + +def write_report( + *, + output_dir: Path, + records: list[BenchCase], + parity: dict[str, dict[str, float]], + e2e_records: list[EndToEndCase], + args: argparse.Namespace, +) -> None: + output_dir.mkdir(parents=True, exist_ok=True) + device_info_line = get_device_info_line() + plot_paths = save_plots(records, output_dir, device_info=device_info_line) + meta = { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "python": platform.python_version(), + "platform": platform.platform(), + "device_info": get_device_info_dict(), + "iters": args.iters, + "warmup": args.warmup, + "profiles": args.profiles, + "num_envs": args.num_envs, + "threads": args.threads, + "requested_threads": args.threads, + "measured_threads": getattr(args, "measured_threads", None), + "skipped_threads": getattr(args, "skipped_threads", None), + "numba_max_threads": getattr(args, "numba_max_threads", None), + "scope": "G1 motion tracking reward+termination hot slice; synthetic arrays", + "e2e_enabled": args.e2e, + "e2e_num_envs": args.e2e_num_envs, + "e2e_case": args.e2e_case, + "e2e_warmup_steps": args.e2e_warmup_steps, + "e2e_measure_steps": args.e2e_measure_steps, + "e2e_numba_threads_source": "best sac_default hot-slice thread per num_env", + } + payload = { + "meta": meta, + "results": [_case_to_dict(record) for record in records], + "end_to_end_results": [_e2e_case_to_dict(record) for record in e2e_records], + "parity": parity, + "plots": plot_paths, + } + json_path = output_dir / "results.json" + json_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + + best_by_case = _best_numba_by_case(records) + summary_lines = [ + "# G1 motion tracking Numba benchmark", + "", + "Scope: reward plus termination for `G1MotionTrackingEnv.update_state`, using", + "deterministic synthetic arrays. Physics stepping, observation assembly, reset,", + "policy inference, learner work, and replay are out of scope for the hot slice.", + "", + "## Numba-specific hot slice", + "", + "Profile meanings:", + "", + "- `ppo_default`: PPO G1 motion-tracking owner-config reward scales.", + "- `sac_default`: SAC G1 motion-tracking owner-config reward scales; this chooses", + " Numba threads for the collector A/B run.", + "- `full_supported`: synthetic stress profile with every reward term supported by", + " `motion_tracking_numba.py` enabled.", + "", + ] + for key in sorted(best_by_case): + best = best_by_case[key] + summary_lines.append( + f"- `{best.profile}` / {best.num_envs} envs: best {best.mean_ms:.3f} ms " + f"at {best.threads} threads, {best.speedup_vs_numpy:.2f}x vs numpy " + f"({best.env_per_s / 1e6:.2f}M env/s)." + ) + if plot_paths: + summary_lines.extend(["", "## Plots", ""]) + for path in plot_paths: + rel_path = Path(path).name + title = rel_path.removesuffix(".png").replace("_", " ") + summary_lines.append(f"![{title}]({rel_path})") + summary_lines.append("") + if e2e_records: + summary_lines.extend( + [ + "", + "## End-to-end collector comparison", + "", + "This section mirrors `benchmark_offpolicy_collector_active.py`: Hydra owner", + "config, `create_env`, actor action sampling, `env.step`, terminal-observation", + "handling, replay writes, and collector bookkeeping are included. Learner", + "updates are not run. The Numba variant uses the best `sac_default` hot-slice", + "thread count found above for each `num_envs`.", + "`other_ms` is the collector active step remainder after subtracting reported", + "`physics_ms` and `update_state_ms`; if the backend does not report physics", + "timing, only `update_state_ms` is subtracted.", + "", + "```text", + _format_e2e_table(e2e_records), + "```", + ] + ) + else: + summary_lines.extend( + [ + "", + "## End-to-end collector comparison", + "", + "Not run. Pass `--e2e` to add a real off-policy collector active-window A/B", + f"comparison for `{args.e2e_case}` with `numba_acceleration=false/true`.", + "This collector comparison does not run learner updates.", + ] + ) + summary_lines.extend( + [ + "", + "## Detailed Results", + "", + "```text", + _format_table(records), + "```", + "", + "## Parity", + "", + "Reward is checked with `rtol=1e-4, atol=1e-5`; termination is exact.", + "", + "```json", + json.dumps(parity, indent=2), + "```", + "", + "## Interpretation", + "", + "- `numba 1 thread` isolates fusion/codegen benefit over numpy reward functions.", + "- Higher thread counts add row-parallel speedup over the same fused kernel.", + "- Hot-slice speedup is an upper bound for collector speedup because collector", + " timing also includes physics, observation assembly, reset/RNG, policy inference,", + " replay, and bookkeeping.", + "- The optional collector comparison still excludes learner updates.", + ] + ) + md_path = output_dir / "report.md" + md_path.write_text("\n".join(summary_lines) + "\n", encoding="utf-8") + print(f"Saved JSON: {json_path.resolve()}") + print(f"Saved report: {md_path.resolve()}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--profiles", + nargs="+", + default=["ppo_default", "sac_default", "full_supported"], + choices=sorted(make_profile_specs()), + ) + parser.add_argument("--num-envs", nargs="+", type=int, default=DEFAULT_NUM_ENVS) + parser.add_argument("--threads", nargs="+", type=int, default=None) + parser.add_argument("--iters", type=int, default=80) + parser.add_argument("--warmup", type=int, default=8) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--e2e", + action="store_true", + help="Also run a real off-policy collector active-window baseline vs numba comparison.", + ) + parser.add_argument("--e2e-num-envs", nargs="+", type=int, default=DEFAULT_E2E_NUM_ENVS) + parser.add_argument("--e2e-case", default=DEFAULT_E2E_CASE) + parser.add_argument("--e2e-warmup-steps", type=int, default=2) + parser.add_argument("--e2e-measure-steps", type=int, default=8) + parser.add_argument( + "--e2e-numba-threads", + type=int, + default=None, + help="Fallback Numba thread count when a requested e2e num_env lacks hot-slice data.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("benchmark/outputs/g1_motion_tracking_numba"), + ) + parser.add_argument( + "--quick", + action="store_true", + help="Short smoke run: sac_default at 512 and 2048 envs.", + ) + args = parser.parse_args() + if args.quick: + args.profiles = ["sac_default"] + args.num_envs = [512, 2048] + args.e2e_num_envs = [1024, 2048] + if args.threads is None: + args.threads = QUICK_THREADS + args.iters = 10 + args.warmup = 2 + elif args.threads is None: + args.threads = DEFAULT_THREADS + return args + + +def main() -> None: + args = parse_args() + specs = make_profile_specs() + all_records: list[BenchCase] = [] + e2e_records: list[EndToEndCase] = [] + parity: dict[str, dict[str, float]] = {} + max_threads = get_num_threads() + args.numba_max_threads = max_threads + args.measured_threads = sorted({1, *(threads for threads in args.threads if threads <= max_threads)}) + args.skipped_threads = sorted({threads for threads in args.threads if threads > max_threads}) + + print("=" * 80) + print("G1 motion tracking Numba benchmark: reward + termination") + print("=" * 80) + print(f"host numba threads: {max_threads}") + print( + f"profiles={args.profiles} num_envs={args.num_envs} " + f"requested_threads={args.threads} measured_threads={args.measured_threads}" + ) + if args.skipped_threads: + print(f"skipped threads above numba max: {args.skipped_threads}") + for profile_name in args.profiles: + spec = specs[profile_name] + for num_envs in args.num_envs: + records, parity_result = bench_one( + profile=spec, + num_envs=num_envs, + thread_counts=args.threads, + iters=args.iters, + warmup=args.warmup, + seed=args.seed, + ) + all_records.extend(records) + parity[f"{profile_name}:{num_envs}"] = parity_result + print() + print(_format_table(records)) + + if args.e2e: + missing_e2e_hot_slice = [ + num_envs + for num_envs in args.e2e_num_envs + if ("sac_default", num_envs) not in _best_numba_by_case(all_records) + ] + if missing_e2e_hot_slice: + print() + print("=" * 80) + print("Completing sac_default hot-slice data for e2e thread selection") + print("=" * 80) + for num_envs in missing_e2e_hot_slice: + records, parity_result = bench_one( + profile=specs["sac_default"], + num_envs=num_envs, + thread_counts=args.threads, + iters=args.iters, + warmup=args.warmup, + seed=args.seed, + ) + all_records.extend(records) + parity[f"sac_default:{num_envs}"] = parity_result + print() + print(_format_table(records)) + selected_threads = _best_threads_for_profile( + all_records, profile="sac_default", num_envs=args.e2e_num_envs + ) + e2e_records = _run_e2e_collector_sweep( + case_name=args.e2e_case, + num_envs=args.e2e_num_envs, + warmup_steps=args.e2e_warmup_steps, + measure_steps=args.e2e_measure_steps, + selected_threads=selected_threads, + fallback_numba_threads=args.e2e_numba_threads, + ) + print() + print(_format_e2e_table(e2e_records)) + + write_report( + output_dir=args.output_dir, + records=all_records, + parity=parity, + e2e_records=e2e_records, + args=args, + ) + + +if __name__ == "__main__": + main() diff --git a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py b/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py new file mode 100644 index 000000000..4ba7ed0d6 --- /dev/null +++ b/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py @@ -0,0 +1,552 @@ +"""Optional Numba hot path for the G1 motion-tracking task. + +This module is deliberately task-owned. It mirrors the reward and termination +math in ``tracking.py`` while keeping the env/backend contracts unchanged. +Importing it is safe when ``numba`` is not installed; constructing the +accelerator is not. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any, Mapping + +import numpy as np + +try: # pragma: no cover - exercised in environments with numba installed + from numba import get_num_threads, get_thread_id, njit, prange, set_num_threads + + NUMBA_AVAILABLE = True +except Exception: # pragma: no cover - default test env may not install numba + get_num_threads = get_thread_id = njit = prange = set_num_threads = None # type: ignore[assignment] + NUMBA_AVAILABLE = False + + +TERM_ORDER: tuple[str, ...] = ( + "motion_global_root_pos", + "motion_global_root_ori", + "motion_body_pos", + "motion_body_ori", + "motion_body_lin_vel", + "motion_body_ang_vel", + "motion_ee_body_pos_z", + "motion_joint_pos", + "motion_joint_vel", + "action_rate_l2", + "joint_limit", + "undesired_contacts", +) +TERM_INDEX = {name: i for i, name in enumerate(TERM_ORDER)} +SUPPORTED_TERMS = frozenset(TERM_ORDER) + + +@dataclass(frozen=True) +class G1MotionTrackingNumbaResult: + reward: np.ndarray + terminated: np.ndarray + log: dict[str, float] + + +def _active_terms(scales: Mapping[str, float]) -> frozenset[str]: + return frozenset(name for name, scale in scales.items() if scale != 0.0) + + +def unsupported_terms(scales: Mapping[str, float]) -> frozenset[str]: + """Return nonzero reward terms this task-specific kernel cannot compute.""" + return _active_terms(scales) - SUPPORTED_TERMS + + +def is_available(scales: Mapping[str, float]) -> bool: + return NUMBA_AVAILABLE and not unsupported_terms(scales) + + +if NUMBA_AVAILABLE: + + def _dev(fn): + return njit(inline="always", fastmath=True, cache=True, nogil=True)(fn) + + @_dev + def _exp_reward(error, std): + return math.exp(error / -(std * std)) + + @_dev + def _quat_angle_sq_i(q1, q2, body_idx, i): + dot = ( + q1[i, body_idx, 0] * q2[i, body_idx, 0] + + q1[i, body_idx, 1] * q2[i, body_idx, 1] + + q1[i, body_idx, 2] * q2[i, body_idx, 2] + + q1[i, body_idx, 3] * q2[i, body_idx, 3] + ) + dot = abs(dot) + if dot > 1.0: + dot = 1.0 + angle = 2.0 * math.acos(dot) + return angle * angle + + @_dev + def motion_global_root_pos_i(motion_pos, robot_pos, anchor, std, i): + dx = motion_pos[i, anchor, 0] - robot_pos[i, anchor, 0] + dy = motion_pos[i, anchor, 1] - robot_pos[i, anchor, 1] + dz = motion_pos[i, anchor, 2] - robot_pos[i, anchor, 2] + return _exp_reward(dx * dx + dy * dy + dz * dz, std) + + @_dev + def motion_global_root_ori_i(motion_quat, robot_quat, anchor, std, i): + return _exp_reward(_quat_angle_sq_i(motion_quat, robot_quat, anchor, i), std) + + @_dev + def _mean_body_xyz_sq_error_i(reference, actual, n_body, i): + acc = 0.0 + for body_idx in range(n_body): + dx = reference[i, body_idx, 0] - actual[i, body_idx, 0] + dy = reference[i, body_idx, 1] - actual[i, body_idx, 1] + dz = reference[i, body_idx, 2] - actual[i, body_idx, 2] + acc += dx * dx + dy * dy + dz * dz + return acc / n_body + + @_dev + def motion_body_pos_i(reference, actual, n_body, std, i): + return _exp_reward(_mean_body_xyz_sq_error_i(reference, actual, n_body, i), std) + + @_dev + def motion_body_ori_i(reference, actual, n_body, std, i): + acc = 0.0 + for body_idx in range(n_body): + acc += _quat_angle_sq_i(reference, actual, body_idx, i) + return _exp_reward(acc / n_body, std) + + @_dev + def motion_body_lin_vel_i(motion_vel, robot_vel, n_body, std, i): + return _exp_reward(_mean_body_xyz_sq_error_i(motion_vel, robot_vel, n_body, i), std) + + @_dev + def motion_body_ang_vel_i(motion_vel, robot_vel, n_body, std, i): + return _exp_reward(_mean_body_xyz_sq_error_i(motion_vel, robot_vel, n_body, i), std) + + @_dev + def motion_ee_body_pos_z_i(reference, actual, ee_indices, std, i): + if ee_indices.shape[0] == 0: + return 0.0 + acc = 0.0 + for idx in range(ee_indices.shape[0]): + body_idx = ee_indices[idx] + dz = reference[i, body_idx, 2] - actual[i, body_idx, 2] + acc += dz * dz + return _exp_reward(acc / ee_indices.shape[0], std) + + @_dev + def _mean_joint_sq_error_i(reference, actual, n_action, i): + acc = 0.0 + for j in range(n_action): + d = reference[i, j] - actual[i, j] + acc += d * d + return acc / n_action + + @_dev + def motion_joint_pos_i(motion_joint_pos, dof_pos, n_action, std, i): + return _exp_reward(_mean_joint_sq_error_i(motion_joint_pos, dof_pos, n_action, i), std) + + @_dev + def motion_joint_vel_i(motion_joint_vel, dof_vel, n_action, std, i): + return _exp_reward(_mean_joint_sq_error_i(motion_joint_vel, dof_vel, n_action, i), std) + + @_dev + def action_rate_l2_i(current_actions, last_actions, n_action, i): + acc = 0.0 + for j in range(n_action): + d = current_actions[i, j] - last_actions[i, j] + acc += d * d + return acc + + @_dev + def joint_limit_i(dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i): + if not has_joint_limits: + return 0.0 + acc = 0.0 + for j in range(n_action): + low = joint_lower[j] - dof_pos[i, j] + if low < 0.0: + low = 0.0 + high = dof_pos[i, j] - joint_upper[j] + if high < 0.0: + high = 0.0 + v = low + high + acc += v * v + return acc + + @_dev + def undesired_contacts_i(robot_body_pos, undesired_indices, undesired_contact_z_threshold, i): + acc = 0.0 + for idx in range(undesired_indices.shape[0]): + if robot_body_pos[i, undesired_indices[idx], 2] < undesired_contact_z_threshold: + acc += 1.0 + return acc + + @_dev + def terminated_i( + motion_pos, + motion_quat, + ref_pos, + robot_pos, + robot_quat, + anchor, + ee_indices, + undesired_indices, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + i, + ): + if abs(motion_pos[i, anchor, 2] - robot_pos[i, anchor, 2]) > anchor_pos_z_threshold: + return True + if anchor_ori_threshold < 2.0: + motion_gravity_z = ( + 2.0 + * ( + motion_quat[i, anchor, 1] * motion_quat[i, anchor, 1] + + motion_quat[i, anchor, 2] * motion_quat[i, anchor, 2] + ) + - 1.0 + ) + robot_gravity_z = ( + 2.0 + * ( + robot_quat[i, anchor, 1] * robot_quat[i, anchor, 1] + + robot_quat[i, anchor, 2] * robot_quat[i, anchor, 2] + ) + - 1.0 + ) + if abs(motion_gravity_z - robot_gravity_z) > anchor_ori_threshold: + return True + for idx in range(ee_indices.shape[0]): + body_idx = ee_indices[idx] + if abs(ref_pos[i, body_idx, 2] - robot_pos[i, body_idx, 2]) > ee_body_pos_z_threshold: + return True + if terminate_on_undesired_contacts: + for idx in range(undesired_indices.shape[0]): + if robot_pos[i, undesired_indices[idx], 2] < undesired_contact_z_threshold: + return True + return False + + @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _compute_reward_termination_kernel( + motion_body_pos_w, + motion_body_quat_w, + motion_body_lin_vel_w, + motion_body_ang_vel_w, + motion_joint_pos, + motion_joint_vel, + ref_body_pos_w, + ref_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + current_actions, + last_actions, + joint_lower, + joint_upper, + scale, + std, + anchor, + ee_indices, + undesired_indices, + ctrl_dt, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + has_joint_limits, + reward, + terminated, + log_scratch, + ): + n = reward.shape[0] + n_body = robot_body_pos_w.shape[1] + n_action = dof_pos.shape[1] + for i in prange(n): + tid = get_thread_id() + total = 0.0 + + w = motion_global_root_pos_i( + motion_body_pos_w, robot_body_pos_w, anchor, std[0], i + ) * scale[0] + total += w + log_scratch[tid, 0] += w + + w = motion_global_root_ori_i( + motion_body_quat_w, robot_body_quat_w, anchor, std[1], i + ) * scale[1] + total += w + log_scratch[tid, 1] += w + + w = motion_body_pos_i(ref_body_pos_w, robot_body_pos_w, n_body, std[2], i) * scale[2] + total += w + log_scratch[tid, 2] += w + + w = motion_body_ori_i(ref_body_quat_w, robot_body_quat_w, n_body, std[3], i) * scale[3] + total += w + log_scratch[tid, 3] += w + + w = motion_body_lin_vel_i( + motion_body_lin_vel_w, robot_body_lin_vel_w, n_body, std[4], i + ) * scale[4] + total += w + log_scratch[tid, 4] += w + + w = motion_body_ang_vel_i( + motion_body_ang_vel_w, robot_body_ang_vel_w, n_body, std[5], i + ) * scale[5] + total += w + log_scratch[tid, 5] += w + + w = motion_ee_body_pos_z_i( + ref_body_pos_w, robot_body_pos_w, ee_indices, std[6], i + ) * scale[6] + total += w + log_scratch[tid, 6] += w + + w = motion_joint_pos_i(motion_joint_pos, dof_pos, n_action, std[7], i) * scale[7] + total += w + log_scratch[tid, 7] += w + + w = motion_joint_vel_i(motion_joint_vel, dof_vel, n_action, std[8], i) * scale[8] + total += w + log_scratch[tid, 8] += w + + w = action_rate_l2_i(current_actions, last_actions, n_action, i) * scale[9] + total += w + log_scratch[tid, 9] += w + + w = joint_limit_i( + dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i + ) * scale[10] + total += w + log_scratch[tid, 10] += w + + w = ( + undesired_contacts_i( + robot_body_pos_w, undesired_indices, undesired_contact_z_threshold, i + ) + * scale[11] + ) + total += w + log_scratch[tid, 11] += w + + reward[i] = total * ctrl_dt + terminated[i] = terminated_i( + motion_body_pos_w, + motion_body_quat_w, + ref_body_pos_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + ee_indices, + undesired_indices, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + i, + ) + + +class G1MotionTrackingNumbaAccelerator: + """Driver that keeps config-derived arrays and calls the fused kernel.""" + + def __init__( + self, + *, + num_envs: int, + num_action: int, + ctrl_dt: float, + reward_config: Any, + anchor_body_idx: int, + ee_body_indices: np.ndarray, + undesired_contact_body_indices: np.ndarray, + joint_lower: np.ndarray | None, + joint_upper: np.ndarray | None, + anchor_pos_z_threshold: float, + anchor_ori_threshold: float, + ee_body_pos_z_threshold: float, + undesired_contact_z_threshold: float, + terminate_on_undesired_contacts: bool, + num_threads: int | None = None, + ) -> None: + self.num_envs = int(num_envs) + self.num_action = int(num_action) + self.ctrl_dt = float(ctrl_dt) + self.anchor_body_idx = int(anchor_body_idx) + self.ee_body_indices = np.asarray(ee_body_indices, dtype=np.int32) + self.undesired_contact_body_indices = np.asarray( + undesired_contact_body_indices, dtype=np.int32 + ) + self.anchor_pos_z_threshold = float(anchor_pos_z_threshold) + self.anchor_ori_threshold = float(anchor_ori_threshold) + self.ee_body_pos_z_threshold = float(ee_body_pos_z_threshold) + self.undesired_contact_z_threshold = float(undesired_contact_z_threshold) + self.terminate_on_undesired_contacts = bool(terminate_on_undesired_contacts) + self.num_threads = num_threads + self.has_joint_limits = joint_lower is not None and joint_upper is not None + if self.has_joint_limits: + self.joint_lower = np.asarray(joint_lower, dtype=np.float64) + self.joint_upper = np.asarray(joint_upper, dtype=np.float64) + else: + self.joint_lower = np.zeros((self.num_action,), dtype=np.float64) + self.joint_upper = np.zeros((self.num_action,), dtype=np.float64) + self.scale = np.zeros((len(TERM_ORDER),), dtype=np.float64) + self.std = self._build_std_vector(reward_config) + self._zero_actions = np.zeros((self.num_envs, self.num_action), dtype=np.float64) + + @classmethod + def from_env( + cls, env: Any, num_threads: int | None = None + ) -> "G1MotionTrackingNumbaAccelerator": + if not NUMBA_AVAILABLE: + raise RuntimeError( + "G1MotionTracking numba_acceleration=True requires numba. Install it or run " + "through `uv run --with numba ...`; disable numba_acceleration to use the " + "numpy path." + ) + unsupported = unsupported_terms(env._cfg.reward_config.scales) + if unsupported: + raise ValueError( + "G1MotionTracking Numba accelerator does not support active reward terms " + f"{sorted(unsupported)}. Disable numba_acceleration or add these terms to " + "src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py." + ) + return cls( + num_envs=env.num_envs, + num_action=env._num_action, + ctrl_dt=env._cfg.ctrl_dt, + reward_config=env._cfg.reward_config, + anchor_body_idx=env.anchor_body_idx, + ee_body_indices=env.ee_body_indices, + undesired_contact_body_indices=env.undesired_contact_body_indices, + joint_lower=env._joint_lower, + joint_upper=env._joint_upper, + anchor_pos_z_threshold=env._cfg.anchor_pos_z_threshold, + anchor_ori_threshold=env._cfg.anchor_ori_threshold, + ee_body_pos_z_threshold=env._cfg.ee_body_pos_z_threshold, + undesired_contact_z_threshold=env._cfg.undesired_contact_z_threshold, + terminate_on_undesired_contacts=env._cfg.terminate_on_undesired_contacts, + num_threads=num_threads, + ) + + def _build_std_vector(self, reward_config: Any) -> np.ndarray: + per_term = { + "motion_global_root_pos": reward_config.std_root_pos, + "motion_global_root_ori": reward_config.std_root_ori, + "motion_body_pos": reward_config.std_body_pos, + "motion_body_ori": reward_config.std_body_ori, + "motion_body_lin_vel": reward_config.std_body_lin_vel, + "motion_body_ang_vel": reward_config.std_body_ang_vel, + "motion_ee_body_pos_z": reward_config.std_body_pos, + "motion_joint_pos": reward_config.std_joint_pos, + "motion_joint_vel": reward_config.std_joint_vel, + "action_rate_l2": 0.0, + "joint_limit": 0.0, + "undesired_contacts": 0.0, + } + return np.array([per_term[name] for name in TERM_ORDER], dtype=np.float64) + + def _sync_scales(self, scales: Mapping[str, float]) -> None: + unsupported = unsupported_terms(scales) + if unsupported: + raise ValueError( + "G1MotionTracking Numba accelerator does not support active reward terms " + f"{sorted(unsupported)}. Disable numba_acceleration or add these terms to " + "src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py." + ) + self.scale.fill(0.0) + for name, value in scales.items(): + idx = TERM_INDEX.get(name) + if idx is not None: + self.scale[idx] = float(value) + + def compute( + self, + *, + info: dict[str, Any], + motion_data: Any, + ref_body_pos_w: np.ndarray, + ref_body_quat_w: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + robot_body_lin_vel_w: np.ndarray, + robot_body_ang_vel_w: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + scales: Mapping[str, float], + enable_log: bool, + ) -> G1MotionTrackingNumbaResult: + if not NUMBA_AVAILABLE: + raise RuntimeError( + "G1MotionTracking Numba accelerator was constructed while numba is " + "unavailable; this indicates an invalid accelerator state." + ) + self._sync_scales(scales) + if self.num_threads is not None: + set_num_threads(self.num_threads) + + dtype = dof_pos.dtype + current_actions = np.asarray(info.get("current_actions", self._zero_actions), dtype=dtype) + last_actions = np.asarray(info.get("last_actions", self._zero_actions), dtype=dtype) + reward = np.empty((dof_pos.shape[0],), dtype=dtype) + terminated = np.empty((dof_pos.shape[0],), dtype=np.bool_) + log_scratch = np.zeros((get_num_threads(), len(TERM_ORDER)), dtype=np.float64) + + _compute_reward_termination_kernel( + motion_data.body_pos_w, + motion_data.body_quat_w, + motion_data.body_lin_vel_w, + motion_data.body_ang_vel_w, + motion_data.joint_pos, + motion_data.joint_vel, + ref_body_pos_w, + ref_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + current_actions, + last_actions, + self.joint_lower, + self.joint_upper, + self.scale, + self.std, + self.anchor_body_idx, + self.ee_body_indices, + self.undesired_contact_body_indices, + self.ctrl_dt, + self.anchor_pos_z_threshold, + self.anchor_ori_threshold, + self.ee_body_pos_z_threshold, + self.undesired_contact_z_threshold, + self.terminate_on_undesired_contacts, + self.has_joint_limits, + reward, + terminated, + log_scratch, + ) + + step_count = info.get("steps") + should_log = enable_log and ( + int(step_count[0]) % 4 == 0 if isinstance(step_count, np.ndarray) else True + ) + log = {} if should_log else info.get("log", {}) + if should_log: + term_sums = log_scratch.sum(axis=0) + for idx, name in enumerate(TERM_ORDER): + if self.scale[idx] != 0.0: + log[f"reward/{name}"] = float(term_sums[idx] / dof_pos.shape[0]) + return G1MotionTrackingNumbaResult(reward=reward, terminated=terminated, log=log) diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index eac620b60..3da83c6fe 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -185,6 +185,8 @@ class G1MotionTrackingCfg(G1BaseCfg): ) undesired_contact_z_threshold: float = 0.05 terminate_on_undesired_contacts: bool = False + numba_acceleration: bool = False + numba_num_threads: int | None = None @registry.envcfg("G1MotionTracking") @@ -633,6 +635,15 @@ def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"): for name, reward_fn in self._reward_fns.items() if self._reward_term_is_active(name) } + self._numba_accelerator = None + if cfg.numba_acceleration: + from unilab.envs.motion_tracking.g1.motion_tracking_numba import ( + G1MotionTrackingNumbaAccelerator, + ) + + self._numba_accelerator = G1MotionTrackingNumbaAccelerator.from_env( + self, num_threads=cfg.numba_num_threads + ) self._clip_end_truncated = np.zeros((num_envs,), dtype=bool) def _effective_default_angles(self, env_ids: np.ndarray | None = None) -> np.ndarray: @@ -834,24 +845,45 @@ def update_state(self, state: NpEnvState) -> NpEnvState: # Compute relative body transforms (for observations and rewards) self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) - # Compute terminations - terminated = self._compute_terminations(motion_data, robot_body_pos_w, robot_body_quat_w) + if self._numba_accelerator is not None: + numba_result = self._numba_accelerator.compute( + info=state.info, + motion_data=motion_data, + ref_body_pos_w=self.body_pos_relative_w, + ref_body_quat_w=self.body_quat_relative_w, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=robot_body_quat_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, + dof_pos=dof_pos, + dof_vel=dof_vel, + scales=self._cfg.reward_config.scales, + enable_log=self._enable_reward_log, + ) + terminated = numba_result.terminated + reward = numba_result.reward + state.info["log"] = numba_result.log + else: + # Compute terminations + terminated = self._compute_terminations( + motion_data, robot_body_pos_w, robot_body_quat_w + ) + + # Compute reward + reward = self._compute_reward( + state.info, + motion_data, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ) # Update failure statistics for adaptive sampling self.motion_sampler.update_failure_stats(terminated) - # Compute reward - reward = self._compute_reward( - state.info, - motion_data, - robot_body_pos_w, - robot_body_quat_w, - robot_body_lin_vel_w, - robot_body_ang_vel_w, - dof_pos, - dof_vel, - ) - # Compute observations obs = self._compute_obs( state.info, diff --git a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py new file mode 100644 index 000000000..9d4f3c097 --- /dev/null +++ b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +from benchmark import benchmark_g1_motion_tracking_numba as bench + + +def test_g1_motion_tracking_numba_benchmark_builds_records_and_matches_numpy() -> None: + spec = bench.make_profile_specs()["sac_default"] + + records, parity = bench.bench_one( + profile=spec, + num_envs=64, + thread_counts=[1], + iters=1, + warmup=0, + seed=0, + ) + + assert parity["termination_mismatch"] == 0.0 + assert parity["max_abs_reward_diff"] < 1.0e-5 + assert {record.path for record in records} == {"numpy_dispatch", "numba_accelerator"} + assert any(record.path == "numba_accelerator" and record.threads == 1 for record in records) + + +def test_g1_motion_tracking_numba_benchmark_formats_end_to_end_records() -> None: + record = bench.EndToEndCase( + case=bench.DEFAULT_E2E_CASE, + path="training_collector_numba", + num_envs=64, + warmup_steps=1, + measure_steps=2, + numba_acceleration=True, + numba_threads=4, + collector_active_steps_per_sec=30_000.0, + total_active_ms=4.2, + collector_step_ms=2.1, + env_step_ms=1.5, + physics_step_ms=0.9, + update_state_ms=0.2, + other_ms=1.0, + speedup_vs_numpy=1.25, + env_step_speedup_vs_numpy=1.5, + update_state_speedup_vs_numpy=2.0, + ) + + payload = bench._e2e_case_to_dict(record) + table = bench._format_e2e_table([record]) + + assert payload["numba_acceleration"] is True + assert payload["numba_threads"] == 4 + assert "training_collector_numba" in table + assert "1.25x" in table + assert "motrixsim" in table + assert "physics ms" in table + + +def test_g1_motion_tracking_numba_benchmark_selects_best_hot_slice_threads() -> None: + records = [ + bench.BenchCase("sac_default", 1024, "numpy_dispatch", None, 1.0, 1.0, 0.0, 1000.0, 1.0), + bench.BenchCase( + "sac_default", 1024, "numba_accelerator", 2, 0.8, 0.8, 0.0, 1250.0, 1.25 + ), + bench.BenchCase( + "sac_default", 1024, "numba_accelerator", 4, 0.6, 0.6, 0.0, 1666.0, 1.67 + ), + bench.BenchCase("ppo_default", 1024, "numba_accelerator", 8, 0.5, 0.5, 0.0, 2000.0, 2.0), + ] + + assert bench._best_threads_for_profile(records, profile="sac_default", num_envs=[1024]) == { + 1024: 4 + } diff --git a/tests/envs/test_g1_motion_tracking_numba.py b/tests/envs/test_g1_motion_tracking_numba.py new file mode 100644 index 000000000..c70c7550e --- /dev/null +++ b/tests/envs/test_g1_motion_tracking_numba.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, cast + +import numpy as np +import pytest + +from unilab.envs.motion_tracking.g1.motion_tracking_numba import ( + NUMBA_AVAILABLE, + G1MotionTrackingNumbaAccelerator, + unsupported_terms, +) +from unilab.envs.motion_tracking.g1.tracking import ( + G1MotionTrackingCfg, + G1MotionTrackingEnv, + RewardConfig, +) + + +def test_unsupported_terms_only_reports_active_unknown_terms(): + assert unsupported_terms({"motion_body_pos": 1.0, "custom": 0.0}) == frozenset() + assert unsupported_terms({"motion_body_pos": 1.0, "custom": 1.0}) == frozenset({"custom"}) + + +def test_numba_accelerator_requires_numba_when_declared(monkeypatch): + import unilab.envs.motion_tracking.g1.motion_tracking_numba as motion_tracking_numba + + monkeypatch.setattr(motion_tracking_numba, "NUMBA_AVAILABLE", False) + with pytest.raises(RuntimeError, match="requires numba"): + G1MotionTrackingNumbaAccelerator.from_env(_make_env(8, include_undesired=False)) + + +def test_numba_accelerator_rejects_unsupported_active_reward_terms(): + env = _make_env(8, include_undesired=False) + env._cfg.reward_config.scales = {"motion_body_pos": 1.0, "custom": 1.0} + if NUMBA_AVAILABLE: + with pytest.raises(ValueError, match="does not support active reward terms"): + G1MotionTrackingNumbaAccelerator.from_env(env) + + +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_motion_tracking_numba_reward_termination_parity(): + n = 512 + env = _make_env(n, include_undesired=True) + motion_data, robot_state, info = _make_batch(n, seed=0) + ( + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ) = robot_state + + env._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + terminated_np = env._compute_terminations(motion_data, robot_body_pos_w, robot_body_quat_w) + reward_np = env._compute_reward( + info, + motion_data, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ) + log_np = dict(info["log"]) + + accel = G1MotionTrackingNumbaAccelerator.from_env(env, num_threads=2) + out = accel.compute( + info={ + "steps": np.zeros((n,), dtype=np.uint32), + "current_actions": info["current_actions"], + "last_actions": info["last_actions"], + }, + motion_data=motion_data, + ref_body_pos_w=env.body_pos_relative_w, + ref_body_quat_w=env.body_quat_relative_w, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=robot_body_quat_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, + dof_pos=dof_pos, + dof_vel=dof_vel, + scales=env._cfg.reward_config.scales, + enable_log=True, + ) + + np.testing.assert_allclose(out.reward, reward_np, rtol=1e-4, atol=1e-5) + np.testing.assert_array_equal(out.terminated, terminated_np) + assert set(out.log) == set(log_np) + for key in log_np: + assert out.log[key] == pytest.approx(log_np[key], rel=1e-3, abs=1e-6) + + +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_motion_tracking_numba_term_py_funcs_match_numpy_math(): + from unilab.envs.motion_tracking.g1 import motion_tracking_numba as T + + n = 4 + env = _make_env(n, include_undesired=True) + motion_data, robot_state, _ = _make_batch(n, seed=1) + robot_body_pos_w, robot_body_quat_w, _, _, dof_pos, dof_vel = robot_state + env._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + i = 2 + + diff = motion_data.body_pos_w[i, env.anchor_body_idx] - robot_body_pos_w[i, env.anchor_body_idx] + expected_root_pos = np.exp( + -np.sum(diff * diff) / (env._cfg.reward_config.std_root_pos**2) + ) + assert T.motion_global_root_pos_i.py_func( + motion_data.body_pos_w, + robot_body_pos_w, + env.anchor_body_idx, + env._cfg.reward_config.std_root_pos, + i, + ) == pytest.approx(expected_root_pos) + + body_diff = env.body_pos_relative_w[i] - robot_body_pos_w[i] + expected_body_pos = np.exp( + -np.sum(body_diff * body_diff) / body_diff.shape[0] / (env._cfg.reward_config.std_body_pos**2) + ) + assert T.motion_body_pos_i.py_func( + env.body_pos_relative_w, + robot_body_pos_w, + env._n_motion_bodies, + env._cfg.reward_config.std_body_pos, + i, + ) == pytest.approx(expected_body_pos) + + expected_action_rate = np.sum((dof_pos[i] - dof_vel[i]) ** 2) + assert T.action_rate_l2_i.py_func(dof_pos, dof_vel, env._num_action, i) == pytest.approx( + expected_action_rate + ) + + expected_undesired = np.sum( + robot_body_pos_w[i, env.undesired_contact_body_indices, 2] + < env._cfg.undesired_contact_z_threshold + ) + assert T.undesired_contacts_i.py_func( + robot_body_pos_w, + env.undesired_contact_body_indices, + env._cfg.undesired_contact_z_threshold, + i, + ) == pytest.approx(float(expected_undesired)) + + +@dataclass +class _MotionData: + body_pos_w: np.ndarray + body_quat_w: np.ndarray + body_lin_vel_w: np.ndarray + body_ang_vel_w: np.ndarray + joint_pos: np.ndarray + joint_vel: np.ndarray + + +@dataclass +class _Cfg: + ctrl_dt: float = 0.02 + reward_config: RewardConfig = field(default_factory=RewardConfig) + body_names: tuple[str, ...] = G1MotionTrackingCfg.body_names + anchor_body_name: str = G1MotionTrackingCfg.anchor_body_name + anchor_pos_z_threshold: float = G1MotionTrackingCfg.anchor_pos_z_threshold + anchor_ori_threshold: float = G1MotionTrackingCfg.anchor_ori_threshold + ee_body_pos_z_threshold: float = G1MotionTrackingCfg.ee_body_pos_z_threshold + ee_body_names: tuple[str, ...] = G1MotionTrackingCfg.ee_body_names + undesired_contact_z_threshold: float = G1MotionTrackingCfg.undesired_contact_z_threshold + terminate_on_undesired_contacts: bool = False + + +def _make_env(n: int, *, include_undesired: bool) -> Any: + env = cast(Any, object.__new__(G1MotionTrackingEnv)) + cfg = _Cfg() + cfg.reward_config.scales = { + "motion_global_root_pos": 0.5, + "motion_global_root_ori": 0.5, + "motion_body_pos": 1.0, + "motion_body_ori": 1.0, + "motion_body_lin_vel": 1.0, + "motion_body_ang_vel": 1.0, + "motion_ee_body_pos_z": 0.3, + "motion_joint_pos": 0.4, + "motion_joint_vel": 0.2, + "action_rate_l2": -0.1, + "joint_limit": -2.0, + "undesired_contacts": -0.1 if include_undesired else 0.0, + } + env._cfg = cfg + env._num_envs = n + env._num_action = 29 + env.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) + env.ee_body_indices = np.array([cfg.body_names.index(name) for name in cfg.ee_body_names]) + ee_set = set(cfg.ee_body_names) + env.undesired_contact_body_indices = np.array( + [idx for idx, name in enumerate(cfg.body_names) if name not in ee_set], dtype=np.int32 + ) + env._has_ee_body_indices = True + env._has_undesired_contact_body_indices = True + env._n_motion_bodies = len(cfg.body_names) + env._joint_range = np.column_stack( + [np.full(env._num_action, -2.0), np.full(env._num_action, 2.0)] + ).astype(np.float32) + env._joint_lower = env._joint_range[:, 0] + env._joint_upper = env._joint_range[:, 1] + env.body_pos_relative_w = np.zeros((n, env._n_motion_bodies, 3), dtype=np.float32) + env.body_quat_relative_w = np.zeros((n, env._n_motion_bodies, 4), dtype=np.float32) + env.body_quat_relative_w[..., 0] = 1.0 + env._delta_pos_w = np.empty((n, 3), dtype=np.float32) + env._delta_ori_w = np.empty((n, 4), dtype=np.float32) + env._body_vec_error = np.empty((n, env._n_motion_bodies, 3), dtype=np.float32) + env._env_error = np.empty((n,), dtype=np.float32) + env._env_error2 = np.empty((n,), dtype=np.float32) + env._reward_term = np.empty((n,), dtype=np.float32) + env._weighted_reward = np.empty((n,), dtype=np.float32) + env._terminated = np.empty((n,), dtype=bool) + env._env_bool = np.empty((n,), dtype=bool) + env._quat_error_w = np.empty((n, env._n_motion_bodies), dtype=np.float32) + env._quat_error_x = np.empty((n, env._n_motion_bodies), dtype=np.float32) + env._joint_error = np.empty((n, env._num_action), dtype=np.float32) + env._joint_error_upper = np.empty((n, env._num_action), dtype=np.float32) + env._ee_pos_error_z = np.empty((n, env.ee_body_indices.size), dtype=np.float32) + env._ee_terminated = np.empty((n, env.ee_body_indices.size), dtype=bool) + env._undesired_contact_mask = np.empty( + (n, env.undesired_contact_body_indices.size), dtype=bool + ) + env._enable_reward_log = True + env._init_reward_functions() + env._active_reward_fns = { + name: reward_fn + for name, reward_fn in env._reward_fns.items() + if env._reward_term_is_active(name) + } + return env + + +def _unit_quats(rng: np.random.Generator, shape: tuple[int, ...]) -> np.ndarray: + q = rng.standard_normal((*shape, 4)).astype(np.float32) + q /= np.linalg.norm(q, axis=-1, keepdims=True) + return q + + +def _perturb_quats(rng: np.random.Generator, q: np.ndarray, sigma: float) -> np.ndarray: + out = q + sigma * rng.standard_normal(q.shape).astype(np.float32) + out /= np.linalg.norm(out, axis=-1, keepdims=True) + return out + + +def _make_batch(n: int, seed: int): + rng = np.random.default_rng(seed) + nb = len(G1MotionTrackingCfg.body_names) + na = 29 + target_pos = rng.uniform(-1.0, 1.0, (n, nb, 3)).astype(np.float32) + target_quat = _unit_quats(rng, (n, nb)) + target_lin_vel = rng.uniform(-2.0, 2.0, (n, nb, 3)).astype(np.float32) + target_ang_vel = rng.uniform(-3.0, 3.0, (n, nb, 3)).astype(np.float32) + target_joint_pos = rng.uniform(-1.0, 1.0, (n, na)).astype(np.float32) + target_joint_vel = rng.uniform(-2.0, 2.0, (n, na)).astype(np.float32) + + motion_data = _MotionData( + body_pos_w=np.ascontiguousarray(target_pos + 0.03 * rng.standard_normal((n, nb, 3))), + body_quat_w=np.ascontiguousarray(_perturb_quats(rng, target_quat, 0.02)), + body_lin_vel_w=np.ascontiguousarray( + target_lin_vel + 0.1 * rng.standard_normal((n, nb, 3)) + ), + body_ang_vel_w=np.ascontiguousarray( + target_ang_vel + 0.1 * rng.standard_normal((n, nb, 3)) + ), + joint_pos=np.ascontiguousarray(target_joint_pos + 0.03 * rng.standard_normal((n, na))), + joint_vel=np.ascontiguousarray(target_joint_vel + 0.05 * rng.standard_normal((n, na))), + ) + robot_body_pos_w = np.ascontiguousarray( + target_pos + 0.05 * rng.standard_normal((n, nb, 3)), dtype=np.float32 + ) + robot_body_quat_w = np.ascontiguousarray(_perturb_quats(rng, target_quat, 0.03)) + robot_body_lin_vel_w = np.ascontiguousarray( + target_lin_vel + 0.1 * rng.standard_normal((n, nb, 3)), dtype=np.float32 + ) + robot_body_ang_vel_w = np.ascontiguousarray( + target_ang_vel + 0.1 * rng.standard_normal((n, nb, 3)), dtype=np.float32 + ) + dof_pos = np.ascontiguousarray( + target_joint_pos + 0.05 * rng.standard_normal((n, na)), dtype=np.float32 + ) + dof_vel = np.ascontiguousarray( + target_joint_vel + 0.1 * rng.standard_normal((n, na)), dtype=np.float32 + ) + current_actions = rng.uniform(-1.0, 1.0, (n, na)).astype(np.float32) + last_actions = (current_actions + 0.1 * rng.standard_normal((n, na))).astype(np.float32) + robot_body_pos_w[: max(1, n // 20), 0, 2] = 0.0 + info = { + "steps": np.zeros((n,), dtype=np.uint32), + "current_actions": current_actions, + "last_actions": last_actions, + } + return ( + motion_data, + ( + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ), + info, + ) From 6f408c598ab2e90a8852b2b7254f8bbabc7f3adf Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 18:12:27 +0800 Subject: [PATCH 09/26] bench: clarify g1 numba speedup accounting --- benchmark/benchmark_g1_joystick_numba.py | 302 +++++++++++--- .../benchmark_g1_motion_tracking_numba.py | 384 +++++++++++++++--- .../test_g1_joystick_numba_benchmark.py | 75 ++++ ...test_g1_motion_tracking_numba_benchmark.py | 101 ++++- 4 files changed, 755 insertions(+), 107 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index e0274abc8..11712c428 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -161,6 +161,8 @@ class BenchCase: env_per_s: float speedup_vs_numpy: float compile_ms: float | None = None + parallel_speedup_vs_numba_1t: float | None = None + parallel_efficiency: float | None = None @dataclass @@ -478,6 +480,16 @@ def bench_one( ) ) set_num_threads(max_threads) + numba_1t = next( + record + for record in records + if record.path == "numba_accelerator" and record.threads == 1 + ) + for record in records: + if record.path != "numba_accelerator" or record.threads is None: + continue + record.parallel_speedup_vs_numba_1t = numba_1t.mean_ms / record.mean_ms + record.parallel_efficiency = record.parallel_speedup_vs_numba_1t / record.threads return records, parity @@ -493,6 +505,8 @@ def _case_to_dict(case: BenchCase) -> dict[str, Any]: "env_per_s": case.env_per_s, "speedup_vs_numpy": case.speedup_vs_numpy, "compile_ms": case.compile_ms, + "parallel_speedup_vs_numba_1t": case.parallel_speedup_vs_numba_1t, + "parallel_efficiency": case.parallel_efficiency, } @@ -662,7 +676,9 @@ def _format_table(records: list[BenchCase]) -> str: "threads", "mean_ms", "min_ms", - "speedup", + "vs numpy", + "vs numba1T", + "parallel eff", "M env/s", ] rows = [] @@ -676,6 +692,12 @@ def _format_table(records: list[BenchCase]) -> str: f"{r.mean_ms:.3f}", f"{r.min_ms:.3f}", f"{r.speedup_vs_numpy:.2f}x", + ( + "-" + if r.parallel_speedup_vs_numba_1t is None + else f"{r.parallel_speedup_vs_numba_1t:.2f}x" + ), + "-" if r.parallel_efficiency is None else f"{100.0 * r.parallel_efficiency:.1f}%", f"{r.env_per_s / 1e6:.2f}", ] ) @@ -692,6 +714,104 @@ def fmt(values: list[str]) -> str: return "\n".join(lines) +def _hot_record_for_threads( + records: list[BenchCase], *, profile: str, num_envs: int, threads: int | None +) -> BenchCase | None: + candidates = [ + record + for record in records + if record.profile == profile + and record.num_envs == num_envs + and record.path == "numba_accelerator" + ] + if not candidates: + return None + if threads is None: + return min(candidates, key=lambda record: record.mean_ms) + return next((record for record in candidates if record.threads == threads), None) + + +def _numpy_record(records: list[BenchCase], *, profile: str, num_envs: int) -> BenchCase | None: + return next( + ( + record + for record in records + if record.profile == profile + and record.num_envs == num_envs + and record.path == "numpy_dispatch" + ), + None, + ) + + +def _format_e2e_reconciliation_table( + *, + hot_records: list[BenchCase], + e2e_records: list[EndToEndCase], + profile: str = "sac_default", +) -> str: + baseline_by_env = { + record.num_envs: record + for record in e2e_records + if not record.numba_acceleration and record.update_state_ms is not None + } + rows = [] + for record in e2e_records: + if not record.numba_acceleration or record.update_state_ms is None: + continue + baseline = baseline_by_env.get(record.num_envs) + numpy_hot = _numpy_record(hot_records, profile=profile, num_envs=record.num_envs) + numba_hot = _hot_record_for_threads( + hot_records, + profile=profile, + num_envs=record.num_envs, + threads=record.numba_threads, + ) + if baseline is None or baseline.update_state_ms is None or numpy_hot is None or numba_hot is None: + continue + hot_saved_ms = numpy_hot.mean_ms - numba_hot.mean_ms + update_saved_ms = baseline.update_state_ms - record.update_state_ms + update_base_ms = baseline.update_state_ms + rows.append( + [ + str(record.num_envs), + "-" if record.numba_threads is None else str(record.numba_threads), + f"{numpy_hot.mean_ms:.3f}", + f"{numba_hot.mean_ms:.3f}", + f"{hot_saved_ms:.3f}", + f"{baseline.update_state_ms:.3f}", + f"{record.update_state_ms:.3f}", + f"{update_saved_ms:.3f}", + f"{100.0 * hot_saved_ms / update_base_ms:.1f}%", + f"{100.0 * update_saved_ms / update_base_ms:.1f}%", + ] + ) + if not rows: + return "" + + headers = [ + "envs", + "threads", + "hot numpy ms", + "hot numba ms", + "hot saved ms", + "e2e update numpy ms", + "e2e update numba ms", + "e2e saved ms", + "hot saved / update base", + "e2e saved / update base", + ] + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + def _format_e2e_table(records: list[EndToEndCase]) -> str: headers = [ "case", @@ -759,7 +879,13 @@ def _best_numba_by_case(records: list[BenchCase]) -> dict[tuple[str, int], Bench return best_by_case -def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) -> list[str]: +def save_plots( + records: list[BenchCase], + e2e_records: list[EndToEndCase], + output_dir: Path, + *, + device_info: str, +) -> list[str]: if plt is None or not records: print("Plotting skipped: matplotlib is not available.") return [] @@ -772,7 +898,6 @@ def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(21, 6)) fig.suptitle(f"G1 joystick Numba reward+termination benchmark\n{device_info}", fontsize=13) - # 1) Best speedup vs env count. ax1 = axes[0] for profile in profiles: x = [] @@ -798,34 +923,17 @@ def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) fontsize=8, ) ax1.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") - ax1.set_title("Best speedup vs numpy dispatch") + ax1.set_title("Reward+termination: best vs numpy") ax1.set_xlabel("num_envs") - ax1.set_ylabel("Speedup") + ax1.set_ylabel("Speedup vs numpy") ax1.set_xscale("log", base=2) ax1.set_xticks(num_envs) ax1.set_xticklabels([str(value) for value in num_envs]) ax1.grid(True, alpha=0.3) ax1.legend(fontsize=8) - # 2) Latency: numpy dispatch vs best numba. ax2 = axes[1] for profile in profiles: - numpy_subset = sorted( - [ - record - for record in records - if record.profile == profile and record.path == "numpy_dispatch" - ], - key=lambda record: record.num_envs, - ) - if numpy_subset: - ax2.plot( - [record.num_envs for record in numpy_subset], - [record.mean_ms for record in numpy_subset], - marker="o", - linestyle="--", - label=f"{profile} numpy", - ) best_subset = [ best_by_case[(profile, env_count)] for env_count in num_envs @@ -834,56 +942,101 @@ def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) if best_subset: ax2.plot( [record.num_envs for record in best_subset], - [record.mean_ms for record in best_subset], + [ + 0.0 + if record.parallel_speedup_vs_numba_1t is None + else record.parallel_speedup_vs_numba_1t + for record in best_subset + ], marker="s", linestyle="-", - label=f"{profile} numba best", + label=profile, ) - ax2.set_title("Latency: numpy dispatch vs best numba") + for record in best_subset: + if record.parallel_efficiency is None: + continue + ax2.annotate( + f"{100.0 * record.parallel_efficiency:.0f}%", + xy=(record.num_envs, record.parallel_speedup_vs_numba_1t or 0.0), + xytext=(0, 6), + textcoords="offset points", + ha="center", + fontsize=8, + ) + ax2.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="1T") + ax2.set_title("Reward+termination: parallel speedup") ax2.set_xlabel("num_envs") - ax2.set_ylabel("Mean latency (ms)") + ax2.set_ylabel("Speedup vs numba 1T") ax2.set_xscale("log", base=2) - ax2.set_yscale("log") ax2.set_xticks(num_envs) ax2.set_xticklabels([str(value) for value in num_envs]) ax2.grid(True, alpha=0.3) ax2.legend(fontsize=7) - # 3) Throughput: numpy dispatch vs best numba. ax3 = axes[2] - for profile in profiles: - numpy_subset = sorted( - [ - record - for record in records - if record.profile == profile and record.path == "numpy_dispatch" - ], - key=lambda record: record.num_envs, + e2e_numba = sorted( + [ + record + for record in e2e_records + if record.numba_acceleration and record.update_state_speedup_vs_numpy is not None + ], + key=lambda record: record.num_envs, + ) + if e2e_numba: + ax3.plot( + [record.num_envs for record in e2e_numba], + [record.update_state_speedup_vs_numpy or 0.0 for record in e2e_numba], + marker="o", + linestyle="-", + label="collector update_state", ) - if numpy_subset: - ax3.plot( - [record.num_envs for record in numpy_subset], - [record.env_per_s / 1e6 for record in numpy_subset], - marker="o", - linestyle="--", - label=f"{profile} numpy", + for record in e2e_numba: + ax3.annotate( + "-" if record.numba_threads is None else f"{record.numba_threads}T", + xy=(record.num_envs, record.update_state_speedup_vs_numpy or 0.0), + xytext=(0, 6), + textcoords="offset points", + ha="center", + fontsize=8, ) - best_subset = [ - best_by_case[(profile, env_count)] - for env_count in num_envs - if (profile, env_count) in best_by_case - ] - if best_subset: - ax3.plot( - [record.num_envs for record in best_subset], - [record.env_per_s / 1e6 for record in best_subset], - marker="s", - linestyle="-", - label=f"{profile} numba best", + ax3.set_title("Collector update_state speedup") + ax3.set_ylabel("Speedup vs numpy collector") + else: + for profile in profiles: + numpy_subset = sorted( + [ + record + for record in records + if record.profile == profile and record.path == "numpy_dispatch" + ], + key=lambda record: record.num_envs, ) - ax3.set_title("Throughput: numpy dispatch vs best numba") + best_subset = [ + best_by_case[(profile, env_count)] + for env_count in num_envs + if (profile, env_count) in best_by_case + ] + if numpy_subset: + ax3.plot( + [record.num_envs for record in numpy_subset], + [record.mean_ms for record in numpy_subset], + marker="o", + linestyle="--", + label=f"{profile} numpy", + ) + if best_subset: + ax3.plot( + [record.num_envs for record in best_subset], + [record.mean_ms for record in best_subset], + marker="s", + linestyle="-", + label=f"{profile} numba best", + ) + ax3.set_title("Latency: numpy vs best numba") + ax3.set_ylabel("Mean latency (ms)") + ax3.set_yscale("log") + ax3.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") ax3.set_xlabel("num_envs") - ax3.set_ylabel("Throughput (M env/s)") ax3.set_xscale("log", base=2) ax3.set_xticks(num_envs) ax3.set_xticklabels([str(value) for value in num_envs]) @@ -908,7 +1061,7 @@ def write_report( ) -> None: output_dir.mkdir(parents=True, exist_ok=True) device_info_line = get_device_info_line() - plot_paths = save_plots(records, output_dir, device_info=device_info_line) + plot_paths = save_plots(records, e2e_records, output_dir, device_info=device_info_line) meta = { "timestamp_utc": datetime.now(timezone.utc).isoformat(), "python": platform.python_version(), @@ -952,6 +1105,12 @@ def write_report( "", "## Numba-specific hot slice", "", + "Definitions:", + "", + "- `vs numpy`: numpy reward+termination time divided by numba reward+termination time.", + "- `vs numba1T`: numba 1-thread time divided by numba N-thread time.", + "- `parallel eff`: `vs numba1T / N`; this is the only parallel-efficiency column.", + "", "This section measures only the part this task-specific Numba backend accelerates:", "`G1WalkEnv` reward dispatch plus termination inside `update_state`. It excludes", "physics, observation assembly, reset/RNG, policy inference, learner work, and replay.", @@ -1000,6 +1159,29 @@ def write_report( "```", ] ) + reconciliation_table = _format_e2e_reconciliation_table( + hot_records=records, + e2e_records=e2e_records, + profile="sac_default", + ) + if reconciliation_table: + summary_lines.extend( + [ + "", + "## E2E Reconciliation", + "", + "This compares the synthetic `sac_default` hot-slice milliseconds saved", + "with the collector-measured `update_state_ms` milliseconds saved at the", + "same `num_envs` and Numba thread count. The hot slice is only the", + "reward+termination work replaced by `joystick_numba.py`; collector", + "`update_state_ms` also includes state reads, observation assembly,", + "resets, state replacement, and bookkeeping around that hot slice.", + "", + "```text", + reconciliation_table, + "```", + ] + ) else: summary_lines.extend( [ @@ -1033,8 +1215,8 @@ def write_report( "- `numba 1 thread` isolates fusion/codegen benefit over Python reward dispatch.", "- Higher thread counts add row-parallel speedup over the same fused kernel.", "- Hot-slice speedup is an upper bound for collector speedup because collector timing", - " also includes physics, observation assembly, reset/RNG, policy inference, replay,", - " and bookkeeping.", + " also includes backend state reads, observation assembly, reset/RNG, policy inference,", + " replay, and bookkeeping.", "- The optional collector comparison still excludes learner updates.", ] ) diff --git a/benchmark/benchmark_g1_motion_tracking_numba.py b/benchmark/benchmark_g1_motion_tracking_numba.py index cf976d222..328a79135 100644 --- a/benchmark/benchmark_g1_motion_tracking_numba.py +++ b/benchmark/benchmark_g1_motion_tracking_numba.py @@ -122,6 +122,21 @@ class BenchCase: env_per_s: float speedup_vs_numpy: float compile_ms: float | None = None + parallel_speedup_vs_numba_1t: float | None = None + parallel_efficiency: float | None = None + + +@dataclass +class ComponentCase: + profile: str + num_envs: int + numba_threads: int + relative_transform_ms: float + numpy_reward_termination_ms: float + numba_reward_termination_ms: float + numpy_total_ms: float + numba_total_ms: float + speedup_vs_numpy: float @dataclass @@ -371,6 +386,12 @@ def compute_numba( return out.reward, out.terminated, out.log +def compute_relative_transforms(batch: SyntheticBatch) -> None: + batch.env._update_relative_transforms( + batch.motion_data, batch.robot_body_pos_w, batch.robot_body_quat_w + ) + + def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: for _ in range(warmup): fn() @@ -406,9 +427,12 @@ def bench_one( iters: int, warmup: int, seed: int, -) -> tuple[list[BenchCase], dict[str, float]]: +) -> tuple[list[BenchCase], ComponentCase, dict[str, float]]: batch = make_batch(num_envs, profile, seed) max_threads = get_num_threads() + relative_mean, _, _ = time_call( + lambda: compute_relative_transforms(batch), iters=iters, warmup=warmup + ) numpy_mean, numpy_min, numpy_std = time_call( lambda: compute_numpy(batch), iters=iters, warmup=warmup ) @@ -454,7 +478,35 @@ def bench_one( ) ) set_num_threads(max_threads) - return records, parity + numba_1t = next( + record + for record in records + if record.path == "numba_accelerator" and record.threads == 1 + ) + for record in records: + if record.path != "numba_accelerator" or record.threads is None: + continue + record.parallel_speedup_vs_numba_1t = numba_1t.mean_ms / record.mean_ms + record.parallel_efficiency = record.parallel_speedup_vs_numba_1t / record.threads + + best_numba = min( + (record for record in records if record.path == "numba_accelerator"), + key=lambda record: record.mean_ms, + ) + numpy_total_ms = relative_mean + numpy_mean + numba_total_ms = relative_mean + best_numba.mean_ms + component = ComponentCase( + profile=profile.name, + num_envs=num_envs, + numba_threads=int(best_numba.threads or 1), + relative_transform_ms=relative_mean, + numpy_reward_termination_ms=numpy_mean, + numba_reward_termination_ms=best_numba.mean_ms, + numpy_total_ms=numpy_total_ms, + numba_total_ms=numba_total_ms, + speedup_vs_numpy=numpy_total_ms / numba_total_ms, + ) + return records, component, parity def _case_to_dict(case: BenchCase) -> dict[str, Any]: @@ -469,6 +521,22 @@ def _case_to_dict(case: BenchCase) -> dict[str, Any]: "env_per_s": case.env_per_s, "speedup_vs_numpy": case.speedup_vs_numpy, "compile_ms": case.compile_ms, + "parallel_speedup_vs_numba_1t": case.parallel_speedup_vs_numba_1t, + "parallel_efficiency": case.parallel_efficiency, + } + + +def _component_case_to_dict(case: ComponentCase) -> dict[str, Any]: + return { + "profile": case.profile, + "num_envs": case.num_envs, + "numba_threads": case.numba_threads, + "relative_transform_ms": case.relative_transform_ms, + "numpy_reward_termination_ms": case.numpy_reward_termination_ms, + "numba_reward_termination_ms": case.numba_reward_termination_ms, + "numpy_total_ms": case.numpy_total_ms, + "numba_total_ms": case.numba_total_ms, + "speedup_vs_numpy": case.speedup_vs_numpy, } @@ -626,7 +694,18 @@ def _run_e2e_collector_sweep( def _format_table(records: list[BenchCase]) -> str: - headers = ["profile", "envs", "path", "threads", "mean_ms", "min_ms", "speedup", "M env/s"] + headers = [ + "profile", + "envs", + "path", + "threads", + "mean_ms", + "min_ms", + "vs numpy", + "vs numba1T", + "parallel eff", + "M env/s", + ] rows = [ [ r.profile, @@ -636,6 +715,12 @@ def _format_table(records: list[BenchCase]) -> str: f"{r.mean_ms:.3f}", f"{r.min_ms:.3f}", f"{r.speedup_vs_numpy:.2f}x", + ( + "-" + if r.parallel_speedup_vs_numba_1t is None + else f"{r.parallel_speedup_vs_numba_1t:.2f}x" + ), + "-" if r.parallel_efficiency is None else f"{100.0 * r.parallel_efficiency:.1f}%", f"{r.env_per_s / 1e6:.2f}", ] for r in records @@ -651,6 +736,43 @@ def fmt(values: list[str]) -> str: return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) +def _format_component_table(records: list[ComponentCase]) -> str: + headers = [ + "profile", + "envs", + "threads", + "rel ms", + "numpy reward+term ms", + "numba reward+term ms", + "numpy total ms", + "numba total ms", + "total speedup", + ] + rows = [ + [ + record.profile, + str(record.num_envs), + str(record.numba_threads), + f"{record.relative_transform_ms:.3f}", + f"{record.numpy_reward_termination_ms:.3f}", + f"{record.numba_reward_termination_ms:.3f}", + f"{record.numpy_total_ms:.3f}", + f"{record.numba_total_ms:.3f}", + f"{record.speedup_vs_numpy:.2f}x", + ] + for record in records + ] + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + def _format_e2e_table(records: list[EndToEndCase]) -> str: headers = [ "case", @@ -705,7 +827,111 @@ def fmt(values: list[str]) -> str: return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) -def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) -> list[str]: +def _hot_record_for_threads( + records: list[BenchCase], *, profile: str, num_envs: int, threads: int | None +) -> BenchCase | None: + candidates = [ + record + for record in records + if record.profile == profile + and record.num_envs == num_envs + and record.path == "numba_accelerator" + ] + if not candidates: + return None + if threads is None: + return min(candidates, key=lambda record: record.mean_ms) + return next((record for record in candidates if record.threads == threads), None) + + +def _numpy_record(records: list[BenchCase], *, profile: str, num_envs: int) -> BenchCase | None: + return next( + ( + record + for record in records + if record.profile == profile + and record.num_envs == num_envs + and record.path == "numpy_dispatch" + ), + None, + ) + + +def _format_e2e_reconciliation_table( + *, + hot_records: list[BenchCase], + e2e_records: list[EndToEndCase], + profile: str = "sac_default", +) -> str: + baseline_by_env = { + record.num_envs: record + for record in e2e_records + if not record.numba_acceleration and record.update_state_ms is not None + } + rows = [] + for record in e2e_records: + if not record.numba_acceleration or record.update_state_ms is None: + continue + baseline = baseline_by_env.get(record.num_envs) + numpy_hot = _numpy_record(hot_records, profile=profile, num_envs=record.num_envs) + numba_hot = _hot_record_for_threads( + hot_records, + profile=profile, + num_envs=record.num_envs, + threads=record.numba_threads, + ) + if baseline is None or baseline.update_state_ms is None or numpy_hot is None or numba_hot is None: + continue + hot_saved_ms = numpy_hot.mean_ms - numba_hot.mean_ms + update_saved_ms = baseline.update_state_ms - record.update_state_ms + update_base_ms = baseline.update_state_ms + rows.append( + [ + str(record.num_envs), + "-" if record.numba_threads is None else str(record.numba_threads), + f"{numpy_hot.mean_ms:.3f}", + f"{numba_hot.mean_ms:.3f}", + f"{hot_saved_ms:.3f}", + f"{baseline.update_state_ms:.3f}", + f"{record.update_state_ms:.3f}", + f"{update_saved_ms:.3f}", + f"{100.0 * hot_saved_ms / update_base_ms:.1f}%", + f"{100.0 * update_saved_ms / update_base_ms:.1f}%", + ] + ) + if not rows: + return "" + + headers = [ + "envs", + "threads", + "hot numpy ms", + "hot numba ms", + "hot saved ms", + "e2e update numpy ms", + "e2e update numba ms", + "e2e saved ms", + "hot saved / update base", + "e2e saved / update base", + ] + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + +def save_plots( + records: list[BenchCase], + component_records: list[ComponentCase], + output_dir: Path, + *, + device_info: str, +) -> list[str]: if plt is None or not records: print("Plotting skipped: matplotlib is not available.") return [] @@ -739,9 +965,9 @@ def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) fontsize=8, ) ax1.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") - ax1.set_title("Best speedup vs numpy") + ax1.set_title("Reward+termination: best vs numpy") ax1.set_xlabel("num_envs") - ax1.set_ylabel("Speedup") + ax1.set_ylabel("Speedup vs numpy") ax1.set_xscale("log", base=2) ax1.set_xticks(num_envs) ax1.set_xticklabels([str(value) for value in num_envs]) @@ -749,18 +975,6 @@ def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) ax1.legend(fontsize=8) for profile in profiles: - numpy_subset = sorted( - [r for r in records if r.profile == profile and r.path == "numpy_dispatch"], - key=lambda r: r.num_envs, - ) - if numpy_subset: - ax2.plot( - [r.num_envs for r in numpy_subset], - [r.mean_ms for r in numpy_subset], - marker="o", - linestyle="--", - label=f"{profile} numpy", - ) best_subset = [ best_by_case[(profile, env_count)] for env_count in num_envs @@ -769,39 +983,63 @@ def save_plots(records: list[BenchCase], output_dir: Path, *, device_info: str) if best_subset: ax2.plot( [r.num_envs for r in best_subset], - [r.mean_ms for r in best_subset], - marker="s", - linestyle="-", - label=f"{profile} numba best", - ) - ax3.plot( - [r.num_envs for r in best_subset], - [r.env_per_s / 1e6 for r in best_subset], + [ + 0.0 + if r.parallel_speedup_vs_numba_1t is None + else r.parallel_speedup_vs_numba_1t + for r in best_subset + ], marker="s", linestyle="-", - label=f"{profile} numba best", - ) - if numpy_subset: - ax3.plot( - [r.num_envs for r in numpy_subset], - [r.env_per_s / 1e6 for r in numpy_subset], - marker="o", - linestyle="--", - label=f"{profile} numpy", + label=profile, ) - ax2.set_title("Latency: numpy vs best numba") + for record in best_subset: + if record.parallel_efficiency is None: + continue + ax2.annotate( + f"{100.0 * record.parallel_efficiency:.0f}%", + xy=(record.num_envs, record.parallel_speedup_vs_numba_1t or 0.0), + xytext=(0, 6), + textcoords="offset points", + ha="center", + fontsize=8, + ) + ax2.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="1T") + ax2.set_title("Reward+termination: parallel speedup") ax2.set_xlabel("num_envs") - ax2.set_ylabel("Mean latency (ms)") + ax2.set_ylabel("Speedup vs numba 1T") ax2.set_xscale("log", base=2) - ax2.set_yscale("log") ax2.set_xticks(num_envs) ax2.set_xticklabels([str(value) for value in num_envs]) ax2.grid(True, alpha=0.3) ax2.legend(fontsize=7) - ax3.set_title("Throughput: numpy vs best numba") + for profile in profiles: + component_subset = sorted( + [record for record in component_records if record.profile == profile], + key=lambda record: record.num_envs, + ) + if component_subset: + ax3.plot( + [record.num_envs for record in component_subset], + [record.speedup_vs_numpy for record in component_subset], + marker="o", + linestyle="-", + label=profile, + ) + for record in component_subset: + ax3.annotate( + f"{record.numba_threads}T", + xy=(record.num_envs, record.speedup_vs_numpy), + xytext=(0, 6), + textcoords="offset points", + ha="center", + fontsize=8, + ) + ax3.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") + ax3.set_title("Relative transform + reward+termination") ax3.set_xlabel("num_envs") - ax3.set_ylabel("Throughput (M env/s)") + ax3.set_ylabel("Total speedup vs numpy") ax3.set_xscale("log", base=2) ax3.set_xticks(num_envs) ax3.set_xticklabels([str(value) for value in num_envs]) @@ -820,13 +1058,14 @@ def write_report( *, output_dir: Path, records: list[BenchCase], + component_records: list[ComponentCase], parity: dict[str, dict[str, float]], e2e_records: list[EndToEndCase], args: argparse.Namespace, ) -> None: output_dir.mkdir(parents=True, exist_ok=True) device_info_line = get_device_info_line() - plot_paths = save_plots(records, output_dir, device_info=device_info_line) + plot_paths = save_plots(records, component_records, output_dir, device_info=device_info_line) meta = { "timestamp_utc": datetime.now(timezone.utc).isoformat(), "python": platform.python_version(), @@ -841,7 +1080,7 @@ def write_report( "measured_threads": getattr(args, "measured_threads", None), "skipped_threads": getattr(args, "skipped_threads", None), "numba_max_threads": getattr(args, "numba_max_threads", None), - "scope": "G1 motion tracking reward+termination hot slice; synthetic arrays", + "scope": "G1 motion tracking reward+termination hot slice plus component reconciliation; synthetic arrays", "e2e_enabled": args.e2e, "e2e_num_envs": args.e2e_num_envs, "e2e_case": args.e2e_case, @@ -852,6 +1091,7 @@ def write_report( payload = { "meta": meta, "results": [_case_to_dict(record) for record in records], + "component_results": [_component_case_to_dict(record) for record in component_records], "end_to_end_results": [_e2e_case_to_dict(record) for record in e2e_records], "parity": parity, "plots": plot_paths, @@ -869,6 +1109,12 @@ def write_report( "", "## Numba-specific hot slice", "", + "Definitions:", + "", + "- `vs numpy`: numpy reward+termination time divided by numba reward+termination time.", + "- `vs numba1T`: numba 1-thread time divided by numba N-thread time.", + "- `parallel eff`: `vs numba1T / N`; this is the only parallel-efficiency column.", + "", "Profile meanings:", "", "- `ppo_default`: PPO G1 motion-tracking owner-config reward scales.", @@ -885,6 +1131,22 @@ def write_report( f"at {best.threads} threads, {best.speedup_vs_numpy:.2f}x vs numpy " f"({best.env_per_s / 1e6:.2f}M env/s)." ) + if component_records: + summary_lines.extend( + [ + "", + "## Component Reconciliation", + "", + "This table adds `_update_relative_transforms` to the reward+termination", + "hot slice. It is still synthetic and still excludes backend state getters,", + "observation assembly, motion sampling, reset/refresh, state replacement,", + "collector bookkeeping, and learner work.", + "", + "```text", + _format_component_table(component_records), + "```", + ] + ) if plot_paths: summary_lines.extend(["", "## Plots", ""]) for path in plot_paths: @@ -912,6 +1174,28 @@ def write_report( "```", ] ) + reconciliation_table = _format_e2e_reconciliation_table( + hot_records=records, + e2e_records=e2e_records, + profile="sac_default", + ) + if reconciliation_table: + summary_lines.extend( + [ + "", + "## E2E Reconciliation", + "", + "This compares the synthetic `sac_default` hot-slice milliseconds saved", + "with the collector-measured `update_state_ms` milliseconds saved at the", + "same `num_envs` and Numba thread count. Large hot-slice speedups can", + "translate to small `update_state` speedups when non-accelerated work", + "dominates the update-state baseline.", + "", + "```text", + reconciliation_table, + "```", + ] + ) else: summary_lines.extend( [ @@ -945,8 +1229,8 @@ def write_report( "- `numba 1 thread` isolates fusion/codegen benefit over numpy reward functions.", "- Higher thread counts add row-parallel speedup over the same fused kernel.", "- Hot-slice speedup is an upper bound for collector speedup because collector", - " timing also includes physics, observation assembly, reset/RNG, policy inference,", - " replay, and bookkeeping.", + " timing also includes backend state getters, relative transforms, observation", + " assembly, motion sampling, reset/RNG, policy inference, replay, and bookkeeping.", "- The optional collector comparison still excludes learner updates.", ] ) @@ -1012,6 +1296,7 @@ def main() -> None: args = parse_args() specs = make_profile_specs() all_records: list[BenchCase] = [] + all_component_records: list[ComponentCase] = [] e2e_records: list[EndToEndCase] = [] parity: dict[str, dict[str, float]] = {} max_threads = get_num_threads() @@ -1032,7 +1317,7 @@ def main() -> None: for profile_name in args.profiles: spec = specs[profile_name] for num_envs in args.num_envs: - records, parity_result = bench_one( + records, component, parity_result = bench_one( profile=spec, num_envs=num_envs, thread_counts=args.threads, @@ -1041,9 +1326,12 @@ def main() -> None: seed=args.seed, ) all_records.extend(records) + all_component_records.append(component) parity[f"{profile_name}:{num_envs}"] = parity_result print() print(_format_table(records)) + print() + print(_format_component_table([component])) if args.e2e: missing_e2e_hot_slice = [ @@ -1057,7 +1345,7 @@ def main() -> None: print("Completing sac_default hot-slice data for e2e thread selection") print("=" * 80) for num_envs in missing_e2e_hot_slice: - records, parity_result = bench_one( + records, component, parity_result = bench_one( profile=specs["sac_default"], num_envs=num_envs, thread_counts=args.threads, @@ -1066,9 +1354,12 @@ def main() -> None: seed=args.seed, ) all_records.extend(records) + all_component_records.append(component) parity[f"sac_default:{num_envs}"] = parity_result print() print(_format_table(records)) + print() + print(_format_component_table([component])) selected_threads = _best_threads_for_profile( all_records, profile="sac_default", num_envs=args.e2e_num_envs ) @@ -1086,6 +1377,7 @@ def main() -> None: write_report( output_dir=args.output_dir, records=all_records, + component_records=all_component_records, parity=parity, e2e_records=e2e_records, args=args, diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index 8ebef5a7a..1d1fdef20 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -19,6 +19,11 @@ def test_g1_joystick_numba_benchmark_builds_records_and_matches_numpy() -> None: assert parity["max_abs_reward_diff"] < 1.0e-5 assert {record.path for record in records} == {"numpy_dispatch", "numba_accelerator"} assert any(record.path == "numba_accelerator" and record.threads == 1 for record in records) + assert all( + record.parallel_speedup_vs_numba_1t is not None + for record in records + if record.path == "numba_accelerator" + ) def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: @@ -53,6 +58,76 @@ def test_g1_joystick_numba_benchmark_formats_end_to_end_records() -> None: assert "physics ms" in table +def test_g1_joystick_numba_benchmark_formats_e2e_reconciliation() -> None: + hot_records = [ + bench.BenchCase( + "sac_default", + 1024, + "numpy_dispatch", + None, + 1.0, + 1.0, + 0.0, + 1024.0, + 1.0, + ), + bench.BenchCase( + "sac_default", + 1024, + "numba_accelerator", + 4, + 0.25, + 0.25, + 0.0, + 4096.0, + 4.0, + ), + ] + e2e_records = [ + bench.EndToEndCase( + case=bench.DEFAULT_E2E_CASE, + path="training_collector_numpy", + num_envs=1024, + warmup_steps=1, + measure_steps=2, + numba_acceleration=False, + numba_threads=None, + collector_active_steps_per_sec=20_000.0, + total_active_ms=8.0, + collector_step_ms=4.0, + env_step_ms=3.0, + physics_step_ms=1.0, + update_state_ms=2.0, + other_ms=1.0, + ), + bench.EndToEndCase( + case=bench.DEFAULT_E2E_CASE, + path="training_collector_numba", + num_envs=1024, + warmup_steps=1, + measure_steps=2, + numba_acceleration=True, + numba_threads=4, + collector_active_steps_per_sec=25_000.0, + total_active_ms=6.4, + collector_step_ms=3.2, + env_step_ms=2.2, + physics_step_ms=1.0, + update_state_ms=1.4, + other_ms=0.8, + ), + ] + + table = bench._format_e2e_reconciliation_table( + hot_records=hot_records, + e2e_records=e2e_records, + ) + + assert "hot saved ms" in table + assert "0.750" in table + assert "30.0%" in table + + def test_g1_joystick_numba_benchmark_selects_best_hot_slice_threads() -> None: records = [ bench.BenchCase("sac_default", 1024, "numpy_dispatch", None, 1.0, 1.0, 0.0, 1000.0, 1.0), diff --git a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py index 9d4f3c097..bc2bc7df2 100644 --- a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py +++ b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py @@ -6,7 +6,7 @@ def test_g1_motion_tracking_numba_benchmark_builds_records_and_matches_numpy() -> None: spec = bench.make_profile_specs()["sac_default"] - records, parity = bench.bench_one( + records, component, parity = bench.bench_one( profile=spec, num_envs=64, thread_counts=[1], @@ -19,6 +19,14 @@ def test_g1_motion_tracking_numba_benchmark_builds_records_and_matches_numpy() - assert parity["max_abs_reward_diff"] < 1.0e-5 assert {record.path for record in records} == {"numpy_dispatch", "numba_accelerator"} assert any(record.path == "numba_accelerator" and record.threads == 1 for record in records) + assert component.profile == "sac_default" + assert component.speedup_vs_numpy > 0.0 + assert component.numpy_total_ms >= component.numpy_reward_termination_ms + assert all( + record.parallel_speedup_vs_numba_1t is not None + for record in records + if record.path == "numba_accelerator" + ) def test_g1_motion_tracking_numba_benchmark_formats_end_to_end_records() -> None: @@ -53,6 +61,97 @@ def test_g1_motion_tracking_numba_benchmark_formats_end_to_end_records() -> None assert "physics ms" in table +def test_g1_motion_tracking_numba_benchmark_formats_component_records() -> None: + record = bench.ComponentCase( + profile="sac_default", + num_envs=1024, + numba_threads=8, + relative_transform_ms=0.5, + numpy_reward_termination_ms=1.0, + numba_reward_termination_ms=0.25, + numpy_total_ms=1.5, + numba_total_ms=0.75, + speedup_vs_numpy=2.0, + ) + + payload = bench._component_case_to_dict(record) + table = bench._format_component_table([record]) + + assert payload["relative_transform_ms"] == 0.5 + assert "rel ms" in table + assert "2.00x" in table + + +def test_g1_motion_tracking_numba_benchmark_formats_e2e_reconciliation() -> None: + hot_records = [ + bench.BenchCase( + "sac_default", + 1024, + "numpy_dispatch", + None, + 1.0, + 1.0, + 0.0, + 1024.0, + 1.0, + ), + bench.BenchCase( + "sac_default", + 1024, + "numba_accelerator", + 4, + 0.25, + 0.25, + 0.0, + 4096.0, + 4.0, + ), + ] + e2e_records = [ + bench.EndToEndCase( + case=bench.DEFAULT_E2E_CASE, + path="training_collector_numpy", + num_envs=1024, + warmup_steps=1, + measure_steps=2, + numba_acceleration=False, + numba_threads=None, + collector_active_steps_per_sec=20_000.0, + total_active_ms=8.0, + collector_step_ms=4.0, + env_step_ms=3.0, + physics_step_ms=1.0, + update_state_ms=2.0, + other_ms=1.0, + ), + bench.EndToEndCase( + case=bench.DEFAULT_E2E_CASE, + path="training_collector_numba", + num_envs=1024, + warmup_steps=1, + measure_steps=2, + numba_acceleration=True, + numba_threads=4, + collector_active_steps_per_sec=25_000.0, + total_active_ms=6.4, + collector_step_ms=3.2, + env_step_ms=2.2, + physics_step_ms=1.0, + update_state_ms=1.4, + other_ms=0.8, + ), + ] + + table = bench._format_e2e_reconciliation_table( + hot_records=hot_records, + e2e_records=e2e_records, + ) + + assert "hot saved ms" in table + assert "0.750" in table + assert "30.0%" in table + + def test_g1_motion_tracking_numba_benchmark_selects_best_hot_slice_threads() -> None: records = [ bench.BenchCase("sac_default", 1024, "numpy_dispatch", None, 1.0, 1.0, 0.0, 1000.0, 1.0), From 60889e7b95999577453d5cc7c666830a7d9ade8a Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 18:26:20 +0800 Subject: [PATCH 10/26] bench: add collector numba ab comparison --- .../benchmark_offpolicy_collector_active.py | 235 +++++++++++++++++- ...st_offpolicy_collector_active_benchmark.py | 54 ++++ 2 files changed, 279 insertions(+), 10 deletions(-) diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 588d11b95..89f6eee30 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -63,6 +63,7 @@ BENCHMARK_BACKENDS = ("mujoco", "motrix") DEFAULT_COLLECTOR_CPU_THREADS = 8 COLLECTOR_CPU_THREADS_ENV = "UNILAB_COLLECTOR_TORCH_THREADS" +NUMBA_ACCELERATED_TASKS = frozenset({"g1_walk_flat", "g1_motion_tracking"}) BACKEND_ALIASES = { # UniLab's registry/backend contract uses "motrix"; "motrixsim" is the package # and benchmark-facing backend family name. @@ -142,6 +143,9 @@ class CollectorCase: algo: str task: str sim: str + variant: str + numba_acceleration: bool + numba_threads: int | None runtime_sim_backend: str command: str training_task_name: str @@ -423,6 +427,7 @@ def _build_case( algo: str, task: str, sim: str, + variant: str, replay_capacity_steps: int, actor_algo_type: str, use_layer_norm: bool, @@ -432,10 +437,14 @@ def _build_case( ) -> CollectorCase: num_envs = int(cfg.algo.num_envs) capacity_steps = max(1, int(replay_capacity_steps)) + numba_threads = OmegaConf.select(cfg, "env.numba_num_threads", default=None) return CollectorCase( algo=algo, task=task, sim=sim, + variant=variant, + numba_acceleration=bool(OmegaConf.select(cfg, "env.numba_acceleration", default=False)), + numba_threads=int(numba_threads) if numba_threads is not None else None, runtime_sim_backend=str(cfg.training.sim_backend), command=f"uv run train --algo {algo} --task {task} --sim {cfg.training.sim_backend}", training_task_name=str(cfg.training.task_name), @@ -739,6 +748,7 @@ def _build_and_run_case( replay_capacity_steps: int, num_envs: int | None, extra_overrides: list[str], + variant: str = "default", ) -> CollectorResult: from unilab.training import BackendAdapter, ensure_registries from unilab.training.seed import apply_training_seed @@ -784,6 +794,7 @@ def _build_and_run_case( algo=algo, task=task, sim=sim, + variant=variant, replay_capacity_steps=replay_capacity_steps, actor_algo_type=actor_algo_type, use_layer_norm=use_layer_norm, @@ -805,6 +816,50 @@ def _build_and_run_case( _cleanup() +def _numba_ab_overrides(*, enabled: bool, numba_threads: int | None) -> list[str]: + overrides = [f"++env.numba_acceleration={'true' if enabled else 'false'}"] + if enabled and numba_threads is not None: + overrides.append(f"++env.numba_num_threads={int(numba_threads)}") + return overrides + + +def _build_and_run_numba_ab_case( + spec: str, + *, + warmup_steps: int, + measure_steps: int, + replay_capacity_steps: int, + num_envs: int | None, + extra_overrides: list[str], + numba_threads: int | None, +) -> list[CollectorResult]: + _algo, task, _sim = _parse_case(spec) + if task not in NUMBA_ACCELERATED_TASKS: + raise ValueError( + f"--numba-ab only supports {sorted(NUMBA_ACCELERATED_TASKS)}, got task={task!r}" + ) + records = [] + for variant, enabled in ( + ("numpy_baseline", False), + ("numba_accelerated", True), + ): + records.append( + _build_and_run_case( + spec, + warmup_steps=warmup_steps, + measure_steps=measure_steps, + replay_capacity_steps=replay_capacity_steps, + num_envs=num_envs, + extra_overrides=[ + *extra_overrides, + *_numba_ab_overrides(enabled=enabled, numba_threads=numba_threads), + ], + variant=variant, + ) + ) + return records + + def _result_to_dict(result: CollectorResult) -> dict[str, Any]: data = asdict(result) data["phase_ms_per_vector_step"] = { @@ -829,6 +884,9 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: "algo", "task", "sim", + "variant", + "numba_acceleration", + "numba_threads", "runtime_sim_backend", "training_task_name", "collector_algo_type", @@ -864,6 +922,11 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: "algo": result.case.algo, "task": result.case.task, "sim": result.case.sim, + "variant": result.case.variant, + "numba_acceleration": result.case.numba_acceleration, + "numba_threads": ( + result.case.numba_threads if result.case.numba_threads is not None else "" + ), "runtime_sim_backend": result.case.runtime_sim_backend, "training_task_name": result.case.training_task_name, "collector_algo_type": result.case.collector_algo_type, @@ -1135,6 +1198,8 @@ def _format_throughput_table(results: list[CollectorResult]) -> str: "Algo", "Task", "Backend", + "Variant", + "Numba", "num_env", "Throughput env/s", "Total active ms", @@ -1161,6 +1226,12 @@ def _format_throughput_table(results: list[CollectorResult]) -> str: result.case.algo, result.case.task, result.case.runtime_sim_backend, + result.case.variant, + ( + f"on/{result.case.numba_threads}T" + if result.case.numba_acceleration + else "off" + ), f"{result.case.num_envs:,}", f"{result.collector_active_steps_per_sec:,.0f}", f"{result.total_active_ms / result.measure_steps:.3f} (100.0%)", @@ -1344,6 +1415,114 @@ def _format_env_step_breakdown_table(results: list[CollectorResult]) -> str: return _format_table(headers, rows) +def _mean_phase_ms(result: CollectorResult, key: str) -> float | None: + stat = result.phase_ms_per_vector_step.get(key) + return stat.mean_ms if stat is not None else None + + +def _mean_np_env_ms(result: CollectorResult, key: str) -> float | None: + stat = result.env_step_timing_ms_per_vector_step.get(key) + return stat.mean_ms if stat is not None else None + + +def _physics_ms(result: CollectorResult) -> float | None: + return ( + result.physics_ms_per_vector_step.mean_ms + if result.physics_ms_per_vector_step is not None + else None + ) + + +def _collector_other_ms(result: CollectorResult) -> float: + total = result.total_active_ms / float(result.measure_steps) + physics = _physics_ms(result) or 0.0 + update_state = _mean_np_env_ms(result, "update_state_ms") or 0.0 + return total - physics - update_state + + +def _format_speedup(new: float, baseline: float) -> str: + return f"{baseline / new:.2f}x" if new > 0.0 else "n/a" + + +def _format_saved(new: float, baseline: float) -> str: + return f"{baseline - new:.3f}" + + +def _format_numba_ab_table(results: list[CollectorResult]) -> str: + baseline_by_case = { + (result.case.algo, result.case.task, result.case.sim, result.case.num_envs): result + for result in results + if result.case.variant == "numpy_baseline" + } + rows = [] + for result in results: + if result.case.variant != "numba_accelerated": + continue + key = (result.case.algo, result.case.task, result.case.sim, result.case.num_envs) + baseline = baseline_by_case.get(key) + if baseline is None: + continue + base_total = baseline.total_active_ms / float(baseline.measure_steps) + new_total = result.total_active_ms / float(result.measure_steps) + base_env = _mean_phase_ms(baseline, "env_step_ms") + new_env = _mean_phase_ms(result, "env_step_ms") + base_update = _mean_np_env_ms(baseline, "update_state_ms") + new_update = _mean_np_env_ms(result, "update_state_ms") + base_physics = _physics_ms(baseline) + new_physics = _physics_ms(result) + base_other = _collector_other_ms(baseline) + new_other = _collector_other_ms(result) + rows.append( + ( + result.case.algo, + result.case.task, + result.case.runtime_sim_backend, + f"{result.case.num_envs:,}", + "-" if result.case.numba_threads is None else str(result.case.numba_threads), + f"{result.collector_active_steps_per_sec / baseline.collector_active_steps_per_sec:.2f}x", + _format_speedup(new_total, base_total), + _format_saved(new_total, base_total), + "n/a" if base_env is None or new_env is None else _format_speedup(new_env, base_env), + "n/a" if base_env is None or new_env is None else _format_saved(new_env, base_env), + ( + "n/a" + if base_update is None or new_update is None + else _format_speedup(new_update, base_update) + ), + ( + "n/a" + if base_update is None or new_update is None + else _format_saved(new_update, base_update) + ), + ( + "n/a" + if base_physics is None or new_physics is None + else _format_saved(new_physics, base_physics) + ), + _format_saved(new_other, base_other), + ) + ) + if not rows: + return "" + headers = ( + "Algo", + "Task", + "Backend", + "num_env", + "Numba T", + "Collector throughput", + "Total step speedup", + "Total saved ms", + "Env step speedup", + "Env saved ms", + "Update speedup", + "Update saved ms", + "Physics saved ms", + "Other saved ms", + ) + return _format_table(headers, rows) + + def _print_result(result: CollectorResult) -> None: case = result.case cpu_str = f"{result.cpu_util_pct:.1f}%" if result.cpu_util_pct == result.cpu_util_pct else "n/a" @@ -1475,6 +1654,20 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default=[], help="Additional Hydra override. Can be passed more than once.", ) + parser.add_argument( + "--numba-ab", + action="store_true", + help=( + "Run numpy baseline and env.numba_acceleration=true variants for supported " + "G1 tasks, then print an A/B bottleneck summary." + ), + ) + parser.add_argument( + "--numba-threads", + type=int, + default=None, + help="Numba thread count used by --numba-ab accelerated variants.", + ) parser.add_argument("--out-json", type=Path, default=DEFAULT_OUTPUT_JSON) parser.add_argument("--out-csv", type=Path, default=None) parser.add_argument("--continue-on-error", action="store_true") @@ -1502,16 +1695,30 @@ def main() -> int: errors: list[dict[str, str]] = [] for spec in specs: try: - result = _build_and_run_case( - spec, - warmup_steps=int(args.warmup_steps), - measure_steps=int(args.measure_steps), - replay_capacity_steps=int(args.replay_capacity_steps), - num_envs=args.num_envs, - extra_overrides=list(args.override), - ) - results.append(result) - _print_result(result) + if args.numba_ab: + case_results = _build_and_run_numba_ab_case( + spec, + warmup_steps=int(args.warmup_steps), + measure_steps=int(args.measure_steps), + replay_capacity_steps=int(args.replay_capacity_steps), + num_envs=args.num_envs, + extra_overrides=list(args.override), + numba_threads=args.numba_threads, + ) + else: + case_results = [ + _build_and_run_case( + spec, + warmup_steps=int(args.warmup_steps), + measure_steps=int(args.measure_steps), + replay_capacity_steps=int(args.replay_capacity_steps), + num_envs=args.num_envs, + extra_overrides=list(args.override), + ) + ] + results.extend(case_results) + for result in case_results: + _print_result(result) except Exception as exc: error = {"case": spec, "type": type(exc).__name__, "message": str(exc)} errors.append(error) @@ -1538,6 +1745,8 @@ def main() -> int: "measure_steps": args.measure_steps, "replay_capacity_steps": args.replay_capacity_steps, "override": args.override, + "numba_ab": args.numba_ab, + "numba_threads": args.numba_threads, }, "results": [_result_to_dict(result) for result in results], "errors": errors, @@ -1552,6 +1761,12 @@ def main() -> int: print("\nTask throughput (active phases; phase percentages add to 100%):") if results: print(_format_throughput_table(results)) + if args.numba_ab: + print( + "\nNumba A/B summary (baseline is env.numba_acceleration=false; positive saved ms means the numba run is faster):" + ) + table = _format_numba_ab_table(results) + print(table if table else "No paired Numba A/B results.") print( "\nEnv step breakdown (subparts of Env step; do not add Env step together with its subparts):" ) diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index 93e7e240f..a0d2a4037 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -24,6 +24,9 @@ def _make_result( algo=algo, task=task, sim=sim, + variant="default", + numba_acceleration=False, + numba_threads=None, runtime_sim_backend=runtime_sim_backend, command=f"uv run train --algo {algo} --task {task} --sim {runtime_sim_backend}", training_task_name="G1WalkFlat", @@ -153,6 +156,7 @@ def test_parse_args_defaults_to_large_env_count_and_longer_measure_window() -> N assert args.measure_steps == 100 assert args.backend == "motrix" assert not args.all_backends + assert not args.numba_ab def test_parse_args_accepts_backend_and_all_backend_modes() -> None: @@ -248,3 +252,53 @@ def test_write_csv_includes_all_phase_columns(tmp_path) -> None: assert "env_step_overhead_pct" in header for key in bench.COLLECTOR_PHASES: assert key in header + + +def test_numba_ab_overrides_enable_and_disable_acceleration() -> None: + assert bench._numba_ab_overrides(enabled=False, numba_threads=8) == [ + "++env.numba_acceleration=false" + ] + assert bench._numba_ab_overrides(enabled=True, numba_threads=8) == [ + "++env.numba_acceleration=true", + "++env.numba_num_threads=8", + ] + + +def test_numba_ab_table_reports_update_state_and_other_saved_ms() -> None: + baseline = _make_result(throughput=1000.0, include_env_step_breakdown=True) + baseline.case = bench.CollectorCase( + **{ + **baseline.case.__dict__, + "variant": "numpy_baseline", + "numba_acceleration": False, + "numba_threads": None, + } + ) + baseline.env_step_timing_ms_per_vector_step["update_state_ms"] = bench.TimingStats( + [0.5], 0.5, 0.5, 0.0, 0.5, 0.5 + ) + + accelerated = _make_result(throughput=1250.0, include_env_step_breakdown=True) + accelerated.case = bench.CollectorCase( + **{ + **accelerated.case.__dict__, + "variant": "numba_accelerated", + "numba_acceleration": True, + "numba_threads": 8, + } + ) + accelerated.total_active_ms = 0.8 + accelerated.phase_ms_per_vector_step["env_step_ms"] = bench.TimingStats( + [0.7], 0.7, 0.7, 0.0, 0.7, 0.7 + ) + accelerated.physics_ms_per_vector_step = bench.TimingStats([0.6], 0.6, 0.6, 0.0, 0.6, 0.6) + accelerated.env_step_timing_ms_per_vector_step["update_state_ms"] = bench.TimingStats( + [0.2], 0.2, 0.2, 0.0, 0.2, 0.2 + ) + + table = bench._format_numba_ab_table([baseline, accelerated]) + + assert "Update saved ms" in table + assert "Other saved ms" in table + assert "0.300" in table + assert "1.25x" in table From b54ac550d54079058a269c815a159131e018b7bc Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 19:39:35 +0800 Subject: [PATCH 11/26] perf: accelerate g1 motion tracking update_state --- .../benchmark_g1_motion_tracking_numba.py | 273 ++++++- .../g1/motion_tracking_numba.py | 767 +++++++++++++++++- .../envs/motion_tracking/g1/tracking.py | 52 +- src/unilab/utils/numba_geometry.py | 290 +++++++ ...test_g1_motion_tracking_numba_benchmark.py | 3 + tests/envs/test_g1_motion_tracking_numba.py | 93 +++ tests/utils/test_numba_geometry.py | 67 ++ tests/utils/test_utils_package_policy.py | 9 +- 8 files changed, 1458 insertions(+), 96 deletions(-) create mode 100644 src/unilab/utils/numba_geometry.py create mode 100644 tests/utils/test_numba_geometry.py diff --git a/benchmark/benchmark_g1_motion_tracking_numba.py b/benchmark/benchmark_g1_motion_tracking_numba.py index 328a79135..365f649b3 100644 --- a/benchmark/benchmark_g1_motion_tracking_numba.py +++ b/benchmark/benchmark_g1_motion_tracking_numba.py @@ -5,14 +5,14 @@ * default hot slice: reward plus termination, using deterministic synthetic arrays and the real ``G1MotionTrackingEnv`` reward/termination methods; -* optional ``--e2e``: collector-side A/B through +* default ``--e2e``: collector-side A/B through ``benchmark_offpolicy_collector_active.py`` without learner updates. Run: uv run python -m benchmark.benchmark_g1_motion_tracking_numba uv run python benchmark/benchmark_g1_motion_tracking_numba.py uv run python -m benchmark.benchmark_g1_motion_tracking_numba --quick - uv run python -m benchmark.benchmark_g1_motion_tracking_numba --quick --e2e + uv run python -m benchmark.benchmark_g1_motion_tracking_numba --no-e2e """ from __future__ import annotations @@ -134,6 +134,8 @@ class ComponentCase: relative_transform_ms: float numpy_reward_termination_ms: float numba_reward_termination_ms: float + numpy_update_state_ms: float + numba_update_state_ms: float numpy_total_ms: float numba_total_ms: float speedup_vs_numpy: float @@ -175,6 +177,8 @@ class SyntheticBatch: env: G1MotionTrackingEnv info: dict[str, Any] motion_data: MotionDataBatch + linvel: np.ndarray + gyro: np.ndarray robot_body_pos_w: np.ndarray robot_body_quat_w: np.ndarray robot_body_lin_vel_w: np.ndarray @@ -197,6 +201,15 @@ class SyntheticCfg: terminate_on_undesired_contacts: bool = G1MotionTrackingCfg.terminate_on_undesired_contacts +@dataclass +class SyntheticNoiseCfg: + level: float = 0.0 + scale_linvel: float = 0.1 + scale_gyro: float = 0.1 + scale_joint_angle: float = 0.02 + scale_joint_vel: float = 0.3 + + def make_profile_specs() -> dict[str, ProfileSpec]: return { "ppo_default": ProfileSpec( @@ -241,12 +254,23 @@ def _make_fake_env(num_envs: int, reward_cfg: RewardConfig) -> G1MotionTrackingE env._joint_upper = env._joint_range[:, 1] dtype = get_global_dtype() n_body = env._n_motion_bodies + env.default_angles = np.linspace(-0.2, 0.2, NUM_ACTION).astype(dtype) + env._actor_obs_width = env._actor_obs_dim(NUM_ACTION) + env._critic_base_obs_width = env._critic_base_obs_dim(NUM_ACTION) + env._critic_obs_width = env._critic_base_obs_width + n_body * 9 + env._cfg.noise_config = SyntheticNoiseCfg() env.body_pos_relative_w = np.zeros((num_envs, n_body, 3), dtype=dtype) env.body_quat_relative_w = np.zeros((num_envs, n_body, 4), dtype=dtype) env.body_quat_relative_w[:, :, 0] = 1.0 env._delta_pos_w = np.empty((num_envs, 3), dtype=dtype) env._delta_ori_w = np.empty((num_envs, 4), dtype=dtype) + env._motion_anchor_pos_b = np.empty((num_envs, 3), dtype=dtype) + env._motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) + env._motion_command = np.empty((num_envs, NUM_ACTION * 2), dtype=dtype) + env._joint_pos_rel = np.empty((num_envs, NUM_ACTION), dtype=dtype) + env._zero_actions = np.zeros((num_envs, NUM_ACTION), dtype=dtype) env._body_vec_error = np.empty((num_envs, n_body, 3), dtype=dtype) + env._body_vec_tmp = np.empty((num_envs, n_body, 3), dtype=dtype) env._env_error = np.empty((num_envs,), dtype=dtype) env._env_error2 = np.empty((num_envs,), dtype=dtype) env._reward_term = np.empty((num_envs,), dtype=dtype) @@ -315,6 +339,8 @@ def f32(value: np.ndarray) -> np.ndarray: robot_body_ang_vel_w = f32( target_ang_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3)) ) + linvel = f32(rng.uniform(-1.0, 1.0, (num_envs, 3))) + gyro = f32(rng.uniform(-1.0, 1.0, (num_envs, 3))) dof_pos = f32(target_joint_pos + 0.05 * rng.standard_normal((num_envs, NUM_ACTION))) dof_vel = f32(target_joint_vel + 0.1 * rng.standard_normal((num_envs, NUM_ACTION))) current_actions = f32(rng.uniform(-1.0, 1.0, (num_envs, NUM_ACTION))) @@ -337,6 +363,8 @@ def f32(value: np.ndarray) -> np.ndarray: env=env, info=info, motion_data=motion_data, + linvel=linvel, + gyro=gyro, robot_body_pos_w=robot_body_pos_w, robot_body_quat_w=robot_body_quat_w, robot_body_lin_vel_w=robot_body_lin_vel_w, @@ -392,6 +420,72 @@ def compute_relative_transforms(batch: SyntheticBatch) -> None: ) +def compute_numpy_update_state( + batch: SyntheticBatch, +) -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray, dict[str, float]]: + env = batch.env + info = {**batch.info, "log": {}} + env._update_relative_transforms( + batch.motion_data, batch.robot_body_pos_w, batch.robot_body_quat_w + ) + terminated = env._compute_terminations( + batch.motion_data, batch.robot_body_pos_w, batch.robot_body_quat_w + ) + reward = env._compute_reward( + info, + batch.motion_data, + batch.robot_body_pos_w, + batch.robot_body_quat_w, + batch.robot_body_lin_vel_w, + batch.robot_body_ang_vel_w, + batch.dof_pos, + batch.dof_vel, + ) + obs = env._compute_obs( + info, + batch.motion_data, + batch.linvel, + batch.gyro, + batch.dof_pos, + batch.dof_vel, + batch.robot_body_pos_w, + batch.robot_body_quat_w, + ) + return obs, reward, terminated, info.get("log", {}) + + +def compute_numba_update_state( + batch: SyntheticBatch, accelerator: G1MotionTrackingNumbaAccelerator +) -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray, dict[str, float]]: + info = {**batch.info, "log": {}} + noise_cfg = batch.env._cfg.noise_config + out = accelerator.compute_update_state( + info=info, + motion_data=batch.motion_data, + linvel=batch.linvel, + gyro=batch.gyro, + dof_pos=batch.dof_pos, + dof_vel=batch.dof_vel, + robot_body_pos_w=batch.robot_body_pos_w, + robot_body_quat_w=batch.robot_body_quat_w, + robot_body_lin_vel_w=batch.robot_body_lin_vel_w, + robot_body_ang_vel_w=batch.robot_body_ang_vel_w, + ref_body_pos_w=batch.env.body_pos_relative_w, + ref_body_quat_w=batch.env.body_quat_relative_w, + motion_anchor_pos_b=batch.env._motion_anchor_pos_b, + motion_anchor_ori_b=batch.env._motion_anchor_ori_b, + joint_pos_rel=batch.env._joint_pos_rel, + scales=batch.env._cfg.reward_config.scales, + enable_log=True, + noise_level=noise_cfg.level, + noise_scale_linvel=noise_cfg.scale_linvel, + noise_scale_gyro=noise_cfg.scale_gyro, + noise_scale_joint_angle=noise_cfg.scale_joint_angle, + noise_scale_joint_vel=noise_cfg.scale_joint_vel, + ) + return out.obs, out.reward, out.terminated, out.log + + def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: for _ in range(warmup): fn() @@ -413,9 +507,23 @@ def check_parity( for key, value in log_np.items(): if key in log_nb: np.testing.assert_allclose(log_nb[key], value, rtol=1e-3, atol=1e-6) + obs_np, full_reward_np, full_terminated_np, _ = compute_numpy_update_state(batch) + obs_nb, full_reward_nb, full_terminated_nb, _ = compute_numba_update_state(batch, accelerator) + np.testing.assert_allclose(full_reward_nb, full_reward_np, rtol=1e-4, atol=1e-5) + np.testing.assert_array_equal(full_terminated_nb, full_terminated_np) + np.testing.assert_allclose(obs_nb["obs"], obs_np["obs"], rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(obs_nb["critic"], obs_np["critic"], rtol=1e-5, atol=1e-5) return { "max_abs_reward_diff": float(np.max(np.abs(reward_nb - reward_np))), "termination_mismatch": float(np.count_nonzero(terminated_nb != terminated_np)), + "max_abs_update_state_reward_diff": float( + np.max(np.abs(full_reward_nb - full_reward_np)) + ), + "update_state_termination_mismatch": float( + np.count_nonzero(full_terminated_nb != full_terminated_np) + ), + "max_abs_actor_obs_diff": float(np.max(np.abs(obs_nb["obs"] - obs_np["obs"]))), + "max_abs_critic_obs_diff": float(np.max(np.abs(obs_nb["critic"] - obs_np["critic"]))), } @@ -433,6 +541,9 @@ def bench_one( relative_mean, _, _ = time_call( lambda: compute_relative_transforms(batch), iters=iters, warmup=warmup ) + numpy_update_state_mean, _, _ = time_call( + lambda: compute_numpy_update_state(batch), iters=iters, warmup=warmup + ) numpy_mean, numpy_min, numpy_std = time_call( lambda: compute_numpy(batch), iters=iters, warmup=warmup ) @@ -453,10 +564,11 @@ def bench_one( compile_driver = G1MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=1) t0 = time.perf_counter() compute_numba(batch, compile_driver) + compute_numba_update_state(batch, compile_driver) compile_ms = (time.perf_counter() - t0) * 1e3 parity = check_parity(batch, compile_driver) - for threads in [1, *thread_counts]: + for threads in dict.fromkeys([1, *thread_counts]): if threads > max_threads: continue accelerator = G1MotionTrackingNumbaAccelerator.from_env(batch.env, num_threads=threads) @@ -493,8 +605,16 @@ def bench_one( (record for record in records if record.path == "numba_accelerator"), key=lambda record: record.mean_ms, ) - numpy_total_ms = relative_mean + numpy_mean - numba_total_ms = relative_mean + best_numba.mean_ms + best_update_accelerator = G1MotionTrackingNumbaAccelerator.from_env( + batch.env, num_threads=best_numba.threads + ) + numba_update_state_mean, _, _ = time_call( + lambda: compute_numba_update_state(batch, best_update_accelerator), + iters=iters, + warmup=warmup, + ) + numpy_total_ms = numpy_update_state_mean + numba_total_ms = numba_update_state_mean component = ComponentCase( profile=profile.name, num_envs=num_envs, @@ -502,6 +622,8 @@ def bench_one( relative_transform_ms=relative_mean, numpy_reward_termination_ms=numpy_mean, numba_reward_termination_ms=best_numba.mean_ms, + numpy_update_state_ms=numpy_update_state_mean, + numba_update_state_ms=numba_update_state_mean, numpy_total_ms=numpy_total_ms, numba_total_ms=numba_total_ms, speedup_vs_numpy=numpy_total_ms / numba_total_ms, @@ -534,6 +656,8 @@ def _component_case_to_dict(case: ComponentCase) -> dict[str, Any]: "relative_transform_ms": case.relative_transform_ms, "numpy_reward_termination_ms": case.numpy_reward_termination_ms, "numba_reward_termination_ms": case.numba_reward_termination_ms, + "numpy_update_state_ms": case.numpy_update_state_ms, + "numba_update_state_ms": case.numba_update_state_ms, "numpy_total_ms": case.numpy_total_ms, "numba_total_ms": case.numba_total_ms, "speedup_vs_numpy": case.speedup_vs_numpy, @@ -744,9 +868,9 @@ def _format_component_table(records: list[ComponentCase]) -> str: "rel ms", "numpy reward+term ms", "numba reward+term ms", - "numpy total ms", - "numba total ms", - "total speedup", + "numpy update_state ms", + "numba update_state ms", + "update_state speedup", ] rows = [ [ @@ -756,8 +880,8 @@ def _format_component_table(records: list[ComponentCase]) -> str: f"{record.relative_transform_ms:.3f}", f"{record.numpy_reward_termination_ms:.3f}", f"{record.numba_reward_termination_ms:.3f}", - f"{record.numpy_total_ms:.3f}", - f"{record.numba_total_ms:.3f}", + f"{record.numpy_update_state_ms:.3f}", + f"{record.numba_update_state_ms:.3f}", f"{record.speedup_vs_numpy:.2f}x", ] for record in records @@ -773,6 +897,72 @@ def fmt(values: list[str]) -> str: return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) +def _format_hot_summary_table(records: list[BenchCase]) -> str: + best_by_case = _best_numba_by_case(records) + rows = [] + for profile in sorted({record.profile for record in records}): + best_records = [ + record for key, record in best_by_case.items() if key[0] == profile + ] + if not best_records: + continue + envs = sorted(record.num_envs for record in best_records) + speedups = [record.speedup_vs_numpy for record in best_records] + throughputs = [record.env_per_s / 1e6 for record in best_records] + rows.append( + [ + profile, + f"{envs[0]}-{envs[-1]}", + ",".join(str(thread) for thread in sorted({record.threads for record in best_records})), + f"{min(speedups):.2f}x-{max(speedups):.2f}x", + f"{max(throughputs):.2f}", + ] + ) + headers = ["profile", "env range", "best threads", "hot speedup range", "peak M env/s"] + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + +def _format_parity_summary_table(parity: dict[str, dict[str, float]]) -> str: + if not parity: + return "" + max_reward_diff = max(item["max_abs_reward_diff"] for item in parity.values()) + max_update_reward_diff = max( + item["max_abs_update_state_reward_diff"] for item in parity.values() + ) + max_actor_obs_diff = max(item["max_abs_actor_obs_diff"] for item in parity.values()) + max_critic_obs_diff = max(item["max_abs_critic_obs_diff"] for item in parity.values()) + termination_mismatch = sum(item["termination_mismatch"] for item in parity.values()) + update_termination_mismatch = sum( + item["update_state_termination_mismatch"] for item in parity.values() + ) + rows = [ + ["reward max abs diff", f"{max_reward_diff:.3e}"], + ["update_state reward max abs diff", f"{max_update_reward_diff:.3e}"], + ["actor obs max abs diff", f"{max_actor_obs_diff:.3e}"], + ["critic obs max abs diff", f"{max_critic_obs_diff:.3e}"], + ["termination mismatches", f"{termination_mismatch:.0f}"], + ["update_state termination mismatches", f"{update_termination_mismatch:.0f}"], + ] + headers = ["check", "value"] + widths = [len(h) for h in headers] + for row in rows: + for idx, value in enumerate(row): + widths[idx] = max(widths[idx], len(value)) + + def fmt(values: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + return "\n".join([fmt(headers), "-+-".join("-" * width for width in widths), *map(fmt, rows)]) + + def _format_e2e_table(records: list[EndToEndCase]) -> str: headers = [ "case", @@ -1037,9 +1227,9 @@ def save_plots( fontsize=8, ) ax3.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") - ax3.set_title("Relative transform + reward+termination") + ax3.set_title("Full update_state array path") ax3.set_xlabel("num_envs") - ax3.set_ylabel("Total speedup vs numpy") + ax3.set_ylabel("Speedup vs numpy") ax3.set_xscale("log", base=2) ax3.set_xticks(num_envs) ax3.set_xticklabels([str(value) for value in num_envs]) @@ -1080,7 +1270,10 @@ def write_report( "measured_threads": getattr(args, "measured_threads", None), "skipped_threads": getattr(args, "skipped_threads", None), "numba_max_threads": getattr(args, "numba_max_threads", None), - "scope": "G1 motion tracking reward+termination hot slice plus component reconciliation; synthetic arrays", + "scope": ( + "G1 motion tracking reward+termination hot slice plus full update_state " + "array-path reconciliation; synthetic arrays" + ), "e2e_enabled": args.e2e, "e2e_num_envs": args.e2e_num_envs, "e2e_case": args.e2e_case, @@ -1099,13 +1292,14 @@ def write_report( json_path = output_dir / "results.json" json_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") - best_by_case = _best_numba_by_case(records) summary_lines = [ "# G1 motion tracking Numba benchmark", "", "Scope: reward plus termination for `G1MotionTrackingEnv.update_state`, using", - "deterministic synthetic arrays. Physics stepping, observation assembly, reset,", - "policy inference, learner work, and replay are out of scope for the hot slice.", + "deterministic synthetic arrays. The component table additionally measures the", + "full array portion of `update_state`: relative transforms, reward, termination,", + "and actor/critic observation assembly. Backend state getters, motion sampling,", + "reset/refresh, policy inference, learner work, and replay remain out of scope.", "", "## Numba-specific hot slice", "", @@ -1123,24 +1317,23 @@ def write_report( "- `full_supported`: synthetic stress profile with every reward term supported by", " `motion_tracking_numba.py` enabled.", "", + "```text", + _format_hot_summary_table(records), + "```", + "", + "Detailed per-env and per-thread data is stored in `results.json`.", ] - for key in sorted(best_by_case): - best = best_by_case[key] - summary_lines.append( - f"- `{best.profile}` / {best.num_envs} envs: best {best.mean_ms:.3f} ms " - f"at {best.threads} threads, {best.speedup_vs_numpy:.2f}x vs numpy " - f"({best.env_per_s / 1e6:.2f}M env/s)." - ) if component_records: summary_lines.extend( [ "", "## Component Reconciliation", "", - "This table adds `_update_relative_transforms` to the reward+termination", - "hot slice. It is still synthetic and still excludes backend state getters,", - "observation assembly, motion sampling, reset/refresh, state replacement,", - "collector bookkeeping, and learner work.", + "This table measures the synthetic full array path of `update_state`:", + "`_update_relative_transforms`, reward, termination, motion-anchor", + "features, joint-relative features, and actor/critic observation assembly.", + "It still excludes backend state getters, motion sampling, reset/refresh,", + "state replacement, collector bookkeeping, and learner work.", "", "```text", _format_component_table(component_records), @@ -1207,31 +1400,24 @@ def write_report( "This collector comparison does not run learner updates.", ] ) + parity_summary = _format_parity_summary_table(parity) summary_lines.extend( [ "", - "## Detailed Results", - "", - "```text", - _format_table(records), - "```", - "", - "## Parity", + "## Parity Summary", "", "Reward is checked with `rtol=1e-4, atol=1e-5`; termination is exact.", - "", - "```json", - json.dumps(parity, indent=2), + "```text", + parity_summary, "```", "", "## Interpretation", "", "- `numba 1 thread` isolates fusion/codegen benefit over numpy reward functions.", "- Higher thread counts add row-parallel speedup over the same fused kernel.", - "- Hot-slice speedup is an upper bound for collector speedup because collector", - " timing also includes backend state getters, relative transforms, observation", - " assembly, motion sampling, reset/RNG, policy inference, replay, and bookkeeping.", - "- The optional collector comparison still excludes learner updates.", + "- The collector comparison excludes learner updates but includes collector-side", + " Hydra setup, env construction, actor sampling, `env.step`, replay writes, and", + " bookkeeping inside the measured active window.", ] ) md_path = output_dir / "report.md" @@ -1255,8 +1441,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=0) parser.add_argument( "--e2e", - action="store_true", - help="Also run a real off-policy collector active-window baseline vs numba comparison.", + action=argparse.BooleanOptionalAction, + default=True, + help="Run a real off-policy collector active-window baseline vs numba comparison.", ) parser.add_argument("--e2e-num-envs", nargs="+", type=int, default=DEFAULT_E2E_NUM_ENVS) parser.add_argument("--e2e-case", default=DEFAULT_E2E_CASE) diff --git a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py b/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py index 4ba7ed0d6..31d23bb2e 100644 --- a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py +++ b/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py @@ -14,6 +14,15 @@ import numpy as np +from unilab.utils.numba_geometry import ( + quat_angle_sq_at, + quat_gravity_z_at, + write_body_pos_relative_to_anchor_at, + write_body_quat_relative_6d_to_anchor_at, + write_relative_anchor_transform_at, + write_yaw_aligned_body_transforms_at, +) + try: # pragma: no cover - exercised in environments with numba installed from numba import get_num_threads, get_thread_id, njit, prange, set_num_threads @@ -48,6 +57,14 @@ class G1MotionTrackingNumbaResult: log: dict[str, float] +@dataclass(frozen=True) +class G1MotionTrackingNumbaUpdateStateResult: + obs: dict[str, np.ndarray] + reward: np.ndarray + terminated: np.ndarray + log: dict[str, float] + + def _active_terms(scales: Mapping[str, float]) -> frozenset[str]: return frozenset(name for name, scale in scales.items() if scale != 0.0) @@ -70,20 +87,6 @@ def _dev(fn): def _exp_reward(error, std): return math.exp(error / -(std * std)) - @_dev - def _quat_angle_sq_i(q1, q2, body_idx, i): - dot = ( - q1[i, body_idx, 0] * q2[i, body_idx, 0] - + q1[i, body_idx, 1] * q2[i, body_idx, 1] - + q1[i, body_idx, 2] * q2[i, body_idx, 2] - + q1[i, body_idx, 3] * q2[i, body_idx, 3] - ) - dot = abs(dot) - if dot > 1.0: - dot = 1.0 - angle = 2.0 * math.acos(dot) - return angle * angle - @_dev def motion_global_root_pos_i(motion_pos, robot_pos, anchor, std, i): dx = motion_pos[i, anchor, 0] - robot_pos[i, anchor, 0] @@ -93,7 +96,7 @@ def motion_global_root_pos_i(motion_pos, robot_pos, anchor, std, i): @_dev def motion_global_root_ori_i(motion_quat, robot_quat, anchor, std, i): - return _exp_reward(_quat_angle_sq_i(motion_quat, robot_quat, anchor, i), std) + return _exp_reward(quat_angle_sq_at(motion_quat, robot_quat, anchor, i), std) @_dev def _mean_body_xyz_sq_error_i(reference, actual, n_body, i): @@ -113,7 +116,7 @@ def motion_body_pos_i(reference, actual, n_body, std, i): def motion_body_ori_i(reference, actual, n_body, std, i): acc = 0.0 for body_idx in range(n_body): - acc += _quat_angle_sq_i(reference, actual, body_idx, i) + acc += quat_angle_sq_at(reference, actual, body_idx, i) return _exp_reward(acc / n_body, std) @_dev @@ -203,22 +206,8 @@ def terminated_i( if abs(motion_pos[i, anchor, 2] - robot_pos[i, anchor, 2]) > anchor_pos_z_threshold: return True if anchor_ori_threshold < 2.0: - motion_gravity_z = ( - 2.0 - * ( - motion_quat[i, anchor, 1] * motion_quat[i, anchor, 1] - + motion_quat[i, anchor, 2] * motion_quat[i, anchor, 2] - ) - - 1.0 - ) - robot_gravity_z = ( - 2.0 - * ( - robot_quat[i, anchor, 1] * robot_quat[i, anchor, 1] - + robot_quat[i, anchor, 2] * robot_quat[i, anchor, 2] - ) - - 1.0 - ) + motion_gravity_z = quat_gravity_z_at(motion_quat, anchor, i) + robot_gravity_z = quat_gravity_z_at(robot_quat, anchor, i) if abs(motion_gravity_z - robot_gravity_z) > anchor_ori_threshold: return True for idx in range(ee_indices.shape[0]): @@ -358,6 +347,553 @@ def _compute_reward_termination_kernel( ) + @_dev + def _write_reference_transforms_i( + motion_body_pos_w, + motion_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + n_body, + ref_body_pos_w, + ref_body_quat_w, + i, + ): + write_yaw_aligned_body_transforms_at( + motion_body_pos_w, + motion_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + n_body, + ref_body_pos_w, + ref_body_quat_w, + i, + ) + + @_dev + def _write_motion_anchor_i( + motion_body_pos_w, + motion_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + motion_anchor_pos_b, + motion_anchor_ori_b, + i, + ): + write_relative_anchor_transform_at( + motion_body_pos_w, + motion_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + motion_anchor_pos_b, + motion_anchor_ori_b, + i, + ) + + @njit(fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _compute_reward_termination_i( + motion_body_pos_w, + motion_body_quat_w, + motion_body_lin_vel_w, + motion_body_ang_vel_w, + motion_joint_pos, + motion_joint_vel, + ref_body_pos_w, + ref_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + current_actions, + last_actions, + joint_lower, + joint_upper, + scale, + std, + anchor, + ee_indices, + undesired_indices, + n_body, + n_action, + ctrl_dt, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + has_joint_limits, + log_scratch, + reward, + terminated, + tid, + i, + ): + total = 0.0 + + w = motion_global_root_pos_i( + motion_body_pos_w, robot_body_pos_w, anchor, std[0], i + ) * scale[0] + total += w + log_scratch[tid, 0] += w + + w = motion_global_root_ori_i( + motion_body_quat_w, robot_body_quat_w, anchor, std[1], i + ) * scale[1] + total += w + log_scratch[tid, 1] += w + + w = motion_body_pos_i(ref_body_pos_w, robot_body_pos_w, n_body, std[2], i) * scale[2] + total += w + log_scratch[tid, 2] += w + + w = motion_body_ori_i(ref_body_quat_w, robot_body_quat_w, n_body, std[3], i) * scale[3] + total += w + log_scratch[tid, 3] += w + + w = motion_body_lin_vel_i( + motion_body_lin_vel_w, robot_body_lin_vel_w, n_body, std[4], i + ) * scale[4] + total += w + log_scratch[tid, 4] += w + + w = motion_body_ang_vel_i( + motion_body_ang_vel_w, robot_body_ang_vel_w, n_body, std[5], i + ) * scale[5] + total += w + log_scratch[tid, 5] += w + + w = motion_ee_body_pos_z_i(ref_body_pos_w, robot_body_pos_w, ee_indices, std[6], i) * scale[6] + total += w + log_scratch[tid, 6] += w + + w = motion_joint_pos_i(motion_joint_pos, dof_pos, n_action, std[7], i) * scale[7] + total += w + log_scratch[tid, 7] += w + + w = motion_joint_vel_i(motion_joint_vel, dof_vel, n_action, std[8], i) * scale[8] + total += w + log_scratch[tid, 8] += w + + w = action_rate_l2_i(current_actions, last_actions, n_action, i) * scale[9] + total += w + log_scratch[tid, 9] += w + + w = joint_limit_i(dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i) * scale[10] + total += w + log_scratch[tid, 10] += w + + w = ( + undesired_contacts_i( + robot_body_pos_w, undesired_indices, undesired_contact_z_threshold, i + ) + * scale[11] + ) + total += w + log_scratch[tid, 11] += w + + reward[i] = total * ctrl_dt + terminated[i] = terminated_i( + motion_body_pos_w, + motion_body_quat_w, + ref_body_pos_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + ee_indices, + undesired_indices, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + i, + ) + + @_dev + def _write_joint_pos_rel_i( + dof_pos, + default_angles, + default_dof_pos_bias, + has_default_dof_pos_bias, + joint_pos_rel, + n_action, + i, + ): + for j in range(n_action): + default = default_angles[j] + if has_default_dof_pos_bias: + default += default_dof_pos_bias[i, j] + joint_pos_rel[i, j] = dof_pos[i, j] - default + + @_dev + def _write_actor_obs_i( + motion_joint_pos, + motion_joint_vel, + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + dof_vel, + current_actions, + joint_pos_rel, + actor_noise_linvel, + actor_noise_gyro, + actor_noise_joint_pos, + actor_noise_dof_vel, + actor_obs, + is_deploy_actor, + n_action, + i, + ): + out = 0 + for j in range(n_action): + actor_obs[i, out + j] = motion_joint_pos[i, j] + out += n_action + for j in range(n_action): + actor_obs[i, out + j] = motion_joint_vel[i, j] + out += n_action + if not is_deploy_actor: + actor_obs[i, out] = motion_anchor_pos_b[i, 0] + actor_obs[i, out + 1] = motion_anchor_pos_b[i, 1] + actor_obs[i, out + 2] = motion_anchor_pos_b[i, 2] + out += 3 + for k in range(6): + actor_obs[i, out + k] = motion_anchor_ori_b[i, k] + out += 6 + if not is_deploy_actor: + actor_obs[i, out] = linvel[i, 0] + actor_noise_linvel[i, 0] + actor_obs[i, out + 1] = linvel[i, 1] + actor_noise_linvel[i, 1] + actor_obs[i, out + 2] = linvel[i, 2] + actor_noise_linvel[i, 2] + out += 3 + actor_obs[i, out] = gyro[i, 0] + actor_noise_gyro[i, 0] + actor_obs[i, out + 1] = gyro[i, 1] + actor_noise_gyro[i, 1] + actor_obs[i, out + 2] = gyro[i, 2] + actor_noise_gyro[i, 2] + out += 3 + for j in range(n_action): + actor_obs[i, out + j] = joint_pos_rel[i, j] + actor_noise_joint_pos[i, j] + out += n_action + for j in range(n_action): + actor_obs[i, out + j] = dof_vel[i, j] + actor_noise_dof_vel[i, j] + out += n_action + for j in range(n_action): + actor_obs[i, out + j] = current_actions[i, j] + + @_dev + def _write_critic_base_obs_i( + motion_joint_pos, + motion_joint_vel, + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + dof_vel, + current_actions, + joint_pos_rel, + critic_obs, + n_action, + i, + ): + out = 0 + for j in range(n_action): + critic_obs[i, out + j] = motion_joint_pos[i, j] + out += n_action + for j in range(n_action): + critic_obs[i, out + j] = motion_joint_vel[i, j] + out += n_action + critic_obs[i, out] = motion_anchor_pos_b[i, 0] + critic_obs[i, out + 1] = motion_anchor_pos_b[i, 1] + critic_obs[i, out + 2] = motion_anchor_pos_b[i, 2] + out += 3 + for k in range(6): + critic_obs[i, out + k] = motion_anchor_ori_b[i, k] + out += 6 + critic_obs[i, out] = linvel[i, 0] + critic_obs[i, out + 1] = linvel[i, 1] + critic_obs[i, out + 2] = linvel[i, 2] + out += 3 + critic_obs[i, out] = gyro[i, 0] + critic_obs[i, out + 1] = gyro[i, 1] + critic_obs[i, out + 2] = gyro[i, 2] + out += 3 + for j in range(n_action): + critic_obs[i, out + j] = joint_pos_rel[i, j] + out += n_action + for j in range(n_action): + critic_obs[i, out + j] = dof_vel[i, j] + out += n_action + for j in range(n_action): + critic_obs[i, out + j] = current_actions[i, j] + + @_dev + def _write_critic_body_pos_i( + robot_body_pos_w, + robot_body_quat_w, + critic_obs, + anchor, + n_body, + out, + i, + ): + write_body_pos_relative_to_anchor_at( + robot_body_pos_w, + robot_body_quat_w, + anchor, + n_body, + critic_obs, + i, + out, + ) + + @_dev + def _write_critic_body_ori_i( + robot_body_quat_w, + critic_obs, + anchor, + n_body, + out, + i, + ): + write_body_quat_relative_6d_to_anchor_at( + robot_body_quat_w, + anchor, + n_body, + critic_obs, + i, + out, + ) + + @_dev + def _write_critic_linvel_tail_i(linvel, critic_obs, out, i): + if critic_obs.shape[1] > out: + critic_obs[i, out] = linvel[i, 0] + critic_obs[i, out + 1] = linvel[i, 1] + critic_obs[i, out + 2] = linvel[i, 2] + + @_dev + def _write_critic_obs_i( + motion_joint_pos, + motion_joint_vel, + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + dof_vel, + current_actions, + joint_pos_rel, + robot_body_pos_w, + robot_body_quat_w, + critic_obs, + anchor, + n_body, + n_action, + i, + ): + _write_critic_base_obs_i( + motion_joint_pos, + motion_joint_vel, + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + dof_vel, + current_actions, + joint_pos_rel, + critic_obs, + n_action, + i, + ) + body_pos_offset = n_action * 5 + 15 + body_ori_offset = body_pos_offset + n_body * 3 + tail_offset = body_ori_offset + n_body * 6 + _write_critic_body_pos_i( + robot_body_pos_w, + robot_body_quat_w, + critic_obs, + anchor, + n_body, + body_pos_offset, + i, + ) + _write_critic_body_ori_i( + robot_body_quat_w, + critic_obs, + anchor, + n_body, + body_ori_offset, + i, + ) + _write_critic_linvel_tail_i(linvel, critic_obs, tail_offset, i) + + + @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _compute_update_state_kernel( + motion_body_pos_w, + motion_body_quat_w, + motion_body_lin_vel_w, + motion_body_ang_vel_w, + motion_joint_pos, + motion_joint_vel, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + linvel, + gyro, + dof_pos, + dof_vel, + default_angles, + default_dof_pos_bias, + current_actions, + last_actions, + joint_lower, + joint_upper, + scale, + std, + anchor, + ee_indices, + undesired_indices, + ctrl_dt, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + has_joint_limits, + has_default_dof_pos_bias, + is_deploy_actor, + actor_noise_linvel, + actor_noise_gyro, + actor_noise_joint_pos, + actor_noise_dof_vel, + ref_body_pos_w, + ref_body_quat_w, + motion_anchor_pos_b, + motion_anchor_ori_b, + joint_pos_rel, + actor_obs, + critic_obs, + reward, + terminated, + log_scratch, + ): + n = reward.shape[0] + n_body = robot_body_pos_w.shape[1] + n_action = dof_pos.shape[1] + + for i in prange(n): + _write_reference_transforms_i( + motion_body_pos_w, + motion_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + n_body, + ref_body_pos_w, + ref_body_quat_w, + i, + ) + _write_motion_anchor_i( + motion_body_pos_w, + motion_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + anchor, + motion_anchor_pos_b, + motion_anchor_ori_b, + i, + ) + _compute_reward_termination_i( + motion_body_pos_w, + motion_body_quat_w, + motion_body_lin_vel_w, + motion_body_ang_vel_w, + motion_joint_pos, + motion_joint_vel, + ref_body_pos_w, + ref_body_quat_w, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + current_actions, + last_actions, + joint_lower, + joint_upper, + scale, + std, + anchor, + ee_indices, + undesired_indices, + n_body, + n_action, + ctrl_dt, + anchor_pos_z_threshold, + anchor_ori_threshold, + ee_body_pos_z_threshold, + undesired_contact_z_threshold, + terminate_on_undesired_contacts, + has_joint_limits, + log_scratch, + reward, + terminated, + get_thread_id(), + i, + ) + _write_joint_pos_rel_i( + dof_pos, + default_angles, + default_dof_pos_bias, + has_default_dof_pos_bias, + joint_pos_rel, + n_action, + i, + ) + _write_actor_obs_i( + motion_joint_pos, + motion_joint_vel, + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + dof_vel, + current_actions, + joint_pos_rel, + actor_noise_linvel, + actor_noise_gyro, + actor_noise_joint_pos, + actor_noise_dof_vel, + actor_obs, + is_deploy_actor, + n_action, + i, + ) + _write_critic_obs_i( + motion_joint_pos, + motion_joint_vel, + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + dof_vel, + current_actions, + joint_pos_rel, + robot_body_pos_w, + robot_body_quat_w, + critic_obs, + anchor, + n_body, + n_action, + i, + ) + + class G1MotionTrackingNumbaAccelerator: """Driver that keeps config-derived arrays and calls the fused kernel.""" @@ -373,6 +909,10 @@ def __init__( undesired_contact_body_indices: np.ndarray, joint_lower: np.ndarray | None, joint_upper: np.ndarray | None, + default_angles: np.ndarray, + actor_obs_width: int, + critic_obs_width: int, + is_deploy_actor: bool, anchor_pos_z_threshold: float, anchor_ori_threshold: float, ee_body_pos_z_threshold: float, @@ -393,6 +933,10 @@ def __init__( self.ee_body_pos_z_threshold = float(ee_body_pos_z_threshold) self.undesired_contact_z_threshold = float(undesired_contact_z_threshold) self.terminate_on_undesired_contacts = bool(terminate_on_undesired_contacts) + self.default_angles = np.asarray(default_angles, dtype=np.float64) + self.actor_obs_width = int(actor_obs_width) + self.critic_obs_width = int(critic_obs_width) + self.is_deploy_actor = bool(is_deploy_actor) self.num_threads = num_threads self.has_joint_limits = joint_lower is not None and joint_upper is not None if self.has_joint_limits: @@ -404,6 +948,12 @@ def __init__( self.scale = np.zeros((len(TERM_ORDER),), dtype=np.float64) self.std = self._build_std_vector(reward_config) self._zero_actions = np.zeros((self.num_envs, self.num_action), dtype=np.float64) + self._zero_default_dof_pos_bias = np.zeros( + (self.num_envs, self.num_action), dtype=np.float64 + ) + self._zero_linvel_noise = np.zeros((self.num_envs, 3), dtype=np.float64) + self._zero_gyro_noise = np.zeros((self.num_envs, 3), dtype=np.float64) + self._zero_joint_noise = np.zeros((self.num_envs, self.num_action), dtype=np.float64) @classmethod def from_env( @@ -422,6 +972,13 @@ def from_env( f"{sorted(unsupported)}. Disable numba_acceleration or add these terms to " "src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py." ) + default_angles = getattr(env, "default_angles", None) + if default_angles is None: + default_angles = np.zeros((env._num_action,), dtype=np.float64) + actor_obs_width = getattr(env, "_actor_obs_width", env._num_action * 5 + 15) + critic_obs_width = env.obs_groups_spec["critic"] if hasattr(env, "obs_groups_spec") else ( + env._num_action * 5 + 15 + len(env._cfg.body_names) * 9 + ) return cls( num_envs=env.num_envs, num_action=env._num_action, @@ -432,6 +989,10 @@ def from_env( undesired_contact_body_indices=env.undesired_contact_body_indices, joint_lower=env._joint_lower, joint_upper=env._joint_upper, + default_angles=default_angles, + actor_obs_width=actor_obs_width, + critic_obs_width=critic_obs_width, + is_deploy_actor=actor_obs_width == env._num_action * 5 + 9, anchor_pos_z_threshold=env._cfg.anchor_pos_z_threshold, anchor_ori_threshold=env._cfg.anchor_ori_threshold, ee_body_pos_z_threshold=env._cfg.ee_body_pos_z_threshold, @@ -550,3 +1111,145 @@ def compute( if self.scale[idx] != 0.0: log[f"reward/{name}"] = float(term_sums[idx] / dof_pos.shape[0]) return G1MotionTrackingNumbaResult(reward=reward, terminated=terminated, log=log) + + def compute_update_state( + self, + *, + info: dict[str, Any], + motion_data: Any, + linvel: np.ndarray, + gyro: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + robot_body_lin_vel_w: np.ndarray, + robot_body_ang_vel_w: np.ndarray, + ref_body_pos_w: np.ndarray, + ref_body_quat_w: np.ndarray, + motion_anchor_pos_b: np.ndarray, + motion_anchor_ori_b: np.ndarray, + joint_pos_rel: np.ndarray, + scales: Mapping[str, float], + enable_log: bool, + noise_level: float, + noise_scale_linvel: float, + noise_scale_gyro: float, + noise_scale_joint_angle: float, + noise_scale_joint_vel: float, + ) -> G1MotionTrackingNumbaUpdateStateResult: + if not NUMBA_AVAILABLE: + raise RuntimeError( + "G1MotionTracking Numba accelerator was constructed while numba is " + "unavailable; this indicates an invalid accelerator state." + ) + self._sync_scales(scales) + if self.num_threads is not None: + set_num_threads(self.num_threads) + + dtype = dof_pos.dtype + n = dof_pos.shape[0] + if n != self.num_envs: + raise ValueError( + "G1MotionTracking Numba update_state only supports full-batch updates; " + f"got {n} rows for configured num_envs={self.num_envs}." + ) + + current_actions = np.asarray(info.get("current_actions", self._zero_actions), dtype=dtype) + last_actions = np.asarray(info.get("last_actions", self._zero_actions), dtype=dtype) + default_dof_pos_bias = info.get("default_dof_pos_bias") + has_default_dof_pos_bias = isinstance(default_dof_pos_bias, np.ndarray) + if has_default_dof_pos_bias: + default_dof_pos_bias_arr = np.asarray(default_dof_pos_bias, dtype=dtype) + else: + default_dof_pos_bias_arr = self._zero_default_dof_pos_bias.astype(dtype, copy=False) + + noise_level = float(noise_level) + if noise_level > 0.0: + actor_noise_linvel = np.random.uniform(-1.0, 1.0, linvel.shape).astype(dtype) + actor_noise_linvel *= noise_level * float(noise_scale_linvel) + actor_noise_gyro = np.random.uniform(-1.0, 1.0, gyro.shape).astype(dtype) + actor_noise_gyro *= noise_level * float(noise_scale_gyro) + actor_noise_joint_pos = np.random.uniform(-1.0, 1.0, dof_pos.shape).astype(dtype) + actor_noise_joint_pos *= noise_level * float(noise_scale_joint_angle) + actor_noise_dof_vel = np.random.uniform(-1.0, 1.0, dof_vel.shape).astype(dtype) + actor_noise_dof_vel *= noise_level * float(noise_scale_joint_vel) + else: + actor_noise_linvel = self._zero_linvel_noise.astype(dtype, copy=False) + actor_noise_gyro = self._zero_gyro_noise.astype(dtype, copy=False) + actor_noise_joint_pos = self._zero_joint_noise.astype(dtype, copy=False) + actor_noise_dof_vel = self._zero_joint_noise.astype(dtype, copy=False) + + actor_obs = np.empty((n, self.actor_obs_width), dtype=dtype) + critic_obs = np.empty((n, self.critic_obs_width), dtype=dtype) + reward = np.empty((n,), dtype=dtype) + terminated = np.empty((n,), dtype=np.bool_) + log_scratch = np.zeros((get_num_threads(), len(TERM_ORDER)), dtype=np.float64) + + _compute_update_state_kernel( + motion_data.body_pos_w, + motion_data.body_quat_w, + motion_data.body_lin_vel_w, + motion_data.body_ang_vel_w, + motion_data.joint_pos, + motion_data.joint_vel, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + linvel, + gyro, + dof_pos, + dof_vel, + self.default_angles, + default_dof_pos_bias_arr, + current_actions, + last_actions, + self.joint_lower, + self.joint_upper, + self.scale, + self.std, + self.anchor_body_idx, + self.ee_body_indices, + self.undesired_contact_body_indices, + self.ctrl_dt, + self.anchor_pos_z_threshold, + self.anchor_ori_threshold, + self.ee_body_pos_z_threshold, + self.undesired_contact_z_threshold, + self.terminate_on_undesired_contacts, + self.has_joint_limits, + has_default_dof_pos_bias, + self.is_deploy_actor, + actor_noise_linvel, + actor_noise_gyro, + actor_noise_joint_pos, + actor_noise_dof_vel, + ref_body_pos_w, + ref_body_quat_w, + motion_anchor_pos_b, + motion_anchor_ori_b, + joint_pos_rel, + actor_obs, + critic_obs, + reward, + terminated, + log_scratch, + ) + + step_count = info.get("steps") + should_log = enable_log and ( + int(step_count[0]) % 4 == 0 if isinstance(step_count, np.ndarray) else True + ) + log = {} if should_log else info.get("log", {}) + if should_log: + term_sums = log_scratch.sum(axis=0) + for idx, name in enumerate(TERM_ORDER): + if self.scale[idx] != 0.0: + log[f"reward/{name}"] = float(term_sums[idx] / n) + return G1MotionTrackingNumbaUpdateStateResult( + obs={"obs": actor_obs, "critic": critic_obs}, + reward=reward, + terminated=terminated, + log=log, + ) diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index 3da83c6fe..2c6bfa0d6 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -842,28 +842,40 @@ def update_state(self, state: NpEnvState) -> NpEnvState: robot_body_ang_vel_w, ) = self._get_body_state_w() - # Compute relative body transforms (for observations and rewards) - self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) - if self._numba_accelerator is not None: - numba_result = self._numba_accelerator.compute( + noise_cfg = self._cfg.noise_config + numba_result = self._numba_accelerator.compute_update_state( info=state.info, motion_data=motion_data, - ref_body_pos_w=self.body_pos_relative_w, - ref_body_quat_w=self.body_quat_relative_w, + linvel=linvel, + gyro=gyro, + dof_pos=dof_pos, + dof_vel=dof_vel, robot_body_pos_w=robot_body_pos_w, robot_body_quat_w=robot_body_quat_w, robot_body_lin_vel_w=robot_body_lin_vel_w, robot_body_ang_vel_w=robot_body_ang_vel_w, - dof_pos=dof_pos, - dof_vel=dof_vel, + ref_body_pos_w=self.body_pos_relative_w, + ref_body_quat_w=self.body_quat_relative_w, + motion_anchor_pos_b=self._motion_anchor_pos_b, + motion_anchor_ori_b=self._motion_anchor_ori_b, + joint_pos_rel=self._joint_pos_rel, scales=self._cfg.reward_config.scales, enable_log=self._enable_reward_log, + noise_level=noise_cfg.level, + noise_scale_linvel=noise_cfg.scale_linvel, + noise_scale_gyro=noise_cfg.scale_gyro, + noise_scale_joint_angle=noise_cfg.scale_joint_angle, + noise_scale_joint_vel=noise_cfg.scale_joint_vel, ) terminated = numba_result.terminated reward = numba_result.reward + obs = numba_result.obs state.info["log"] = numba_result.log else: + # Compute relative body transforms (for observations and rewards) + self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + # Compute terminations terminated = self._compute_terminations( motion_data, robot_body_pos_w, robot_body_quat_w @@ -881,21 +893,21 @@ def update_state(self, state: NpEnvState) -> NpEnvState: dof_vel, ) + # Compute observations + obs = self._compute_obs( + state.info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ) + # Update failure statistics for adaptive sampling self.motion_sampler.update_failure_stats(terminated) - # Compute observations - obs = self._compute_obs( - state.info, - motion_data, - linvel, - gyro, - dof_pos, - dof_vel, - robot_body_pos_w, - robot_body_quat_w, - ) - # Advance motion frames done_env_ids = self.motion_sampler.step() if len(done_env_ids) > 0: diff --git a/src/unilab/utils/numba_geometry.py b/src/unilab/utils/numba_geometry.py new file mode 100644 index 000000000..d08e5bffc --- /dev/null +++ b/src/unilab/utils/numba_geometry.py @@ -0,0 +1,290 @@ +"""Reusable Numba geometry helpers. + +The helpers in this module are intentionally low level: quaternion math, +body-frame vector transforms, and compact rotation representations. Task-owned +Numba kernels can compose them without moving reward or observation layout +knowledge into shared utilities. +""" + +from __future__ import annotations + +import math +from typing import Any + +try: # pragma: no cover - exercised when optional numba dependency is installed + from numba import njit + + NUMBA_GEOMETRY_AVAILABLE = True +except Exception: # pragma: no cover + njit = None # type: ignore[assignment] + NUMBA_GEOMETRY_AVAILABLE = False + + +def _missing_numba(*_args: Any, **_kwargs: Any) -> None: + raise RuntimeError("numba_geometry helpers require numba to be installed") + + +if NUMBA_GEOMETRY_AVAILABLE: + + def _dev(fn): + return njit(inline="always", fastmath=True, cache=True, nogil=True)(fn) + + @_dev + def quat_angle_sq_at(q1, q2, item_idx, i): + dot = ( + q1[i, item_idx, 0] * q2[i, item_idx, 0] + + q1[i, item_idx, 1] * q2[i, item_idx, 1] + + q1[i, item_idx, 2] * q2[i, item_idx, 2] + + q1[i, item_idx, 3] * q2[i, item_idx, 3] + ) + dot = abs(dot) + if dot > 1.0: + dot = 1.0 + angle = 2.0 * math.acos(dot) + return angle * angle + + @_dev + def quat_gravity_z_at(quat, item_idx, i): + return ( + 2.0 + * ( + quat[i, item_idx, 1] * quat[i, item_idx, 1] + + quat[i, item_idx, 2] * quat[i, item_idx, 2] + ) + - 1.0 + ) + + @_dev + def quat_yaw_from_components(qw, qx, qy, qz): + return math.atan2( + 2.0 * (qw * qz + qx * qy), + 1.0 - 2.0 * (qy * qy + qz * qz), + ) + + @_dev + def rotate_vec_by_inv_quat_components(qw, qx, qy, qz, vx, vy, vz, out, i, offset): + ix = -qx + iy = -qy + iz = -qz + tx = 2.0 * (iy * vz - iz * vy) + ty = 2.0 * (iz * vx - ix * vz) + tz = 2.0 * (ix * vy - iy * vx) + out[i, offset] = vx + qw * tx + iy * tz - iz * ty + out[i, offset + 1] = vy + qw * ty + iz * tx - ix * tz + out[i, offset + 2] = vz + qw * tz + ix * ty - iy * tx + + @_dev + def write_quat_first_two_matrix_cols_from_components(qw, qx, qy, qz, out, i, offset): + xx = qx * qx + yy = qy * qy + zz = qz * qz + xy = qx * qy + xz = qx * qz + yz = qy * qz + wx = qw * qx + wy = qw * qy + wz = qw * qz + out[i, offset] = 1.0 - 2.0 * (yy + zz) + out[i, offset + 1] = 2.0 * (xy - wz) + out[i, offset + 2] = 2.0 * (xy + wz) + out[i, offset + 3] = 1.0 - 2.0 * (xx + zz) + out[i, offset + 4] = 2.0 * (xz - wy) + out[i, offset + 5] = 2.0 * (yz + wx) + + @_dev + def write_yaw_aligned_body_transforms_at( + source_body_pos_w, + source_body_quat_w, + target_body_pos_w, + target_body_quat_w, + anchor, + n_body, + out_body_pos_w, + out_body_quat_w, + i, + ): + anchor_px = source_body_pos_w[i, anchor, 0] + anchor_py = source_body_pos_w[i, anchor, 1] + anchor_pz = source_body_pos_w[i, anchor, 2] + anchor_qw = source_body_quat_w[i, anchor, 0] + anchor_qx = source_body_quat_w[i, anchor, 1] + anchor_qy = source_body_quat_w[i, anchor, 2] + anchor_qz = source_body_quat_w[i, anchor, 3] + + target_anchor_px = target_body_pos_w[i, anchor, 0] + target_anchor_py = target_body_pos_w[i, anchor, 1] + target_anchor_qw = target_body_quat_w[i, anchor, 0] + target_anchor_qx = target_body_quat_w[i, anchor, 1] + target_anchor_qy = target_body_quat_w[i, anchor, 2] + target_anchor_qz = target_body_quat_w[i, anchor, 3] + + qw = ( + target_anchor_qw * anchor_qw + + target_anchor_qx * anchor_qx + + target_anchor_qy * anchor_qy + + target_anchor_qz * anchor_qz + ) + qx = ( + -target_anchor_qw * anchor_qx + + target_anchor_qx * anchor_qw + - target_anchor_qy * anchor_qz + + target_anchor_qz * anchor_qy + ) + qy = ( + -target_anchor_qw * anchor_qy + + target_anchor_qx * anchor_qz + + target_anchor_qy * anchor_qw + - target_anchor_qz * anchor_qx + ) + qz = ( + -target_anchor_qw * anchor_qz + - target_anchor_qx * anchor_qy + + target_anchor_qy * anchor_qx + + target_anchor_qz * anchor_qw + ) + half_yaw = 0.5 * quat_yaw_from_components(qw, qx, qy, qz) + delta_qw = math.cos(half_yaw) + delta_qz = math.sin(half_yaw) + yaw_cross = 2.0 * delta_qw * delta_qz + yaw_z2 = 2.0 * delta_qz * delta_qz + + for body_idx in range(n_body): + mw = source_body_quat_w[i, body_idx, 0] + mx = source_body_quat_w[i, body_idx, 1] + my = source_body_quat_w[i, body_idx, 2] + mz = source_body_quat_w[i, body_idx, 3] + out_body_quat_w[i, body_idx, 0] = delta_qw * mw - delta_qz * mz + out_body_quat_w[i, body_idx, 1] = delta_qw * mx - delta_qz * my + out_body_quat_w[i, body_idx, 2] = delta_qw * my + delta_qz * mx + out_body_quat_w[i, body_idx, 3] = delta_qw * mz + delta_qz * mw + + vx = source_body_pos_w[i, body_idx, 0] - anchor_px + vy = source_body_pos_w[i, body_idx, 1] - anchor_py + vz = source_body_pos_w[i, body_idx, 2] - anchor_pz + out_body_pos_w[i, body_idx, 0] = ( + vx - yaw_cross * vy - yaw_z2 * vx + target_anchor_px + ) + out_body_pos_w[i, body_idx, 1] = ( + vy + yaw_cross * vx - yaw_z2 * vy + target_anchor_py + ) + out_body_pos_w[i, body_idx, 2] = vz + anchor_pz + + @_dev + def write_relative_anchor_transform_at( + source_body_pos_w, + source_body_quat_w, + target_body_pos_w, + target_body_quat_w, + anchor, + out_pos_b, + out_rot6d_b, + i, + ): + anchor_px = source_body_pos_w[i, anchor, 0] + anchor_py = source_body_pos_w[i, anchor, 1] + anchor_pz = source_body_pos_w[i, anchor, 2] + anchor_qw = source_body_quat_w[i, anchor, 0] + anchor_qx = source_body_quat_w[i, anchor, 1] + anchor_qy = source_body_quat_w[i, anchor, 2] + anchor_qz = source_body_quat_w[i, anchor, 3] + + target_anchor_px = target_body_pos_w[i, anchor, 0] + target_anchor_py = target_body_pos_w[i, anchor, 1] + target_anchor_pz = target_body_pos_w[i, anchor, 2] + target_anchor_qw = target_body_quat_w[i, anchor, 0] + target_anchor_qx = target_body_quat_w[i, anchor, 1] + target_anchor_qy = target_body_quat_w[i, anchor, 2] + target_anchor_qz = target_body_quat_w[i, anchor, 3] + + rotate_vec_by_inv_quat_components( + target_anchor_qw, + target_anchor_qx, + target_anchor_qy, + target_anchor_qz, + anchor_px - target_anchor_px, + anchor_py - target_anchor_py, + anchor_pz - target_anchor_pz, + out_pos_b, + i, + 0, + ) + + rw = ( + target_anchor_qw * anchor_qw + + target_anchor_qx * anchor_qx + + target_anchor_qy * anchor_qy + + target_anchor_qz * anchor_qz + ) + rx = ( + target_anchor_qw * anchor_qx + - target_anchor_qx * anchor_qw + - target_anchor_qy * anchor_qz + + target_anchor_qz * anchor_qy + ) + ry = ( + target_anchor_qw * anchor_qy + + target_anchor_qx * anchor_qz + - target_anchor_qy * anchor_qw + - target_anchor_qz * anchor_qx + ) + rz = ( + target_anchor_qw * anchor_qz + - target_anchor_qx * anchor_qy + + target_anchor_qy * anchor_qx + - target_anchor_qz * anchor_qw + ) + write_quat_first_two_matrix_cols_from_components(rw, rx, ry, rz, out_rot6d_b, i, 0) + + @_dev + def write_body_pos_relative_to_anchor_at(body_pos_w, anchor_quat_w, anchor, n_body, out, i, offset): + anchor_px = body_pos_w[i, anchor, 0] + anchor_py = body_pos_w[i, anchor, 1] + anchor_pz = body_pos_w[i, anchor, 2] + aw = anchor_quat_w[i, anchor, 0] + ax = anchor_quat_w[i, anchor, 1] + ay = anchor_quat_w[i, anchor, 2] + az = anchor_quat_w[i, anchor, 3] + + for body_idx in range(n_body): + rotate_vec_by_inv_quat_components( + aw, + ax, + ay, + az, + body_pos_w[i, body_idx, 0] - anchor_px, + body_pos_w[i, body_idx, 1] - anchor_py, + body_pos_w[i, body_idx, 2] - anchor_pz, + out, + i, + offset + body_idx * 3, + ) + + @_dev + def write_body_quat_relative_6d_to_anchor_at(body_quat_w, anchor, n_body, out, i, offset): + aw = body_quat_w[i, anchor, 0] + ax = body_quat_w[i, anchor, 1] + ay = body_quat_w[i, anchor, 2] + az = body_quat_w[i, anchor, 3] + for body_idx in range(n_body): + bw = body_quat_w[i, body_idx, 0] + bx = body_quat_w[i, body_idx, 1] + by = body_quat_w[i, body_idx, 2] + bz = body_quat_w[i, body_idx, 3] + rw = aw * bw + ax * bx + ay * by + az * bz + rx = aw * bx - ax * bw - ay * bz + az * by + ry = aw * by + ax * bz - ay * bw - az * bx + rz = aw * bz - ax * by + ay * bx - az * bw + write_quat_first_two_matrix_cols_from_components( + rw, rx, ry, rz, out, i, offset + body_idx * 6 + ) + +else: + quat_angle_sq_at = _missing_numba + quat_gravity_z_at = _missing_numba + quat_yaw_from_components = _missing_numba + rotate_vec_by_inv_quat_components = _missing_numba + write_quat_first_two_matrix_cols_from_components = _missing_numba + write_yaw_aligned_body_transforms_at = _missing_numba + write_relative_anchor_transform_at = _missing_numba + write_body_pos_relative_to_anchor_at = _missing_numba + write_body_quat_relative_6d_to_anchor_at = _missing_numba diff --git a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py index bc2bc7df2..722c551c8 100644 --- a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py +++ b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py @@ -69,6 +69,8 @@ def test_g1_motion_tracking_numba_benchmark_formats_component_records() -> None: relative_transform_ms=0.5, numpy_reward_termination_ms=1.0, numba_reward_termination_ms=0.25, + numpy_update_state_ms=3.0, + numba_update_state_ms=1.5, numpy_total_ms=1.5, numba_total_ms=0.75, speedup_vs_numpy=2.0, @@ -79,6 +81,7 @@ def test_g1_motion_tracking_numba_benchmark_formats_component_records() -> None: assert payload["relative_transform_ms"] == 0.5 assert "rel ms" in table + assert "numpy update_state ms" in table assert "2.00x" in table diff --git a/tests/envs/test_g1_motion_tracking_numba.py b/tests/envs/test_g1_motion_tracking_numba.py index c70c7550e..5c3a00a81 100644 --- a/tests/envs/test_g1_motion_tracking_numba.py +++ b/tests/envs/test_g1_motion_tracking_numba.py @@ -94,6 +94,90 @@ def test_g1_motion_tracking_numba_reward_termination_parity(): assert out.log[key] == pytest.approx(log_np[key], rel=1e-3, abs=1e-6) +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_motion_tracking_numba_update_state_parity_without_noise(): + n = 512 + env = _make_env(n, include_undesired=True) + env.default_angles = np.linspace(-0.2, 0.2, env._num_action).astype(np.float32) + env._actor_obs_width = env._actor_obs_dim(env._num_action) + env._critic_base_obs_width = env._critic_base_obs_dim(env._num_action) + env._critic_obs_width = env._critic_base_obs_width + env._n_motion_bodies * 9 + env._motion_anchor_pos_b = np.empty((n, 3), dtype=np.float32) + env._motion_anchor_ori_b = np.empty((n, 6), dtype=np.float32) + env._joint_pos_rel = np.empty((n, env._num_action), dtype=np.float32) + env._motion_command = np.empty((n, env._num_action * 2), dtype=np.float32) + env._zero_actions = np.zeros((n, env._num_action), dtype=np.float32) + env._body_vec_tmp = np.empty((n, env._n_motion_bodies, 3), dtype=np.float32) + env._cfg.noise_config = _NoiseCfg() + + motion_data, robot_state, info = _make_batch(n, seed=2) + ( + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ) = robot_state + rng = np.random.default_rng(22) + linvel = rng.uniform(-1.0, 1.0, (n, 3)).astype(np.float32) + gyro = rng.uniform(-1.0, 1.0, (n, 3)).astype(np.float32) + + env._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + reward_np = env._compute_reward( + {**info, "log": {}}, + motion_data, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + dof_vel, + ) + terminated_np = env._compute_terminations(motion_data, robot_body_pos_w, robot_body_quat_w) + obs_np = env._compute_obs( + info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ) + + accel = G1MotionTrackingNumbaAccelerator.from_env(env, num_threads=2) + out = accel.compute_update_state( + info=info, + motion_data=motion_data, + linvel=linvel, + gyro=gyro, + dof_pos=dof_pos, + dof_vel=dof_vel, + robot_body_pos_w=robot_body_pos_w, + robot_body_quat_w=robot_body_quat_w, + robot_body_lin_vel_w=robot_body_lin_vel_w, + robot_body_ang_vel_w=robot_body_ang_vel_w, + ref_body_pos_w=env.body_pos_relative_w, + ref_body_quat_w=env.body_quat_relative_w, + motion_anchor_pos_b=env._motion_anchor_pos_b, + motion_anchor_ori_b=env._motion_anchor_ori_b, + joint_pos_rel=env._joint_pos_rel, + scales=env._cfg.reward_config.scales, + enable_log=True, + noise_level=0.0, + noise_scale_linvel=0.0, + noise_scale_gyro=0.0, + noise_scale_joint_angle=0.0, + noise_scale_joint_vel=0.0, + ) + + np.testing.assert_allclose(out.reward, reward_np, rtol=1e-4, atol=1e-5) + np.testing.assert_array_equal(out.terminated, terminated_np) + np.testing.assert_allclose(out.obs["obs"], obs_np["obs"], rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(out.obs["critic"], obs_np["critic"], rtol=1e-5, atol=1e-5) + + @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") def test_g1_motion_tracking_numba_term_py_funcs_match_numpy_math(): from unilab.envs.motion_tracking.g1 import motion_tracking_numba as T @@ -170,6 +254,15 @@ class _Cfg: terminate_on_undesired_contacts: bool = False +@dataclass +class _NoiseCfg: + level: float = 0.0 + scale_linvel: float = 0.1 + scale_gyro: float = 0.1 + scale_joint_angle: float = 0.02 + scale_joint_vel: float = 0.3 + + def _make_env(n: int, *, include_undesired: bool) -> Any: env = cast(Any, object.__new__(G1MotionTrackingEnv)) cfg = _Cfg() diff --git a/tests/utils/test_numba_geometry.py b/tests/utils/test_numba_geometry.py new file mode 100644 index 000000000..c0e54ef33 --- /dev/null +++ b/tests/utils/test_numba_geometry.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import numpy as np +import pytest + +from unilab.envs.common.rotation import ( + np_matrix_first_two_cols_from_quat, + np_subtract_frame_transforms, +) +from unilab.utils.numba_geometry import ( + NUMBA_GEOMETRY_AVAILABLE, + quat_angle_sq_at, + write_relative_anchor_transform_at, +) + + +@pytest.mark.skipif(not NUMBA_GEOMETRY_AVAILABLE, reason="numba is optional") +def test_numba_geometry_quat_angle_sq_matches_expected_angles(): + from numba import njit + + @njit + def run(q1, q2): + return quat_angle_sq_at(q1, q2, 0, 0) + + q1 = np.array([[[1.0, 0.0, 0.0, 0.0]]], dtype=np.float64) + q2 = np.array([[[0.0, 1.0, 0.0, 0.0]]], dtype=np.float64) + + assert run(q1, q1) == pytest.approx(0.0) + assert run(q1, q2) == pytest.approx(np.pi * np.pi) + + +@pytest.mark.skipif(not NUMBA_GEOMETRY_AVAILABLE, reason="numba is optional") +def test_numba_geometry_relative_anchor_transform_matches_numpy_rotation_helpers(): + from numba import njit + + @njit + def run(source_pos, source_quat, target_pos, target_quat, out_pos, out_rot6d): + write_relative_anchor_transform_at( + source_pos, + source_quat, + target_pos, + target_quat, + 0, + out_pos, + out_rot6d, + 0, + ) + + source_pos = np.array([[[0.7, -0.2, 1.3]]], dtype=np.float64) + source_quat = np.array([[[0.9238795325, 0.0, 0.3826834324, 0.0]]], dtype=np.float64) + target_pos = np.array([[[-0.1, 0.4, 0.8]]], dtype=np.float64) + target_quat = np.array([[[0.9659258263, 0.0, 0.0, 0.2588190451]]], dtype=np.float64) + out_pos = np.empty((1, 3), dtype=np.float64) + out_rot6d = np.empty((1, 6), dtype=np.float64) + + run(source_pos, source_quat, target_pos, target_quat, out_pos, out_rot6d) + + expected_pos, expected_quat = np_subtract_frame_transforms( + target_pos[0, 0], + target_quat[0, 0], + source_pos[0, 0], + source_quat[0, 0], + ) + expected_rot6d = np_matrix_first_two_cols_from_quat(expected_quat) + + np.testing.assert_allclose(out_pos[0], expected_pos, rtol=1e-9, atol=1e-9) + np.testing.assert_allclose(out_rot6d[0], expected_rot6d, rtol=1e-9, atol=1e-9) diff --git a/tests/utils/test_utils_package_policy.py b/tests/utils/test_utils_package_policy.py index b424cc414..19a5bf0ec 100644 --- a/tests/utils/test_utils_package_policy.py +++ b/tests/utils/test_utils_package_policy.py @@ -4,7 +4,14 @@ import unilab.utils ALLOWED_UTILS_API = {"get_default_device", "to_numpy", "to_torch"} -ALLOWED_UTILS_MODULES = {"__init__", "device", "nan_guard", "support_matrix", "tensor"} +ALLOWED_UTILS_MODULES = { + "__init__", + "device", + "nan_guard", + "numba_geometry", + "support_matrix", + "tensor", +} REMOVED_UTILS_SHIMS = { "algo_utils", "device_utils", From b813b0ef956e3f6431b2f6376f4a44f1bd803f64 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 20:40:46 +0800 Subject: [PATCH 12/26] refactor: consolidate numpy geometry helpers under unilab.utils Move the shared numeric helpers out of unilab.envs.common into unilab.utils and extract per-env private geometry/coordinate helpers (tracking, stewart, allegro, sharpa, go2_arm) into a new unilab.utils.geometry module. Each extraction preserves the caller-visible dtype cast (e.g. get_global_dtype for allegro rotation axis) so refactored code is bit-identical to the pre-refactor implementation across random inputs. Notable moves: - envs/common/{rotation,math}.py -> utils/{rotation,math}.py, then math.py's lone np_sample_uniform folded into utils/geometry.py so utils/ no longer hosts a single-function module. - utils/geometry.py picks up the frame-transform, roll/pitch, gravity-in-body, quat-normalize, orientation-error, angvel-from-pair, uniform-quat-sample, and sphere/cart conversions that were duplicated across envs. - Dead helpers _normalize_quat and _orientation_error_local removed from go2_arm/base.py; call sites now hit the shared util directly. - All 30+ import sites (src/tests/scripts/docs/AGENTS.md) updated; utils package whitelist and sphinx toctree adjusted. --- AGENTS.md | 2 +- .../source/api_reference/envs/common.md | 13 -- .../sphinx/source/api_reference/envs/index.md | 1 - .../source/api_reference/utils/index.md | 3 +- scripts/deploy/sim_prototype.py | 2 +- scripts/manip_loco/diagnose_go2_arm_ik.py | 2 +- scripts/motion/bones_seed_csv_to_npz.py | 2 +- scripts/motion/csv_to_npz.py | 2 +- scripts/motion/x2_csv_to_tracking_npz.py | 2 +- src/unilab/envs/common/__init__.py | 1 - src/unilab/envs/common/math.py | 13 -- src/unilab/envs/locomotion/common/commands.py | 2 +- .../envs/locomotion/common/dr_provider.py | 2 +- src/unilab/envs/locomotion/go1/joystick.py | 2 +- src/unilab/envs/locomotion/go1/rough.py | 2 +- src/unilab/envs/locomotion/go2/footstand.py | 2 +- src/unilab/envs/locomotion/go2/rough.py | 2 +- src/unilab/envs/locomotion/go2_arm/base.py | 27 +-- .../envs/locomotion/go2_arm/manip_loco.py | 27 +-- src/unilab/envs/locomotion/go2w/joystick.py | 2 +- src/unilab/envs/locomotion/go2w/rough.py | 2 +- .../manipulation/allegro_inhand/rotation.py | 15 +- .../envs/manipulation/sharpa_inhand/base.py | 2 +- .../manipulation/sharpa_inhand/grasp_gen.py | 2 +- .../manipulation/sharpa_inhand/rotation.py | 16 +- .../envs/manipulation/stewart/balance.py | 12 +- .../envs/motion_tracking/g1/box_tracking.py | 4 +- .../envs/motion_tracking/g1/tracking.py | 74 +----- src/unilab/utils/geometry.py | 216 ++++++++++++++++++ src/unilab/{envs/common => utils}/rotation.py | 0 tests/envs/test_env_configs.py | 6 +- tests/test_sharpa.py | 2 +- tests/utils/test_math_utils.py | 4 +- tests/utils/test_utils_package_policy.py | 2 + 34 files changed, 281 insertions(+), 187 deletions(-) delete mode 100644 docs/sphinx/source/api_reference/envs/common.md delete mode 100644 src/unilab/envs/common/__init__.py delete mode 100644 src/unilab/envs/common/math.py create mode 100644 src/unilab/utils/geometry.py rename src/unilab/{envs/common => utils}/rotation.py (100%) diff --git a/AGENTS.md b/AGENTS.md index c96a0e92c..501dd26ca 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -46,7 +46,7 @@ UniLab 是一个 **高性能、模块化、contract 驱动** 的 RL infrastructu - backend contract: `src/unilab/base/backend/base.py` - training run helpers: `src/unilab/training/run.py` - visualization helpers: `src/unilab/visualization/` -- env shared numeric helpers: `src/unilab/envs/common/rotation.py`, `src/unilab/envs/common/math.py` +- shared numeric helpers: `src/unilab/utils/rotation.py`, `src/unilab/utils/geometry.py` - MLX rotation helpers: `src/unilab/algos/mlx/common/rotation.py` - config schema: `src/unilab/structured_configs.py` - async runner: `src/unilab/ipc/async_runner.py` diff --git a/docs/sphinx/source/api_reference/envs/common.md b/docs/sphinx/source/api_reference/envs/common.md deleted file mode 100644 index 35ff7a466..000000000 --- a/docs/sphinx/source/api_reference/envs/common.md +++ /dev/null @@ -1,13 +0,0 @@ -# `unilab.envs.common` - -Rotation / math helpers used by every env. - -```{eval-rst} -.. automodule:: unilab.envs.common.math - :members: -``` - -```{eval-rst} -.. automodule:: unilab.envs.common.rotation - :members: -``` diff --git a/docs/sphinx/source/api_reference/envs/index.md b/docs/sphinx/source/api_reference/envs/index.md index 2b1babc66..6665cd124 100644 --- a/docs/sphinx/source/api_reference/envs/index.md +++ b/docs/sphinx/source/api_reference/envs/index.md @@ -15,7 +15,6 @@ can be selected via `uv run train --algo --task --sim `. locomotion manipulation motion_tracking -common ``` ```{eval-rst} diff --git a/docs/sphinx/source/api_reference/utils/index.md b/docs/sphinx/source/api_reference/utils/index.md index 3b5c461da..99e6e4270 100644 --- a/docs/sphinx/source/api_reference/utils/index.md +++ b/docs/sphinx/source/api_reference/utils/index.md @@ -1,6 +1,7 @@ # `unilab.utils` — Utilities -Device probing, tensor helpers, support-matrix bookkeeping, NaN guards. +Device probing, tensor helpers, support-matrix bookkeeping, NaN guards, +and pure-numpy geometry/rotation helpers shared across envs and scripts. ```{eval-rst} .. autosummary:: diff --git a/scripts/deploy/sim_prototype.py b/scripts/deploy/sim_prototype.py index e8db273f6..7043b9bd2 100644 --- a/scripts/deploy/sim_prototype.py +++ b/scripts/deploy/sim_prototype.py @@ -42,7 +42,7 @@ REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT / "src")) -from unilab.envs.common.rotation import ( # noqa: E402 +from unilab.utils.rotation import ( # noqa: E402 np_matrix_from_quat, np_subtract_frame_transforms, ) diff --git a/scripts/manip_loco/diagnose_go2_arm_ik.py b/scripts/manip_loco/diagnose_go2_arm_ik.py index a1a03d821..21e26e3b7 100644 --- a/scripts/manip_loco/diagnose_go2_arm_ik.py +++ b/scripts/manip_loco/diagnose_go2_arm_ik.py @@ -16,7 +16,7 @@ from unilab.base import registry from unilab.base.registry import ensure_registries -from unilab.envs.common.rotation import np_matrix_from_quat +from unilab.utils.rotation import np_matrix_from_quat from unilab.envs.locomotion.go2_arm.manip_loco import RewardConfig diff --git a/scripts/motion/bones_seed_csv_to_npz.py b/scripts/motion/bones_seed_csv_to_npz.py index 3fc10ecc0..3031cce4d 100644 --- a/scripts/motion/bones_seed_csv_to_npz.py +++ b/scripts/motion/bones_seed_csv_to_npz.py @@ -42,7 +42,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors -from unilab.envs.common.rotation import np_quat_angular_velocity, np_quat_ensure_continuity +from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity ROOT_COLUMNS = [ "Frame", diff --git a/scripts/motion/csv_to_npz.py b/scripts/motion/csv_to_npz.py index e3d2f5ddc..2794a7114 100644 --- a/scripts/motion/csv_to_npz.py +++ b/scripts/motion/csv_to_npz.py @@ -30,7 +30,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors -from unilab.envs.common.rotation import np_quat_angular_velocity, np_quat_ensure_continuity +from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity def quat_slerp(q1: np.ndarray, q2: np.ndarray, t: float) -> np.ndarray: diff --git a/scripts/motion/x2_csv_to_tracking_npz.py b/scripts/motion/x2_csv_to_tracking_npz.py index ba2515349..3892eaba2 100644 --- a/scripts/motion/x2_csv_to_tracking_npz.py +++ b/scripts/motion/x2_csv_to_tracking_npz.py @@ -20,7 +20,7 @@ from unilab.assets import ASSETS_ROOT_PATH from unilab.base.backend.mujoco.xml import inject_mujoco_tracking_sensors -from unilab.envs.common.rotation import np_quat_angular_velocity, np_quat_ensure_continuity +from unilab.utils.rotation import np_quat_angular_velocity, np_quat_ensure_continuity ROOT_QPOS_DIM = 7 ROOT_QVEL_DIM = 6 diff --git a/src/unilab/envs/common/__init__.py b/src/unilab/envs/common/__init__.py deleted file mode 100644 index a05e8faba..000000000 --- a/src/unilab/envs/common/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared environment-owned numeric helpers.""" diff --git a/src/unilab/envs/common/math.py b/src/unilab/envs/common/math.py deleted file mode 100644 index e96d7f105..000000000 --- a/src/unilab/envs/common/math.py +++ /dev/null @@ -1,13 +0,0 @@ -from __future__ import annotations - -import numpy as np - - -def np_sample_uniform( - lower: float | np.ndarray, - upper: float | np.ndarray, - size: tuple[int, ...], - dtype=np.float32, -) -> np.ndarray: - """Sample uniformly from [lower, upper] with output dtype.""" - return np.random.uniform(lower, upper, size).astype(dtype) diff --git a/src/unilab/envs/locomotion/common/commands.py b/src/unilab/envs/locomotion/common/commands.py index 9a2069d2f..e14f32252 100644 --- a/src/unilab/envs/locomotion/common/commands.py +++ b/src/unilab/envs/locomotion/common/commands.py @@ -6,7 +6,7 @@ import numpy as np from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_wrap_to_pi, np_yaw_from_quat +from unilab.utils.rotation import np_wrap_to_pi, np_yaw_from_quat @dataclass diff --git a/src/unilab/envs/locomotion/common/dr_provider.py b/src/unilab/envs/locomotion/common/dr_provider.py index 2c9db268c..24f75ab54 100644 --- a/src/unilab/envs/locomotion/common/dr_provider.py +++ b/src/unilab/envs/locomotion/common/dr_provider.py @@ -26,7 +26,7 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_quat_mul, np_yaw_to_quat +from unilab.utils.rotation import np_quat_mul, np_yaw_to_quat class LocomotionDRProvider(DomainRandomizationProvider): diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index 9a1734a0b..f085c832f 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -12,7 +12,7 @@ from unilab.base.np_env import NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_quat_mul, np_yaw_to_quat +from unilab.utils.rotation import np_quat_mul, np_yaw_to_quat from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig diff --git a/src/unilab/envs/locomotion/go1/rough.py b/src/unilab/envs/locomotion/go1/rough.py index 1bbb084bc..2124de2a5 100644 --- a/src/unilab/envs/locomotion/go1/rough.py +++ b/src/unilab/envs/locomotion/go1/rough.py @@ -14,7 +14,7 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_quat_apply_inverse, np_quat_from_euler_xyz, np_quat_mul, diff --git a/src/unilab/envs/locomotion/go2/footstand.py b/src/unilab/envs/locomotion/go2/footstand.py index 956233403..3e5f6572c 100644 --- a/src/unilab/envs/locomotion/go2/footstand.py +++ b/src/unilab/envs/locomotion/go2/footstand.py @@ -13,7 +13,7 @@ from unilab.base.scene import SceneCfg from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_quat_apply, np_quat_apply_inverse +from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig diff --git a/src/unilab/envs/locomotion/go2/rough.py b/src/unilab/envs/locomotion/go2/rough.py index f37952ce8..b7bee9b4f 100644 --- a/src/unilab/envs/locomotion/go2/rough.py +++ b/src/unilab/envs/locomotion/go2/rough.py @@ -12,7 +12,7 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_quat_apply_inverse, np_quat_from_euler_xyz, np_quat_mul, diff --git a/src/unilab/envs/locomotion/go2_arm/base.py b/src/unilab/envs/locomotion/go2_arm/base.py index 087923395..ec76bb539 100644 --- a/src/unilab/envs/locomotion/go2_arm/base.py +++ b/src/unilab/envs/locomotion/go2_arm/base.py @@ -5,18 +5,14 @@ import numpy as np from unilab.base.backend import SimBackend -from unilab.envs.common.rotation import ( - np_matrix_from_quat, - np_quat_canonicalize, - np_quat_inv, - np_quat_mul, -) +from unilab.utils.rotation import np_matrix_from_quat from unilab.envs.locomotion.common.base import ( ControlConfigBase, LocomotionBaseCfg, LocomotionBaseEnv, Sensor, ) +from unilab.utils.geometry import np_quat_orientation_error_local DEFAULT_LEG_ANGLES = np.asarray( [ @@ -241,7 +237,7 @@ def compute_arm_ik_delta( "goal_local_quat and curr_local_quat are required when " "ik.use_orientation=true and ik.orientation_mode='target'" ) - orn_err = _orientation_error_local(goal_local_quat, curr_local_quat) + orn_err = np_quat_orientation_error_local(goal_local_quat, curr_local_quat) elif cfg.orientation_mode == "zero_error": orn_err = np.zeros_like(pos_err) else: @@ -264,20 +260,3 @@ def compute_arm_ik_delta( if cfg.dq_clip > 0.0: dq = np.clip(dq, -cfg.dq_clip, cfg.dq_clip) return dq - - -def _normalize_quat(q: np.ndarray) -> np.ndarray: - q = np.asarray(q) - if q.ndim == 1: - q = q[None, :] - norm = np.linalg.norm(q, axis=1, keepdims=True) - return q / np.clip(norm, 1.0e-8, None) - - -def _orientation_error_local(goal_quat: np.ndarray, curr_quat: np.ndarray) -> np.ndarray: - goal = _normalize_quat(goal_quat) - curr = _normalize_quat(curr_quat) - rel = np_quat_mul(goal, np_quat_inv(curr)) - rel = np_quat_canonicalize(rel) - sign = np.where(rel[:, 0:1] < 0.0, -1.0, 1.0) - return rel[:, 1:] * sign diff --git a/src/unilab/envs/locomotion/go2_arm/manip_loco.py b/src/unilab/envs/locomotion/go2_arm/manip_loco.py index f93a48440..44d1729d1 100644 --- a/src/unilab/envs/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/envs/locomotion/go2_arm/manip_loco.py @@ -12,7 +12,7 @@ from unilab.base.scene import SceneCfg from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_matrix_from_quat, np_quat_from_euler_xyz +from unilab.utils.rotation import np_matrix_from_quat, np_quat_from_euler_xyz from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig @@ -25,6 +25,10 @@ Go2ArmSensor, build_go2_arm_position_gains, ) +from unilab.utils.geometry import ( + np_cartesian_to_spherical as _cart2sphere, + np_spherical_to_cartesian as _sphere2cart, +) def _default_go2_arm_model_file() -> str: @@ -50,27 +54,6 @@ def _resolve_go2_arm_scene(cfg: "Go2ArmManipLocoCfg") -> SceneCfg: return scene -def _sphere2cart(sphere: np.ndarray) -> np.ndarray: - """Convert (..., 3)[l, phi, theta] to (..., 3)[x, y, z].""" - l = sphere[..., 0] - phi = sphere[..., 1] - theta = sphere[..., 2] - x = l * np.cos(phi) * np.cos(theta) - y = l * np.sin(theta) - z = l * np.sin(phi) * np.cos(theta) - return np.stack([x, y, z], axis=-1) - - -def _cart2sphere(cart: np.ndarray) -> np.ndarray: - """Convert (..., 3)[x, y, z] to (..., 3)[l, phi, theta].""" - cart = np.asarray(cart) - l_sq = np.sum(cart**2, axis=-1, keepdims=True) - l = np.sqrt(np.maximum(l_sq, 1e-12)) - phi = np.arctan2(cart[..., 2:3], cart[..., 0:1]) - theta = np.arcsin(np.clip(cart[..., 1:2] / l, -1.0, 1.0)) - return np.concatenate([l, phi, theta], axis=-1) - - @dataclass class InitState: pos: list[float] = field(default_factory=lambda: [0.0, 0.0, 0.42]) diff --git a/src/unilab/envs/locomotion/go2w/joystick.py b/src/unilab/envs/locomotion/go2w/joystick.py index 1b7c2082f..a36786810 100644 --- a/src/unilab/envs/locomotion/go2w/joystick.py +++ b/src/unilab/envs/locomotion/go2w/joystick.py @@ -17,7 +17,7 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_quat_mul, np_yaw_to_quat, ) diff --git a/src/unilab/envs/locomotion/go2w/rough.py b/src/unilab/envs/locomotion/go2w/rough.py index f27506801..77866c99b 100644 --- a/src/unilab/envs/locomotion/go2w/rough.py +++ b/src/unilab/envs/locomotion/go2w/rough.py @@ -12,7 +12,7 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_quat_from_euler_xyz, np_quat_mul, ) diff --git a/src/unilab/envs/manipulation/allegro_inhand/rotation.py b/src/unilab/envs/manipulation/allegro_inhand/rotation.py index d1597434e..74173a807 100644 --- a/src/unilab/envs/manipulation/allegro_inhand/rotation.py +++ b/src/unilab/envs/manipulation/allegro_inhand/rotation.py @@ -27,7 +27,10 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_quat_conjugate, np_quat_mul, np_quat_to_axis_angle +from unilab.utils.geometry import ( + np_normalize_axis, + np_quat_angular_velocity_from_pair, +) from .base import AllegroBaseCfg, AllegroBaseEnv @@ -41,15 +44,19 @@ def resolve_grasp_cache_path(cache_path: str) -> epath.Path: def normalize_rotation_axis(rotation_axis: tuple[float, float, float]) -> np.ndarray: + # Cast to the training dtype first so the norm and division happen at that + # precision, matching the pre-refactor bit-exact behavior for float32 runs. axis = np.asarray(rotation_axis, dtype=get_global_dtype()) - return np.asarray(axis / np.linalg.norm(axis), dtype=get_global_dtype()) + return np.asarray(np_normalize_axis(axis), dtype=get_global_dtype()) def compute_ball_angvel( ball_quat: np.ndarray, prev_ball_quat: np.ndarray, ctrl_dt: float ) -> np.ndarray: - rel_quat = np_quat_mul(ball_quat, np_quat_conjugate(prev_ball_quat)) - return np.asarray(np_quat_to_axis_angle(rel_quat) / ctrl_dt, dtype=get_global_dtype()) + return np.asarray( + np_quat_angular_velocity_from_pair(ball_quat, prev_ball_quat, ctrl_dt), + dtype=get_global_dtype(), + ) def compute_pd_torques( diff --git a/src/unilab/envs/manipulation/sharpa_inhand/base.py b/src/unilab/envs/manipulation/sharpa_inhand/base.py index 6b330ec4e..bdace32ee 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/base.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/base.py @@ -13,7 +13,7 @@ from unilab.base.np_env import NpEnv, NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import np_quat_apply, np_quat_mul +from unilab.utils.rotation import np_quat_apply, np_quat_mul DEFAULT_ACTUATED_JOINT_NAMES: list[str] = [ "right_thumb_CMC_FE", diff --git a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py b/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py index bdbc5c81f..b36c07fb1 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py @@ -10,7 +10,7 @@ from unilab.base.np_env import NpEnvState from unilab.dr import ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization -from unilab.envs.common.rotation import np_quat_error_magnitude +from unilab.utils.rotation import np_quat_error_magnitude from unilab.envs.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, SharpaDomainRandConfig, diff --git a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py b/src/unilab/envs/manipulation/sharpa_inhand/rotation.py index 542298ef9..030816786 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/rotation.py @@ -30,7 +30,7 @@ ResetRandomizationPayload, ) from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_quat_apply, np_quat_conjugate, np_quat_mul, @@ -43,6 +43,7 @@ resolve_grasp_cache_file, sample_scale_grasp_caches, ) +from unilab.utils.geometry import np_sample_uniform_quaternion @dataclass @@ -76,18 +77,9 @@ def sample_random_quaternion(num_envs: int) -> np.ndarray: num_envs: Number of quaternions to sample. Returns: - Quaternion array with shape ``(num_envs, 4)``. + Quaternion array with shape ``(num_envs, 4)`` in float64. """ - u1 = np.random.rand(num_envs) - u2 = np.random.rand(num_envs) * 2.0 * np.pi - u3 = np.random.rand(num_envs) * 2.0 * np.pi - - q1 = np.sqrt(1.0 - u1) * np.sin(u2) - q2 = np.sqrt(1.0 - u1) * np.cos(u2) - q3 = np.sqrt(u1) * np.sin(u3) - q4 = np.sqrt(u1) * np.cos(u3) - - return np.stack([q4, q1, q2, q3], axis=1).astype(np.float64) + return np_sample_uniform_quaternion(num_envs).astype(np.float64) class SharpaInhandRotationDRProvider(DomainRandomizationProvider): diff --git a/src/unilab/envs/manipulation/stewart/balance.py b/src/unilab/envs/manipulation/stewart/balance.py index 63e91f57e..2db2acaaa 100644 --- a/src/unilab/envs/manipulation/stewart/balance.py +++ b/src/unilab/envs/manipulation/stewart/balance.py @@ -23,7 +23,7 @@ from unilab.dr.provider import DomainRandomizationProvider from unilab.dr.types import DomainRandomizationCapabilities, ResetPlan from unilab.dtype_config import get_global_dtype -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_quat_apply_batched, np_quat_apply_inverse, np_quat_conjugate_batched, @@ -31,6 +31,7 @@ np_quat_mul_batched, np_quat_to_axis_angle, ) +from unilab.utils.geometry import np_roll_pitch_from_quat # Leg base / top-connect bodies in actuator order (a0..a5 -> slide00,slide10, # slide01,slide11,slide02,slide12), produced by the `replicate count=3` in the XML. @@ -120,13 +121,8 @@ def _ball_home_z(cfg: StewartBalanceCfg) -> float: def _roll_pitch_from_quat(quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Extract roll/pitch (rad) from a wxyz quaternion via its rotation matrix.""" - w, x, y, z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3] - r20 = 2.0 * (x * z - w * y) - r21 = 2.0 * (y * z + w * x) - r22 = 1.0 - 2.0 * (x * x + y * y) - roll = np.arctan2(r21, r22) - pitch = np.arctan2(-r20, np.sqrt(np.clip(r21 * r21 + r22 * r22, 0.0, None))) + """Extract roll/pitch (rad) from a wxyz quaternion, cast to float32.""" + roll, pitch = np_roll_pitch_from_quat(quat) return roll.astype(np.float32), pitch.astype(np.float32) diff --git a/src/unilab/envs/motion_tracking/g1/box_tracking.py b/src/unilab/envs/motion_tracking/g1/box_tracking.py index ebf78600e..5024323bb 100644 --- a/src/unilab/envs/motion_tracking/g1/box_tracking.py +++ b/src/unilab/envs/motion_tracking/g1/box_tracking.py @@ -13,8 +13,8 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.envs.common.math import np_sample_uniform -from unilab.envs.common.rotation import ( +from unilab.utils.geometry import np_sample_uniform +from unilab.utils.rotation import ( np_matrix_from_quat, np_quat_apply, np_quat_error_magnitude, diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index 2c6bfa0d6..c8cef8d4f 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -28,14 +28,18 @@ ) from unilab.dr.types import RESET_TERM_GEOM_FRICTION, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.envs.common.math import np_sample_uniform -from unilab.envs.common.rotation import ( +from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv +from unilab.utils.geometry import ( + np_gravity_z_in_body_from_quat, + np_sample_uniform, + np_write_relative_anchor_transform_pos_rot6d, +) +from unilab.utils.rotation import ( np_quat_apply, np_quat_from_euler_xyz, np_quat_inv, np_quat_mul, ) -from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv from .motion_loader import MotionData, MotionLoader, MotionSampler @@ -275,64 +279,6 @@ def _build_motion_reference_state( return qpos, qvel -def _gravity_z_in_body_from_quat_w(quat_w: np.ndarray) -> np.ndarray: - """Z component of world gravity ``[0, 0, -1]`` expressed in body frame.""" - return 2.0 * (quat_w[:, 1] * quat_w[:, 1] + quat_w[:, 2] * quat_w[:, 2]) - 1.0 - - -def _write_motion_anchor_transform( - robot_anchor_pos_w: np.ndarray, - robot_anchor_quat_w: np.ndarray, - anchor_pos_w: np.ndarray, - anchor_quat_w: np.ndarray, - out_pos: np.ndarray, - out_ori6: np.ndarray, -) -> None: - aw = robot_anchor_quat_w[:, 0] - ax = robot_anchor_quat_w[:, 1] - ay = robot_anchor_quat_w[:, 2] - az = robot_anchor_quat_w[:, 3] - - vx = anchor_pos_w[:, 0] - robot_anchor_pos_w[:, 0] - vy = anchor_pos_w[:, 1] - robot_anchor_pos_w[:, 1] - vz = anchor_pos_w[:, 2] - robot_anchor_pos_w[:, 2] - - qx = -ax - qy = -ay - qz = -az - tx = 2 * (qy * vz - qz * vy) - ty = 2 * (qz * vx - qx * vz) - tz = 2 * (qx * vy - qy * vx) - out_pos[:, 0] = vx + aw * tx + qy * tz - qz * ty - out_pos[:, 1] = vy + aw * ty + qz * tx - qx * tz - out_pos[:, 2] = vz + aw * tz + qx * ty - qy * tx - - bw = anchor_quat_w[:, 0] - bx = anchor_quat_w[:, 1] - by = anchor_quat_w[:, 2] - bz = anchor_quat_w[:, 3] - rw = aw * bw + ax * bx + ay * by + az * bz - rx = aw * bx - ax * bw - ay * bz + az * by - ry = aw * by + ax * bz - ay * bw - az * bx - rz = aw * bz - ax * by + ay * bx - az * bw - - xx = rx * rx - yy = ry * ry - zz = rz * rz - xy = rx * ry - xz = rx * rz - yz = ry * rz - wx = rw * rx - wy = rw * ry - wz = rw * rz - out_ori6[:, 0] = 1 - 2 * (yy + zz) - out_ori6[:, 1] = 2 * (xy - wz) - out_ori6[:, 2] = 2 * (xy + wz) - out_ori6[:, 3] = 1 - 2 * (xx + zz) - out_ori6[:, 4] = 2 * (xz - wy) - out_ori6[:, 5] = 2 * (yz + wx) - - class G1MotionTrackingDomainRandomizationProvider(DomainRandomizationProvider): def __init__( self, @@ -1045,8 +991,8 @@ def _compute_terminations( if self._cfg.anchor_ori_threshold < 2.0: anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] - motion_gravity_z_b = _gravity_z_in_body_from_quat_w(anchor_quat_w) - robot_gravity_z_b = _gravity_z_in_body_from_quat_w(robot_anchor_quat_w) + motion_gravity_z_b = np_gravity_z_in_body_from_quat(anchor_quat_w) + robot_gravity_z_b = np_gravity_z_in_body_from_quat(robot_anchor_quat_w) np.subtract(motion_gravity_z_b, robot_gravity_z_b, out=self._env_error) np.abs(self._env_error, out=self._env_error) np.greater(self._env_error, self._cfg.anchor_ori_threshold, out=self._env_bool) @@ -1171,7 +1117,7 @@ def _compute_obs( motion_anchor_ori_b = np.empty((num_envs, 6), dtype=dtype) joint_pos_rel = np.empty((num_envs, n_action), dtype=dtype) zero_actions = np.zeros((num_envs, n_action), dtype=dtype) - _write_motion_anchor_transform( + np_write_relative_anchor_transform_pos_rot6d( robot_anchor_pos_w, robot_anchor_quat_w, anchor_pos_w, diff --git a/src/unilab/utils/geometry.py b/src/unilab/utils/geometry.py new file mode 100644 index 000000000..0c9f22489 --- /dev/null +++ b/src/unilab/utils/geometry.py @@ -0,0 +1,216 @@ +"""Standalone numpy helpers for coordinate frames and geometry. + +Reusable pure-numpy helpers for rotations, coordinate-frame transforms, +and geometric conversions. Kept dtype-agnostic and side-effect-free so +that env / task / backend code can compose them without carrying task +policy inside a shared module. Complements the vectorized quaternion +primitives in :mod:`unilab.utils.rotation`. +""" + +from __future__ import annotations + +import numpy as np + +from unilab.utils.rotation import ( + np_quat_canonicalize, + np_quat_conjugate, + np_quat_inv, + np_quat_mul, + np_quat_to_axis_angle, +) + + +def np_sample_uniform( + lower: float | np.ndarray, + upper: float | np.ndarray, + size: tuple[int, ...], + dtype=np.float32, +) -> np.ndarray: + """Sample uniformly from ``[lower, upper]`` and cast to ``dtype``.""" + return np.random.uniform(lower, upper, size).astype(dtype) + + +def np_quat_normalize(q: np.ndarray) -> np.ndarray: + """L2-normalize quaternion(s), clamping tiny norms to 1e-8 to avoid divide-by-zero.""" + q = np.asarray(q) + if q.ndim == 1: + norm = float(np.linalg.norm(q)) + return q / max(norm, 1.0e-8) + norm = np.linalg.norm(q, axis=-1, keepdims=True) + return q / np.clip(norm, 1.0e-8, None) + + +def np_normalize_axis(axis: np.ndarray | tuple[float, ...] | list[float]) -> np.ndarray: + """Return a unit-length copy of a rotation axis vector. Raises on zero norm.""" + axis = np.asarray(axis) + norm = float(np.linalg.norm(axis)) + if norm <= 0.0: + raise ValueError(f"axis must be non-zero, got {axis!r}") + return axis / norm + + +def np_roll_pitch_from_quat(quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Roll and pitch (rad) from a w-first quaternion, computed via rotation-matrix rows. + + Supports either ``(4,)`` or ``(..., 4)`` inputs. Returned arrays match + the caller's leading shape and dtype (no upcast). + """ + w = quat[..., 0] + x = quat[..., 1] + y = quat[..., 2] + z = quat[..., 3] + r20 = 2.0 * (x * z - w * y) + r21 = 2.0 * (y * z + w * x) + r22 = 1.0 - 2.0 * (x * x + y * y) + roll = np.arctan2(r21, r22) + pitch = np.arctan2(-r20, np.sqrt(np.clip(r21 * r21 + r22 * r22, 0.0, None))) + return roll, pitch + + +def np_gravity_z_in_body_from_quat(quat_w: np.ndarray) -> np.ndarray: + """Z component of world gravity ``[0, 0, -1]`` expressed in body frame. + + Equivalent to ``np_quat_apply_inverse(quat_w, [0, 0, -1])[..., 2]`` but + computed directly from quaternion components to skip the intermediate. + """ + return 2.0 * (quat_w[..., 1] * quat_w[..., 1] + quat_w[..., 2] * quat_w[..., 2]) - 1.0 + + +def np_quat_orientation_error_local( + goal_quat: np.ndarray, curr_quat: np.ndarray +) -> np.ndarray: + """Local-frame orientation error as the signed xyz of the relative quaternion. + + Both inputs are re-normalized to unit length. Returns the imaginary + part of ``goal * inv(curr)`` after canonicalization, giving a + signed-axis-scaled error suitable for PD-style feedback terms. + """ + goal = np_quat_normalize(goal_quat) + curr = np_quat_normalize(curr_quat) + if goal.ndim == 1: + goal = goal[None, :] + if curr.ndim == 1: + curr = curr[None, :] + rel = np_quat_mul(goal, np_quat_inv(curr)) + rel = np_quat_canonicalize(rel) + sign = np.where(rel[:, 0:1] < 0.0, -1.0, 1.0) + return rel[:, 1:] * sign + + +def np_quat_angular_velocity_from_pair( + quat: np.ndarray, prev_quat: np.ndarray, dt: float +) -> np.ndarray: + """Angular velocity from two consecutive quaternions via axis-angle diff / dt.""" + rel = np_quat_mul(quat, np_quat_conjugate(prev_quat)) + return np_quat_to_axis_angle(rel) / dt + + +def np_sample_uniform_quaternion(num_samples: int) -> np.ndarray: + """Sample uniformly random unit quaternions (w-first) via Shoemake (1992). + + Returns an ``(num_samples, 4)`` array in float64 (leaves any downstream + dtype conversion to the caller). + """ + u1 = np.random.rand(num_samples) + u2 = np.random.rand(num_samples) * 2.0 * np.pi + u3 = np.random.rand(num_samples) * 2.0 * np.pi + + r1 = np.sqrt(1.0 - u1) + r2 = np.sqrt(u1) + q1 = r1 * np.sin(u2) + q2 = r1 * np.cos(u2) + q3 = r2 * np.sin(u3) + q4 = r2 * np.cos(u3) + + return np.stack([q4, q1, q2, q3], axis=1) + + +def np_spherical_to_cartesian(sphere: np.ndarray) -> np.ndarray: + """Convert ``(..., 3)[l, phi, theta]`` to ``(..., 3)[x, y, z]``. + + Uses the go2-arm spherical convention: ``phi`` sweeps in the x-z plane + from the positive x axis, and ``theta`` measures elevation toward + positive y. + """ + length = sphere[..., 0] + phi = sphere[..., 1] + theta = sphere[..., 2] + x = length * np.cos(phi) * np.cos(theta) + y = length * np.sin(theta) + z = length * np.sin(phi) * np.cos(theta) + return np.stack([x, y, z], axis=-1) + + +def np_cartesian_to_spherical(cart: np.ndarray) -> np.ndarray: + """Convert ``(..., 3)[x, y, z]`` to ``(..., 3)[l, phi, theta]`` (inverse of + :func:`np_spherical_to_cartesian`).""" + cart = np.asarray(cart) + l_sq = np.sum(cart**2, axis=-1, keepdims=True) + length = np.sqrt(np.maximum(l_sq, 1e-12)) + phi = np.arctan2(cart[..., 2:3], cart[..., 0:1]) + theta = np.arcsin(np.clip(cart[..., 1:2] / length, -1.0, 1.0)) + return np.concatenate([length, phi, theta], axis=-1) + + +def np_write_relative_anchor_transform_pos_rot6d( + source_anchor_pos_w: np.ndarray, + source_anchor_quat_w: np.ndarray, + target_anchor_pos_w: np.ndarray, + target_anchor_quat_w: np.ndarray, + out_pos: np.ndarray, + out_rot6d: np.ndarray, +) -> None: + """Fused frame-transform + 6D rotation flatten, writing in place. + + Computes the position of ``target_anchor`` in ``source_anchor``'s frame + (written to ``out_pos`` of shape ``(N, 3)``) and the relative rotation + ``conj(source_anchor) * target_anchor`` expressed as the flattened + first two columns of its rotation matrix (written to ``out_rot6d`` of + shape ``(N, 6)``). No intermediate quaternion arrays are allocated. + + This is the numpy analogue of the numba kernel + :func:`unilab.utils.numba_geometry.write_relative_anchor_transform_at`. + """ + aw = source_anchor_quat_w[:, 0] + ax = source_anchor_quat_w[:, 1] + ay = source_anchor_quat_w[:, 2] + az = source_anchor_quat_w[:, 3] + + vx = target_anchor_pos_w[:, 0] - source_anchor_pos_w[:, 0] + vy = target_anchor_pos_w[:, 1] - source_anchor_pos_w[:, 1] + vz = target_anchor_pos_w[:, 2] - source_anchor_pos_w[:, 2] + + qx = -ax + qy = -ay + qz = -az + tx = 2 * (qy * vz - qz * vy) + ty = 2 * (qz * vx - qx * vz) + tz = 2 * (qx * vy - qy * vx) + out_pos[:, 0] = vx + aw * tx + qy * tz - qz * ty + out_pos[:, 1] = vy + aw * ty + qz * tx - qx * tz + out_pos[:, 2] = vz + aw * tz + qx * ty - qy * tx + + bw = target_anchor_quat_w[:, 0] + bx = target_anchor_quat_w[:, 1] + by = target_anchor_quat_w[:, 2] + bz = target_anchor_quat_w[:, 3] + rw = aw * bw + ax * bx + ay * by + az * bz + rx = aw * bx - ax * bw - ay * bz + az * by + ry = aw * by + ax * bz - ay * bw - az * bx + rz = aw * bz - ax * by + ay * bx - az * bw + + xx = rx * rx + yy = ry * ry + zz = rz * rz + xy = rx * ry + xz = rx * rz + yz = ry * rz + wx = rw * rx + wy = rw * ry + wz = rw * rz + out_rot6d[:, 0] = 1 - 2 * (yy + zz) + out_rot6d[:, 1] = 2 * (xy - wz) + out_rot6d[:, 2] = 2 * (xy + wz) + out_rot6d[:, 3] = 1 - 2 * (xx + zz) + out_rot6d[:, 4] = 2 * (xz - wy) + out_rot6d[:, 5] = 2 * (yz + wx) diff --git a/src/unilab/envs/common/rotation.py b/src/unilab/utils/rotation.py similarity index 100% rename from src/unilab/envs/common/rotation.py rename to src/unilab/utils/rotation.py diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index be0eaa9a8..d3221079d 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -654,7 +654,7 @@ def test_g1_motion_tracking_critic_uses_clean_beyondmimic_aligned_terms(): def test_g1_motion_tracking_anchor_frame_writers_match_reference(): - from unilab.envs.common.rotation import ( + from unilab.utils.rotation import ( np_matrix_from_quat, np_quat_apply, np_quat_inv, @@ -704,7 +704,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_relative_transform_fast_path_matches_reference(): - from unilab.envs.common.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat + from unilab.utils.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv rng = np.random.default_rng(321) @@ -761,7 +761,7 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_reward_fast_path_matches_reference(): - from unilab.envs.common.rotation import np_quat_error_magnitude + from unilab.utils.rotation import np_quat_error_magnitude from unilab.envs.motion_tracking.g1.motion_loader import MotionData from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig diff --git a/tests/test_sharpa.py b/tests/test_sharpa.py index 4f5059305..47ef1da18 100644 --- a/tests/test_sharpa.py +++ b/tests/test_sharpa.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from unilab.envs.common.rotation import np_quat_apply +from unilab.utils.rotation import np_quat_apply from unilab.envs.manipulation.sharpa_inhand.base import SharpaDomainRandConfig from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, diff --git a/tests/utils/test_math_utils.py b/tests/utils/test_math_utils.py index 2193f1aed..da9c4effa 100644 --- a/tests/utils/test_math_utils.py +++ b/tests/utils/test_math_utils.py @@ -1,10 +1,10 @@ -"""Tests for quaternion helpers in unilab.envs.common.rotation.""" +"""Tests for quaternion helpers in unilab.utils.rotation.""" from __future__ import annotations import numpy as np -from unilab.envs.common.rotation import ( +from unilab.utils.rotation import ( np_matrix_first_two_cols_from_quat, np_matrix_from_quat, np_quat_angular_velocity, diff --git a/tests/utils/test_utils_package_policy.py b/tests/utils/test_utils_package_policy.py index 19a5bf0ec..e2b3f3b1d 100644 --- a/tests/utils/test_utils_package_policy.py +++ b/tests/utils/test_utils_package_policy.py @@ -7,8 +7,10 @@ ALLOWED_UTILS_MODULES = { "__init__", "device", + "geometry", "nan_guard", "numba_geometry", + "rotation", "support_matrix", "tensor", } From 898473bc61b7b87313709e63752f419918e70447 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 20:42:58 +0800 Subject: [PATCH 13/26] style: apply ruff import-order and line-length fixups from make check --- scripts/manip_loco/diagnose_go2_arm_ik.py | 2 +- src/unilab/envs/locomotion/go1/joystick.py | 2 +- src/unilab/envs/locomotion/go1/rough.py | 10 +++++----- src/unilab/envs/locomotion/go2/footstand.py | 2 +- src/unilab/envs/locomotion/go2/rough.py | 10 +++++----- src/unilab/envs/locomotion/go2_arm/base.py | 2 +- src/unilab/envs/locomotion/go2_arm/manip_loco.py | 4 +++- src/unilab/envs/locomotion/go2w/joystick.py | 8 ++++---- src/unilab/envs/locomotion/go2w/rough.py | 8 ++++---- .../envs/manipulation/sharpa_inhand/grasp_gen.py | 2 +- .../envs/manipulation/sharpa_inhand/rotation.py | 12 ++++++------ src/unilab/envs/manipulation/stewart/balance.py | 2 +- src/unilab/utils/geometry.py | 4 +--- tests/envs/test_env_configs.py | 6 +++--- tests/test_sharpa.py | 2 +- 15 files changed, 38 insertions(+), 38 deletions(-) diff --git a/scripts/manip_loco/diagnose_go2_arm_ik.py b/scripts/manip_loco/diagnose_go2_arm_ik.py index 21e26e3b7..601eb70b7 100644 --- a/scripts/manip_loco/diagnose_go2_arm_ik.py +++ b/scripts/manip_loco/diagnose_go2_arm_ik.py @@ -16,8 +16,8 @@ from unilab.base import registry from unilab.base.registry import ensure_registries -from unilab.utils.rotation import np_matrix_from_quat from unilab.envs.locomotion.go2_arm.manip_loco import RewardConfig +from unilab.utils.rotation import np_matrix_from_quat def _parse_vec3(raw: list[float], *, name: str) -> list[float]: diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index f085c832f..6a5985fd9 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -12,7 +12,6 @@ from unilab.base.np_env import NpEnvState from unilab.base.scene import SceneCfg from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import np_quat_mul, np_yaw_to_quat from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig @@ -23,6 +22,7 @@ TerrainSpawnManager, ) from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv +from unilab.utils.rotation import np_quat_mul, np_yaw_to_quat @dataclass diff --git a/src/unilab/envs/locomotion/go1/rough.py b/src/unilab/envs/locomotion/go1/rough.py index 2124de2a5..9f67c2305 100644 --- a/src/unilab/envs/locomotion/go1/rough.py +++ b/src/unilab/envs/locomotion/go1/rough.py @@ -14,11 +14,6 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import ( - np_quat_apply_inverse, - np_quat_from_euler_xyz, - np_quat_mul, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import ( Commands, @@ -58,6 +53,11 @@ random_rough, wave_terrain, ) +from unilab.utils.rotation import ( + np_quat_apply_inverse, + np_quat_from_euler_xyz, + np_quat_mul, +) # pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false diff --git a/src/unilab/envs/locomotion/go2/footstand.py b/src/unilab/envs/locomotion/go2/footstand.py index 3e5f6572c..24d876279 100644 --- a/src/unilab/envs/locomotion/go2/footstand.py +++ b/src/unilab/envs/locomotion/go2/footstand.py @@ -13,13 +13,13 @@ from unilab.base.scene import SceneCfg from unilab.dr import ResetPlan, ResetRandomizationPayload from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig from unilab.envs.locomotion.common.dr_provider import LocomotionDRProvider from unilab.envs.locomotion.common.rewards import RewardContext from unilab.envs.locomotion.go2.base import ControlConfig, Go2BaseCfg, Go2BaseEnv, NoiseConfig +from unilab.utils.rotation import np_quat_apply, np_quat_apply_inverse @dataclass diff --git a/src/unilab/envs/locomotion/go2/rough.py b/src/unilab/envs/locomotion/go2/rough.py index b7bee9b4f..0d7d5975e 100644 --- a/src/unilab/envs/locomotion/go2/rough.py +++ b/src/unilab/envs/locomotion/go2/rough.py @@ -12,11 +12,6 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization, zero_actions from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import ( - np_quat_apply_inverse, - np_quat_from_euler_xyz, - np_quat_mul, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import ( apply_heading_yaw_feedback, @@ -54,6 +49,11 @@ random_rough, wave_terrain, ) +from unilab.utils.rotation import ( + np_quat_apply_inverse, + np_quat_from_euler_xyz, + np_quat_mul, +) # pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false diff --git a/src/unilab/envs/locomotion/go2_arm/base.py b/src/unilab/envs/locomotion/go2_arm/base.py index ec76bb539..c54e13718 100644 --- a/src/unilab/envs/locomotion/go2_arm/base.py +++ b/src/unilab/envs/locomotion/go2_arm/base.py @@ -5,7 +5,6 @@ import numpy as np from unilab.base.backend import SimBackend -from unilab.utils.rotation import np_matrix_from_quat from unilab.envs.locomotion.common.base import ( ControlConfigBase, LocomotionBaseCfg, @@ -13,6 +12,7 @@ Sensor, ) from unilab.utils.geometry import np_quat_orientation_error_local +from unilab.utils.rotation import np_matrix_from_quat DEFAULT_LEG_ANGLES = np.asarray( [ diff --git a/src/unilab/envs/locomotion/go2_arm/manip_loco.py b/src/unilab/envs/locomotion/go2_arm/manip_loco.py index 44d1729d1..08194b4c4 100644 --- a/src/unilab/envs/locomotion/go2_arm/manip_loco.py +++ b/src/unilab/envs/locomotion/go2_arm/manip_loco.py @@ -12,7 +12,6 @@ from unilab.base.scene import SceneCfg from unilab.dr.types import ResetPlan from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import np_matrix_from_quat, np_quat_from_euler_xyz from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import Commands from unilab.envs.locomotion.common.domain_rand import DomainRandConfig @@ -27,8 +26,11 @@ ) from unilab.utils.geometry import ( np_cartesian_to_spherical as _cart2sphere, +) +from unilab.utils.geometry import ( np_spherical_to_cartesian as _sphere2cart, ) +from unilab.utils.rotation import np_matrix_from_quat, np_quat_from_euler_xyz def _default_go2_arm_model_file() -> str: diff --git a/src/unilab/envs/locomotion/go2w/joystick.py b/src/unilab/envs/locomotion/go2w/joystick.py index a36786810..43059e83e 100644 --- a/src/unilab/envs/locomotion/go2w/joystick.py +++ b/src/unilab/envs/locomotion/go2w/joystick.py @@ -17,10 +17,6 @@ zero_actions, ) from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import ( - np_quat_mul, - np_yaw_to_quat, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import ( Commands, @@ -43,6 +39,10 @@ compute_go2w_motor_ctrl, stack_joint_sensors, ) +from unilab.utils.rotation import ( + np_quat_mul, + np_yaw_to_quat, +) GO2W_HIP_INDICES = np.asarray([0, 3, 6, 9], dtype=np.int32) diff --git a/src/unilab/envs/locomotion/go2w/rough.py b/src/unilab/envs/locomotion/go2w/rough.py index 77866c99b..ab96b61ba 100644 --- a/src/unilab/envs/locomotion/go2w/rough.py +++ b/src/unilab/envs/locomotion/go2w/rough.py @@ -12,10 +12,6 @@ from unilab.dr import DomainRandomizationManager, ResetPlan from unilab.dr.dr_utils import zero_actions from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import ( - np_quat_from_euler_xyz, - np_quat_mul, -) from unilab.envs.locomotion.common import rewards from unilab.envs.locomotion.common.commands import ( Commands, @@ -54,6 +50,10 @@ random_rough, wave_terrain, ) +from unilab.utils.rotation import ( + np_quat_from_euler_xyz, + np_quat_mul, +) # pyright: reportIncompatibleVariableOverride=false, reportAttributeAccessIssue=false, reportCallIssue=false diff --git a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py b/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py index b36c07fb1..2081bf95e 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/grasp_gen.py @@ -10,7 +10,6 @@ from unilab.base.np_env import NpEnvState from unilab.dr import ResetPlan from unilab.dr.dr_utils import build_common_reset_randomization -from unilab.utils.rotation import np_quat_error_magnitude from unilab.envs.manipulation.sharpa_inhand.base import ( SOURCE_DEFAULT_HAND_JOINT_POS_DEG, SharpaDomainRandConfig, @@ -22,6 +21,7 @@ SharpaInhandRotationDRProvider, SharpaInhandRotationEnv, ) +from unilab.utils.rotation import np_quat_error_magnitude def _default_sharpa_grasp_domain_rand() -> SharpaDomainRandConfig: diff --git a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py b/src/unilab/envs/manipulation/sharpa_inhand/rotation.py index 030816786..04285d40b 100644 --- a/src/unilab/envs/manipulation/sharpa_inhand/rotation.py +++ b/src/unilab/envs/manipulation/sharpa_inhand/rotation.py @@ -30,12 +30,6 @@ ResetRandomizationPayload, ) from unilab.dtype_config import get_global_dtype -from unilab.utils.rotation import ( - np_quat_apply, - np_quat_conjugate, - np_quat_mul, - np_quat_to_axis_angle, -) from unilab.envs.manipulation.sharpa_inhand.base import ( SharpaInhandBaseCfg, SharpaInhandBaseEnv, @@ -44,6 +38,12 @@ sample_scale_grasp_caches, ) from unilab.utils.geometry import np_sample_uniform_quaternion +from unilab.utils.rotation import ( + np_quat_apply, + np_quat_conjugate, + np_quat_mul, + np_quat_to_axis_angle, +) @dataclass diff --git a/src/unilab/envs/manipulation/stewart/balance.py b/src/unilab/envs/manipulation/stewart/balance.py index 2db2acaaa..75552146f 100644 --- a/src/unilab/envs/manipulation/stewart/balance.py +++ b/src/unilab/envs/manipulation/stewart/balance.py @@ -23,6 +23,7 @@ from unilab.dr.provider import DomainRandomizationProvider from unilab.dr.types import DomainRandomizationCapabilities, ResetPlan from unilab.dtype_config import get_global_dtype +from unilab.utils.geometry import np_roll_pitch_from_quat from unilab.utils.rotation import ( np_quat_apply_batched, np_quat_apply_inverse, @@ -31,7 +32,6 @@ np_quat_mul_batched, np_quat_to_axis_angle, ) -from unilab.utils.geometry import np_roll_pitch_from_quat # Leg base / top-connect bodies in actuator order (a0..a5 -> slide00,slide10, # slide01,slide11,slide02,slide12), produced by the `replicate count=3` in the XML. diff --git a/src/unilab/utils/geometry.py b/src/unilab/utils/geometry.py index 0c9f22489..019360092 100644 --- a/src/unilab/utils/geometry.py +++ b/src/unilab/utils/geometry.py @@ -76,9 +76,7 @@ def np_gravity_z_in_body_from_quat(quat_w: np.ndarray) -> np.ndarray: return 2.0 * (quat_w[..., 1] * quat_w[..., 1] + quat_w[..., 2] * quat_w[..., 2]) - 1.0 -def np_quat_orientation_error_local( - goal_quat: np.ndarray, curr_quat: np.ndarray -) -> np.ndarray: +def np_quat_orientation_error_local(goal_quat: np.ndarray, curr_quat: np.ndarray) -> np.ndarray: """Local-frame orientation error as the signed xyz of the relative quaternion. Both inputs are re-normalized to unit length. Returns the imaginary diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index d3221079d..3d9b0009f 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -654,13 +654,13 @@ def test_g1_motion_tracking_critic_uses_clean_beyondmimic_aligned_terms(): def test_g1_motion_tracking_anchor_frame_writers_match_reference(): + from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv from unilab.utils.rotation import ( np_matrix_from_quat, np_quat_apply, np_quat_inv, np_quat_mul, ) - from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv rng = np.random.default_rng(123) num_envs = 4 @@ -704,8 +704,8 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_relative_transform_fast_path_matches_reference(): - from unilab.utils.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv + from unilab.utils.rotation import np_quat_apply, np_quat_inv, np_quat_mul, np_yaw_quat rng = np.random.default_rng(321) num_envs = 4 @@ -761,9 +761,9 @@ def random_quat(shape: tuple[int, ...]) -> np.ndarray: def test_g1_motion_tracking_reward_fast_path_matches_reference(): - from unilab.utils.rotation import np_quat_error_magnitude from unilab.envs.motion_tracking.g1.motion_loader import MotionData from unilab.envs.motion_tracking.g1.tracking import G1MotionTrackingEnv, RewardConfig + from unilab.utils.rotation import np_quat_error_magnitude rng = np.random.default_rng(456) num_envs = 3 diff --git a/tests/test_sharpa.py b/tests/test_sharpa.py index 47ef1da18..5b4618f02 100644 --- a/tests/test_sharpa.py +++ b/tests/test_sharpa.py @@ -5,7 +5,6 @@ import numpy as np import pytest -from unilab.utils.rotation import np_quat_apply from unilab.envs.manipulation.sharpa_inhand.base import SharpaDomainRandConfig from unilab.envs.manipulation.sharpa_inhand.grasp_gen import ( SharpaInhandRotationGraspCfg, @@ -15,6 +14,7 @@ SharpaInhandRotationEnv, sample_random_quaternion, ) +from unilab.utils.rotation import np_quat_apply def test_sharpa_gravity_direction_randomization_matches_rotated_gravity( From 34072808a8024332f20f9768309a24d19b175817 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 21:22:52 +0800 Subject: [PATCH 14/26] perf: accelerate g1 joystick update_state with numba --- benchmark/benchmark_g1_joystick_numba.py | 67 +- src/unilab/envs/locomotion/g1/joystick.py | 6 +- .../envs/locomotion/g1/joystick_numba.py | 706 ++++++++++++++---- .../test_g1_joystick_numba_benchmark.py | 2 + .../envs/locomotion/g1/test_joystick_numba.py | 136 ++++ 5 files changed, 763 insertions(+), 154 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index 11712c428..8442dc21f 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -1,10 +1,10 @@ -"""Benchmark the task-specific Numba backend for G1 joystick rewards. +"""Benchmark the task-specific Numba backend for G1 joystick update_state. This benchmark intentionally avoids MuJoCo/Motrix construction. It measures the hot slice owned by ``src/unilab/envs/locomotion/g1/joystick.py``: -* baseline: real ``G1WalkEnv._compute_reward`` reward dispatch + termination; -* accelerated: ``G1WalkNumbaAccelerator.compute``. +* baseline: real ``G1WalkEnv`` reward, termination, and observation assembly; +* accelerated: ``G1WalkNumbaAccelerator.compute_update_state``. Synthetic backend arrays keep the benchmark deterministic while still using the same reward functions, reward config fields, sensor names, and accelerator entry @@ -209,8 +209,17 @@ class SyntheticBatch: dof_vel: np.ndarray +class _NoiseCfg: + level = 0.0 + scale_gyro = 0.1 + scale_gravity = 0.0 + scale_joint_angle = 0.02 + scale_joint_vel = 0.3 + + class _Cfg: ctrl_dt = 0.02 + noise_config = _NoiseCfg() def make_profile_specs() -> dict[str, ProfileSpec]: @@ -360,7 +369,9 @@ def make_batch(num_envs: int, spec: ProfileSpec, seed: int) -> SyntheticBatch: ) -def compute_numpy(batch: SyntheticBatch) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: +def compute_numpy( + batch: SyntheticBatch, +) -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray, dict[str, float]]: env = batch.env info = {**batch.info, "log": {}} max_tilt_rad = np.deg2rad(env._reward_cfg.max_tilt_deg) @@ -377,14 +388,15 @@ def compute_numpy(batch: SyntheticBatch) -> tuple[np.ndarray, np.ndarray, dict[s batch.dof_pos, batch.dof_vel, ) - return reward, terminated, info.get("log", {}) + obs = env._compute_obs(info, batch.linvel, batch.gyro, batch.gravity, batch.dof_pos, batch.dof_vel) + return obs, reward, terminated, info.get("log", {}) def compute_numba( batch: SyntheticBatch, accelerator: G1WalkNumbaAccelerator -) -> tuple[np.ndarray, np.ndarray, dict[str, float]]: +) -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray, dict[str, float]]: info = {**batch.info, "log": {}} - out = accelerator.compute( + out = accelerator.compute_update_state( env=batch.env, info=info, linvel=batch.linvel, @@ -394,8 +406,9 @@ def compute_numba( dof_vel=batch.dof_vel, scales=batch.env._reward_cfg.scales, enable_log=True, + noise_level=0.0, ) - return out.reward, out.terminated, out.log + return out.obs, out.reward, out.terminated, out.log def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: @@ -410,16 +423,20 @@ def time_call(fn, *, iters: int, warmup: int) -> tuple[float, float, float]: def check_parity(batch: SyntheticBatch, accelerator: G1WalkNumbaAccelerator) -> dict[str, float]: - reward_np, terminated_np, log_np = compute_numpy(batch) - reward_nb, terminated_nb, log_nb = compute_numba(batch, accelerator) + obs_np, reward_np, terminated_np, log_np = compute_numpy(batch) + obs_nb, reward_nb, terminated_nb, log_nb = compute_numba(batch, accelerator) np.testing.assert_allclose(reward_nb, reward_np, rtol=5e-5, atol=5e-5) np.testing.assert_array_equal(terminated_nb, terminated_np) + np.testing.assert_allclose(obs_nb["obs"], obs_np["obs"], rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(obs_nb["critic"], obs_np["critic"], rtol=1e-6, atol=1e-6) for key, value in log_np.items(): if key in log_nb: np.testing.assert_allclose(log_nb[key], value, rtol=5e-5, atol=5e-5) return { "max_abs_reward_diff": float(np.max(np.abs(reward_nb - reward_np))), "termination_mismatch": float(np.count_nonzero(terminated_nb != terminated_np)), + "max_abs_actor_obs_diff": float(np.max(np.abs(obs_nb["obs"] - obs_np["obs"]))), + "max_abs_critic_obs_diff": float(np.max(np.abs(obs_nb["critic"] - obs_np["critic"]))), } @@ -551,7 +568,7 @@ def _run_e2e_collector_pair( create_env, actor action sampling, env.step, terminal-observation handling, replay writes, and bookkeeping are all included. Learner updates are excluded. It is intentionally optional because it constructs a real MuJoCo env and is - much heavier than the synthetic reward+termination hot-slice benchmark above. + much heavier than the synthetic update_state hot-slice benchmark above. """ from benchmark.benchmark_offpolicy_collector_active import _build_and_run_case @@ -896,7 +913,7 @@ def save_plots( best_by_case = _best_numba_by_case(records) fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(21, 6)) - fig.suptitle(f"G1 joystick Numba reward+termination benchmark\n{device_info}", fontsize=13) + fig.suptitle(f"G1 joystick Numba update_state benchmark\n{device_info}", fontsize=13) ax1 = axes[0] for profile in profiles: @@ -923,7 +940,7 @@ def save_plots( fontsize=8, ) ax1.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="break-even") - ax1.set_title("Reward+termination: best vs numpy") + ax1.set_title("Update_state: best vs numpy") ax1.set_xlabel("num_envs") ax1.set_ylabel("Speedup vs numpy") ax1.set_xscale("log", base=2) @@ -964,7 +981,7 @@ def save_plots( fontsize=8, ) ax2.axhline(1.0, color="grey", linestyle=":", linewidth=0.9, label="1T") - ax2.set_title("Reward+termination: parallel speedup") + ax2.set_title("Update_state: parallel speedup") ax2.set_xlabel("num_envs") ax2.set_ylabel("Speedup vs numba 1T") ax2.set_xscale("log", base=2) @@ -1076,7 +1093,7 @@ def write_report( "measured_threads": getattr(args, "measured_threads", None), "skipped_threads": getattr(args, "skipped_threads", None), "numba_max_threads": getattr(args, "numba_max_threads", None), - "scope": "G1 joystick reward+termination hot slice; synthetic backend arrays", + "scope": "G1 joystick update_state hot slice; synthetic backend arrays", "e2e_enabled": args.e2e, "e2e_num_envs": args.e2e_num_envs, "e2e_case": args.e2e_case, @@ -1099,21 +1116,21 @@ def write_report( summary_lines = [ "# G1 joystick Numba benchmark", "", - "Scope: reward dispatch plus termination for `G1WalkEnv.update_state`, using", - "deterministic synthetic backend arrays. Physics stepping, obs assembly, reset,", + "Scope: reward, termination, and observation assembly for `G1WalkEnv.update_state`, using", + "deterministic synthetic backend arrays. Physics stepping, reset,", "and policy inference are intentionally out of scope.", "", "## Numba-specific hot slice", "", "Definitions:", "", - "- `vs numpy`: numpy reward+termination time divided by numba reward+termination time.", + "- `vs numpy`: numpy update_state hot-slice time divided by numba update_state time.", "- `vs numba1T`: numba 1-thread time divided by numba N-thread time.", "- `parallel eff`: `vs numba1T / N`; this is the only parallel-efficiency column.", "", "This section measures only the part this task-specific Numba backend accelerates:", - "`G1WalkEnv` reward dispatch plus termination inside `update_state`. It excludes", - "physics, observation assembly, reset/RNG, policy inference, learner work, and replay.", + "`G1WalkEnv` reward, termination, and observation assembly inside `update_state`. It excludes", + "physics, reset/RNG, policy inference, learner work, and replay.", "", "Profile meanings:", "", @@ -1172,10 +1189,10 @@ def write_report( "", "This compares the synthetic `sac_default` hot-slice milliseconds saved", "with the collector-measured `update_state_ms` milliseconds saved at the", - "same `num_envs` and Numba thread count. The hot slice is only the", - "reward+termination work replaced by `joystick_numba.py`; collector", - "`update_state_ms` also includes state reads, observation assembly,", - "resets, state replacement, and bookkeeping around that hot slice.", + "same `num_envs` and Numba thread count. The hot slice is the synthetic", + "full update_state array work replaced by `joystick_numba.py`; collector", + "`update_state_ms` also includes backend state reads, resets, state", + "replacement, and bookkeeping around that hot slice.", "", "```text", reconciliation_table, @@ -1295,7 +1312,7 @@ def main() -> None: args.skipped_threads = sorted({threads for threads in args.threads if threads > max_threads}) print("=" * 80) - print("G1 joystick Numba benchmark: reward dispatch + termination") + print("G1 joystick Numba benchmark: update_state hot slice") print("=" * 80) print(f"host numba threads: {max_threads}") print( diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 90201dbd2..7c19939f2 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -362,7 +362,7 @@ def update_state(self, state: NpEnvState) -> NpEnvState: dof_vel = self.get_dof_vel() if self._numba_accelerator is not None: - accel_result = self._numba_accelerator.compute( + accel_result = self._numba_accelerator.compute_update_state( env=self, info=state.info, linvel=linvel, @@ -372,9 +372,11 @@ def update_state(self, state: NpEnvState) -> NpEnvState: dof_vel=dof_vel, scales=self._reward_cfg.scales, enable_log=self._enable_reward_log, + noise_level=self._cfg.noise_config.level, ) terminated = accel_result.terminated reward = accel_result.reward + obs = accel_result.obs state.info["log"] = accel_result.log else: max_tilt_rad = np.deg2rad(self._reward_cfg.max_tilt_deg) @@ -384,8 +386,8 @@ def update_state(self, state: NpEnvState) -> NpEnvState: self._terrain_relative_base_height() < self._reward_cfg.min_base_height, ) reward = self._compute_reward(state.info, linvel, gyro, gravity, dof_pos, dof_vel) + obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel) - obs = self._compute_obs(state.info, linvel, gyro, gravity, dof_pos, dof_vel) state = state.replace(obs=obs, reward=reward, terminated=terminated) done = state.terminated | state.truncated diff --git a/src/unilab/envs/locomotion/g1/joystick_numba.py b/src/unilab/envs/locomotion/g1/joystick_numba.py index e9eec25e5..57dbec099 100644 --- a/src/unilab/envs/locomotion/g1/joystick_numba.py +++ b/src/unilab/envs/locomotion/g1/joystick_numba.py @@ -63,6 +63,14 @@ class G1WalkNumbaResult: log: dict[str, float] +@dataclass(frozen=True) +class G1WalkNumbaUpdateStateResult: + obs: dict[str, np.ndarray] + reward: np.ndarray + terminated: np.ndarray + log: dict[str, float] + + def _active_terms(scales: Mapping[str, float]) -> frozenset[str]: return frozenset(name for name, scale in scales.items() if scale != 0.0) @@ -276,6 +284,175 @@ def terminated_i(gravity, base_height, max_tilt_rad, min_base_height, i): gz = 1.0 return math.acos(gz) > max_tilt_rad or base_height[i] < min_base_height + @njit(fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _compute_reward_termination_i( + linvel, + gyro, + gravity, + dof_pos, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + default_angles, + pose_weights, + upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + scale, + ctrl_dt, + tracking_sigma, + base_height_target, + min_base_height, + max_tilt_rad, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + close_feet_threshold, + reward, + terminated, + log_scratch, + tid, + i, + ): + n_action = dof_pos.shape[1] + r = 0.0 + + w = tracking_lin_vel_i(linvel, commands, tracking_sigma, i) * scale[0] + r += w + log_scratch[tid, 0] += w + + w = tracking_ang_vel_i(gyro, commands, tracking_sigma, i) * scale[1] + r += w + log_scratch[tid, 1] += w + + w = forward_progress_i(linvel, commands, i) * scale[2] + r += w + log_scratch[tid, 2] += w + + w = under_speed_i(linvel, commands, i) * scale[3] + r += w + log_scratch[tid, 3] += w + + w = lin_vel_z_i(linvel, i) * scale[4] + r += w + log_scratch[tid, 4] += w + + orientation = orientation_i(gravity, i) + w = orientation * scale[5] + r += w + log_scratch[tid, 5] += w + w = orientation * scale[6] + r += w + log_scratch[tid, 6] += w + + ang_vel_xy = ang_vel_xy_i(gyro, i) + w = ang_vel_xy * scale[7] + r += w + log_scratch[tid, 7] += w + w = ang_vel_xy * scale[8] + r += w + log_scratch[tid, 8] += w + + action_rate = action_rate_i(current_actions, last_actions, n_action, i) + w = action_rate * scale[9] + r += w + log_scratch[tid, 9] += w + w = action_rate * scale[10] + r += w + log_scratch[tid, 10] += w + + w = base_height_i(base_height, base_height_target, i) * scale[11] + r += w + log_scratch[tid, 11] += w + + pose = weighted_pose_i(dof_pos, default_angles, pose_weights, n_action, i) + w = pose * scale[12] + r += w + log_scratch[tid, 12] += w + upper_body_pose = weighted_pose_i( + dof_pos, default_angles, upper_body_pose_weights, n_action, i + ) + w = upper_body_pose * scale[13] + r += w + log_scratch[tid, 13] += w + + w = close_feet_xy_i(left_foot_pos, right_foot_pos, close_feet_threshold, i) * scale[14] + r += w + log_scratch[tid, 14] += w + + w = feet_ori_i(left_foot_quat, right_foot_quat, i) * scale[15] + r += w + log_scratch[tid, 15] += w + + w = ( + feet_phase_i( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + i, + ) + * scale[16] + ) + r += w + log_scratch[tid, 16] += w + + w = ( + feet_phase_contrast_i( + linvel, + gait_phase, + left_foot_pos, + right_foot_pos, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + i, + ) + * scale[17] + ) + r += w + log_scratch[tid, 17] += w + + w = ( + feet_phase_contact_i( + linvel, + gait_phase, + left_contact, + right_contact, + feet_phase_swing_height, + min_forward_speed_for_gait_reward, + i, + ) + * scale[18] + ) + r += w + log_scratch[tid, 18] += w + + w = feet_double_stance_i(commands, left_contact, right_contact, i) * scale[19] + r += w + log_scratch[tid, 19] += w + + w = feet_air_time_i(feet_air_time, i) * scale[20] + r += w + log_scratch[tid, 20] += w + + w = scale[21] + r += w + log_scratch[tid, 21] += w + + reward[i] = r * ctrl_dt + terminated[i] = terminated_i(gravity, base_height, max_tilt_rad, min_base_height, i) + @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] def _compute_reward_termination_kernel( linvel, @@ -311,141 +488,249 @@ def _compute_reward_termination_kernel( reward, terminated, log_scratch, + ): + n = reward.shape[0] + for i in prange(n): + _compute_reward_termination_i( + linvel, + gyro, + gravity, + dof_pos, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + default_angles, + pose_weights, + upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + scale, + ctrl_dt, + tracking_sigma, + base_height_target, + min_base_height, + max_tilt_rad, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + close_feet_threshold, + reward, + terminated, + log_scratch, + get_thread_id(), + i, + ) + + @_dev + def _write_actor_obs_i( + gyro, + gravity, + dof_pos, + dof_vel, + current_actions, + commands, + gait_phase, + default_angles, + actor_gyro_scale, + actor_dof_vel_scale, + actor_obs, + n_action, + i, + ): + out = 0 + actor_obs[i, out] = gyro[i, 0] * actor_gyro_scale + actor_obs[i, out + 1] = gyro[i, 1] * actor_gyro_scale + actor_obs[i, out + 2] = gyro[i, 2] * actor_gyro_scale + out += 3 + actor_obs[i, out] = -gravity[i, 0] + actor_obs[i, out + 1] = -gravity[i, 1] + actor_obs[i, out + 2] = -gravity[i, 2] + out += 3 + for j in range(n_action): + actor_obs[i, out + j] = dof_pos[i, j] - default_angles[j] + out += n_action + for j in range(n_action): + actor_obs[i, out + j] = dof_vel[i, j] * actor_dof_vel_scale + out += n_action + for j in range(n_action): + actor_obs[i, out + j] = current_actions[i, j] + out += n_action + actor_obs[i, out] = commands[i, 0] + actor_obs[i, out + 1] = commands[i, 1] + actor_obs[i, out + 2] = commands[i, 2] + out += 3 + actor_obs[i, out] = gait_phase[i, 0] + actor_obs[i, out + 1] = gait_phase[i, 1] + + @_dev + def _write_critic_obs_i( + linvel, + gyro, + gravity, + dof_pos, + dof_vel, + current_actions, + commands, + gait_phase, + default_angles, + critic_gyro_scale, + critic_dof_vel_scale, + critic_linvel_scale, + critic_obs, + n_action, + i, + ): + out = 0 + critic_obs[i, out] = gyro[i, 0] * critic_gyro_scale + critic_obs[i, out + 1] = gyro[i, 1] * critic_gyro_scale + critic_obs[i, out + 2] = gyro[i, 2] * critic_gyro_scale + out += 3 + critic_obs[i, out] = -gravity[i, 0] + critic_obs[i, out + 1] = -gravity[i, 1] + critic_obs[i, out + 2] = -gravity[i, 2] + out += 3 + for j in range(n_action): + critic_obs[i, out + j] = dof_pos[i, j] - default_angles[j] + out += n_action + for j in range(n_action): + critic_obs[i, out + j] = dof_vel[i, j] * critic_dof_vel_scale + out += n_action + for j in range(n_action): + critic_obs[i, out + j] = current_actions[i, j] + out += n_action + critic_obs[i, out] = commands[i, 0] + critic_obs[i, out + 1] = commands[i, 1] + critic_obs[i, out + 2] = commands[i, 2] + out += 3 + critic_obs[i, out] = gait_phase[i, 0] + critic_obs[i, out + 1] = gait_phase[i, 1] + out += 2 + critic_obs[i, out] = linvel[i, 0] * critic_linvel_scale + critic_obs[i, out + 1] = linvel[i, 1] * critic_linvel_scale + critic_obs[i, out + 2] = linvel[i, 2] * critic_linvel_scale + + @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _compute_update_state_kernel( + linvel, + gyro, + gravity, + dof_pos, + dof_vel, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + default_angles, + pose_weights, + upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + scale, + ctrl_dt, + tracking_sigma, + base_height_target, + min_base_height, + max_tilt_rad, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + close_feet_threshold, + actor_gyro_scale, + actor_dof_vel_scale, + critic_gyro_scale, + critic_dof_vel_scale, + critic_linvel_scale, + actor_obs, + critic_obs, + reward, + terminated, + log_scratch, ): n = reward.shape[0] n_action = dof_pos.shape[1] for i in prange(n): tid = get_thread_id() - r = 0.0 - - w = tracking_lin_vel_i(linvel, commands, tracking_sigma, i) * scale[0] - r += w - log_scratch[tid, 0] += w - - w = tracking_ang_vel_i(gyro, commands, tracking_sigma, i) * scale[1] - r += w - log_scratch[tid, 1] += w - - w = forward_progress_i(linvel, commands, i) * scale[2] - r += w - log_scratch[tid, 2] += w - - w = under_speed_i(linvel, commands, i) * scale[3] - r += w - log_scratch[tid, 3] += w - - w = lin_vel_z_i(linvel, i) * scale[4] - r += w - log_scratch[tid, 4] += w - - orientation = orientation_i(gravity, i) - w = orientation * scale[5] - r += w - log_scratch[tid, 5] += w - w = orientation * scale[6] - r += w - log_scratch[tid, 6] += w - - ang_vel_xy = ang_vel_xy_i(gyro, i) - w = ang_vel_xy * scale[7] - r += w - log_scratch[tid, 7] += w - w = ang_vel_xy * scale[8] - r += w - log_scratch[tid, 8] += w - - action_rate = action_rate_i(current_actions, last_actions, n_action, i) - w = action_rate * scale[9] - r += w - log_scratch[tid, 9] += w - w = action_rate * scale[10] - r += w - log_scratch[tid, 10] += w - - w = base_height_i(base_height, base_height_target, i) * scale[11] - r += w - log_scratch[tid, 11] += w - - pose = weighted_pose_i(dof_pos, default_angles, pose_weights, n_action, i) - w = pose * scale[12] - r += w - log_scratch[tid, 12] += w - upper_body_pose = weighted_pose_i( - dof_pos, default_angles, upper_body_pose_weights, n_action, i - ) - w = upper_body_pose * scale[13] - r += w - log_scratch[tid, 13] += w - - w = close_feet_xy_i(left_foot_pos, right_foot_pos, close_feet_threshold, i) * scale[14] - r += w - log_scratch[tid, 14] += w - - w = feet_ori_i(left_foot_quat, right_foot_quat, i) * scale[15] - r += w - log_scratch[tid, 15] += w - - w = ( - feet_phase_i( - linvel, - gait_phase, - left_foot_pos, - right_foot_pos, - feet_phase_swing_height, - feet_phase_tracking_sigma, - min_forward_speed_for_gait_reward, - i, - ) - * scale[16] + _compute_reward_termination_i( + linvel, + gyro, + gravity, + dof_pos, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + default_angles, + pose_weights, + upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + scale, + ctrl_dt, + tracking_sigma, + base_height_target, + min_base_height, + max_tilt_rad, + feet_phase_swing_height, + feet_phase_tracking_sigma, + min_forward_speed_for_gait_reward, + close_feet_threshold, + reward, + terminated, + log_scratch, + tid, + i, ) - r += w - log_scratch[tid, 16] += w - - w = ( - feet_phase_contrast_i( - linvel, - gait_phase, - left_foot_pos, - right_foot_pos, - feet_phase_swing_height, - feet_phase_tracking_sigma, - min_forward_speed_for_gait_reward, - i, - ) - * scale[17] + _write_actor_obs_i( + gyro, + gravity, + dof_pos, + dof_vel, + current_actions, + commands, + gait_phase, + default_angles, + actor_gyro_scale, + actor_dof_vel_scale, + actor_obs, + n_action, + i, ) - r += w - log_scratch[tid, 17] += w - - w = ( - feet_phase_contact_i( - linvel, - gait_phase, - left_contact, - right_contact, - feet_phase_swing_height, - min_forward_speed_for_gait_reward, - i, - ) - * scale[18] + _write_critic_obs_i( + linvel, + gyro, + gravity, + dof_pos, + dof_vel, + current_actions, + commands, + gait_phase, + default_angles, + critic_gyro_scale, + critic_dof_vel_scale, + critic_linvel_scale, + critic_obs, + n_action, + i, ) - r += w - log_scratch[tid, 18] += w - - w = feet_double_stance_i(commands, left_contact, right_contact, i) * scale[19] - r += w - log_scratch[tid, 19] += w - - w = feet_air_time_i(feet_air_time, i) * scale[20] - r += w - log_scratch[tid, 20] += w - - w = scale[21] - r += w - log_scratch[tid, 21] += w - - reward[i] = r * ctrl_dt - terminated[i] = terminated_i(gravity, base_height, max_tilt_rad, min_base_height, i) class G1WalkNumbaAccelerator: @@ -468,6 +753,9 @@ def __init__( default_angles: np.ndarray, pose_weights: np.ndarray, upper_body_pose_weights: np.ndarray, + actor_obs_width: int, + critic_obs_width: int, + walk_observation_profile: bool, num_threads: int | None = None, ) -> None: self.num_envs = int(num_envs) @@ -484,6 +772,9 @@ def __init__( self.default_angles = np.asarray(default_angles, dtype=np.float64) self.pose_weights = np.asarray(pose_weights, dtype=np.float64) self.upper_body_pose_weights = np.asarray(upper_body_pose_weights, dtype=np.float64) + self.actor_obs_width = int(actor_obs_width) + self.critic_obs_width = int(critic_obs_width) + self.walk_observation_profile = bool(walk_observation_profile) self.num_threads = num_threads self.scale = np.zeros((len(TERM_ORDER),), dtype=np.float64) self._zero_vec2 = np.zeros((self.num_envs, 2), dtype=np.float64) @@ -515,6 +806,9 @@ def from_env(cls, env: Any, num_threads: int | None = None) -> "G1WalkNumbaAccel default_angles=env.default_angles, pose_weights=env._pose_weights, upper_body_pose_weights=env._upper_body_pose_weights, + actor_obs_width=env.obs_groups_spec["obs"], + critic_obs_width=env.obs_groups_spec["critic"], + walk_observation_profile=env._uses_walk_observation_profile(), num_threads=num_threads, ) @@ -654,3 +948,161 @@ def compute( if self.scale[idx] != 0.0: log[f"reward/{name}"] = float(term_sums[idx] / linvel.shape[0]) return G1WalkNumbaResult(reward=reward, terminated=terminated, log=log) + + def compute_update_state( + self, + *, + env: Any, + info: dict[str, Any], + linvel: np.ndarray, + gyro: np.ndarray, + gravity: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + scales: Mapping[str, float], + enable_log: bool, + noise_level: float, + ) -> G1WalkNumbaUpdateStateResult: + if not NUMBA_AVAILABLE: + raise RuntimeError( + "G1Walk Numba accelerator was constructed while numba is unavailable; " + "this indicates an invalid accelerator state." + ) + if float(noise_level) > 0.0: + raise RuntimeError( + "G1Walk numba_acceleration=True does not support observation noise yet. " + "Set env.noise_config.level=0.0 or disable numba_acceleration." + ) + self._sync_scales(scales) + + n = linvel.shape[0] + if n != self.num_envs: + raise ValueError( + "G1Walk Numba update_state only supports full-batch updates; " + f"got {n} rows for configured num_envs={self.num_envs}." + ) + + active = _active_terms(scales) + backend = env._backend + dtype = linvel.dtype + base_height = np.asarray(backend.get_base_pos()[:, 2], dtype=dtype) + commands = np.asarray(info["commands"], dtype=dtype) + current_actions = np.asarray( + info.get("current_actions", np.zeros((self.num_envs, self.num_action), dtype=dtype)), + dtype=dtype, + ) + last_actions = np.asarray( + info.get("last_actions", np.zeros((self.num_envs, self.num_action), dtype=dtype)), + dtype=dtype, + ) + gait_phase = np.asarray(info.get("gait_phase", self._zero_vec2), dtype=dtype) + feet_air_time = np.asarray(info.get("feet_air_time", self._zero_vec2), dtype=dtype) + + if active & FOOT_POSITION_TERMS: + left_foot_pos = np.asarray(backend.get_sensor_data("left_foot_pos"), dtype=dtype) + right_foot_pos = np.asarray(backend.get_sensor_data("right_foot_pos"), dtype=dtype) + else: + left_foot_pos = right_foot_pos = self._zero_vec3.astype(dtype, copy=False) + + if active & FOOT_QUAT_TERMS: + left_foot_quat = np.asarray(backend.get_sensor_data("left_foot_quat"), dtype=dtype) + right_foot_quat = np.asarray(backend.get_sensor_data("right_foot_quat"), dtype=dtype) + else: + left_foot_quat = right_foot_quat = self._zero_vec4.astype(dtype, copy=False) + + if active & FOOT_CONTACT_TERMS: + left_contact = _aggregated_contact( + backend, + ( + "left_foot_contact_0", + "left_foot_contact_1", + "left_foot_contact_2", + "left_foot_contact_3", + ), + ) + right_contact = _aggregated_contact( + backend, + ( + "right_foot_contact_0", + "right_foot_contact_1", + "right_foot_contact_2", + "right_foot_contact_3", + ), + ) + else: + left_contact = right_contact = self._zero_bool + + if self.num_threads is not None: + set_num_threads(self.num_threads) + nthreads = get_num_threads() + actor_obs = np.empty((n, self.actor_obs_width), dtype=dtype) + critic_obs = np.empty((n, self.critic_obs_width), dtype=dtype) + reward = np.empty((n,), dtype=dtype) + terminated = np.empty((n,), dtype=np.bool_) + log_scratch = np.zeros((nthreads, len(TERM_ORDER)), dtype=np.float64) + + actor_gyro_scale = 0.25 if self.walk_observation_profile else 1.0 + actor_dof_vel_scale = 0.05 if self.walk_observation_profile else 1.0 + critic_gyro_scale = 0.25 if self.walk_observation_profile else 1.0 + critic_dof_vel_scale = 0.05 if self.walk_observation_profile else 1.0 + critic_linvel_scale = 2.0 if self.walk_observation_profile else 1.0 + + _compute_update_state_kernel( + linvel, + gyro, + gravity, + dof_pos, + dof_vel, + base_height, + commands, + current_actions, + last_actions, + gait_phase, + self.default_angles, + self.pose_weights, + self.upper_body_pose_weights, + left_foot_pos, + right_foot_pos, + left_foot_quat, + right_foot_quat, + left_contact, + right_contact, + feet_air_time, + self.scale, + self.ctrl_dt, + self.tracking_sigma, + self.base_height_target, + self.min_base_height, + self.max_tilt_rad, + self.feet_phase_swing_height, + self.feet_phase_tracking_sigma, + self.min_forward_speed_for_gait_reward, + self.close_feet_threshold, + actor_gyro_scale, + actor_dof_vel_scale, + critic_gyro_scale, + critic_dof_vel_scale, + critic_linvel_scale, + actor_obs, + critic_obs, + reward, + terminated, + log_scratch, + ) + + step_count = info.get("steps") + should_log = enable_log and ( + int(step_count[0]) % 4 == 0 if isinstance(step_count, np.ndarray) else True + ) + log = {} if should_log else info.get("log", {}) + if should_log: + term_sums = log_scratch.sum(axis=0) + for idx, name in enumerate(TERM_ORDER): + if self.scale[idx] != 0.0: + log[f"reward/{name}"] = float(term_sums[idx] / n) + return G1WalkNumbaUpdateStateResult( + obs={"obs": actor_obs, "critic": critic_obs}, + reward=reward, + terminated=terminated, + log=log, + ) diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index 1d1fdef20..a62632472 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -17,6 +17,8 @@ def test_g1_joystick_numba_benchmark_builds_records_and_matches_numpy() -> None: assert parity["termination_mismatch"] == 0.0 assert parity["max_abs_reward_diff"] < 1.0e-5 + assert parity["max_abs_actor_obs_diff"] < 1.0e-6 + assert parity["max_abs_critic_obs_diff"] < 1.0e-6 assert {record.path for record in records} == {"numpy_dispatch", "numba_accelerator"} assert any(record.path == "numba_accelerator" and record.threads == 1 for record in records) assert all( diff --git a/tests/envs/locomotion/g1/test_joystick_numba.py b/tests/envs/locomotion/g1/test_joystick_numba.py index dc3402303..b094430cc 100644 --- a/tests/envs/locomotion/g1/test_joystick_numba.py +++ b/tests/envs/locomotion/g1/test_joystick_numba.py @@ -49,6 +49,15 @@ def get_sensor_data(self, name: str): class _Cfg: ctrl_dt = 0.02 + noise_config = None + + +class _NoiseCfg: + level = 0.0 + scale_gyro = 0.1 + scale_gravity = 0.0 + scale_joint_angle = 0.02 + scale_joint_vel = 0.3 class _RewardCfg: @@ -65,14 +74,24 @@ class _RewardCfg: class _Env: def __init__(self, n: int, n_action: int): self.num_envs = n + self._num_envs = n self._num_action = n_action self._cfg = _Cfg() + self._cfg.noise_config = _NoiseCfg() self._reward_cfg = _RewardCfg() + self._enable_reward_log = True self.default_angles = np.zeros((n_action,), dtype=np.float32) self._pose_weights = np.ones((n_action,), dtype=np.float32) self._upper_body_pose_weights = np.ones((n_action,), dtype=np.float32) self._backend = _Backend(n) + @property + def obs_groups_spec(self): + return {"obs": 98, "critic": 101} + + def _uses_walk_observation_profile(self): + return True + @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") def test_g1_walk_numba_unsupported_active_term_raises(): @@ -164,6 +183,123 @@ def test_g1_walk_numba_basic_reward_parity(): assert set(out.log) == {f"reward/{name}" for name in scales} +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_walk_numba_update_state_parity_without_noise(): + from unilab.envs.locomotion.g1.joystick import G1WalkEnv + + n = 512 + n_action = 29 + rng = np.random.default_rng(3) + env = _Env(n, n_action) + env.default_angles = rng.normal(scale=0.1, size=(n_action,)).astype(np.float32) + env._pose_weights = rng.uniform(0.1, 2.0, size=(n_action,)).astype(np.float32) + env._upper_body_pose_weights = env._pose_weights.copy() + env._upper_body_pose_weights[:12] = 0.0 + env._reward_cfg.scales = { + "tracking_lin_vel": 2.0, + "tracking_ang_vel": 1.5, + "penalty_orientation": -10.0, + "penalty_action_rate": -4.0, + "pose": -0.5, + "penalty_feet_ori": -20.0, + "feet_phase": 5.0, + "alive": 10.0, + } + env._init_reward_functions = G1WalkEnv._init_reward_functions.__get__(env, G1WalkEnv) + env._build_reward_context = G1WalkEnv._build_reward_context.__get__(env, G1WalkEnv) + env._compute_reward = G1WalkEnv._compute_reward.__get__(env, G1WalkEnv) + env._compute_obs = G1WalkEnv._compute_obs.__get__(env, G1WalkEnv) + env._gait_reward_gate = G1WalkEnv._gait_reward_gate.__get__(env, G1WalkEnv) + env._obs_noise = lambda data, scale: data + env._reward_fns = { + "tracking_lin_vel": __import__( + "unilab.envs.locomotion.common.rewards", fromlist=["tracking_lin_vel"] + ).tracking_lin_vel, + "tracking_ang_vel": __import__( + "unilab.envs.locomotion.common.rewards", fromlist=["tracking_ang_vel"] + ).tracking_ang_vel, + "penalty_orientation": __import__( + "unilab.envs.locomotion.common.rewards", fromlist=["orientation"] + ).orientation, + "penalty_action_rate": __import__( + "unilab.envs.locomotion.common.rewards", fromlist=["action_rate"] + ).action_rate, + "pose": __import__( + "unilab.envs.locomotion.common.rewards", fromlist=["weighted_pose"] + ).weighted_pose, + "penalty_feet_ori": G1WalkEnv._reward_feet_ori.__get__(env, G1WalkEnv), + "feet_phase": G1WalkEnv._reward_feet_phase.__get__(env, G1WalkEnv), + "alive": __import__("unilab.envs.locomotion.common.rewards", fromlist=["alive"]).alive, + } + + info = { + "steps": np.zeros((n,), dtype=np.uint32), + "commands": rng.normal(size=(n, 3)).astype(np.float32), + "current_actions": rng.normal(size=(n, n_action)).astype(np.float32), + "last_actions": rng.normal(size=(n, n_action)).astype(np.float32), + "gait_phase": rng.uniform(0.0, 2 * np.pi, size=(n, 2)).astype(np.float32), + "log": {}, + } + linvel = rng.normal(size=(n, 3)).astype(np.float32) + gyro = rng.normal(size=(n, 3)).astype(np.float32) + gravity = np.tile(np.array([0.01, -0.02, 0.99], dtype=np.float32), (n, 1)) + dof_pos = rng.normal(scale=0.01, size=(n, n_action)).astype(np.float32) + dof_vel = rng.normal(size=(n, n_action)).astype(np.float32) + + max_tilt_rad = np.deg2rad(env._reward_cfg.max_tilt_deg) + terminated_np = ( + np.arccos(np.clip(gravity[:, 2], -1.0, 1.0)) > max_tilt_rad + ) | (env._backend.get_base_pos()[:, 2] < env._reward_cfg.min_base_height) + reward_np = env._compute_reward( + {**info, "log": {}}, linvel, gyro, gravity, dof_pos, dof_vel + ) + obs_np = env._compute_obs(info, linvel, gyro, gravity, dof_pos, dof_vel) + + accel = G1WalkNumbaAccelerator.from_env(env, num_threads=2) + out = accel.compute_update_state( + env=env, + info=info, + linvel=linvel, + gyro=gyro, + gravity=gravity, + dof_pos=dof_pos, + dof_vel=dof_vel, + scales=env._reward_cfg.scales, + enable_log=True, + noise_level=0.0, + ) + + np.testing.assert_allclose(out.reward, reward_np, rtol=5e-5, atol=5e-5) + np.testing.assert_array_equal(out.terminated, terminated_np) + np.testing.assert_allclose(out.obs["obs"], obs_np["obs"], rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(out.obs["critic"], obs_np["critic"], rtol=1e-6, atol=1e-6) + + +@pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +def test_g1_walk_numba_update_state_rejects_noise(): + env = _Env(8, 29) + accel = G1WalkNumbaAccelerator.from_env(env, num_threads=1) + arrays = { + "linvel": np.zeros((8, 3), dtype=np.float32), + "gyro": np.zeros((8, 3), dtype=np.float32), + "gravity": np.tile(np.array([0.0, 0.0, 1.0], dtype=np.float32), (8, 1)), + "dof_pos": np.zeros((8, 29), dtype=np.float32), + "dof_vel": np.zeros((8, 29), dtype=np.float32), + } + with pytest.raises(RuntimeError, match="does not support observation noise"): + accel.compute_update_state( + env=env, + info={ + "steps": np.zeros((8,), dtype=np.uint32), + "commands": np.zeros((8, 3), dtype=np.float32), + }, + scales={"tracking_lin_vel": 1.0}, + enable_log=False, + noise_level=1.0, + **arrays, + ) + + @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") def test_g1_walk_numba_term_py_funcs_match_numpy_math(): from unilab.envs.locomotion.g1 import joystick_numba as T From f7a5bfc3d90acd4de48cd97208bb743565f97bb3 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Tue, 7 Jul 2026 23:45:21 +0800 Subject: [PATCH 15/26] perf: profile and optimize Motrix reset_done / set_state hot path Extend SimBackend.set_state contract to return an optional {"timing": {...}} dict so the DR manager can surface per-substep timings next to the outer dr_reset_set_state_ms wall-clock. Both Motrix and MuJoCo backends emit the same 16-key schema for column stability. Instrument Motrix set_state into: qpos_convert, mask, data_slice, data_reset, clear_forces, geom_overrides, reset_rand, set_dof_vel, set_dof_pos, actuator_ctrl, forward_kinematic, refresh_pose_cache, invalidate_velocity, internal_gap. MuJoCo populates pool_reset + state_scatter; the motrix-only keys stay at 0.0 for schema parity. Apply owner-layer Motrix set_state optimizations, all driven by the new sub-timing report: - Reusable scratch buffers for the bool mask and the qpos motrix conversion, sized once at backend init; no per-reset np.zeros / np.array(copy=True) allocation. - Hoist env_ids_intp = np.asarray(env_indices, dtype=np.intp) to the top of set_state and pass it into _clear_applied_body_forces, _apply_init_geom_size_overrides, _apply_reset_randomization, and _refresh_link_pose_cache so the helpers stop re-converting. - In-place _mujoco_qpos_to_motrix_into(dst, qpos) for the hot path; original _mujoco_qpos_to_motrix stays for cold paths. - Skip np.ascontiguousarray on actuator ctrl when the slice is already c-contiguous. Benchmark harness (benchmark_offpolicy_collector_active.py) exposes the new keys via NP_ENV_STEP_TIMING_KEYS + CSV fields, plus two new formatters _format_set_state_detail_table (motrix keyset) and _format_set_state_mujoco_table. New tests: schema + parity contract for both backends (tests/base/test_sim_backend_set_state_timing.py), DR manager merge + malformed-timing tolerance (tests/dr/test_manager.py), benchmark CSV schema + table formatters (tests/benchmark/test_offpolicy_collector_active_benchmark.py). New microbenchmark: benchmark/benchmark_motrix_set_state_ab.py drives set_state under baseline vs optimized variants to isolate the Python- side wins from Rust-boundary noise. Local A/B (Apple M5 Max, 8192 env, motrix, 100 measured steps): - dr_reset_set_state_ms: 3.274 ms (walk_flat numba), fully broken down by sub-key (Data reset 37% + FK 20% + Data slice 16% dominant, all in Rust boundary). - reset_done_ms: g1_walk_flat -0.102 ms, g1_motion_tracking -0.157 ms vs b96224a6 baseline (issue #663 comment 4904777445). - set_state microbenchmark: mask 1.6-2.3x, qpos_convert 1.05-1.35x, actuator_ctrl 1.07-1.17x; outer set_state 1.08-1.16x at low reset ratios (which is the collector's normal regime). Refs: #679, #665 --- benchmark/benchmark_motrix_set_state_ab.py | 368 ++++++++++++++++++ .../benchmark_offpolicy_collector_active.py | 153 +++++++- src/unilab/base/backend/base.py | 10 +- src/unilab/base/backend/motrix/backend.py | 172 +++++++- src/unilab/base/backend/mujoco/backend.py | 39 +- src/unilab/base/np_env.py | 40 ++ src/unilab/dr/manager.py | 13 +- .../base/test_sim_backend_set_state_timing.py | 190 +++++++++ ...st_offpolicy_collector_active_benchmark.py | 81 ++++ tests/dr/test_manager.py | 90 +++++ 10 files changed, 1128 insertions(+), 28 deletions(-) create mode 100644 benchmark/benchmark_motrix_set_state_ab.py create mode 100644 tests/base/test_sim_backend_set_state_timing.py diff --git a/benchmark/benchmark_motrix_set_state_ab.py b/benchmark/benchmark_motrix_set_state_ab.py new file mode 100644 index 000000000..d9bf5c4dd --- /dev/null +++ b/benchmark/benchmark_motrix_set_state_ab.py @@ -0,0 +1,368 @@ +"""Local A/B microbenchmark for Motrix ``set_state`` (issue #679). + +Two variants run against the same ``MotrixBackend`` instance and the same +sequence of reset requests: + +* ``baseline`` — the pre-#679 body: fresh ``np.zeros(num_envs, bool)`` mask each + call, full ``np.array(qpos, copy=True)`` conversion, and helpers receive raw + ``env_indices`` so they redo ``np.asarray(env_indices, dtype=np.intp)`` + internally, plus an unconditional ``np.ascontiguousarray`` on actuator ctrl. +* ``optimized`` — the current backend method: reusable scratch buffers, single + hoisted ``env_ids_intp``, and a c-contiguous fast path for actuator ctrl. + +Both variants instrument the same 16-key schema documented in +:mod:`unilab.base.np_env`. Only the internals differ. + +The script does NOT touch the collector loop / physics — it only drives the +``set_state`` method — so it isolates the change under study. + +Usage:: + + uv run python benchmark/benchmark_motrix_set_state_ab.py \ + --num-envs 1024 --iters 30 --warmup 5 + uv run python benchmark/benchmark_motrix_set_state_ab.py --out benchmark/outputs/set_state_ab.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time +from pathlib import Path +from typing import Any + +import numpy as np + +ROOT_DIR = Path(__file__).resolve().parents[1] +if str(ROOT_DIR) not in sys.path: + sys.path.append(str(ROOT_DIR)) + +from unilab.assets import ASSETS_ROOT_PATH # noqa: E402 +from unilab.base.scene import SceneCfg # noqa: E402 +from unilab.dr.types import ResetRandomizationPayload # noqa: E402 + +_G1_MODEL_FILE = str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") + +# Schema mirrors BACKEND_SET_STATE_DETAIL_TIMING_KEYS in np_env.py; keys that +# don't apply to motrix stay at 0.0. +_TIMING_KEYS = ( + "set_state_mask_ms", + "set_state_data_slice_ms", + "set_state_data_reset_ms", + "set_state_clear_forces_ms", + "set_state_geom_overrides_ms", + "set_state_reset_rand_ms", + "set_state_set_dof_vel_ms", + "set_state_set_dof_pos_ms", + "set_state_actuator_ctrl_ms", + "set_state_forward_kinematic_ms", + "set_state_refresh_pose_cache_ms", + "set_state_invalidate_velocity_ms", + "set_state_qpos_convert_ms", + "set_state_pool_reset_ms", + "set_state_state_scatter_ms", + "set_state_internal_gap_ms", +) + + +def _baseline_set_state( + self: Any, + env_indices: np.ndarray, + qpos: np.ndarray, + qvel: np.ndarray, + randomization: ResetRandomizationPayload | None = None, +) -> dict: + """Inline copy of the pre-#679 body, monkey-patched onto the backend. + + Kept in the benchmark rather than the backend so the shipping code stays + single-path. Uses only the helper kwargs that existed before the refactor + (which are still accepted with defaults), so it stays a faithful baseline. + """ + timing: dict[str, float] = {k: 0.0 for k in _TIMING_KEYS} + outer_t0 = time.perf_counter() + + t0 = time.perf_counter() + # NOTE: allocates a fresh (num_envs, qpos_dim) array every call. + qpos_motrix = self._mujoco_qpos_to_motrix(qpos) + timing["set_state_qpos_convert_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + # NOTE: allocates a fresh bool mask every call. + mask = np.zeros(self._num_envs, dtype=bool) + mask[env_indices] = True + timing["set_state_mask_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + data_slice = self._data[mask] + timing["set_state_data_slice_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + data_slice.reset(self._model) + timing["set_state_data_reset_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + # NOTE: helper redoes np.asarray(env_indices, dtype=np.intp) internally + # because we do NOT pass env_ids_intp. + self._clear_applied_body_forces(env_indices) + timing["set_state_clear_forces_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._apply_init_geom_size_overrides(data_slice, env_indices) + timing["set_state_geom_overrides_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._apply_reset_randomization(data_slice, env_indices, randomization) + timing["set_state_reset_rand_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + data_slice.set_dof_vel(qvel) + timing["set_state_set_dof_vel_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + data_slice.set_dof_pos(qpos_motrix, self._model) + timing["set_state_set_dof_pos_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + if self._supports_position_actuator_gains and len(self._joint_dof_pos_indices) == int( + self.num_actuators + ): + if self._joint_dof_pos_slice is not None: + ctrl = qpos_motrix[:, self._joint_dof_pos_slice] + else: + ctrl = qpos_motrix[:, self._joint_dof_pos_indices] + elif self._actuator_joint_pos_indices is not None: + if self._actuator_joint_pos_slice is not None: + ctrl = qpos_motrix[:, self._actuator_joint_pos_slice] + else: + ctrl = qpos_motrix[:, self._actuator_joint_pos_indices] + else: + ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype) + # NOTE: unconditional ascontiguousarray copy. + data_slice.actuator_ctrls = np.ascontiguousarray(ctrl) + timing["set_state_actuator_ctrl_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._model.forward_kinematic(data_slice) + timing["set_state_forward_kinematic_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._refresh_link_pose_cache(env_indices, data_slice=data_slice) + timing["set_state_refresh_pose_cache_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._invalidate_link_velocity_cache() + timing["set_state_invalidate_velocity_ms"] = (time.perf_counter() - t0) * 1000.0 + + outer_total_ms = (time.perf_counter() - outer_t0) * 1000.0 + measured_ms = sum( + v + for k, v in timing.items() + if k + not in ( + "set_state_pool_reset_ms", + "set_state_state_scatter_ms", + "set_state_internal_gap_ms", + ) + ) + timing["set_state_internal_gap_ms"] = outer_total_ms - measured_ms + return {"timing": timing} + + +def _make_backend(num_envs: int): + from unilab.base.backend.motrix.backend import MotrixBackend + + return MotrixBackend( + SceneCfg(model_file=_G1_MODEL_FILE), + num_envs, + sim_dt=0.005, + base_name="pelvis", + ) + + +def _identity_qpos(nq: int, num_envs: int, xyz=(0.0, 0.0, 0.8)) -> np.ndarray: + q = np.zeros((num_envs, nq), dtype=np.float32) + q[:, :3] = xyz + q[:, 3] = 1.0 + return q + + +def _run_variant( + backend, + *, + variant: str, + num_envs: int, + reset_ratio: float, + warmup: int, + iters: int, + rng: np.random.Generator, + randomization: ResetRandomizationPayload | None, +) -> dict[str, Any]: + """Drive set_state ``warmup+iters`` times; report per-key mean/median ms.""" + nq = backend.get_dof_pos().shape[-1] + 7 + nv = backend.get_dof_vel().shape[-1] + 6 + num_reset = max(1, int(num_envs * reset_ratio)) + + samples: dict[str, list[float]] = {k: [] for k in _TIMING_KEYS} + outer_samples: list[float] = [] + + # New env_indices per call (mimics real reset_done pattern). + def _sample_indices() -> np.ndarray: + return rng.choice(num_envs, size=num_reset, replace=False).astype(np.int32) + + fn = ( + backend.set_state + if variant == "optimized" + else lambda *a, **kw: _baseline_set_state(backend, *a, **kw) + ) + + for step in range(warmup + iters): + env_indices = _sample_indices() + qpos = _identity_qpos(nq, num_reset) + qvel = np.zeros((num_reset, nv), dtype=np.float32) + t0 = time.perf_counter() + result = fn(env_indices, qpos, qvel, randomization=randomization) + outer_ms = (time.perf_counter() - t0) * 1000.0 + if step < warmup: + continue + outer_samples.append(outer_ms) + assert isinstance(result, dict) and "timing" in result, "variant must return timing dict" + for k in _TIMING_KEYS: + samples[k].append(float(result["timing"].get(k, 0.0))) + + summary = { + "variant": variant, + "num_envs": num_envs, + "num_reset": num_reset, + "reset_ratio": reset_ratio, + "warmup": warmup, + "iters": iters, + "outer_ms": { + "mean": statistics.mean(outer_samples), + "median": statistics.median(outer_samples), + "min": min(outer_samples), + "max": max(outer_samples), + }, + "sub_ms": { + k: { + "mean": statistics.mean(v), + "median": statistics.median(v), + } + for k, v in samples.items() + }, + } + return summary + + +def _format_ab_table(baseline: dict[str, Any], optimized: dict[str, Any]) -> str: + """Render an ASCII A/B table with per-key mean ms and delta.""" + header = f"{'sub-timing':40s} {'baseline ms':>12s} {'optimized ms':>13s} {'delta ms':>10s} {'speedup':>9s}" + rows = [header, "-" * len(header)] + total_saved = 0.0 + for k in _TIMING_KEYS: + b = baseline["sub_ms"][k]["mean"] + o = optimized["sub_ms"][k]["mean"] + delta = b - o + total_saved += delta + speedup = (b / o) if o > 1e-9 else float("inf") + rows.append(f"{k:40s} {b:12.4f} {o:13.4f} {delta:10.4f} {speedup:9.2f}x") + rows.append("-" * len(header)) + b_outer = baseline["outer_ms"]["mean"] + o_outer = optimized["outer_ms"]["mean"] + rows.append( + f"{'outer set_state wall-clock':40s} {b_outer:12.4f} {o_outer:13.4f}" + f" {b_outer - o_outer:10.4f} {(b_outer / o_outer) if o_outer > 1e-9 else float('inf'):9.2f}x" + ) + rows.append(f"{'sum of sub-key deltas':40s} {'':12s} {'':13s} {total_saved:10.4f}") + return "\n".join(rows) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--num-envs", type=int, default=1024) + parser.add_argument( + "--reset-ratios", + type=str, + default="0.05,0.25,0.50", + help="Comma-separated reset ratios (fraction of envs reset per call).", + ) + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--out", + type=Path, + default=None, + help="Optional path for JSON output.", + ) + parser.add_argument( + "--randomization", + action="store_true", + help="Include a small DR payload (gravity) to exercise set_state_reset_rand_ms.", + ) + args = parser.parse_args(argv) + + ratios = tuple(float(x) for x in args.reset_ratios.split(",")) + backend = _make_backend(args.num_envs) + print(f"MotrixBackend ready: num_envs={backend.num_envs}, dt=0.005") + + output: dict[str, Any] = { + "num_envs": args.num_envs, + "iters": args.iters, + "warmup": args.warmup, + "results": [], + } + + for ratio in ratios: + num_reset = max(1, int(args.num_envs * ratio)) + payload = ( + ResetRandomizationPayload( + gravity=np.tile(np.asarray([0.0, 0.0, -9.81], dtype=np.float32), (num_reset, 1)), + ) + if args.randomization + else None + ) + rng_a = np.random.default_rng(args.seed) + rng_b = np.random.default_rng(args.seed) + baseline = _run_variant( + backend, + variant="baseline", + num_envs=args.num_envs, + reset_ratio=ratio, + warmup=args.warmup, + iters=args.iters, + rng=rng_a, + randomization=payload, + ) + optimized = _run_variant( + backend, + variant="optimized", + num_envs=args.num_envs, + reset_ratio=ratio, + warmup=args.warmup, + iters=args.iters, + rng=rng_b, + randomization=payload, + ) + print(f"\n=== reset_ratio={ratio:.2f} ({num_reset}/{args.num_envs} envs per set_state) ===") + print(_format_ab_table(baseline, optimized)) + output["results"].append( + { + "reset_ratio": ratio, + "num_reset": num_reset, + "baseline": baseline, + "optimized": optimized, + } + ) + + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(output, indent=2), encoding="utf-8") + print(f"\nWrote JSON: {args.out}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 89f6eee30..7b002bac1 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -104,6 +104,25 @@ "dr_reset_obs_get_body_pose_ms", "dr_reset_observation_compute_obs_ms", "dr_reset_observation_internal_gap_ms", + # Backend set_state internals (see BACKEND_SET_STATE_DETAIL_TIMING_KEYS in + # src/unilab/base/np_env.py). All backends emit the same key set for column + # stability; sub-keys that don't apply report 0.0. + "set_state_mask_ms", + "set_state_data_slice_ms", + "set_state_data_reset_ms", + "set_state_clear_forces_ms", + "set_state_geom_overrides_ms", + "set_state_reset_rand_ms", + "set_state_set_dof_vel_ms", + "set_state_set_dof_pos_ms", + "set_state_actuator_ctrl_ms", + "set_state_forward_kinematic_ms", + "set_state_refresh_pose_cache_ms", + "set_state_invalidate_velocity_ms", + "set_state_qpos_convert_ms", + "set_state_pool_reset_ms", + "set_state_state_scatter_ms", + "set_state_internal_gap_ms", ) NP_ENV_STEP_COUNT_KEYS = ("reset_done_count",) NP_ENV_STEP_SAMPLE_KEYS = (*NP_ENV_STEP_TIMING_KEYS, *NP_ENV_STEP_COUNT_KEYS) @@ -135,6 +154,22 @@ ("dr_reset_obs_get_body_pose_ms", "dr_reset_obs_get_body_pose_ms"), ("dr_reset_observation_compute_obs_ms", "dr_reset_observation_compute_obs_ms"), ("dr_reset_observation_internal_gap_ms", "dr_reset_observation_internal_gap_ms"), + ("set_state_mask_ms", "set_state_mask_ms"), + ("set_state_data_slice_ms", "set_state_data_slice_ms"), + ("set_state_data_reset_ms", "set_state_data_reset_ms"), + ("set_state_clear_forces_ms", "set_state_clear_forces_ms"), + ("set_state_geom_overrides_ms", "set_state_geom_overrides_ms"), + ("set_state_reset_rand_ms", "set_state_reset_rand_ms"), + ("set_state_set_dof_vel_ms", "set_state_set_dof_vel_ms"), + ("set_state_set_dof_pos_ms", "set_state_set_dof_pos_ms"), + ("set_state_actuator_ctrl_ms", "set_state_actuator_ctrl_ms"), + ("set_state_forward_kinematic_ms", "set_state_forward_kinematic_ms"), + ("set_state_refresh_pose_cache_ms", "set_state_refresh_pose_cache_ms"), + ("set_state_invalidate_velocity_ms", "set_state_invalidate_velocity_ms"), + ("set_state_qpos_convert_ms", "set_state_qpos_convert_ms"), + ("set_state_pool_reset_ms", "set_state_pool_reset_ms"), + ("set_state_state_scatter_ms", "set_state_state_scatter_ms"), + ("set_state_internal_gap_ms", "set_state_internal_gap_ms"), ) @@ -1193,6 +1228,23 @@ def _format_np_env_value(result: CollectorResult, key: str, *, digits: int = 1) return f"{stat.mean_ms:.{digits}f}" +def _format_set_state_sub_ms(result: CollectorResult, key: str) -> str: + """Format a backend set_state sub-timing as ``ms (%of set_state)``. + + The percentage is relative to ``dr_reset_set_state_ms`` (the outer + wall-clock measurement in DomainRandomizationManager), not to env_step_ms, + so a reader can see which sub-step dominates set_state. + """ + stat = result.env_step_timing_ms_per_vector_step.get(key) + if stat is None: + return "n/a" + outer = result.env_step_timing_ms_per_vector_step.get("dr_reset_set_state_ms") + if outer is None or outer.mean_ms <= 0.0: + return f"{stat.mean_ms:.3f} (n/a)" + pct = 100.0 * stat.mean_ms / outer.mean_ms + return f"{stat.mean_ms:.3f} ({pct:.1f}%)" + + def _format_throughput_table(results: list[CollectorResult]) -> str: headers = ( "Algo", @@ -1227,11 +1279,7 @@ def _format_throughput_table(results: list[CollectorResult]) -> str: result.case.task, result.case.runtime_sim_backend, result.case.variant, - ( - f"on/{result.case.numba_threads}T" - if result.case.numba_acceleration - else "off" - ), + (f"on/{result.case.numba_threads}T" if result.case.numba_acceleration else "off"), f"{result.case.num_envs:,}", f"{result.collector_active_steps_per_sec:,.0f}", f"{result.total_active_ms / result.measure_steps:.3f} (100.0%)", @@ -1322,6 +1370,90 @@ def _format_dr_reset_timing_table(results: list[CollectorResult]) -> str: return _format_table(headers, rows) +_SET_STATE_MOTRIX_KEYS = ( + ("set_state_qpos_convert_ms", "QPos convert"), + ("set_state_mask_ms", "Mask"), + ("set_state_data_slice_ms", "Data slice"), + ("set_state_data_reset_ms", "Data reset"), + ("set_state_clear_forces_ms", "Clear forces"), + ("set_state_geom_overrides_ms", "Geom overrides"), + ("set_state_reset_rand_ms", "Reset rand"), + ("set_state_set_dof_vel_ms", "Set dof vel"), + ("set_state_set_dof_pos_ms", "Set dof pos"), + ("set_state_actuator_ctrl_ms", "Actuator ctrl"), + ("set_state_forward_kinematic_ms", "FK"), + ("set_state_refresh_pose_cache_ms", "Refresh pose"), + ("set_state_invalidate_velocity_ms", "Inval vel"), + ("set_state_internal_gap_ms", "Gap"), +) + +_SET_STATE_MUJOCO_KEYS = ( + ("set_state_qpos_convert_ms", "QPos convert"), + ("set_state_pool_reset_ms", "Pool reset"), + ("set_state_state_scatter_ms", "State scatter"), + ("set_state_internal_gap_ms", "Gap"), +) + + +def _format_set_state_detail_table(results: list[CollectorResult]) -> str: + """Backend set_state sub-timing table (motrix keyset). + + Renders the 14 motrix-oriented sub-keys next to the outer + ``dr_reset_set_state_ms``. Backends that don't populate a key emit 0.0 so + columns stay stable across backends. MuJoCo runs will show 0.0 for the + motrix-only sub-keys; use :func:`_format_set_state_mujoco_table` for the + MuJoCo-oriented view instead. + """ + headers = ( + "Algo", + "Task", + "Backend", + "Set state ms (%env, %active)", + *(label for _, label in _SET_STATE_MOTRIX_KEYS), + ) + rows = [] + for result in results: + env_step = result.phase_ms_per_vector_step.get("env_step_ms") + if env_step is None: + continue + rows.append( + ( + result.case.algo, + result.case.task, + result.case.runtime_sim_backend, + _format_np_env_timing(result, "dr_reset_set_state_ms"), + *(_format_set_state_sub_ms(result, key) for key, _ in _SET_STATE_MOTRIX_KEYS), + ) + ) + return _format_table(headers, rows) + + +def _format_set_state_mujoco_table(results: list[CollectorResult]) -> str: + """Backend set_state sub-timing table (mujoco keyset).""" + headers = ( + "Algo", + "Task", + "Backend", + "Set state ms (%env, %active)", + *(label for _, label in _SET_STATE_MUJOCO_KEYS), + ) + rows = [] + for result in results: + env_step = result.phase_ms_per_vector_step.get("env_step_ms") + if env_step is None: + continue + rows.append( + ( + result.case.algo, + result.case.task, + result.case.runtime_sim_backend, + _format_np_env_timing(result, "dr_reset_set_state_ms"), + *(_format_set_state_sub_ms(result, key) for key, _ in _SET_STATE_MUJOCO_KEYS), + ) + ) + return _format_table(headers, rows) + + def _format_np_env_step_timing_table(results: list[CollectorResult]) -> str: headers = ( "Algo", @@ -1482,7 +1614,9 @@ def _format_numba_ab_table(results: list[CollectorResult]) -> str: f"{result.collector_active_steps_per_sec / baseline.collector_active_steps_per_sec:.2f}x", _format_speedup(new_total, base_total), _format_saved(new_total, base_total), - "n/a" if base_env is None or new_env is None else _format_speedup(new_env, base_env), + "n/a" + if base_env is None or new_env is None + else _format_speedup(new_env, base_env), "n/a" if base_env is None or new_env is None else _format_saved(new_env, base_env), ( "n/a" @@ -1781,6 +1915,13 @@ def main() -> int: "\nDR reset timing (subparts of reset call; reset obs getters currently read full batch):" ) print(_format_dr_reset_timing_table(results)) + print( + "\nBackend set_state detail — motrix keyset " + "(sub-timings sum to dr_reset_set_state_ms; % is share of that outer wall-clock):" + ) + print(_format_set_state_detail_table(results)) + print("\nBackend set_state detail — mujoco keyset:") + print(_format_set_state_mujoco_table(results)) else: print("No successful benchmark cases.") return 0 if not errors else 1 diff --git a/src/unilab/base/backend/base.py b/src/unilab/base/backend/base.py index 96cc56f53..6367de989 100644 --- a/src/unilab/base/backend/base.py +++ b/src/unilab/base/backend/base.py @@ -271,7 +271,7 @@ def set_state( qpos: np.ndarray, qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, - ) -> None: + ) -> dict | None: """Set physics state for selected environments. Args: @@ -279,6 +279,14 @@ def set_state( qpos: Position state. qvel: Velocity state. randomization: Optional backend randomization payload. + + Returns: + Optional dictionary. Backends MAY include a ``"timing"`` key with + per-substep timings in milliseconds (e.g. ``set_state_mask_ms``, + ``set_state_data_slice_ms``, ...). Callers MUST treat ``None`` or + missing keys as "not reported" — the outer wall-clock measurement in + ``DomainRandomizationManager.reset`` (``dr_reset_set_state_ms``) + remains authoritative for total set_state time. """ @abc.abstractmethod diff --git a/src/unilab/base/backend/motrix/backend.py b/src/unilab/base/backend/motrix/backend.py index 49aaa417f..d2d5ff266 100644 --- a/src/unilab/base/backend/motrix/backend.py +++ b/src/unilab/base/backend/motrix/backend.py @@ -307,6 +307,11 @@ def __init__( self._link_velocity_cache_valid = False self._refresh_link_pose_cache() + # Scratch buffers reused by set_state() to avoid per-reset allocations. + # Sized to the full env count and rewritten in place each call. + self._set_state_mask_scratch: np.ndarray = np.zeros(self._num_envs, dtype=bool) + self._set_state_qpos_motrix_scratch: np.ndarray | None = None + def get_motion_body_ids(self, names: Sequence[str]) -> np.ndarray: ids: list[int] = [] for name in names: @@ -595,22 +600,83 @@ def set_state( qpos: np.ndarray, qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, - ) -> None: - qpos_motrix = self._mujoco_qpos_to_motrix(qpos) + ) -> dict | None: + timing: dict[str, float] = { + "set_state_mask_ms": 0.0, + "set_state_data_slice_ms": 0.0, + "set_state_data_reset_ms": 0.0, + "set_state_clear_forces_ms": 0.0, + "set_state_geom_overrides_ms": 0.0, + "set_state_reset_rand_ms": 0.0, + "set_state_set_dof_vel_ms": 0.0, + "set_state_set_dof_pos_ms": 0.0, + "set_state_actuator_ctrl_ms": 0.0, + "set_state_forward_kinematic_ms": 0.0, + "set_state_refresh_pose_cache_ms": 0.0, + "set_state_invalidate_velocity_ms": 0.0, + "set_state_qpos_convert_ms": 0.0, + "set_state_pool_reset_ms": 0.0, + "set_state_state_scatter_ms": 0.0, + "set_state_internal_gap_ms": 0.0, + } + outer_t0 = time.perf_counter() - # Create mask for batch operation - mask = np.zeros(self._num_envs, dtype=bool) - mask[env_indices] = True + # Pre-convert env_indices once; every downstream helper reuses this. + env_ids_intp = np.asarray(env_indices, dtype=np.intp) + + t0 = time.perf_counter() + # Reuse the scratch qpos buffer when its shape matches; the reset path + # feeds a fixed-shape (num_envs, qpos_dim) array so this holds after + # the first call. + scratch = self._set_state_qpos_motrix_scratch + if scratch is None or scratch.shape != qpos.shape or scratch.dtype != qpos.dtype: + scratch = np.empty_like(qpos) + self._set_state_qpos_motrix_scratch = scratch + qpos_motrix = self._mujoco_qpos_to_motrix_into(scratch, qpos) + timing["set_state_qpos_convert_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + # Reuse the persistent bool-mask scratch; only the touched entries need + # rewriting each call. + mask = self._set_state_mask_scratch + if mask.shape[0] != self._num_envs: + mask = np.zeros(self._num_envs, dtype=bool) + self._set_state_mask_scratch = mask + mask.fill(False) + mask[env_ids_intp] = True + timing["set_state_mask_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() data_slice = self._data[mask] + timing["set_state_data_slice_ms"] = (time.perf_counter() - t0) * 1000.0 - # Batch set state + t0 = time.perf_counter() data_slice.reset(self._model) - self._clear_applied_body_forces(env_indices) - self._apply_init_geom_size_overrides(data_slice, env_indices) - self._apply_reset_randomization(data_slice, env_indices, randomization) + timing["set_state_data_reset_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._clear_applied_body_forces(env_indices, env_ids_intp=env_ids_intp) + timing["set_state_clear_forces_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._apply_init_geom_size_overrides(data_slice, env_indices, env_ids_intp=env_ids_intp) + timing["set_state_geom_overrides_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._apply_reset_randomization( + data_slice, env_indices, randomization, env_ids_intp=env_ids_intp + ) + timing["set_state_reset_rand_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() data_slice.set_dof_vel(qvel) + timing["set_state_set_dof_vel_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() data_slice.set_dof_pos(qpos_motrix, self._model) + timing["set_state_set_dof_pos_ms"] = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() if self._supports_position_actuator_gains and len(self._joint_dof_pos_indices) == int( self.num_actuators ): @@ -627,11 +693,44 @@ def set_state( ctrl = qpos_motrix[:, self._actuator_joint_pos_indices] else: ctrl = np.zeros((len(env_indices), self.num_actuators), dtype=self._np_dtype) - data_slice.actuator_ctrls = np.ascontiguousarray(ctrl) + # Only pay the copy when the underlying slice is non-contiguous; view + # slices of a contiguous scratch buffer already satisfy the motrixsim + # contiguous requirement. + if not ctrl.flags.c_contiguous: + ctrl = np.ascontiguousarray(ctrl) + data_slice.actuator_ctrls = ctrl + timing["set_state_actuator_ctrl_ms"] = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() self._model.forward_kinematic(data_slice) - self._refresh_link_pose_cache(env_indices, data_slice=data_slice) + timing["set_state_forward_kinematic_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() + self._refresh_link_pose_cache(env_indices, data_slice=data_slice, env_ids_intp=env_ids_intp) + timing["set_state_refresh_pose_cache_ms"] = (time.perf_counter() - t0) * 1000.0 + + t0 = time.perf_counter() self._invalidate_link_velocity_cache() + timing["set_state_invalidate_velocity_ms"] = (time.perf_counter() - t0) * 1000.0 + + outer_total_ms = (time.perf_counter() - outer_t0) * 1000.0 + measured_ms = ( + timing["set_state_qpos_convert_ms"] + + timing["set_state_mask_ms"] + + timing["set_state_data_slice_ms"] + + timing["set_state_data_reset_ms"] + + timing["set_state_clear_forces_ms"] + + timing["set_state_geom_overrides_ms"] + + timing["set_state_reset_rand_ms"] + + timing["set_state_set_dof_vel_ms"] + + timing["set_state_set_dof_pos_ms"] + + timing["set_state_actuator_ctrl_ms"] + + timing["set_state_forward_kinematic_ms"] + + timing["set_state_refresh_pose_cache_ms"] + + timing["set_state_invalidate_velocity_ms"] + ) + timing["set_state_internal_gap_ms"] = outer_total_ms - measured_ms + return {"timing": timing} def get_dr_capabilities(self) -> DomainRandomizationCapabilities: supported_reset_terms = { @@ -1000,6 +1099,20 @@ def _mujoco_qpos_to_motrix(self, qpos: np.ndarray) -> np.ndarray: qpos_motrix[..., quat_indices] = qpos[..., quat_indices[[1, 2, 3, 0]]] return qpos_motrix + def _mujoco_qpos_to_motrix_into(self, dst: np.ndarray, qpos: np.ndarray) -> np.ndarray: + """Hot-path variant that writes into ``dst`` in place. + + Same conversion as :meth:`_mujoco_qpos_to_motrix` but avoids the per-call + ``np.array(qpos, copy=True)`` allocation. Returns ``dst`` so callers can + chain: ``qpos_motrix = self._mujoco_qpos_to_motrix_into(scratch, qpos)``. + ``dst`` must have the same shape as ``qpos``; caller is responsible for + sizing the scratch buffer. + """ + np.copyto(dst, qpos, casting="same_kind") + for quat_indices in self._floating_base_quat_indices: + dst[..., quat_indices] = qpos[..., quat_indices[[1, 2, 3, 0]]] + return dst + def _motrix_qpos_to_mujoco(self, qpos: np.ndarray) -> np.ndarray: """Convert every Motrix freejoint quaternion slice from xyzw to wxyz.""" qpos_mujoco = np.array(qpos, copy=True) @@ -1008,7 +1121,10 @@ def _motrix_qpos_to_mujoco(self, qpos: np.ndarray) -> np.ndarray: return qpos_mujoco def _refresh_link_pose_cache( - self, env_indices: np.ndarray | None = None, data_slice: Any | None = None + self, + env_indices: np.ndarray | None = None, + data_slice: Any | None = None, + env_ids_intp: np.ndarray | None = None, ) -> None: if env_indices is None: self._link_poses = self._model.get_link_poses(self._data) @@ -1017,7 +1133,11 @@ def _refresh_link_pose_cache( mask = np.zeros(self._num_envs, dtype=bool) mask[env_indices] = True data_slice = self._data[mask] - self._link_poses[env_indices] = self._model.get_link_poses(data_slice) + # Indexing with intp is a hair faster than the arbitrary-dtype + # ndarray path; use the pre-converted array when set_state hands + # it down. + idx = env_indices if env_ids_intp is None else env_ids_intp + self._link_poses[idx] = self._model.get_link_poses(data_slice) def _refresh_link_velocity_cache(self, env_indices: np.ndarray | None = None) -> None: if env_indices is None: @@ -1058,10 +1178,17 @@ def _coerce_reset_field( return arr.reshape(shaped).copy() raise ValueError(f"{name} must have shape {shaped} or {flat_shape}, got {arr.shape}") - def _apply_init_geom_size_overrides(self, data_slice, env_indices: np.ndarray) -> None: + def _apply_init_geom_size_overrides( + self, + data_slice, + env_indices: np.ndarray, + env_ids_intp: np.ndarray | None = None, + ) -> None: if not self._init_geom_size_overrides: return - env_ids = np.asarray(env_indices, dtype=np.intp) + env_ids = ( + env_ids_intp if env_ids_intp is not None else np.asarray(env_indices, dtype=np.intp) + ) for geom_id, values in self._init_geom_size_overrides.items(): geom = _require_not_none( self._model.get_geom(int(geom_id)), @@ -1106,10 +1233,16 @@ def _set_geom_friction_overrides(self, data_slice, geom_friction: np.ndarray) -> np.ascontiguousarray(np.asarray(geom_friction[:, geom_id, :], dtype=np.float32)), ) - def _clear_applied_body_forces(self, env_indices: np.ndarray) -> None: + def _clear_applied_body_forces( + self, + env_indices: np.ndarray, + env_ids_intp: np.ndarray | None = None, + ) -> None: if not self._applied_body_forces: return - env_ids = np.asarray(env_indices, dtype=np.intp) + env_ids = ( + env_ids_intp if env_ids_intp is not None else np.asarray(env_indices, dtype=np.intp) + ) for applied_force in self._applied_body_forces.values(): applied_force[env_ids, :] = 0.0 @@ -1339,6 +1472,7 @@ def _apply_reset_randomization( data_slice, env_indices: np.ndarray, randomization: ResetRandomizationPayload | None, + env_ids_intp: np.ndarray | None = None, ) -> None: if randomization is None or randomization.is_empty(): return @@ -1351,7 +1485,9 @@ def _apply_reset_randomization( f"{self.backend_type} backend does not support reset randomization terms: {terms}" ) - env_ids = np.asarray(env_indices, dtype=np.intp) + env_ids = ( + env_ids_intp if env_ids_intp is not None else np.asarray(env_indices, dtype=np.intp) + ) num_reset = len(env_ids) body_mass = None if randomization.body_mass is not None: diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index d2c9e4013..dd45f2487 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -877,23 +877,58 @@ def set_state( qpos: np.ndarray, qvel: np.ndarray, randomization: ResetRandomizationPayload | None = None, - ) -> None: + ) -> dict | None: + timing: dict[str, float] = { + "set_state_mask_ms": 0.0, + "set_state_data_slice_ms": 0.0, + "set_state_data_reset_ms": 0.0, + "set_state_clear_forces_ms": 0.0, + "set_state_geom_overrides_ms": 0.0, + "set_state_reset_rand_ms": 0.0, + "set_state_set_dof_vel_ms": 0.0, + "set_state_set_dof_pos_ms": 0.0, + "set_state_actuator_ctrl_ms": 0.0, + "set_state_forward_kinematic_ms": 0.0, + "set_state_refresh_pose_cache_ms": 0.0, + "set_state_invalidate_velocity_ms": 0.0, + "set_state_qpos_convert_ms": 0.0, + "set_state_pool_reset_ms": 0.0, + "set_state_state_scatter_ms": 0.0, + "set_state_internal_gap_ms": 0.0, + } if len(env_indices) == 0: - return + return {"timing": timing} + outer_t0 = time.perf_counter() + + t0 = time.perf_counter() num_reset = len(env_indices) state_np = np.zeros((num_reset, self._physics_state.shape[1]), dtype=np.float64) state_np[:, self._idx_qpos : self._idx_qpos + self.nq] = qpos state_np[:, self._idx_qvel : self._idx_qvel + self.nv] = qvel + timing["set_state_qpos_convert_ms"] = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() state_out, sensor_np = self._pool.reset( # type: ignore[union-attr] env_ids=np.asarray(env_indices, dtype=np.int32), initial_state=state_np, randomization=self._translate_reset_randomization(randomization, num_reset), ) + timing["set_state_pool_reset_ms"] = (time.perf_counter() - t0) * 1000.0 + t0 = time.perf_counter() self._physics_state[env_indices] = state_out.astype(self._np_dtype) self._sensor_data[env_indices] = sensor_np.astype(self._np_dtype) + timing["set_state_state_scatter_ms"] = (time.perf_counter() - t0) * 1000.0 + + outer_total_ms = (time.perf_counter() - outer_t0) * 1000.0 + measured_ms = ( + timing["set_state_qpos_convert_ms"] + + timing["set_state_pool_reset_ms"] + + timing["set_state_state_scatter_ms"] + ) + timing["set_state_internal_gap_ms"] = outer_total_ms - measured_ms + return {"timing": timing} def get_dr_capabilities(self) -> DomainRandomizationCapabilities: return DomainRandomizationCapabilities( diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 65408d330..3b5777fa0 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -46,6 +46,46 @@ "dr_reset_obs_get_body_pose_ms", "dr_reset_observation_compute_obs_ms", "dr_reset_observation_internal_gap_ms", + # Backend-internal set_state sub-timings. All backends report the same key + # set for column stability; sub-keys that don't apply report 0.0. + "set_state_mask_ms", + "set_state_data_slice_ms", + "set_state_data_reset_ms", + "set_state_clear_forces_ms", + "set_state_geom_overrides_ms", + "set_state_reset_rand_ms", + "set_state_set_dof_vel_ms", + "set_state_set_dof_pos_ms", + "set_state_actuator_ctrl_ms", + "set_state_forward_kinematic_ms", + "set_state_refresh_pose_cache_ms", + "set_state_invalidate_velocity_ms", + "set_state_qpos_convert_ms", + "set_state_pool_reset_ms", + "set_state_state_scatter_ms", + "set_state_internal_gap_ms", +) + +# Subset of RESET_DONE_DETAIL_TIMING_KEYS that comes from the backend's +# set_state() timing dict. DR manager merges these keys 1-to-1 from the +# backend return value. +BACKEND_SET_STATE_DETAIL_TIMING_KEYS = ( + "set_state_mask_ms", + "set_state_data_slice_ms", + "set_state_data_reset_ms", + "set_state_clear_forces_ms", + "set_state_geom_overrides_ms", + "set_state_reset_rand_ms", + "set_state_set_dof_vel_ms", + "set_state_set_dof_pos_ms", + "set_state_actuator_ctrl_ms", + "set_state_forward_kinematic_ms", + "set_state_refresh_pose_cache_ms", + "set_state_invalidate_velocity_ms", + "set_state_qpos_convert_ms", + "set_state_pool_reset_ms", + "set_state_state_scatter_ms", + "set_state_internal_gap_ms", ) diff --git a/src/unilab/dr/manager.py b/src/unilab/dr/manager.py index 6beddb281..1c233085f 100644 --- a/src/unilab/dr/manager.py +++ b/src/unilab/dr/manager.py @@ -49,13 +49,22 @@ def reset(self, env_ids: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: payload_filter_ms = (time.perf_counter() - t0) * 1000.0 t0 = time.perf_counter() - self._env._backend.set_state( + set_state_result = self._env._backend.set_state( plan.env_ids, plan.qpos, plan.qvel, randomization=payload, ) set_state_ms = (time.perf_counter() - t0) * 1000.0 + backend_set_state_timing: dict[str, float] = {} + if isinstance(set_state_result, dict): + backend_timing = set_state_result.get("timing") + if isinstance(backend_timing, dict): + for key, value in backend_timing.items(): + try: + backend_set_state_timing[str(key)] = float(value) + except (TypeError, ValueError): + continue t0 = time.perf_counter() obs = self._provider.build_reset_observation(self._env, plan.env_ids, plan.info_updates) @@ -71,6 +80,8 @@ def reset(self, env_ids: np.ndarray) -> tuple[dict[str, np.ndarray], dict]: "dr_reset_build_observation_ms": build_observation_ms, "dr_reset_internal_gap_ms": total_ms - measured_ms, } + if backend_set_state_timing: + timing.update(backend_set_state_timing) provider_timing = getattr(self._provider, "last_reset_observation_timing_ms", {}) if isinstance(provider_timing, dict): timing.update(provider_timing) diff --git a/tests/base/test_sim_backend_set_state_timing.py b/tests/base/test_sim_backend_set_state_timing.py new file mode 100644 index 000000000..987213297 --- /dev/null +++ b/tests/base/test_sim_backend_set_state_timing.py @@ -0,0 +1,190 @@ +"""Contract + parity tests for the extended ``SimBackend.set_state`` timing dict. + +Issue #679 extends the ``SimBackend.set_state`` return type from ``None`` to +``dict | None`` so backends can surface per-substep timing to the DR manager. This +module locks in three invariants: + +1. Both MuJoCo and Motrix ``set_state`` return a dict with a ``"timing"`` sub-dict + whose keys are the schema documented in + ``unilab.base.np_env.BACKEND_SET_STATE_DETAIL_TIMING_KEYS``. +2. Every reported sub-timing is a non-negative float. +3. The reported sub-timings sum to within ~5 ms of the outer wall-clock cost + (``set_state_internal_gap_ms`` catches the residual). This is the same + consistency check the benchmark harness relies on when computing % share. + +The tests double as parity anchors for the Step 5 Motrix refactor: pre-refactor +they capture the current behavior (all schema keys populated, gap tiny); post- +refactor they must still hold, so any regression that skips a sub-key or blows +up the gap is caught before the benchmark run. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import pytest + +from unilab.assets import ASSETS_ROOT_PATH +from unilab.base.np_env import BACKEND_SET_STATE_DETAIL_TIMING_KEYS +from unilab.base.scene import SceneCfg +from unilab.dr.types import ResetRandomizationPayload + +pytest.importorskip("mujoco", reason="mujoco not installed") + + +NUM_ENVS = 2 +SIM_DT = 0.005 +_G1_MODEL_FILE = str(ASSETS_ROOT_PATH / "robots" / "g1" / "scene_flat.xml") + + +def _identity_qpos_mujoco(nq: int, xyz=(0.0, 0.0, 0.8)) -> np.ndarray: + q = np.zeros((1, nq)) + q[0, :3] = xyz + q[0, 3] = 1.0 + return q + + +def _assert_timing_dict_shape(result: Any) -> dict[str, float]: + """Common assertions across backends.""" + assert isinstance(result, dict), f"set_state must return a dict, got {type(result)!r}" + timing = result.get("timing") + assert isinstance(timing, dict), "set_state result must contain a 'timing' sub-dict" + + missing = [key for key in BACKEND_SET_STATE_DETAIL_TIMING_KEYS if key not in timing] + assert not missing, f"missing keys in set_state timing: {missing}" + + for key in BACKEND_SET_STATE_DETAIL_TIMING_KEYS: + value = timing[key] + assert isinstance(value, float), f"{key} must be float, got {type(value)!r}" + # Gap can be very slightly negative due to perf_counter noise; the + # sum-based check below is what we really care about. + if not key.endswith("internal_gap_ms"): + assert value >= 0.0, f"{key} must be non-negative, got {value}" + + return timing + + +def _assert_gap_bounded(timing: dict[str, float]) -> None: + """Sub-timings should sum to ~ the outer wall-clock cost; gap is the residual. + + We just assert the reported ``set_state_internal_gap_ms`` stays within a + generous absolute bound so a refactor that forgets to track a new sub-step + (making the gap balloon) is caught. + """ + gap = timing["set_state_internal_gap_ms"] + assert abs(gap) < 5.0, f"set_state_internal_gap_ms out of bounds: {gap}" + + +def test_mujoco_set_state_returns_schema_conformant_timing() -> None: + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + backend = MuJoCoBackend( + SceneCfg(model_file=_G1_MODEL_FILE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + backend.materialize() + qpos = _identity_qpos_mujoco(backend.model.nq, xyz=(1.0, 2.0, 0.8)) + qvel = np.zeros((1, backend.model.nv)) + + result = backend.set_state(np.array([0]), qpos, qvel) + + timing = _assert_timing_dict_shape(result) + _assert_gap_bounded(timing) + # MuJoCo-specific sub-keys must be populated on the mujoco path. + assert timing["set_state_pool_reset_ms"] > 0.0 + assert timing["set_state_state_scatter_ms"] > 0.0 + # Motrix-only sub-keys report 0.0 on the mujoco backend. + assert timing["set_state_mask_ms"] == 0.0 + assert timing["set_state_data_slice_ms"] == 0.0 + assert timing["set_state_forward_kinematic_ms"] == 0.0 + assert timing["set_state_refresh_pose_cache_ms"] == 0.0 + + +def test_mujoco_set_state_empty_env_indices_still_returns_timing_dict() -> None: + """Backends must never return None even on the empty-reset fast path.""" + from unilab.base.backend.mujoco.backend import MuJoCoBackend + + backend = MuJoCoBackend( + SceneCfg(model_file=_G1_MODEL_FILE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + backend.materialize() + + result = backend.set_state( + np.array([], dtype=np.int32), + np.zeros((0, backend.model.nq)), + np.zeros((0, backend.model.nv)), + ) + + timing = _assert_timing_dict_shape(result) + # Empty-reset path skips all substeps. + assert timing["set_state_pool_reset_ms"] == 0.0 + assert timing["set_state_state_scatter_ms"] == 0.0 + assert timing["set_state_internal_gap_ms"] == 0.0 + + +def test_motrix_set_state_returns_schema_conformant_timing() -> None: + pytest.importorskip("motrixsim") + + from unilab.base.backend.motrix.backend import MotrixBackend + + backend = MotrixBackend( + SceneCfg(model_file=_G1_MODEL_FILE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + nq = backend.get_dof_pos().shape[-1] + 7 + nv = backend.get_dof_vel().shape[-1] + 6 + qpos = _identity_qpos_mujoco(nq, xyz=(1.0, 2.0, 0.8)) + qvel = np.zeros((1, nv)) + randomization = ResetRandomizationPayload( + gravity=np.asarray([[0.5, 0.5, -3.0]], dtype=np.float64) + ) + + result = backend.set_state(np.array([0]), qpos, qvel, randomization=randomization) + + timing = _assert_timing_dict_shape(result) + _assert_gap_bounded(timing) + # Motrix path must populate the sub-steps it actually runs. + assert timing["set_state_mask_ms"] > 0.0 + assert timing["set_state_data_slice_ms"] > 0.0 + assert timing["set_state_data_reset_ms"] > 0.0 + assert timing["set_state_set_dof_vel_ms"] > 0.0 + assert timing["set_state_set_dof_pos_ms"] > 0.0 + assert timing["set_state_forward_kinematic_ms"] > 0.0 + assert timing["set_state_refresh_pose_cache_ms"] >= 0.0 + # A gravity payload was provided, so the DR sub-step should be non-zero. + assert timing["set_state_reset_rand_ms"] > 0.0 + # MuJoCo-only sub-keys report 0.0 on motrix. + assert timing["set_state_pool_reset_ms"] == 0.0 + assert timing["set_state_state_scatter_ms"] == 0.0 + + +def test_motrix_set_state_without_randomization_leaves_reset_rand_zero() -> None: + pytest.importorskip("motrixsim") + + from unilab.base.backend.motrix.backend import MotrixBackend + + backend = MotrixBackend( + SceneCfg(model_file=_G1_MODEL_FILE), + NUM_ENVS, + SIM_DT, + base_name="pelvis", + ) + nq = backend.get_dof_pos().shape[-1] + 7 + nv = backend.get_dof_vel().shape[-1] + 6 + qpos = _identity_qpos_mujoco(nq, xyz=(1.0, 2.0, 0.8)) + qvel = np.zeros((1, nv)) + + result = backend.set_state(np.array([0]), qpos, qvel) + + timing = _assert_timing_dict_shape(result) + # No payload → the fast-return branch in _apply_reset_randomization keeps + # this sub-step at essentially zero. + assert timing["set_state_reset_rand_ms"] < 1.0 diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index a0d2a4037..73aa5f836 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -254,6 +254,87 @@ def test_write_csv_includes_all_phase_columns(tmp_path) -> None: assert key in header +def test_write_csv_includes_backend_set_state_sub_timing_columns(tmp_path) -> None: + """CSV headers must expose every backend set_state sub-key so downstream + reporting can diff pre/post-optimization runs.""" + result = _make_result(num_envs=2, throughput=2000.0, include_env_step_breakdown=True) + out_csv = tmp_path / "collector.csv" + + bench._write_csv(out_csv, [result]) + + header = out_csv.read_text(encoding="utf-8").splitlines()[0] + expected_keys = ( + "set_state_mask_ms", + "set_state_data_slice_ms", + "set_state_data_reset_ms", + "set_state_clear_forces_ms", + "set_state_geom_overrides_ms", + "set_state_reset_rand_ms", + "set_state_set_dof_vel_ms", + "set_state_set_dof_pos_ms", + "set_state_actuator_ctrl_ms", + "set_state_forward_kinematic_ms", + "set_state_refresh_pose_cache_ms", + "set_state_invalidate_velocity_ms", + "set_state_qpos_convert_ms", + "set_state_pool_reset_ms", + "set_state_state_scatter_ms", + "set_state_internal_gap_ms", + ) + for key in expected_keys: + assert key in header, f"CSV header missing {key!r}" + + +def test_format_set_state_detail_table_reports_sub_key_percentages() -> None: + """The detail table shows each sub-key next to its share of + ``dr_reset_set_state_ms`` (percentages must appear).""" + result = _make_result(num_envs=2, throughput=2000.0, include_env_step_breakdown=True) + result.env_step_timing_ms_per_vector_step["dr_reset_set_state_ms"] = bench.TimingStats( + [4.0], 4.0, 4.0, 0.0, 4.0, 4.0 + ) + result.env_step_timing_ms_per_vector_step["set_state_forward_kinematic_ms"] = bench.TimingStats( + [2.0], 2.0, 2.0, 0.0, 2.0, 2.0 + ) + result.env_step_timing_ms_per_vector_step["set_state_mask_ms"] = bench.TimingStats( + [0.1], 0.1, 0.1, 0.0, 0.1, 0.1 + ) + + table = bench._format_set_state_detail_table([result]) + + # Header exposes the "FK" and "Mask" sub-columns. + assert "FK" in table + assert "Mask" in table + # 2.0 ms / 4.0 ms = 50%; 0.1 ms / 4.0 ms = 2.5%. + assert "2.000 (50.0%)" in table + assert "0.100 (2.5%)" in table + + +def test_format_set_state_mujoco_table_covers_mujoco_only_keys() -> None: + """The mujoco keyset table exposes pool_reset / state_scatter columns.""" + result = _make_result( + runtime_sim_backend="mujoco", + num_envs=2, + throughput=2000.0, + include_env_step_breakdown=True, + ) + result.env_step_timing_ms_per_vector_step["dr_reset_set_state_ms"] = bench.TimingStats( + [5.0], 5.0, 5.0, 0.0, 5.0, 5.0 + ) + result.env_step_timing_ms_per_vector_step["set_state_pool_reset_ms"] = bench.TimingStats( + [4.0], 4.0, 4.0, 0.0, 4.0, 4.0 + ) + result.env_step_timing_ms_per_vector_step["set_state_state_scatter_ms"] = bench.TimingStats( + [0.75], 0.75, 0.75, 0.0, 0.75, 0.75 + ) + + table = bench._format_set_state_mujoco_table([result]) + + assert "Pool reset" in table + assert "State scatter" in table + assert "4.000 (80.0%)" in table + assert "0.750 (15.0%)" in table + + def test_numba_ab_overrides_enable_and_disable_acceleration() -> None: assert bench._numba_ab_overrides(enabled=False, numba_threads=8) == [ "++env.numba_acceleration=false" diff --git a/tests/dr/test_manager.py b/tests/dr/test_manager.py index c8150b30a..c8c4eac12 100644 --- a/tests/dr/test_manager.py +++ b/tests/dr/test_manager.py @@ -6,6 +6,7 @@ from typing import Any import numpy as np +import pytest from unilab.dr import ( DomainRandomizationCapabilities, @@ -135,6 +136,33 @@ def set_state( self.last_randomization = randomization +@dataclass +class _FakeTimedBackend: + """Backend that reports set_state sub-timings via the extended contract.""" + + capabilities: DomainRandomizationCapabilities + timing: dict[str, float] + backend_type: str = "motrix" + + def __post_init__(self) -> None: + self.last_randomization: ResetRandomizationPayload | None = None + self.call_count = 0 + + def get_dr_capabilities(self) -> DomainRandomizationCapabilities: + return self.capabilities + + def set_state( + self, + env_indices: np.ndarray, + qpos: np.ndarray, + qvel: np.ndarray, + randomization: ResetRandomizationPayload | None = None, + ) -> dict: + self.last_randomization = randomization + self.call_count += 1 + return {"timing": dict(self.timing)} + + class _FakeProvider(DomainRandomizationProvider): def validate(self, env: Any, capabilities: DomainRandomizationCapabilities) -> None: return None @@ -199,3 +227,65 @@ def test_manager_keeps_supported_reset_terms_without_warning(caplog): assert backend.last_randomization.base_mass_delta is not None assert backend.last_randomization.kp is not None assert "skipping them" not in caplog.text + + +def test_manager_merges_backend_set_state_sub_timings(): + """Backend-reported ``{"timing": {...}}`` from set_state flows into + ``last_reset_timing_ms`` next to ``dr_reset_set_state_ms``.""" + reported = { + "set_state_mask_ms": 0.12, + "set_state_data_slice_ms": 0.34, + "set_state_forward_kinematic_ms": 2.1, + "set_state_internal_gap_ms": 0.05, + } + backend = _FakeTimedBackend( + capabilities=DomainRandomizationCapabilities( + supported_reset_terms=frozenset({RESET_TERM_BASE_MASS, RESET_TERM_KP}) + ), + timing=reported, + ) + env = SimpleNamespace(_backend=backend) + manager = DomainRandomizationManager(env, _FakeProvider()) + + obs, _ = manager.reset(np.array([0, 1, 2], dtype=np.int32)) + + assert obs["obs"].shape == (3, 1) + assert backend.call_count == 1 + timings = manager.last_reset_timing_ms + # Outer wall-clock measurement still present. + assert "dr_reset_set_state_ms" in timings + # Every reported backend sub-key is merged in. + for key, expected in reported.items(): + assert key in timings + assert timings[key] == pytest.approx(expected) + + +def test_manager_tolerates_missing_or_malformed_backend_timing(): + """Backends may return ``None`` (unchanged behavior) or a dict with + non-numeric values; the manager must not crash and must not add spurious + sub-keys in either case.""" + plain_backend = _FakeBackend( + capabilities=DomainRandomizationCapabilities( + supported_reset_terms=frozenset({RESET_TERM_BASE_MASS, RESET_TERM_KP}) + ) + ) + env = SimpleNamespace(_backend=plain_backend) + manager = DomainRandomizationManager(env, _FakeProvider()) + manager.reset(np.array([0], dtype=np.int32)) + plain_keys = set(manager.last_reset_timing_ms) + assert "dr_reset_set_state_ms" in plain_keys + assert not any(k.startswith("set_state_") for k in plain_keys) + + malformed_backend = _FakeTimedBackend( + capabilities=DomainRandomizationCapabilities( + supported_reset_terms=frozenset({RESET_TERM_BASE_MASS, RESET_TERM_KP}) + ), + timing={"set_state_mask_ms": "not a number", "set_state_data_slice_ms": 0.5}, + ) + env = SimpleNamespace(_backend=malformed_backend) + manager = DomainRandomizationManager(env, _FakeProvider()) + manager.reset(np.array([0], dtype=np.int32)) + timings = manager.last_reset_timing_ms + # Malformed value dropped, well-formed one merged. + assert "set_state_mask_ms" not in timings + assert timings["set_state_data_slice_ms"] == pytest.approx(0.5) From dd2228987a8318599f309e8b8bdc7db64892c89f Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Wed, 8 Jul 2026 01:19:07 +0800 Subject: [PATCH 16/26] bench: profile RNG noise buffer cost --- benchmark/benchmark_numba_random_noise.py | 614 ++++++++++++++++++ .../benchmark_offpolicy_collector_active.py | 395 +++++++---- .../test_numba_random_noise_benchmark.py | 90 +++ 3 files changed, 966 insertions(+), 133 deletions(-) create mode 100644 benchmark/benchmark_numba_random_noise.py create mode 100644 tests/benchmark/test_numba_random_noise_benchmark.py diff --git a/benchmark/benchmark_numba_random_noise.py b/benchmark/benchmark_numba_random_noise.py new file mode 100644 index 000000000..6bc669132 --- /dev/null +++ b/benchmark/benchmark_numba_random_noise.py @@ -0,0 +1,614 @@ +#!/usr/bin/env python3 +"""Benchmark Numba RNG against NumPy RNG for UniLab noise-buffer workloads. + +Issue #680 asks for an isolated measurement before changing env hot paths. This +script uses owner-config shapes from G1 motion tracking and G1 walk/joystick +training profiles, then compares: + +* ``numpy_random_uniform_alloc``: current-style ``np.random.uniform(...).astype``; +* ``numpy_generator_random_out``: NumPy with preallocated output buffers; +* ``numba_random_prange``: ``np.random.random`` inside an ``njit(parallel=True)`` + kernel, also writing into preallocated buffers. + +Run: + uv run python -m benchmark.benchmark_numba_random_noise + uv run python benchmark/benchmark_numba_random_noise.py --quick + uv run python -m benchmark.benchmark_numba_random_noise --profiles sac_g1_motion_tracking_mujoco --threads 1 4 8 +""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import time +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from statistics import mean, stdev +from typing import Any + +REPO_ROOT = Path(__file__).resolve().parents[1] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import numpy as np + +try: # pragma: no cover - exercised when numba is installed + from numba import get_num_threads, njit, prange, set_num_threads + + NUMBA_AVAILABLE = True +except Exception: # pragma: no cover - benchmark still reports NumPy baselines + get_num_threads = njit = prange = set_num_threads = None # type: ignore[assignment] + NUMBA_AVAILABLE = False + +from benchmark.core.device_info import get_device_info_dict, get_device_info_line + +DEFAULT_THREADS = [1, 2, 4, 8, 16, 32, 64] +QUICK_THREADS = [1, 2] +DEFAULT_NUM_ENVS = [1024, 2048, 4096, 8192, 16384, 32768] +QUICK_NUM_ENVS = [1024, 2048] +NUM_ACTION = 29 + + +@dataclass(frozen=True) +class NoiseField: + name: str + width: int + scale: float + + +@dataclass(frozen=True) +class NoiseProfile: + name: str + owner_config: str + default_num_envs: int + fields: tuple[NoiseField, ...] + note: str + + @property + def total_width(self) -> int: + return sum(field.width for field in self.fields) + + @property + def values_per_step(self) -> int: + return self.default_num_envs * self.total_width + + +@dataclass +class BenchCase: + profile: str + owner_config: str + num_envs: int + dtype: str + path: str + threads: int | None + values: int + mean_ms: float + min_ms: float + std_ms: float + values_per_s: float + speedup_vs_numpy_alloc: float + compile_ms: float | None = None + deterministic_same_seed: bool | None = None + distribution_mean: float | None = None + distribution_std: float | None = None + distribution_min: float | None = None + distribution_max: float | None = None + + +def make_profiles() -> dict[str, NoiseProfile]: + """Profiles mirror current env noise calls and owner-config scales. + + G1Walk/G1Joystick still calls ``_obs_noise`` for gyro and gravity even when a + scale is zero, so those zero-scale fields are intentionally kept to measure + the current hot-path RNG cost. + """ + + return { + "ppo_g1_motion_tracking_motrix": NoiseProfile( + name="ppo_g1_motion_tracking_motrix", + owner_config="conf/ppo/task/g1_motion_tracking/motrix.yaml", + default_num_envs=1024, + fields=( + NoiseField("linvel", 3, 0.1), + NoiseField("gyro", 3, 0.2), + NoiseField("joint_pos", NUM_ACTION, 0.01), + NoiseField("dof_vel", NUM_ACTION, 1.5), + ), + note="G1MotionTracking actor observation noise.", + ), + "sac_g1_motion_tracking_mujoco": NoiseProfile( + name="sac_g1_motion_tracking_mujoco", + owner_config="conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml", + default_num_envs=2048, + fields=( + NoiseField("linvel", 3, 0.1), + NoiseField("gyro", 3, 0.2), + NoiseField("joint_pos", NUM_ACTION, 0.01), + NoiseField("dof_vel", NUM_ACTION, 1.5), + ), + note="G1MotionTrackingSAC actor observation noise.", + ), + "ppo_g1_walk_flat_motrix": NoiseProfile( + name="ppo_g1_walk_flat_motrix", + owner_config="conf/ppo/task/g1_walk_flat/motrix.yaml", + default_num_envs=2048, + fields=( + NoiseField("gyro", 3, 0.2), + NoiseField("gravity", 3, 0.0), + NoiseField("joint_pos", NUM_ACTION, 0.01), + NoiseField("dof_vel", NUM_ACTION, 1.5), + ), + note="G1WalkFlat joystick actor observation noise.", + ), + "sac_g1_walk_flat_mujoco": NoiseProfile( + name="sac_g1_walk_flat_mujoco", + owner_config="conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml", + default_num_envs=2048, + fields=( + NoiseField("gyro", 3, 0.0), + NoiseField("gravity", 3, 0.0), + NoiseField("joint_pos", NUM_ACTION, 0.01), + NoiseField("dof_vel", NUM_ACTION, 0.1), + ), + note="G1WalkFlat SAC actor observation noise.", + ), + } + + +def _dtype(name: str) -> np.dtype: + if name == "float32": + return np.dtype(np.float32) + if name == "float64": + return np.dtype(np.float64) + raise ValueError(f"unsupported dtype: {name}") + + +def _allocate_buffers(profile: NoiseProfile, num_envs: int, dtype: np.dtype) -> list[np.ndarray]: + return [np.empty((num_envs, field.width), dtype=dtype) for field in profile.fields] + + +def _scale_array(profile: NoiseProfile, noise_level: float) -> np.ndarray: + return np.asarray([noise_level * field.scale for field in profile.fields], dtype=np.float64) + + +def _numpy_random_uniform_alloc( + profile: NoiseProfile, + num_envs: int, + dtype: np.dtype, + noise_level: float, +) -> list[np.ndarray]: + out = [] + for field in profile.fields: + noise = np.random.uniform(-1.0, 1.0, size=(num_envs, field.width)).astype(dtype) + noise *= noise_level * field.scale + out.append(noise) + return out + + +def _numpy_generator_random_out( + rng: np.random.Generator, + buffers: list[np.ndarray], + profile: NoiseProfile, + noise_level: float, +) -> None: + for buffer, field in zip(buffers, profile.fields): + rng.random(out=buffer, dtype=buffer.dtype) + buffer *= 2.0 + buffer -= 1.0 + buffer *= noise_level * field.scale + + +if NUMBA_AVAILABLE: + + @njit(cache=True) # type: ignore[misc] + def _numba_seed(seed: int) -> None: + np.random.seed(seed) + + @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] + def _numba_fill_four( + buffer0: np.ndarray, + buffer1: np.ndarray, + buffer2: np.ndarray, + buffer3: np.ndarray, + scales: np.ndarray, + ) -> None: + n = buffer0.shape[0] + for i in prange(n): + for j in range(buffer0.shape[1]): + buffer0[i, j] = (np.random.random() * 2.0 - 1.0) * scales[0] + for j in range(buffer1.shape[1]): + buffer1[i, j] = (np.random.random() * 2.0 - 1.0) * scales[1] + for j in range(buffer2.shape[1]): + buffer2[i, j] = (np.random.random() * 2.0 - 1.0) * scales[2] + for j in range(buffer3.shape[1]): + buffer3[i, j] = (np.random.random() * 2.0 - 1.0) * scales[3] + + +def _numba_random_prange(buffers: list[np.ndarray], scales: np.ndarray) -> None: + if not NUMBA_AVAILABLE: + raise RuntimeError("numba is not available") + if len(buffers) != 4: + raise ValueError("numba benchmark expects exactly four noise fields") + _numba_fill_four(buffers[0], buffers[1], buffers[2], buffers[3], scales) + + +def _time_ms(fn: Any, *, iters: int, warmup: int) -> tuple[float, float, float]: + for _ in range(warmup): + fn() + samples = [] + for _ in range(iters): + t0 = time.perf_counter() + fn() + samples.append((time.perf_counter() - t0) * 1000.0) + return mean(samples), min(samples), stdev(samples) if len(samples) > 1 else 0.0 + + +def _flatten(buffers: list[np.ndarray]) -> np.ndarray: + return np.concatenate([buffer.ravel() for buffer in buffers]) + + +def _distribution(buffers: list[np.ndarray]) -> dict[str, float]: + flat = _flatten(buffers) + if flat.size == 0: + return {"mean": 0.0, "std": 0.0, "min": 0.0, "max": 0.0} + return { + "mean": float(np.mean(flat)), + "std": float(np.std(flat)), + "min": float(np.min(flat)), + "max": float(np.max(flat)), + } + + +def _numba_same_seed_is_deterministic( + profile: NoiseProfile, + num_envs: int, + dtype: np.dtype, + scales: np.ndarray, + threads: int, + seed: int, +) -> bool: + if not NUMBA_AVAILABLE: + return False + previous_threads = get_num_threads() + set_num_threads(threads) + first = _allocate_buffers(profile, num_envs, dtype) + second = _allocate_buffers(profile, num_envs, dtype) + _numba_seed(seed) + _numba_random_prange(first, scales) + _numba_seed(seed) + _numba_random_prange(second, scales) + set_num_threads(previous_threads) + return all(np.array_equal(a, b) for a, b in zip(first, second)) + + +def bench_one( + *, + profile: NoiseProfile, + num_envs: int, + dtype_name: str = "float32", + thread_counts: list[int] | None = None, + iters: int = 20, + warmup: int = 5, + seed: int = 0, + noise_level: float = 1.0, +) -> list[BenchCase]: + dtype = _dtype(dtype_name) + values = num_envs * profile.total_width + records: list[BenchCase] = [] + + np.random.seed(seed) + numpy_buffers = _numpy_random_uniform_alloc(profile, num_envs, dtype, noise_level) + numpy_dist = _distribution(numpy_buffers) + + def numpy_alloc_call() -> None: + _numpy_random_uniform_alloc(profile, num_envs, dtype, noise_level) + + numpy_ms, numpy_min, numpy_std = _time_ms( + numpy_alloc_call, + iters=iters, + warmup=warmup, + ) + records.append( + BenchCase( + profile=profile.name, + owner_config=profile.owner_config, + num_envs=num_envs, + dtype=dtype_name, + path="numpy_random_uniform_alloc", + threads=None, + values=values, + mean_ms=numpy_ms, + min_ms=numpy_min, + std_ms=numpy_std, + values_per_s=values / (numpy_ms * 1.0e-3), + speedup_vs_numpy_alloc=1.0, + distribution_mean=numpy_dist["mean"], + distribution_std=numpy_dist["std"], + distribution_min=numpy_dist["min"], + distribution_max=numpy_dist["max"], + ) + ) + + generator_buffers = _allocate_buffers(profile, num_envs, dtype) + generator = np.random.default_rng(seed) + _numpy_generator_random_out(generator, generator_buffers, profile, noise_level) + generator_dist = _distribution(generator_buffers) + + def numpy_out_call() -> None: + _numpy_generator_random_out(generator, generator_buffers, profile, noise_level) + + numpy_out_ms, numpy_out_min, numpy_out_std = _time_ms( + numpy_out_call, + iters=iters, + warmup=warmup, + ) + records.append( + BenchCase( + profile=profile.name, + owner_config=profile.owner_config, + num_envs=num_envs, + dtype=dtype_name, + path="numpy_generator_random_out", + threads=None, + values=values, + mean_ms=numpy_out_ms, + min_ms=numpy_out_min, + std_ms=numpy_out_std, + values_per_s=values / (numpy_out_ms * 1.0e-3), + speedup_vs_numpy_alloc=numpy_ms / numpy_out_ms, + distribution_mean=generator_dist["mean"], + distribution_std=generator_dist["std"], + distribution_min=generator_dist["min"], + distribution_max=generator_dist["max"], + ) + ) + + if not NUMBA_AVAILABLE: + return records + + assert thread_counts is not None + max_threads = get_num_threads() + scales = _scale_array(profile, noise_level) + previous_threads = max_threads + try: + for threads in thread_counts: + if threads > max_threads: + continue + set_num_threads(threads) + numba_buffers = _allocate_buffers(profile, num_envs, dtype) + _numba_seed(seed) + t0 = time.perf_counter() + _numba_random_prange(numba_buffers, scales) + compile_ms = (time.perf_counter() - t0) * 1000.0 + numba_dist = _distribution(numba_buffers) + + def numba_call() -> None: + _numba_random_prange(numba_buffers, scales) + + numba_ms, numba_min, numba_std = _time_ms( + numba_call, + iters=iters, + warmup=warmup, + ) + same_seed = _numba_same_seed_is_deterministic( + profile=profile, + num_envs=min(num_envs, 256), + dtype=dtype, + scales=scales, + threads=threads, + seed=seed, + ) + records.append( + BenchCase( + profile=profile.name, + owner_config=profile.owner_config, + num_envs=num_envs, + dtype=dtype_name, + path="numba_random_prange", + threads=threads, + values=values, + mean_ms=numba_ms, + min_ms=numba_min, + std_ms=numba_std, + values_per_s=values / (numba_ms * 1.0e-3), + speedup_vs_numpy_alloc=numpy_ms / numba_ms, + compile_ms=compile_ms, + deterministic_same_seed=same_seed, + distribution_mean=numba_dist["mean"], + distribution_std=numba_dist["std"], + distribution_min=numba_dist["min"], + distribution_max=numba_dist["max"], + ) + ) + finally: + set_num_threads(previous_threads) + + return records + + +def _case_to_dict(case: BenchCase) -> dict[str, Any]: + return { + "profile": case.profile, + "owner_config": case.owner_config, + "num_envs": case.num_envs, + "dtype": case.dtype, + "path": case.path, + "threads": case.threads, + "values": case.values, + "mean_ms": case.mean_ms, + "min_ms": case.min_ms, + "std_ms": case.std_ms, + "values_per_s": case.values_per_s, + "mvalues_per_s": case.values_per_s / 1.0e6, + "speedup_vs_numpy_alloc": case.speedup_vs_numpy_alloc, + "compile_ms": case.compile_ms, + "deterministic_same_seed": case.deterministic_same_seed, + "distribution_mean": case.distribution_mean, + "distribution_std": case.distribution_std, + "distribution_min": case.distribution_min, + "distribution_max": case.distribution_max, + } + + +def _format_table(records: list[BenchCase]) -> str: + headers = [ + "profile", + "envs", + "path", + "thr", + "mean ms", + "min ms", + "M vals/s", + "speedup", + "seed", + ] + rows = [] + for record in records: + rows.append( + [ + record.profile, + str(record.num_envs), + record.path, + "-" if record.threads is None else str(record.threads), + f"{record.mean_ms:.3f}", + f"{record.min_ms:.3f}", + f"{record.values_per_s / 1.0e6:.1f}", + f"{record.speedup_vs_numpy_alloc:.2f}x", + "-" if record.deterministic_same_seed is None else str(record.deterministic_same_seed), + ] + ) + 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: list[str]) -> str: + return " | ".join(value.ljust(widths[idx]) for idx, value in enumerate(values)) + + out = [fmt(headers), "-+-".join("-" * width for width in widths)] + out.extend(fmt(row) for row in rows) + return "\n".join(out) + + +def _format_profile_notes(profiles: list[NoiseProfile]) -> str: + lines = ["Profiles:"] + for profile in profiles: + fields = ", ".join(f"{field.name}[{field.width}]x{field.scale:g}" for field in profile.fields) + lines.append(f"- {profile.name}: {profile.owner_config}; {fields}") + return "\n".join(lines) + + +def _select_profiles(names: list[str]) -> list[NoiseProfile]: + profiles = make_profiles() + if names == ["all"]: + return list(profiles.values()) + missing = [name for name in names if name not in profiles] + if missing: + raise ValueError(f"unknown profiles: {missing}; available: {sorted(profiles)}") + return [profiles[name] for name in names] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--profiles", + nargs="+", + default=["all"], + help="Profile names to run, or 'all'.", + ) + parser.add_argument( + "--num-envs", + nargs="+", + type=int, + default=None, + help="Override env counts. Defaults to each profile's owner-config count plus benchmark sizes.", + ) + parser.add_argument("--threads", nargs="+", type=int, default=None) + parser.add_argument("--dtype", choices=["float32", "float64"], default="float32") + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--noise-level", type=float, default=1.0) + parser.add_argument("--quick", action="store_true") + parser.add_argument("--output", type=Path, default=None, help="Optional JSON output path.") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + selected_profiles = _select_profiles(args.profiles) + if args.quick: + selected_profiles = selected_profiles[:2] + args.iters = min(args.iters, 3) + args.warmup = min(args.warmup, 1) + + thread_counts = args.threads or (QUICK_THREADS if args.quick else DEFAULT_THREADS) + num_envs_default = QUICK_NUM_ENVS if args.quick else DEFAULT_NUM_ENVS + + if NUMBA_AVAILABLE: + max_threads = get_num_threads() + skipped_threads = [thread for thread in thread_counts if thread > max_threads] + thread_counts = [thread for thread in thread_counts if thread <= max_threads] + else: + max_threads = 0 + skipped_threads = thread_counts + thread_counts = [] + + print("=" * 88) + print("Numba RNG vs NumPy RNG: UniLab observation-noise buffer benchmark") + print("=" * 88) + print(get_device_info_line()) + print(f"python={platform.python_version()} numpy={np.__version__} numba_available={NUMBA_AVAILABLE}") + print(f"host numba threads={max_threads}") + if skipped_threads: + print(f"skipped numba threads: {skipped_threads}") + print(_format_profile_notes(selected_profiles)) + print( + "Seed note: NumPy and Numba RNG streams are not expected to be bit-identical; " + "the benchmark only checks same-seed repeatability within each numba thread count." + ) + + all_records: list[BenchCase] = [] + for profile in selected_profiles: + sizes = args.num_envs or sorted({profile.default_num_envs, *num_envs_default}) + for num_envs in sizes: + print() + print(f"profile={profile.name} num_envs={num_envs} values={num_envs * profile.total_width}") + records = bench_one( + profile=profile, + num_envs=num_envs, + dtype_name=args.dtype, + thread_counts=thread_counts, + iters=args.iters, + warmup=args.warmup, + seed=args.seed, + noise_level=args.noise_level, + ) + all_records.extend(records) + print(_format_table(records)) + + if args.output is not None: + payload = { + "meta": { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "device": get_device_info_dict(), + "numpy_version": np.__version__, + "numba_available": NUMBA_AVAILABLE, + "dtype": args.dtype, + "iters": args.iters, + "warmup": args.warmup, + "seed": args.seed, + "noise_level": args.noise_level, + }, + "results": [_case_to_dict(record) for record in all_records], + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"\nSaved JSON: {args.output.resolve()}") + + +if __name__ == "__main__": + main() diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 7b002bac1..657b4b8cf 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -126,6 +126,15 @@ ) NP_ENV_STEP_COUNT_KEYS = ("reset_done_count",) NP_ENV_STEP_SAMPLE_KEYS = (*NP_ENV_STEP_TIMING_KEYS, *NP_ENV_STEP_COUNT_KEYS) +NP_RANDOM_PROFILE_FUNCTIONS = ( + "uniform", + "randint", + "random", + "rand", + "randn", + "normal", + "choice", +) NP_ENV_STEP_TIMING_CSV_FIELDS = ( ("env_step_total_ms", "np_env_step_total_ms"), ("apply_action_ms", "np_env_apply_action_ms"), @@ -226,6 +235,10 @@ class CollectorResult: env_step_overhead_ms_per_vector_step: TimingStats | None = None # System-wide CPU utilization (%) over the measured window. NaN if unknown. cpu_util_pct: float = float("nan") + # NumPy random time measured inside collector env.step() only. This is an + # optional benchmark-side monkeypatch used for RNG/noise-buffer profiling. + numpy_random_ms_per_vector_step: TimingStats | None = None + numpy_random_calls_per_vector_step: TimingStats | None = None def _stats(samples_ms: list[float]) -> TimingStats: @@ -302,6 +315,54 @@ def _cleanup() -> None: torch.mps.empty_cache() +class _NumpyRandomProfiler: + """Measure selected ``np.random`` module calls inside a benchmark window.""" + + def __init__(self) -> None: + self.enabled = False + self.step_ms = 0.0 + self.step_calls = 0 + self._originals: dict[str, Any] = {} + + def install(self) -> None: + if self._originals: + return + for name in NP_RANDOM_PROFILE_FUNCTIONS: + original = getattr(np.random, name, None) + if original is None: + continue + self._originals[name] = original + setattr(np.random, name, self._wrap(original)) + + def uninstall(self) -> None: + for name, original in self._originals.items(): + setattr(np.random, name, original) + self._originals.clear() + self.enabled = False + + def begin_step(self) -> None: + self.step_ms = 0.0 + self.step_calls = 0 + self.enabled = True + + def end_step(self) -> tuple[float, int]: + self.enabled = False + return self.step_ms, self.step_calls + + def _wrap(self, fn: Any) -> Any: + def wrapped(*args: Any, **kwargs: Any) -> Any: + if not self.enabled: + return fn(*args, **kwargs) + t0 = time.perf_counter_ns() + try: + return fn(*args, **kwargs) + finally: + self.step_ms += (time.perf_counter_ns() - t0) / 1e6 + self.step_calls += 1 + + return wrapped + + def _configure_collector_cpu_threads(cap: int | None = None) -> int: desired = DEFAULT_COLLECTOR_CPU_THREADS if cap is None else int(cap) override = os.environ.get(COLLECTOR_CPU_THREADS_ENV) @@ -552,6 +613,7 @@ def _run_active_window_case( actor, warmup_steps: int, measure_steps: int, + profile_numpy_random: bool = False, ) -> CollectorResult: from unilab.algos.torch.offpolicy.worker import ( resolve_offpolicy_actor_priv_info, @@ -590,6 +652,9 @@ def _run_active_window_case( "env_step_overhead_ms": [], } env_step_timing_samples: dict[str, list[float]] = {key: [] for key in NP_ENV_STEP_SAMPLE_KEYS} + numpy_random_ms_samples: list[float] = [] + numpy_random_call_samples: list[float] = [] + random_profiler = _NumpyRandomProfiler() if profile_numpy_random else None total_active_ns = 0 total_steps = int(warmup_steps) + int(measure_steps) if total_steps <= 0 or measure_steps <= 0: @@ -601,142 +666,160 @@ def _run_active_window_case( # "160 cores but only 25% utilized" scaling problem. cpu_probe_start: tuple[float, float] | None = None - for step_idx in range(total_steps): - record = step_idx >= warmup_steps - if record and cpu_probe_start is None: - cpu_probe_start = _read_cpu_times() # begin measured-window CPU sampling - - phase_start_ns = time.perf_counter_ns() - # No learner exists in this benchmark, so weight sync is intentionally a - # zero-work phase. Keeping it explicit preserves the training collector's - # phase schema. - weight_sync_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 - - phase_start_ns = time.perf_counter_ns() - with torch.no_grad(): - obs_torch = torch.from_numpy(obs_np) - dones_torch = torch.from_numpy(prev_dones_np) - priv_info_np = resolve_offpolicy_actor_priv_info( - algo_type=case.collector_algo_type, - obs_np=obs_np, - critic_np=critic_np, - info=info_dict, + if random_profiler is not None: + random_profiler.install() + try: + for step_idx in range(total_steps): + record = step_idx >= warmup_steps + if record and cpu_probe_start is None: + cpu_probe_start = _read_cpu_times() # begin measured-window CPU sampling + + phase_start_ns = time.perf_counter_ns() + # No learner exists in this benchmark, so weight sync is intentionally a + # zero-work phase. Keeping it explicit preserves the training collector's + # phase schema. + weight_sync_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 + + phase_start_ns = time.perf_counter_ns() + with torch.no_grad(): + obs_torch = torch.from_numpy(obs_np) + dones_torch = torch.from_numpy(prev_dones_np) + priv_info_np = resolve_offpolicy_actor_priv_info( + algo_type=case.collector_algo_type, + obs_np=obs_np, + critic_np=critic_np, + info=info_dict, + ) + priv_info_torch = ( + torch.from_numpy(priv_info_np) if priv_info_np is not None else None + ) + actions_torch = sample_offpolicy_actions( + actor=actor, + algo_type=case.collector_algo_type, + obs_torch=obs_torch, + prev_dones_torch=dones_torch, + priv_info_torch=priv_info_torch, + ) + actions_np = actions_torch.numpy() + action_select_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 + + phase_start_ns = time.perf_counter_ns() + if random_profiler is not None and record: + random_profiler.begin_step() + try: + state = env.step(actions_np) + finally: + if random_profiler is not None and record: + random_ms, random_calls = random_profiler.end_step() + numpy_random_ms_samples.append(random_ms) + numpy_random_call_samples.append(float(random_calls)) + env_step_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 + # Backend-internal physics time (sub-part of env_step_ms). Present for + # MuJoCo (via backend.step timing); absent for backends that don't + # report it -> NaN, rendered as "n/a". + _timing = state.info.get("timing", {}) if isinstance(state.info, dict) else {} + physics_ms = _optional_timing_ms(_timing, "backend_physics_ms") + env_step_timing_values = { + key: _optional_timing_ms(_timing, key) + for key in NP_ENV_STEP_SAMPLE_KEYS + if key != "env_step_internal_gap_ms" + } + internal_children = ( + env_step_timing_values["apply_action_ms"], + env_step_timing_values["step_core_ms"], + env_step_timing_values["update_state_ms"], + env_step_timing_values["reset_done_ms"], ) - priv_info_torch = torch.from_numpy(priv_info_np) if priv_info_np is not None else None - actions_torch = sample_offpolicy_actions( - actor=actor, - algo_type=case.collector_algo_type, - obs_torch=obs_torch, - prev_dones_torch=dones_torch, - priv_info_torch=priv_info_torch, + if all(value is not None for value in internal_children): + env_step_timing_values["env_step_internal_gap_ms"] = env_step_ms - sum( + cast(float, value) for value in internal_children + ) + else: + env_step_timing_values["env_step_internal_gap_ms"] = None + + phase_start_ns = time.perf_counter_ns() + next_obs_np, next_critic_np = split_obs_dict(state.obs) + next_obs_np = np.asarray(next_obs_np, dtype=np.float32) + next_critic_np = np.asarray(next_critic_np, dtype=np.float32) + rewards_np = np.asarray(state.reward, dtype=np.float32).ravel() + truncated_np = state.truncated.astype(np.float32, copy=False).ravel() + combined_dones = (state.terminated | state.truncated).astype( + np.float32, copy=False + ).ravel() + terminal_contract = resolve_terminal_observation_contract( + next_obs_batch_size=next_obs_np.shape[0], + final_observation=state.final_observation, + done=combined_dones > 0.5, + info=state.info, + truncated=truncated_np, ) - actions_np = actions_torch.numpy() - action_select_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 - - phase_start_ns = time.perf_counter_ns() - state = env.step(actions_np) - env_step_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 - # Backend-internal physics time (sub-part of env_step_ms). Present for - # MuJoCo (via backend.step timing); absent for backends that don't - # report it -> NaN, rendered as "n/a". - _timing = state.info.get("timing", {}) if isinstance(state.info, dict) else {} - physics_ms = _optional_timing_ms(_timing, "backend_physics_ms") - env_step_timing_values = { - key: _optional_timing_ms(_timing, key) - for key in NP_ENV_STEP_SAMPLE_KEYS - if key != "env_step_internal_gap_ms" - } - internal_children = ( - env_step_timing_values["apply_action_ms"], - env_step_timing_values["step_core_ms"], - env_step_timing_values["update_state_ms"], - env_step_timing_values["reset_done_ms"], - ) - if all(value is not None for value in internal_children): - env_step_timing_values["env_step_internal_gap_ms"] = env_step_ms - sum( - cast(float, value) for value in internal_children + replay_buffer.add( + torch.from_numpy(obs_np), + torch.from_numpy(actions_np), + torch.from_numpy(rewards_np), + torch.from_numpy(next_obs_np), + torch.from_numpy(combined_dones), + torch.from_numpy(truncated_np), + terminal_mask=torch.from_numpy(terminal_contract.terminal_mask), + terminal_next_obs=( + torch.from_numpy(terminal_contract.terminal_obs) + if terminal_contract.terminal_obs is not None + else None + ), + critic=torch.from_numpy(critic_np), + next_critic=torch.from_numpy(next_critic_np), + terminal_next_critic=( + torch.from_numpy(terminal_contract.terminal_critic) + if terminal_contract.terminal_critic is not None + else None + ), ) - else: - env_step_timing_values["env_step_internal_gap_ms"] = None - - phase_start_ns = time.perf_counter_ns() - next_obs_np, next_critic_np = split_obs_dict(state.obs) - next_obs_np = np.asarray(next_obs_np, dtype=np.float32) - next_critic_np = np.asarray(next_critic_np, dtype=np.float32) - rewards_np = np.asarray(state.reward, dtype=np.float32).ravel() - truncated_np = state.truncated.astype(np.float32, copy=False).ravel() - combined_dones = (state.terminated | state.truncated).astype(np.float32, copy=False).ravel() - terminal_contract = resolve_terminal_observation_contract( - next_obs_batch_size=next_obs_np.shape[0], - final_observation=state.final_observation, - done=combined_dones > 0.5, - info=state.info, - truncated=truncated_np, - ) - replay_buffer.add( - torch.from_numpy(obs_np), - torch.from_numpy(actions_np), - torch.from_numpy(rewards_np), - torch.from_numpy(next_obs_np), - torch.from_numpy(combined_dones), - torch.from_numpy(truncated_np), - terminal_mask=torch.from_numpy(terminal_contract.terminal_mask), - terminal_next_obs=( - torch.from_numpy(terminal_contract.terminal_obs) - if terminal_contract.terminal_obs is not None - else None - ), - critic=torch.from_numpy(critic_np), - next_critic=torch.from_numpy(next_critic_np), - terminal_next_critic=( - torch.from_numpy(terminal_contract.terminal_critic) - if terminal_contract.terminal_critic is not None - else None - ), - ) - replay_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 - - phase_start_ns = time.perf_counter_ns() - current_ep_rewards += rewards_np - current_ep_lengths += 1 - reset_indices = np.where(combined_dones > 0.5)[0] - if len(reset_indices) > 0: - ep_rewards.extend(current_ep_rewards[reset_indices].tolist()) - ep_lengths.extend(current_ep_lengths[reset_indices].tolist()) - current_ep_rewards[reset_indices] = 0.0 - current_ep_lengths[reset_indices] = 0 - - log_info = state.info.get("log", {}) - if log_info: - for key, value in log_info.items(): - if key.startswith("reward/"): - ep_reward_components[key].append(value) - bookkeeping_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 - - obs_np = next_obs_np - critic_np = next_critic_np - info_dict = state.info - prev_dones_np = combined_dones - - if record: - phase_values = { - "weight_sync_ms": weight_sync_ms, - "action_select_ms": action_select_ms, - "env_step_ms": env_step_ms, - "replay_ms": replay_ms, - "bookkeeping_ms": bookkeeping_ms, - } - step_active_ns = int(sum(phase_values.values()) * 1e6) - total_active_ns += step_active_ns - for key, value in phase_values.items(): - samples[key].append(value) - # Env-step breakdown is aux only; env_step_ms is the additive phase. - if physics_ms is not None: - aux_samples["physics_ms"].append(physics_ms) - aux_samples["env_step_overhead_ms"].append(env_step_ms - physics_ms) - for key, value in env_step_timing_values.items(): - if value is not None: - env_step_timing_samples[key].append(value) + replay_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 + + phase_start_ns = time.perf_counter_ns() + current_ep_rewards += rewards_np + current_ep_lengths += 1 + reset_indices = np.where(combined_dones > 0.5)[0] + if len(reset_indices) > 0: + ep_rewards.extend(current_ep_rewards[reset_indices].tolist()) + ep_lengths.extend(current_ep_lengths[reset_indices].tolist()) + current_ep_rewards[reset_indices] = 0.0 + current_ep_lengths[reset_indices] = 0 + + log_info = state.info.get("log", {}) + if log_info: + for key, value in log_info.items(): + if key.startswith("reward/"): + ep_reward_components[key].append(value) + bookkeeping_ms = (time.perf_counter_ns() - phase_start_ns) / 1e6 + + obs_np = next_obs_np + critic_np = next_critic_np + info_dict = state.info + prev_dones_np = combined_dones + + if record: + phase_values = { + "weight_sync_ms": weight_sync_ms, + "action_select_ms": action_select_ms, + "env_step_ms": env_step_ms, + "replay_ms": replay_ms, + "bookkeeping_ms": bookkeeping_ms, + } + step_active_ns = int(sum(phase_values.values()) * 1e6) + total_active_ns += step_active_ns + for key, value in phase_values.items(): + samples[key].append(value) + # Env-step breakdown is aux only; env_step_ms is the additive phase. + if physics_ms is not None: + aux_samples["physics_ms"].append(physics_ms) + aux_samples["env_step_overhead_ms"].append(env_step_ms - physics_ms) + for key, value in env_step_timing_values.items(): + if value is not None: + env_step_timing_samples[key].append(value) + finally: + if random_profiler is not None: + random_profiler.uninstall() # Measured-window system CPU utilization (see cpu_probe_start comment). cpu_util_pct = _cpu_util_pct(cpu_probe_start, _read_cpu_times()) @@ -772,6 +855,12 @@ def _run_active_window_case( physics_ms_per_vector_step=physics_stats, env_step_overhead_ms_per_vector_step=env_step_overhead_stats, cpu_util_pct=cpu_util_pct, + numpy_random_ms_per_vector_step=( + _stats(numpy_random_ms_samples) if numpy_random_ms_samples else None + ), + numpy_random_calls_per_vector_step=( + _stats(numpy_random_call_samples) if numpy_random_call_samples else None + ), ) @@ -784,6 +873,7 @@ def _build_and_run_case( num_envs: int | None, extra_overrides: list[str], variant: str = "default", + profile_numpy_random: bool = False, ) -> CollectorResult: from unilab.training import BackendAdapter, ensure_registries from unilab.training.seed import apply_training_seed @@ -844,6 +934,7 @@ def _build_and_run_case( actor=actor, warmup_steps=warmup_steps, measure_steps=measure_steps, + profile_numpy_random=profile_numpy_random, ) finally: if env is not None: @@ -867,6 +958,7 @@ def _build_and_run_numba_ab_case( num_envs: int | None, extra_overrides: list[str], numba_threads: int | None, + profile_numpy_random: bool = False, ) -> list[CollectorResult]: _algo, task, _sim = _parse_case(spec) if task not in NUMBA_ACCELERATED_TASKS: @@ -890,6 +982,7 @@ def _build_and_run_numba_ab_case( *_numba_ab_overrides(enabled=enabled, numba_threads=numba_threads), ], variant=variant, + profile_numpy_random=profile_numpy_random, ) ) return records @@ -935,6 +1028,8 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: "replay_ms", "weight_sync_ms", "bookkeeping_ms", + "numpy_random_ms", + "numpy_random_calls", "physics_ms", "env_step_overhead_ms", "reset_done_count", @@ -980,6 +1075,16 @@ def _write_csv(path: Path, results: list[CollectorResult]) -> None: if result.physics_ms_per_vector_step is not None else "" ) + row["numpy_random_ms"] = ( + result.numpy_random_ms_per_vector_step.mean_ms + if result.numpy_random_ms_per_vector_step is not None + else "" + ) + row["numpy_random_calls"] = ( + result.numpy_random_calls_per_vector_step.mean_ms + if result.numpy_random_calls_per_vector_step is not None + else "" + ) row["env_step_overhead_ms"] = ( result.env_step_overhead_ms_per_vector_step.mean_ms if result.env_step_overhead_ms_per_vector_step is not None @@ -1692,6 +1797,19 @@ def _print_result(result: CollectorResult) -> None: f"pct_env={_env_step_child_env_pct(result, upper):5.1f}% " f"pct_active={_env_step_child_pct(result, upper):5.1f}%" ) + if result.numpy_random_ms_per_vector_step is not None: + random_ms = result.numpy_random_ms_per_vector_step.mean_ms + random_calls = ( + result.numpy_random_calls_per_vector_step.mean_ms + if result.numpy_random_calls_per_vector_step is not None + else 0.0 + ) + print( + f" {'numpy_random_ms':<18} mean={random_ms:8.3f} ms " + f"calls={random_calls:5.1f} " + f"pct_env={_env_step_child_env_pct(result, random_ms):5.1f}% " + f"pct_active={_env_step_child_pct(result, random_ms):5.1f}%" + ) if result.env_step_timing_ms_per_vector_step: reset_count = result.env_step_timing_ms_per_vector_step.get("reset_done_count") if reset_count is not None: @@ -1802,6 +1920,14 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: default=None, help="Numba thread count used by --numba-ab accelerated variants.", ) + parser.add_argument( + "--profile-numpy-random", + action="store_true", + help=( + "Measure selected np.random module calls during measured env.step() " + "windows and report their collector active share." + ), + ) parser.add_argument("--out-json", type=Path, default=DEFAULT_OUTPUT_JSON) parser.add_argument("--out-csv", type=Path, default=None) parser.add_argument("--continue-on-error", action="store_true") @@ -1838,6 +1964,7 @@ def main() -> int: num_envs=args.num_envs, extra_overrides=list(args.override), numba_threads=args.numba_threads, + profile_numpy_random=bool(args.profile_numpy_random), ) else: case_results = [ @@ -1848,6 +1975,7 @@ def main() -> int: replay_capacity_steps=int(args.replay_capacity_steps), num_envs=args.num_envs, extra_overrides=list(args.override), + profile_numpy_random=bool(args.profile_numpy_random), ) ] results.extend(case_results) @@ -1881,6 +2009,7 @@ def main() -> int: "override": args.override, "numba_ab": args.numba_ab, "numba_threads": args.numba_threads, + "profile_numpy_random": args.profile_numpy_random, }, "results": [_result_to_dict(result) for result in results], "errors": errors, diff --git a/tests/benchmark/test_numba_random_noise_benchmark.py b/tests/benchmark/test_numba_random_noise_benchmark.py new file mode 100644 index 000000000..b87ee7bb0 --- /dev/null +++ b/tests/benchmark/test_numba_random_noise_benchmark.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import math + +import numpy as np +from benchmark import benchmark_numba_random_noise as bench + + +def test_numba_random_noise_profiles_match_training_shapes() -> None: + profiles = bench.make_profiles() + + motion = profiles["sac_g1_motion_tracking_mujoco"] + joystick = profiles["sac_g1_walk_flat_mujoco"] + + assert motion.default_num_envs == 2048 + assert motion.total_width == 64 + assert [field.name for field in motion.fields] == [ + "linvel", + "gyro", + "joint_pos", + "dof_vel", + ] + assert joystick.total_width == 64 + assert joystick.fields[0].scale == 0.0 + assert joystick.fields[1].scale == 0.0 + + +def test_numba_random_noise_benchmark_builds_records() -> None: + profile = bench.make_profiles()["ppo_g1_motion_tracking_motrix"] + threads = [1] if bench.NUMBA_AVAILABLE else [] + + records = bench.bench_one( + profile=profile, + num_envs=32, + dtype_name="float32", + thread_counts=threads, + iters=1, + warmup=0, + seed=0, + ) + + paths = {record.path for record in records} + assert "numpy_random_uniform_alloc" in paths + assert "numpy_generator_random_out" in paths + if bench.NUMBA_AVAILABLE: + assert "numba_random_prange" in paths + assert all(record.values == 32 * profile.total_width for record in records) + assert all(record.mean_ms >= 0.0 for record in records) + assert all(math.isfinite(record.speedup_vs_numpy_alloc) for record in records) + + +def test_numba_random_noise_zero_scale_fields_are_zeroed() -> None: + profile = bench.make_profiles()["sac_g1_walk_flat_mujoco"] + buffers = bench._numpy_random_uniform_alloc( + profile, + num_envs=8, + dtype=np.dtype(np.float32), + noise_level=1.0, + ) + + assert np.max(np.abs(buffers[0])) == 0.0 + assert np.max(np.abs(buffers[1])) == 0.0 + assert np.max(np.abs(buffers[2])) > 0.0 + + +def test_numba_random_noise_formats_records() -> None: + record = bench.BenchCase( + profile="sac_g1_motion_tracking_mujoco", + owner_config="conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml", + num_envs=64, + dtype="float32", + path="numba_random_prange", + threads=4, + values=4096, + mean_ms=0.25, + min_ms=0.2, + std_ms=0.01, + values_per_s=16_384_000.0, + speedup_vs_numpy_alloc=2.0, + compile_ms=10.0, + deterministic_same_seed=True, + ) + + payload = bench._case_to_dict(record) + table = bench._format_table([record]) + + assert payload["mvalues_per_s"] == 16.384 + assert "numba_random_prange" in table + assert "2.00x" in table + assert "True" in table From f539e5564be88822dd2f0b1d3ee01db60cdb2e6c Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Wed, 8 Jul 2026 02:16:04 +0800 Subject: [PATCH 17/26] fix: satisfy numba accel gate checks --- benchmark/benchmark_g1_joystick_numba.py | 19 ++-- .../benchmark_g1_motion_tracking_numba.py | 39 ++++---- benchmark/benchmark_numba_random_noise.py | 16 +++- .../benchmark_offpolicy_collector_active.py | 6 +- .../g1/motion_tracking_numba.py | 92 +++++++++++-------- .../envs/motion_tracking/g1/tracking.py | 5 +- src/unilab/utils/numba_geometry.py | 32 +++++-- .../test_g1_joystick_numba_benchmark.py | 8 +- ...test_g1_motion_tracking_numba_benchmark.py | 8 +- .../envs/locomotion/g1/test_joystick_numba.py | 8 +- tests/envs/test_g1_motion_tracking_numba.py | 20 ++-- tests/utils/test_numba_geometry.py | 8 +- 12 files changed, 147 insertions(+), 114 deletions(-) diff --git a/benchmark/benchmark_g1_joystick_numba.py b/benchmark/benchmark_g1_joystick_numba.py index 8442dc21f..43a2ff965 100644 --- a/benchmark/benchmark_g1_joystick_numba.py +++ b/benchmark/benchmark_g1_joystick_numba.py @@ -388,7 +388,9 @@ def compute_numpy( batch.dof_pos, batch.dof_vel, ) - obs = env._compute_obs(info, batch.linvel, batch.gyro, batch.gravity, batch.dof_pos, batch.dof_vel) + obs = env._compute_obs( + info, batch.linvel, batch.gyro, batch.gravity, batch.dof_pos, batch.dof_vel + ) return obs, reward, terminated, info.get("log", {}) @@ -498,9 +500,7 @@ def bench_one( ) set_num_threads(max_threads) numba_1t = next( - record - for record in records - if record.path == "numba_accelerator" and record.threads == 1 + record for record in records if record.path == "numba_accelerator" and record.threads == 1 ) for record in records: if record.path != "numba_accelerator" or record.threads is None: @@ -784,7 +784,12 @@ def _format_e2e_reconciliation_table( num_envs=record.num_envs, threads=record.numba_threads, ) - if baseline is None or baseline.update_state_ms is None or numpy_hot is None or numba_hot is None: + if ( + baseline is None + or baseline.update_state_ms is None + or numpy_hot is None + or numba_hot is None + ): continue hot_saved_ms = numpy_hot.mean_ms - numba_hot.mean_ms update_saved_ms = baseline.update_state_ms - record.update_state_ms @@ -1308,7 +1313,9 @@ def main() -> None: parity: dict[str, dict[str, float]] = {} max_threads = get_num_threads() args.numba_max_threads = max_threads - args.measured_threads = sorted({1, *(threads for threads in args.threads if threads <= max_threads)}) + args.measured_threads = sorted( + {1, *(threads for threads in args.threads if threads <= max_threads)} + ) args.skipped_threads = sorted({threads for threads in args.threads if threads > max_threads}) print("=" * 80) diff --git a/benchmark/benchmark_g1_motion_tracking_numba.py b/benchmark/benchmark_g1_motion_tracking_numba.py index 365f649b3..5db0f5beb 100644 --- a/benchmark/benchmark_g1_motion_tracking_numba.py +++ b/benchmark/benchmark_g1_motion_tracking_numba.py @@ -333,12 +333,8 @@ def f32(value: np.ndarray) -> np.ndarray: ) robot_body_pos_w = f32(target_pos + 0.06 * rng.standard_normal((num_envs, n_body, 3))) robot_body_quat_w = f32(_perturb_quats(rng, target_quat, 0.03)) - robot_body_lin_vel_w = f32( - target_lin_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3)) - ) - robot_body_ang_vel_w = f32( - target_ang_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3)) - ) + robot_body_lin_vel_w = f32(target_lin_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3))) + robot_body_ang_vel_w = f32(target_ang_vel + 0.1 * rng.standard_normal((num_envs, n_body, 3))) linvel = f32(rng.uniform(-1.0, 1.0, (num_envs, 3))) gyro = f32(rng.uniform(-1.0, 1.0, (num_envs, 3))) dof_pos = f32(target_joint_pos + 0.05 * rng.standard_normal((num_envs, NUM_ACTION))) @@ -516,9 +512,7 @@ def check_parity( return { "max_abs_reward_diff": float(np.max(np.abs(reward_nb - reward_np))), "termination_mismatch": float(np.count_nonzero(terminated_nb != terminated_np)), - "max_abs_update_state_reward_diff": float( - np.max(np.abs(full_reward_nb - full_reward_np)) - ), + "max_abs_update_state_reward_diff": float(np.max(np.abs(full_reward_nb - full_reward_np))), "update_state_termination_mismatch": float( np.count_nonzero(full_terminated_nb != full_terminated_np) ), @@ -591,9 +585,7 @@ def bench_one( ) set_num_threads(max_threads) numba_1t = next( - record - for record in records - if record.path == "numba_accelerator" and record.threads == 1 + record for record in records if record.path == "numba_accelerator" and record.threads == 1 ) for record in records: if record.path != "numba_accelerator" or record.threads is None: @@ -901,9 +893,7 @@ def _format_hot_summary_table(records: list[BenchCase]) -> str: best_by_case = _best_numba_by_case(records) rows = [] for profile in sorted({record.profile for record in records}): - best_records = [ - record for key, record in best_by_case.items() if key[0] == profile - ] + best_records = [record for key, record in best_by_case.items() if key[0] == profile] if not best_records: continue envs = sorted(record.num_envs for record in best_records) @@ -913,7 +903,9 @@ def _format_hot_summary_table(records: list[BenchCase]) -> str: [ profile, f"{envs[0]}-{envs[-1]}", - ",".join(str(thread) for thread in sorted({record.threads for record in best_records})), + ",".join( + str(thread) for thread in sorted({record.threads for record in best_records}) + ), f"{min(speedups):.2f}x-{max(speedups):.2f}x", f"{max(throughputs):.2f}", ] @@ -1070,7 +1062,12 @@ def _format_e2e_reconciliation_table( num_envs=record.num_envs, threads=record.numba_threads, ) - if baseline is None or baseline.update_state_ms is None or numpy_hot is None or numba_hot is None: + if ( + baseline is None + or baseline.update_state_ms is None + or numpy_hot is None + or numba_hot is None + ): continue hot_saved_ms = numpy_hot.mean_ms - numba_hot.mean_ms update_saved_ms = baseline.update_state_ms - record.update_state_ms @@ -1131,7 +1128,9 @@ def save_plots( best_by_case = _best_numba_by_case(records) fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(21, 6)) - fig.suptitle(f"G1 motion tracking Numba reward+termination benchmark\n{device_info}", fontsize=13) + fig.suptitle( + f"G1 motion tracking Numba reward+termination benchmark\n{device_info}", fontsize=13 + ) ax1, ax2, ax3 = axes for profile in profiles: @@ -1488,7 +1487,9 @@ def main() -> None: parity: dict[str, dict[str, float]] = {} max_threads = get_num_threads() args.numba_max_threads = max_threads - args.measured_threads = sorted({1, *(threads for threads in args.threads if threads <= max_threads)}) + args.measured_threads = sorted( + {1, *(threads for threads in args.threads if threads <= max_threads)} + ) args.skipped_threads = sorted({threads for threads in args.threads if threads > max_threads}) print("=" * 80) diff --git a/benchmark/benchmark_numba_random_noise.py b/benchmark/benchmark_numba_random_noise.py index 6bc669132..188a3f39d 100644 --- a/benchmark/benchmark_numba_random_noise.py +++ b/benchmark/benchmark_numba_random_noise.py @@ -477,7 +477,9 @@ def _format_table(records: list[BenchCase]) -> str: f"{record.min_ms:.3f}", f"{record.values_per_s / 1.0e6:.1f}", f"{record.speedup_vs_numpy_alloc:.2f}x", - "-" if record.deterministic_same_seed is None else str(record.deterministic_same_seed), + "-" + if record.deterministic_same_seed is None + else str(record.deterministic_same_seed), ] ) widths = [len(header) for header in headers] @@ -496,7 +498,9 @@ def fmt(values: list[str]) -> str: def _format_profile_notes(profiles: list[NoiseProfile]) -> str: lines = ["Profiles:"] for profile in profiles: - fields = ", ".join(f"{field.name}[{field.width}]x{field.scale:g}" for field in profile.fields) + fields = ", ".join( + f"{field.name}[{field.width}]x{field.scale:g}" for field in profile.fields + ) lines.append(f"- {profile.name}: {profile.owner_config}; {fields}") return "\n".join(lines) @@ -561,7 +565,9 @@ def main() -> None: print("Numba RNG vs NumPy RNG: UniLab observation-noise buffer benchmark") print("=" * 88) print(get_device_info_line()) - print(f"python={platform.python_version()} numpy={np.__version__} numba_available={NUMBA_AVAILABLE}") + print( + f"python={platform.python_version()} numpy={np.__version__} numba_available={NUMBA_AVAILABLE}" + ) print(f"host numba threads={max_threads}") if skipped_threads: print(f"skipped numba threads: {skipped_threads}") @@ -576,7 +582,9 @@ def main() -> None: sizes = args.num_envs or sorted({profile.default_num_envs, *num_envs_default}) for num_envs in sizes: print() - print(f"profile={profile.name} num_envs={num_envs} values={num_envs * profile.total_width}") + print( + f"profile={profile.name} num_envs={num_envs} values={num_envs * profile.total_width}" + ) records = bench_one( profile=profile, num_envs=num_envs, diff --git a/benchmark/benchmark_offpolicy_collector_active.py b/benchmark/benchmark_offpolicy_collector_active.py index 657b4b8cf..7795df25d 100644 --- a/benchmark/benchmark_offpolicy_collector_active.py +++ b/benchmark/benchmark_offpolicy_collector_active.py @@ -743,9 +743,9 @@ def _run_active_window_case( next_critic_np = np.asarray(next_critic_np, dtype=np.float32) rewards_np = np.asarray(state.reward, dtype=np.float32).ravel() truncated_np = state.truncated.astype(np.float32, copy=False).ravel() - combined_dones = (state.terminated | state.truncated).astype( - np.float32, copy=False - ).ravel() + combined_dones = ( + (state.terminated | state.truncated).astype(np.float32, copy=False).ravel() + ) terminal_contract = resolve_terminal_observation_contract( next_obs_batch_size=next_obs_np.shape[0], final_observation=state.final_observation, diff --git a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py b/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py index 31d23bb2e..e8b988655 100644 --- a/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py +++ b/src/unilab/envs/motion_tracking/g1/motion_tracking_numba.py @@ -263,15 +263,17 @@ def _compute_reward_termination_kernel( tid = get_thread_id() total = 0.0 - w = motion_global_root_pos_i( - motion_body_pos_w, robot_body_pos_w, anchor, std[0], i - ) * scale[0] + w = ( + motion_global_root_pos_i(motion_body_pos_w, robot_body_pos_w, anchor, std[0], i) + * scale[0] + ) total += w log_scratch[tid, 0] += w - w = motion_global_root_ori_i( - motion_body_quat_w, robot_body_quat_w, anchor, std[1], i - ) * scale[1] + w = ( + motion_global_root_ori_i(motion_body_quat_w, robot_body_quat_w, anchor, std[1], i) + * scale[1] + ) total += w log_scratch[tid, 1] += w @@ -283,21 +285,28 @@ def _compute_reward_termination_kernel( total += w log_scratch[tid, 3] += w - w = motion_body_lin_vel_i( - motion_body_lin_vel_w, robot_body_lin_vel_w, n_body, std[4], i - ) * scale[4] + w = ( + motion_body_lin_vel_i( + motion_body_lin_vel_w, robot_body_lin_vel_w, n_body, std[4], i + ) + * scale[4] + ) total += w log_scratch[tid, 4] += w - w = motion_body_ang_vel_i( - motion_body_ang_vel_w, robot_body_ang_vel_w, n_body, std[5], i - ) * scale[5] + w = ( + motion_body_ang_vel_i( + motion_body_ang_vel_w, robot_body_ang_vel_w, n_body, std[5], i + ) + * scale[5] + ) total += w log_scratch[tid, 5] += w - w = motion_ee_body_pos_z_i( - ref_body_pos_w, robot_body_pos_w, ee_indices, std[6], i - ) * scale[6] + w = ( + motion_ee_body_pos_z_i(ref_body_pos_w, robot_body_pos_w, ee_indices, std[6], i) + * scale[6] + ) total += w log_scratch[tid, 6] += w @@ -313,9 +322,10 @@ def _compute_reward_termination_kernel( total += w log_scratch[tid, 9] += w - w = joint_limit_i( - dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i - ) * scale[10] + w = ( + joint_limit_i(dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i) + * scale[10] + ) total += w log_scratch[tid, 10] += w @@ -346,7 +356,6 @@ def _compute_reward_termination_kernel( i, ) - @_dev def _write_reference_transforms_i( motion_body_pos_w, @@ -435,15 +444,17 @@ def _compute_reward_termination_i( ): total = 0.0 - w = motion_global_root_pos_i( - motion_body_pos_w, robot_body_pos_w, anchor, std[0], i - ) * scale[0] + w = ( + motion_global_root_pos_i(motion_body_pos_w, robot_body_pos_w, anchor, std[0], i) + * scale[0] + ) total += w log_scratch[tid, 0] += w - w = motion_global_root_ori_i( - motion_body_quat_w, robot_body_quat_w, anchor, std[1], i - ) * scale[1] + w = ( + motion_global_root_ori_i(motion_body_quat_w, robot_body_quat_w, anchor, std[1], i) + * scale[1] + ) total += w log_scratch[tid, 1] += w @@ -455,19 +466,24 @@ def _compute_reward_termination_i( total += w log_scratch[tid, 3] += w - w = motion_body_lin_vel_i( - motion_body_lin_vel_w, robot_body_lin_vel_w, n_body, std[4], i - ) * scale[4] + w = ( + motion_body_lin_vel_i(motion_body_lin_vel_w, robot_body_lin_vel_w, n_body, std[4], i) + * scale[4] + ) total += w log_scratch[tid, 4] += w - w = motion_body_ang_vel_i( - motion_body_ang_vel_w, robot_body_ang_vel_w, n_body, std[5], i - ) * scale[5] + w = ( + motion_body_ang_vel_i(motion_body_ang_vel_w, robot_body_ang_vel_w, n_body, std[5], i) + * scale[5] + ) total += w log_scratch[tid, 5] += w - w = motion_ee_body_pos_z_i(ref_body_pos_w, robot_body_pos_w, ee_indices, std[6], i) * scale[6] + w = ( + motion_ee_body_pos_z_i(ref_body_pos_w, robot_body_pos_w, ee_indices, std[6], i) + * scale[6] + ) total += w log_scratch[tid, 6] += w @@ -483,7 +499,10 @@ def _compute_reward_termination_i( total += w log_scratch[tid, 9] += w - w = joint_limit_i(dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i) * scale[10] + w = ( + joint_limit_i(dof_pos, joint_lower, joint_upper, n_action, has_joint_limits, i) + * scale[10] + ) total += w log_scratch[tid, 10] += w @@ -729,7 +748,6 @@ def _write_critic_obs_i( ) _write_critic_linvel_tail_i(linvel, critic_obs, tail_offset, i) - @njit(parallel=True, fastmath=True, cache=True, nogil=True) # type: ignore[misc] def _compute_update_state_kernel( motion_body_pos_w, @@ -976,8 +994,10 @@ def from_env( if default_angles is None: default_angles = np.zeros((env._num_action,), dtype=np.float64) actor_obs_width = getattr(env, "_actor_obs_width", env._num_action * 5 + 15) - critic_obs_width = env.obs_groups_spec["critic"] if hasattr(env, "obs_groups_spec") else ( - env._num_action * 5 + 15 + len(env._cfg.body_names) * 9 + critic_obs_width = ( + env.obs_groups_spec["critic"] + if hasattr(env, "obs_groups_spec") + else (env._num_action * 5 + 15 + len(env._cfg.body_names) * 9) ) return cls( num_envs=env.num_envs, diff --git a/src/unilab/envs/motion_tracking/g1/tracking.py b/src/unilab/envs/motion_tracking/g1/tracking.py index c8cef8d4f..4a9b4df98 100644 --- a/src/unilab/envs/motion_tracking/g1/tracking.py +++ b/src/unilab/envs/motion_tracking/g1/tracking.py @@ -788,9 +788,10 @@ def update_state(self, state: NpEnvState) -> NpEnvState: robot_body_ang_vel_w, ) = self._get_body_state_w() - if self._numba_accelerator is not None: + numba_accelerator = getattr(self, "_numba_accelerator", None) + if numba_accelerator is not None: noise_cfg = self._cfg.noise_config - numba_result = self._numba_accelerator.compute_update_state( + numba_result = numba_accelerator.compute_update_state( info=state.info, motion_data=motion_data, linvel=linvel, diff --git a/src/unilab/utils/numba_geometry.py b/src/unilab/utils/numba_geometry.py index d08e5bffc..bfd6b9fd9 100644 --- a/src/unilab/utils/numba_geometry.py +++ b/src/unilab/utils/numba_geometry.py @@ -9,7 +9,8 @@ from __future__ import annotations import math -from typing import Any +from collections.abc import Callable +from typing import Any, TypeAlias, cast try: # pragma: no cover - exercised when optional numba dependency is installed from numba import njit @@ -24,10 +25,23 @@ def _missing_numba(*_args: Any, **_kwargs: Any) -> None: raise RuntimeError("numba_geometry helpers require numba to be installed") +NumbaHelper: TypeAlias = Callable[..., Any] +quat_angle_sq_at: NumbaHelper +quat_gravity_z_at: NumbaHelper +quat_yaw_from_components: NumbaHelper +rotate_vec_by_inv_quat_components: NumbaHelper +write_quat_first_two_matrix_cols_from_components: NumbaHelper +write_yaw_aligned_body_transforms_at: NumbaHelper +write_relative_anchor_transform_at: NumbaHelper +write_body_pos_relative_to_anchor_at: NumbaHelper +write_body_quat_relative_6d_to_anchor_at: NumbaHelper + + if NUMBA_GEOMETRY_AVAILABLE: + _njit = cast(Any, njit) - def _dev(fn): - return njit(inline="always", fastmath=True, cache=True, nogil=True)(fn) + def _dev(fn: NumbaHelper) -> NumbaHelper: + return cast(NumbaHelper, _njit(inline="always", fastmath=True, cache=True, nogil=True)(fn)) @_dev def quat_angle_sq_at(q1, q2, item_idx, i): @@ -161,12 +175,8 @@ def write_yaw_aligned_body_transforms_at( vx = source_body_pos_w[i, body_idx, 0] - anchor_px vy = source_body_pos_w[i, body_idx, 1] - anchor_py vz = source_body_pos_w[i, body_idx, 2] - anchor_pz - out_body_pos_w[i, body_idx, 0] = ( - vx - yaw_cross * vy - yaw_z2 * vx + target_anchor_px - ) - out_body_pos_w[i, body_idx, 1] = ( - vy + yaw_cross * vx - yaw_z2 * vy + target_anchor_py - ) + out_body_pos_w[i, body_idx, 0] = vx - yaw_cross * vy - yaw_z2 * vx + target_anchor_px + out_body_pos_w[i, body_idx, 1] = vy + yaw_cross * vx - yaw_z2 * vy + target_anchor_py out_body_pos_w[i, body_idx, 2] = vz + anchor_pz @_dev @@ -236,7 +246,9 @@ def write_relative_anchor_transform_at( write_quat_first_two_matrix_cols_from_components(rw, rx, ry, rz, out_rot6d_b, i, 0) @_dev - def write_body_pos_relative_to_anchor_at(body_pos_w, anchor_quat_w, anchor, n_body, out, i, offset): + def write_body_pos_relative_to_anchor_at( + body_pos_w, anchor_quat_w, anchor, n_body, out, i, offset + ): anchor_px = body_pos_w[i, anchor, 0] anchor_py = body_pos_w[i, anchor, 1] anchor_pz = body_pos_w[i, anchor, 2] diff --git a/tests/benchmark/test_g1_joystick_numba_benchmark.py b/tests/benchmark/test_g1_joystick_numba_benchmark.py index a62632472..f2a09f037 100644 --- a/tests/benchmark/test_g1_joystick_numba_benchmark.py +++ b/tests/benchmark/test_g1_joystick_numba_benchmark.py @@ -133,12 +133,8 @@ def test_g1_joystick_numba_benchmark_formats_e2e_reconciliation() -> None: def test_g1_joystick_numba_benchmark_selects_best_hot_slice_threads() -> None: records = [ bench.BenchCase("sac_default", 1024, "numpy_dispatch", None, 1.0, 1.0, 0.0, 1000.0, 1.0), - bench.BenchCase( - "sac_default", 1024, "numba_accelerator", 2, 0.8, 0.8, 0.0, 1250.0, 1.25 - ), - bench.BenchCase( - "sac_default", 1024, "numba_accelerator", 4, 0.6, 0.6, 0.0, 1666.0, 1.67 - ), + bench.BenchCase("sac_default", 1024, "numba_accelerator", 2, 0.8, 0.8, 0.0, 1250.0, 1.25), + bench.BenchCase("sac_default", 1024, "numba_accelerator", 4, 0.6, 0.6, 0.0, 1666.0, 1.67), bench.BenchCase("ppo_default", 1024, "numba_accelerator", 8, 0.5, 0.5, 0.0, 2000.0, 2.0), ] diff --git a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py index 722c551c8..cdba5e232 100644 --- a/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py +++ b/tests/benchmark/test_g1_motion_tracking_numba_benchmark.py @@ -158,12 +158,8 @@ def test_g1_motion_tracking_numba_benchmark_formats_e2e_reconciliation() -> None def test_g1_motion_tracking_numba_benchmark_selects_best_hot_slice_threads() -> None: records = [ bench.BenchCase("sac_default", 1024, "numpy_dispatch", None, 1.0, 1.0, 0.0, 1000.0, 1.0), - bench.BenchCase( - "sac_default", 1024, "numba_accelerator", 2, 0.8, 0.8, 0.0, 1250.0, 1.25 - ), - bench.BenchCase( - "sac_default", 1024, "numba_accelerator", 4, 0.6, 0.6, 0.0, 1666.0, 1.67 - ), + bench.BenchCase("sac_default", 1024, "numba_accelerator", 2, 0.8, 0.8, 0.0, 1250.0, 1.25), + bench.BenchCase("sac_default", 1024, "numba_accelerator", 4, 0.6, 0.6, 0.0, 1666.0, 1.67), bench.BenchCase("ppo_default", 1024, "numba_accelerator", 8, 0.5, 0.5, 0.0, 2000.0, 2.0), ] diff --git a/tests/envs/locomotion/g1/test_joystick_numba.py b/tests/envs/locomotion/g1/test_joystick_numba.py index b094430cc..de0beb7a4 100644 --- a/tests/envs/locomotion/g1/test_joystick_numba.py +++ b/tests/envs/locomotion/g1/test_joystick_numba.py @@ -247,12 +247,10 @@ def test_g1_walk_numba_update_state_parity_without_noise(): dof_vel = rng.normal(size=(n, n_action)).astype(np.float32) max_tilt_rad = np.deg2rad(env._reward_cfg.max_tilt_deg) - terminated_np = ( - np.arccos(np.clip(gravity[:, 2], -1.0, 1.0)) > max_tilt_rad - ) | (env._backend.get_base_pos()[:, 2] < env._reward_cfg.min_base_height) - reward_np = env._compute_reward( - {**info, "log": {}}, linvel, gyro, gravity, dof_pos, dof_vel + terminated_np = (np.arccos(np.clip(gravity[:, 2], -1.0, 1.0)) > max_tilt_rad) | ( + env._backend.get_base_pos()[:, 2] < env._reward_cfg.min_base_height ) + reward_np = env._compute_reward({**info, "log": {}}, linvel, gyro, gravity, dof_pos, dof_vel) obs_np = env._compute_obs(info, linvel, gyro, gravity, dof_pos, dof_vel) accel = G1WalkNumbaAccelerator.from_env(env, num_threads=2) diff --git a/tests/envs/test_g1_motion_tracking_numba.py b/tests/envs/test_g1_motion_tracking_numba.py index 5c3a00a81..8a16adec7 100644 --- a/tests/envs/test_g1_motion_tracking_numba.py +++ b/tests/envs/test_g1_motion_tracking_numba.py @@ -190,9 +190,7 @@ def test_g1_motion_tracking_numba_term_py_funcs_match_numpy_math(): i = 2 diff = motion_data.body_pos_w[i, env.anchor_body_idx] - robot_body_pos_w[i, env.anchor_body_idx] - expected_root_pos = np.exp( - -np.sum(diff * diff) / (env._cfg.reward_config.std_root_pos**2) - ) + expected_root_pos = np.exp(-np.sum(diff * diff) / (env._cfg.reward_config.std_root_pos**2)) assert T.motion_global_root_pos_i.py_func( motion_data.body_pos_w, robot_body_pos_w, @@ -203,7 +201,9 @@ def test_g1_motion_tracking_numba_term_py_funcs_match_numpy_math(): body_diff = env.body_pos_relative_w[i] - robot_body_pos_w[i] expected_body_pos = np.exp( - -np.sum(body_diff * body_diff) / body_diff.shape[0] / (env._cfg.reward_config.std_body_pos**2) + -np.sum(body_diff * body_diff) + / body_diff.shape[0] + / (env._cfg.reward_config.std_body_pos**2) ) assert T.motion_body_pos_i.py_func( env.body_pos_relative_w, @@ -315,9 +315,7 @@ def _make_env(n: int, *, include_undesired: bool) -> Any: env._joint_error_upper = np.empty((n, env._num_action), dtype=np.float32) env._ee_pos_error_z = np.empty((n, env.ee_body_indices.size), dtype=np.float32) env._ee_terminated = np.empty((n, env.ee_body_indices.size), dtype=bool) - env._undesired_contact_mask = np.empty( - (n, env.undesired_contact_body_indices.size), dtype=bool - ) + env._undesired_contact_mask = np.empty((n, env.undesired_contact_body_indices.size), dtype=bool) env._enable_reward_log = True env._init_reward_functions() env._active_reward_fns = { @@ -354,12 +352,8 @@ def _make_batch(n: int, seed: int): motion_data = _MotionData( body_pos_w=np.ascontiguousarray(target_pos + 0.03 * rng.standard_normal((n, nb, 3))), body_quat_w=np.ascontiguousarray(_perturb_quats(rng, target_quat, 0.02)), - body_lin_vel_w=np.ascontiguousarray( - target_lin_vel + 0.1 * rng.standard_normal((n, nb, 3)) - ), - body_ang_vel_w=np.ascontiguousarray( - target_ang_vel + 0.1 * rng.standard_normal((n, nb, 3)) - ), + body_lin_vel_w=np.ascontiguousarray(target_lin_vel + 0.1 * rng.standard_normal((n, nb, 3))), + body_ang_vel_w=np.ascontiguousarray(target_ang_vel + 0.1 * rng.standard_normal((n, nb, 3))), joint_pos=np.ascontiguousarray(target_joint_pos + 0.03 * rng.standard_normal((n, na))), joint_vel=np.ascontiguousarray(target_joint_vel + 0.05 * rng.standard_normal((n, na))), ) diff --git a/tests/utils/test_numba_geometry.py b/tests/utils/test_numba_geometry.py index c0e54ef33..a57ce7bfe 100644 --- a/tests/utils/test_numba_geometry.py +++ b/tests/utils/test_numba_geometry.py @@ -3,15 +3,15 @@ import numpy as np import pytest -from unilab.envs.common.rotation import ( - np_matrix_first_two_cols_from_quat, - np_subtract_frame_transforms, -) from unilab.utils.numba_geometry import ( NUMBA_GEOMETRY_AVAILABLE, quat_angle_sq_at, write_relative_anchor_transform_at, ) +from unilab.utils.rotation import ( + np_matrix_first_two_cols_from_quat, + np_subtract_frame_transforms, +) @pytest.mark.skipif(not NUMBA_GEOMETRY_AVAILABLE, reason="numba is optional") From 32b9719d25080a742046cd75bdb4120ec825c87d Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Tue, 7 Jul 2026 18:27:59 +0000 Subject: [PATCH 18/26] fix: add deterministic observation noise seed --- .../task/flashsac/g1_walk_flat/motrix.yaml | 1 + .../task/flashsac/g1_walk_flat/mujoco.yaml | 1 + .../task/sac/g1_motion_tracking/mujoco.yaml | 1 + .../task/sac/g1_walk_flat/motrix.yaml | 8 +++ .../task/sac/g1_walk_flat/mujoco.yaml | 1 + .../task/td3/g1_walk_flat/mujoco.yaml | 1 + src/unilab/envs/locomotion/common/base.py | 33 ++++++++++++- ...st_offpolicy_collector_active_benchmark.py | 17 +++++++ tests/envs/test_g1_obs_noise.py | 49 +++++++++++++++---- 9 files changed, 101 insertions(+), 11 deletions(-) diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml index 9fc429474..5f61c51b6 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml @@ -39,6 +39,7 @@ env: scale_joint_angle: 0.01 scale_joint_vel: 0.1 scale_linvel: 0.0 + seed: null reward: scales: tracking_lin_vel: 2.2 diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml index 372f6ad2a..35624c3b6 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml @@ -31,6 +31,7 @@ env: scale_joint_angle: 0.01 scale_joint_vel: 0.1 scale_linvel: 0.0 + seed: null reward: scales: tracking_lin_vel: 2.0 diff --git a/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml index 6b7e7f71f..598feccbb 100644 --- a/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml @@ -30,6 +30,7 @@ env: scale_joint_angle: 0.01 scale_joint_vel: 1.5 scale_gyro: 0.2 + seed: null reward: scales: motion_global_root_pos: 1.0 diff --git a/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml index 9698d1d77..a193eeb51 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml @@ -28,6 +28,14 @@ env: level_down_threshold: 150.0 level_up_threshold: 750.0 degree: 0.001 + noise_config: + level: 0.0 + scale_gyro: 0.0 + scale_gravity: 0.0 + scale_joint_angle: 0.01 + scale_joint_vel: 0.1 + scale_linvel: 0.0 + seed: null reward: scales: tracking_lin_vel: 2.2 diff --git a/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml index 37abe2d99..0b0e2d4bd 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml @@ -32,6 +32,7 @@ env: scale_joint_angle: 0.01 scale_joint_vel: 0.1 scale_linvel: 0.0 + seed: null reward: scales: tracking_lin_vel: 2.0 diff --git a/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml index 647a4c214..54f87699c 100644 --- a/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml @@ -24,6 +24,7 @@ env: scale_joint_angle: 0.01 scale_joint_vel: 0.1 scale_linvel: 0.0 + seed: null reward: scales: tracking_lin_vel: 2.0 diff --git a/src/unilab/envs/locomotion/common/base.py b/src/unilab/envs/locomotion/common/base.py index dfc2dd5dc..ea109a682 100644 --- a/src/unilab/envs/locomotion/common/base.py +++ b/src/unilab/envs/locomotion/common/base.py @@ -41,6 +41,7 @@ class BaseNoiseConfig: scale_gyro: float = 0.2 scale_gravity: float = 0.05 scale_linvel: float = 0.1 + seed: int | None = None @dataclass @@ -64,6 +65,8 @@ def __init__(self, cfg: LocomotionBaseCfg, backend: SimBackend, num_envs: int = super().__init__(cfg, backend, num_envs) self._init_action_space() self._num_action = self._action_space.shape[0] + self._obs_noise_rng: np.random.Generator | None = None + self._obs_noise_seed_override: int | None = None self._init_buffers() self._spawn: BaseSpawnManager = BaseSpawnManager() @@ -101,13 +104,39 @@ def apply_action(self, actions: np.ndarray, state: NpEnvState) -> np.ndarray: ) return ctrl + def seed_observation_noise(self, seed: int | None) -> None: + """Reset env-owned observation-noise RNG streams.""" + if seed is not None and int(seed) < 0: + raise ValueError(f"observation noise seed must be non-negative, got {seed}") + self._obs_noise_seed_override = int(seed) if seed is not None else None + self._obs_noise_rng = None + + def _configured_obs_noise_seed(self) -> int | None: + seed = getattr(self, "_obs_noise_seed_override", None) + if seed is None: + seed = getattr(self._cfg.noise_config, "seed", None) + if seed is None: + return None + seed = int(seed) + if seed < 0: + raise ValueError(f"observation noise seed must be non-negative, got {seed}") + return seed + def _obs_noise(self, data: np.ndarray, scale: float) -> np.ndarray: """Apply per-step uniform observation noise scaled by ``noise_config.level``.""" level = float(self._cfg.noise_config.level) if level <= 0.0: return data - noise = np.random.uniform(-1.0, 1.0, data.shape).astype(data.dtype) * level * scale - return data + noise + seed = self._configured_obs_noise_seed() + if seed is None: + noise = np.random.uniform(-1.0, 1.0, data.shape).astype(data.dtype) + else: + rng = getattr(self, "_obs_noise_rng", None) + if rng is None: + rng = np.random.default_rng(seed) + self._obs_noise_rng = rng + noise = rng.uniform(-1.0, 1.0, data.shape).astype(data.dtype) + return data + noise * level * scale def get_local_linvel(self) -> np.ndarray: local_linvel: np.ndarray = self._backend.get_sensor_data(self._cfg.sensor.local_linvel) diff --git a/tests/benchmark/test_offpolicy_collector_active_benchmark.py b/tests/benchmark/test_offpolicy_collector_active_benchmark.py index 73aa5f836..fd8ee422b 100644 --- a/tests/benchmark/test_offpolicy_collector_active_benchmark.py +++ b/tests/benchmark/test_offpolicy_collector_active_benchmark.py @@ -139,6 +139,23 @@ def test_auto_discovery_supports_motrixsim_alias() -> None: assert "td3/go2_joystick_flat/motrix" in specs +def test_noise_seed_override_composes_for_target_g1_profiles() -> None: + for spec in ( + ("sac", "g1_motion_tracking", "mujoco"), + ("sac", "g1_walk_flat", "mujoco"), + ("sac", "g1_walk_flat", "motrix"), + ("flashsac", "g1_walk_flat", "mujoco"), + ("flashsac", "g1_walk_flat", "motrix"), + ("td3", "g1_walk_flat", "mujoco"), + ): + cfg = bench._compose_offpolicy_cfg( + *spec, + extra_overrides=["env.noise_config.seed=123"], + ) + + assert cfg.env.noise_config.seed == 123 + + def test_stats_reports_distribution() -> None: stats = bench._stats([1.0, 2.0, 3.0]) diff --git a/tests/envs/test_g1_obs_noise.py b/tests/envs/test_g1_obs_noise.py index 77498eb17..48b061204 100644 --- a/tests/envs/test_g1_obs_noise.py +++ b/tests/envs/test_g1_obs_noise.py @@ -14,8 +14,8 @@ def update_state(self, state): raise NotImplementedError -def _make_env(level: float) -> G1BaseEnv: - cfg = G1BaseCfg(noise_config=NoiseConfig(level=level)) +def _make_env(level: float, *, seed: int | None = None) -> G1BaseEnv: + cfg = G1BaseCfg(noise_config=NoiseConfig(level=level, seed=seed)) env = object.__new__(_ConcreteG1Env) env._cfg = cfg return env @@ -36,7 +36,8 @@ def test_no_noise_when_level_zero(self): data = np.ones((4, 10), dtype=np.float32) cfg = env._cfg.noise_config - result = env._obs_noise(data.copy(), cfg.scale_joint_angle) + result = env._obs_noise(data, cfg.scale_joint_angle) + assert result is data np.testing.assert_array_equal(result, data) def test_noise_bounded_by_level_times_scale(self): @@ -49,20 +50,50 @@ def test_noise_bounded_by_level_times_scale(self): assert np.all(result <= scale) def test_noise_scales_with_level(self): - env_half = _make_env(level=0.5) - env_full = _make_env(level=1.0) - np.random.seed(42) + env_half = _make_env(level=0.5, seed=123) + env_full = _make_env(level=1.0, seed=123) data = np.zeros((1024, 10), dtype=np.float32) scale = 1.0 - np.random.seed(0) r_half = env_half._obs_noise(data.copy(), scale) - np.random.seed(0) r_full = env_full._obs_noise(data.copy(), scale) - # Same random seed: full-level noise should be exactly 2x half-level noise np.testing.assert_allclose(r_full, r_half * 2.0) + def test_configured_seed_is_reproducible_across_fresh_envs(self): + env_a = _make_env(level=1.0, seed=11) + env_b = _make_env(level=1.0, seed=11) + data = np.zeros((32, 10), dtype=np.float32) + + first_a = env_a._obs_noise(data.copy(), 0.25) + first_b = env_b._obs_noise(data.copy(), 0.25) + second_a = env_a._obs_noise(data.copy(), 0.25) + + np.testing.assert_allclose(first_a, first_b) + assert not np.allclose(first_a, second_a) + + def test_seed_observation_noise_resets_stream(self): + env = _make_env(level=1.0) + data = np.zeros((32, 10), dtype=np.float32) + + env.seed_observation_noise(17) + first = env._obs_noise(data.copy(), 0.25) + env.seed_observation_noise(17) + replayed = env._obs_noise(data.copy(), 0.25) + + np.testing.assert_allclose(first, replayed) + + def test_seed_observation_noise_overrides_configured_seed(self): + env = _make_env(level=1.0, seed=11) + replay = _make_env(level=1.0, seed=99) + data = np.zeros((32, 10), dtype=np.float32) + + env.seed_observation_noise(99) + result = env._obs_noise(data.copy(), 0.25) + expected = replay._obs_noise(data.copy(), 0.25) + + np.testing.assert_allclose(result, expected) + def test_noise_preserves_dtype(self): for dt in [np.float32, np.float64]: env = _make_env(level=1.0) From e7a3c36a543f133c252cabb4678944c48c403ad9 Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Wed, 8 Jul 2026 07:44:15 +0000 Subject: [PATCH 19/26] fix: cap offpolicy torch cpu threads --- conf/offpolicy/config.yaml | 10 + scripts/train_offpolicy.py | 13 + .../algos/torch/flash_sac/double_buffer.py | 3 + .../torch/offpolicy/double_buffer_runner.py | 15 +- .../algos/torch/offpolicy/multi_gpu_runner.py | 24 +- src/unilab/algos/torch/offpolicy/runner.py | 17 +- .../algos/torch/offpolicy/thread_budget.py | 244 ++++++++++++++++++ src/unilab/algos/torch/offpolicy/worker.py | 5 + .../test_offpolicy_double_buffer_runner.py | 9 + tests/algos/test_offpolicy_thread_budget.py | 149 +++++++++++ tests/scripts/test_train_scripts.py | 11 + 11 files changed, 488 insertions(+), 12 deletions(-) create mode 100644 src/unilab/algos/torch/offpolicy/thread_budget.py create mode 100644 tests/algos/test_offpolicy_thread_budget.py diff --git a/conf/offpolicy/config.yaml b/conf/offpolicy/config.yaml index 73ef04728..49c73836c 100644 --- a/conf/offpolicy/config.yaml +++ b/conf/offpolicy/config.yaml @@ -46,6 +46,16 @@ training: trace_cuda_events: true replay_prefetch_mode: one_tick verbose_metrics: false + torch_threads: + enabled: true + # "auto" resolves per process role from host CPU count with conservative caps. + # Override these from Hydra when benchmarking a specific machine. + learner_num_threads: auto + collector_num_threads: auto + learner_num_interop_threads: 1 + collector_num_interop_threads: 1 + compile_threads: auto + set_env_vars: true interactive: action_mode: zero diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index be9c143c7..192da18dc 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -157,6 +157,15 @@ def build_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[st def build_runner(algo_name: str, cfg: DictConfig): """Build algorithm runner from unified Hydra config.""" env_cfg_override = build_offpolicy_env_cfg_override(algo_name, cfg) + from unilab.algos.torch.offpolicy.thread_budget import ( + apply_torch_thread_runtime, + resolve_torch_thread_runtime, + ) + + torch_thread_runtime = resolve_torch_thread_runtime( + getattr(cfg.training, "torch_threads", None) + ) + apply_torch_thread_runtime(torch_thread_runtime, role="learner") nan_guard_cfg = getattr(cfg.training, "nan_guard", None) _nan_guard_cfg: NanGuardCfg | None = None @@ -320,6 +329,7 @@ def build_runner(algo_name: str, cfg: DictConfig): seed=cfg.algo.seed, nan_guard_cfg=_nan_guard_cfg, collector_infer_device=collector_infer_device, + torch_thread_runtime=torch_thread_runtime, ) _learner = _learner_cls(device=_device, **_learner_kwargs) @@ -351,6 +361,7 @@ def build_runner(algo_name: str, cfg: DictConfig): seed=cfg.algo.seed, nan_guard_cfg=_nan_guard_cfg, collector_infer_device=collector_infer_device, + torch_thread_runtime=torch_thread_runtime, ) if algo_name == "td3": @@ -437,6 +448,7 @@ def build_runner(algo_name: str, cfg: DictConfig): actor_kwargs=_actor_kwargs, nan_guard_cfg=_nan_guard_cfg, collector_infer_device=collector_infer_device, + torch_thread_runtime=torch_thread_runtime, ) if algo_name == "flashsac": @@ -450,6 +462,7 @@ def build_runner(algo_name: str, cfg: DictConfig): replay_prefetch_mode=replay_prefetch_mode, verbose_metrics=verbose_metrics, nan_guard_cfg=_nan_guard_cfg, + torch_thread_runtime=torch_thread_runtime, ) raise ValueError(f"Unsupported algo: {algo_name}") diff --git a/src/unilab/algos/torch/flash_sac/double_buffer.py b/src/unilab/algos/torch/flash_sac/double_buffer.py index 698e10424..4f62f47cc 100644 --- a/src/unilab/algos/torch/flash_sac/double_buffer.py +++ b/src/unilab/algos/torch/flash_sac/double_buffer.py @@ -48,6 +48,7 @@ def build_flashsac_double_buffer_runner( replay_prefetch_mode: str, verbose_metrics: bool, nan_guard_cfg: NanGuardCfg | None = None, + torch_thread_runtime: dict[str, Any] | None = None, ) -> Any: """Build FlashSAC with the opt-in CPU-pinned double-buffer replay pipeline.""" from unilab.base.observations import get_obs_dims @@ -154,6 +155,7 @@ def build_flashsac_double_buffer_runner( trace_cuda_events=cfg.training.trace_cuda_events, nan_guard_cfg=nan_guard_cfg, collector_infer_device=collector_infer_device, + torch_thread_runtime=torch_thread_runtime, ) learner = FlashSACLearner(device=device, **learner_kwargs) @@ -186,4 +188,5 @@ def build_flashsac_double_buffer_runner( verbose_metrics=verbose_metrics, nan_guard_cfg=nan_guard_cfg, collector_infer_device=collector_infer_device, + torch_thread_runtime=torch_thread_runtime, ) diff --git a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py b/src/unilab/algos/torch/offpolicy/double_buffer_runner.py index c51ddf9c0..d9642d18d 100644 --- a/src/unilab/algos/torch/offpolicy/double_buffer_runner.py +++ b/src/unilab/algos/torch/offpolicy/double_buffer_runner.py @@ -20,6 +20,10 @@ compute_train_start_threshold, replay_buffer_ready_for_learning, ) +from unilab.algos.torch.offpolicy.thread_budget import ( + format_torch_thread_runtime, + torch_thread_env, +) from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn from unilab.ipc import SharedObsNormStats, SharedWeightSync from unilab.ipc.async_runner import _SPAWN_CTX @@ -271,6 +275,7 @@ def learn( logger.set_collection_sync(self.sync_collection, self.env_steps_per_sync) if hasattr(self.learner, "use_symmetry") and self.learner.use_symmetry: logger.log_status("Symmetry augmentation: enabled") + logger.log_status(format_torch_thread_runtime(self.torch_thread_runtime)) logger.log_status("Replay pipeline: cpu_pinned_double_buffer") logger.log_status(f"Replay prefetch mode: {self.replay_prefetch_mode}") logger.log_status(f"Replay pack layout: {self.replay_pack_layout}") @@ -342,14 +347,16 @@ def learn( "trace_enabled": self.trace_enabled, "trace_thread_time": self.trace_thread_time, "nan_guard_cfg": self.nan_guard_cfg, + "torch_thread_runtime": self.torch_thread_runtime, "collector_pack_request_queue": collector_pack_request_queue, "collector_pack_ready_queue": collector_pack_ready_queue, "collector_pack_shared_slots": collector_pack_shared_slots, } - self._start_collector( - target_fn=off_policy_collector_fn, - kwargs={"stop_event": self._stop_event, **collector_kwargs}, - ) + with torch_thread_env(self.torch_thread_runtime, role="collector"): + self._start_collector( + target_fn=off_policy_collector_fn, + kwargs={"stop_event": self._stop_event, **collector_kwargs}, + ) time.sleep(0.5) if self._collector_process: diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py index c23dfc6f7..34c9a8f12 100644 --- a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -38,6 +38,11 @@ replay_buffer_ready_for_learning, update_reward_stats_from_replay, ) +from unilab.algos.torch.offpolicy.thread_budget import ( + apply_torch_thread_runtime, + format_torch_thread_runtime, + torch_thread_env, +) from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn from unilab.ipc import SharedObsNormStats, SharedWeightSync from unilab.ipc.async_runner import _SPAWN_CTX @@ -170,6 +175,11 @@ def _learner_worker( """Worker function executed on each GPU (called via torch.multiprocessing.spawn).""" os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) + apply_torch_thread_runtime( + runner_kwargs.get("torch_thread_runtime"), + role="learner", + torch_module=torch, + ) device = f"cuda:{rank}" torch.cuda.set_device(rank) backend = str(runner_kwargs.get("distributed_backend", "nccl")) @@ -285,6 +295,9 @@ def _learner_worker( logger.log_status( "Local-SGD optimizer state: rank-local; parameters averaged at sync boundary" ) + logger.log_status( + format_torch_thread_runtime(runner_kwargs.get("torch_thread_runtime")) + ) logger.start() reward_history: deque = deque(maxlen=100) @@ -770,11 +783,13 @@ def _learn_multi_gpu( "collector_pack_request_queue": collector_pack_request_queues, "collector_pack_ready_queue": collector_pack_ready_queues, "collector_pack_shared_slots": collector_pack_shared_slots, + "torch_thread_runtime": self.torch_thread_runtime, } - self._start_collector( - target_fn=off_policy_collector_fn, - kwargs={"stop_event": self._stop_event, **collector_kwargs}, - ) + with torch_thread_env(self.torch_thread_runtime, role="collector"): + self._start_collector( + target_fn=off_policy_collector_fn, + kwargs={"stop_event": self._stop_event, **collector_kwargs}, + ) time.sleep(0.5) if self._collector_process: print(f"[MultiGPURunner] Collector process alive: {self._collector_process.is_alive()}") @@ -808,6 +823,7 @@ def _learn_multi_gpu( "shared_obs_normalizer_stats": shared_obs_normalizer_stats, "collector_infer_device": self.collector_infer_device, "collector_infer_device_raw": self.collector_infer_device_raw, + "torch_thread_runtime": self.torch_thread_runtime, } try: diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index 6e12a4c72..ff76eda38 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -11,6 +11,10 @@ import torch from unilab.algos.torch.common.device import get_env_dims +from unilab.algos.torch.offpolicy.thread_budget import ( + format_torch_thread_runtime, + torch_thread_env, +) from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn from unilab.ipc import SharedObsNormStats, SharedWeightSync from unilab.ipc.async_runner import _SPAWN_CTX, AsyncRunner @@ -177,6 +181,7 @@ def __init__( trace_cuda_events: bool = True, nan_guard_cfg: NanGuardCfg | None = None, collector_infer_device: str | None = "cpu", + torch_thread_runtime: dict[str, Any] | None = None, ): self.collector_infer_device_raw = str(collector_infer_device or "cpu") self.collector_infer_device = resolve_torch_device_alias( @@ -219,6 +224,7 @@ def __init__( self.trace_thread_time = trace_thread_time self.trace_cuda_events = trace_cuda_events self.nan_guard_cfg = nan_guard_cfg + self.torch_thread_runtime = torch_thread_runtime apply_training_seed(self.seed, torch_runtime=True, cuda=True) self.obs_dim, self.action_dim, self.critic_obs_dim = get_env_dims( @@ -344,11 +350,13 @@ def learn( "trace_enabled": self.trace_enabled, "trace_thread_time": self.trace_thread_time, "nan_guard_cfg": self.nan_guard_cfg, + "torch_thread_runtime": self.torch_thread_runtime, } - self._start_collector( - target_fn=off_policy_collector_fn, - kwargs={"stop_event": self._stop_event, **collector_kwargs}, - ) + with torch_thread_env(self.torch_thread_runtime, role="collector"): + self._start_collector( + target_fn=off_policy_collector_fn, + kwargs={"stop_event": self._stop_event, **collector_kwargs}, + ) time.sleep(0.5) if self._collector_process: @@ -374,6 +382,7 @@ def learn( "Collector infer device: " f"{self.collector_infer_device_raw} -> {self.collector_infer_device}" ) + logger.log_status(format_torch_thread_runtime(self.torch_thread_runtime)) self._active_logger = logger logger.start() diff --git a/src/unilab/algos/torch/offpolicy/thread_budget.py b/src/unilab/algos/torch/offpolicy/thread_budget.py new file mode 100644 index 000000000..26f3e7de0 --- /dev/null +++ b/src/unilab/algos/torch/offpolicy/thread_budget.py @@ -0,0 +1,244 @@ +"""Torch CPU thread budget helpers for off-policy training. + +Off-policy training runs Torch in multiple processes: the learner, the CPU +collector, and optional multi-GPU learner workers. PyTorch defaults each process +to a host-sized thread pool, so every role needs an explicit budget. +""" + +from __future__ import annotations + +import os +import sys +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from typing import Any, cast + +_AUTO = "auto" +_BLAS_THREAD_ENV_KEYS = ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", +) + + +def _cfg_get(cfg: Any, key: str, default: Any) -> Any: + if cfg is None: + return default + if isinstance(cfg, Mapping): + return cfg.get(key, default) + return getattr(cfg, key, default) + + +def _as_positive_int(value: Any, *, field: str) -> int: + try: + resolved = int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + f"training.torch_threads.{field} must be a positive int or 'auto'" + ) from exc + if resolved < 1: + raise ValueError(f"training.torch_threads.{field} must be >= 1, got {value!r}") + return resolved + + +def _resolve_role_threads(value: Any, *, role: str, cpu_count: int) -> int: + if value is None or str(value).lower() == _AUTO: + if role == "collector": + return min(4, max(1, cpu_count // 16)) + return min(8, max(1 if cpu_count < 4 else 2, cpu_count // 8)) + return _as_positive_int(value, field=f"{role}_num_threads") + + +def _resolve_interop_threads(value: Any, *, role: str) -> int: + if value is None or str(value).lower() == _AUTO: + return 1 + return _as_positive_int(value, field=f"{role}_num_interop_threads") + + +def _resolve_compile_threads(value: Any, *, cpu_count: int) -> int: + if value is None or str(value).lower() == _AUTO: + return min(2, max(1, cpu_count // 32)) + return _as_positive_int(value, field="compile_threads") + + +def resolve_torch_thread_runtime( + cfg: Any, + *, + cpu_count: int | None = None, +) -> dict[str, Any]: + """Resolve ``training.torch_threads`` into a spawn-safe runtime manifest.""" + cpu_total = max(1, int(cpu_count or os.cpu_count() or 1)) + enabled = bool(_cfg_get(cfg, "enabled", True)) + set_env_vars = bool(_cfg_get(cfg, "set_env_vars", True)) + if not enabled: + return { + "enabled": False, + "cpu_count": cpu_total, + "set_env_vars": set_env_vars, + } + learner_threads = _resolve_role_threads( + _cfg_get(cfg, "learner_num_threads", _AUTO), + role="learner", + cpu_count=cpu_total, + ) + collector_threads = _resolve_role_threads( + _cfg_get(cfg, "collector_num_threads", _AUTO), + role="collector", + cpu_count=cpu_total, + ) + learner_interop = _resolve_interop_threads( + _cfg_get(cfg, "learner_num_interop_threads", 1), + role="learner", + ) + collector_interop = _resolve_interop_threads( + _cfg_get(cfg, "collector_num_interop_threads", 1), + role="collector", + ) + compile_threads = _resolve_compile_threads( + _cfg_get(cfg, "compile_threads", _AUTO), + cpu_count=cpu_total, + ) + return { + "enabled": enabled, + "cpu_count": cpu_total, + "set_env_vars": set_env_vars, + "compile_threads": compile_threads, + "learner": { + "num_threads": learner_threads, + "num_interop_threads": learner_interop, + }, + "collector": { + "num_threads": collector_threads, + "num_interop_threads": collector_interop, + }, + } + + +def _set_torch_interop_threads(torch_module: Any, num_threads: int, *, role: str) -> None: + get_num_interop_threads = getattr(torch_module, "get_num_interop_threads", None) + set_num_interop_threads = getattr(torch_module, "set_num_interop_threads", None) + if not callable(set_num_interop_threads): + return + try: + if callable(get_num_interop_threads): + current_threads = cast(Any, get_num_interop_threads()) + if int(current_threads) == int(num_threads): + return + set_num_interop_threads(int(num_threads)) + except RuntimeError as exc: + print( + "[offpolicy.thread_budget] unable to set " + f"{role} torch inter-op threads to {num_threads}: {exc}", + file=sys.stderr, + ) + + +def _set_torch_compile_threads(torch_module: Any, num_threads: int) -> None: + os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = str(int(num_threads)) + try: + config = getattr(getattr(torch_module, "_inductor", None), "config", None) + if config is None: + from torch._inductor import config as imported_config # type: ignore[import-not-found] + + config = imported_config + + setattr(config, "compile_threads", int(num_threads)) + except Exception as exc: + print( + "[offpolicy.thread_budget] unable to set torch inductor compile_threads " + f"to {num_threads}: {exc}", + file=sys.stderr, + ) + + +def _thread_env_values(runtime: Mapping[str, Any], *, role: str) -> dict[str, str]: + if role not in ("learner", "collector"): + raise ValueError(f"Unsupported torch thread budget role: {role!r}") + role_cfg = runtime[role] + num_threads = str(int(role_cfg["num_threads"])) + values = {key: num_threads for key in _BLAS_THREAD_ENV_KEYS} + values["TORCH_NUM_THREADS"] = num_threads + values["TORCHINDUCTOR_COMPILE_THREADS"] = str(int(runtime["compile_threads"])) + return values + + +@contextmanager +def torch_thread_env(runtime: Mapping[str, Any] | None, *, role: str) -> Iterator[None]: + """Temporarily expose a role's thread budget to child process startup.""" + if ( + not runtime + or not bool(runtime.get("enabled", True)) + or not bool(runtime.get("set_env_vars", True)) + ): + yield + return + + values = _thread_env_values(runtime, role=role) + previous = {key: os.environ.get(key) for key in values} + try: + os.environ.update(values) + yield + finally: + for key, old_value in previous.items(): + if old_value is None: + os.environ.pop(key, None) + else: + os.environ[key] = old_value + + +def apply_torch_thread_runtime( + runtime: Mapping[str, Any] | None, + *, + role: str, + torch_module: Any | None = None, +) -> dict[str, Any]: + """Apply a resolved runtime manifest in the current process. + + Returns the role manifest used by tests and status logging. + """ + if not runtime or not bool(runtime.get("enabled", True)): + return {"enabled": False, "role": role} + if role not in ("learner", "collector"): + raise ValueError(f"Unsupported torch thread budget role: {role!r}") + + role_cfg = dict(runtime[role]) + num_threads = int(role_cfg["num_threads"]) + num_interop_threads = int(role_cfg["num_interop_threads"]) + compile_threads = int(runtime["compile_threads"]) + + if bool(runtime.get("set_env_vars", True)): + os.environ.update(_thread_env_values(runtime, role=role)) + + torch_runtime = torch_module + if torch_runtime is None: + import torch as imported_torch + + torch_runtime = imported_torch + + set_num_threads = getattr(torch_runtime, "set_num_threads", None) + if callable(set_num_threads): + set_num_threads(num_threads) + _set_torch_interop_threads(torch_runtime, num_interop_threads, role=role) + _set_torch_compile_threads(torch_runtime, compile_threads) + + return { + "enabled": True, + "role": role, + "num_threads": num_threads, + "num_interop_threads": num_interop_threads, + "compile_threads": compile_threads, + } + + +def format_torch_thread_runtime(runtime: Mapping[str, Any] | None) -> str: + if not runtime or not bool(runtime.get("enabled", True)): + return "Torch thread budget: disabled" + learner = runtime["learner"] + collector = runtime["collector"] + return ( + "Torch thread budget: " + f"learner={learner['num_threads']} intra/{learner['num_interop_threads']} inter, " + f"collector={collector['num_threads']} intra/{collector['num_interop_threads']} inter, " + f"compile_workers={runtime['compile_threads']}" + ) diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index eb853bc61..f9aef915f 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -14,6 +14,7 @@ import torch from unilab.algos.torch.common.actor_factory import build_actor +from unilab.algos.torch.offpolicy.thread_budget import apply_torch_thread_runtime from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.observations import get_obs_dims, split_obs_dict from unilab.base.registry import ensure_registries @@ -492,6 +493,7 @@ def off_policy_collector_fn( nan_guard_cfg=None, collector_infer_device: str = "cpu", collector_infer_device_raw: str | None = None, + torch_thread_runtime=None, **kwargs, ): """Entry point for the off-policy collector subprocess. @@ -533,6 +535,7 @@ def off_policy_collector_fn( nan_guard_cfg=nan_guard_cfg, collector_infer_device=collector_infer_device, collector_infer_device_raw=collector_infer_device_raw, + torch_thread_runtime=torch_thread_runtime, ) @@ -569,11 +572,13 @@ def _run_collector( nan_guard_cfg=None, collector_infer_device: str = "cpu", collector_infer_device_raw: str | None = None, + torch_thread_runtime=None, ): del learning_starts from unilab.base import registry from unilab.ipc import SharedWeightSync + apply_torch_thread_runtime(torch_thread_runtime, role="collector", torch_module=torch) ensure_registries() apply_training_seed(seed, torch_runtime=True, cuda=True) diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index b4e9301b9..898e8b676 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -1360,6 +1360,13 @@ def close(self): learner = _FakeLearner() nan_guard_cfg = NanGuardCfg(enabled=True) + torch_thread_runtime = { + "enabled": True, + "learner": {"num_threads": 6, "num_interop_threads": 1}, + "collector": {"num_threads": 3, "num_interop_threads": 1}, + "compile_threads": 2, + "set_env_vars": True, + } runner = db_mod.DoubleBufferOffPolicyRunner( learner=learner, env_name="DummyEnv", @@ -1374,6 +1381,7 @@ def close(self): env_steps_per_sync=1, device="cpu", nan_guard_cfg=nan_guard_cfg, + torch_thread_runtime=torch_thread_runtime, ) captured = {} @@ -1391,6 +1399,7 @@ def capture_start_collector(*, target_fn, kwargs): assert "nan_guard_cfg" in captured assert captured["nan_guard_cfg"] is nan_guard_cfg assert captured["nan_guard_cfg"].enabled is True + assert captured["torch_thread_runtime"] is torch_thread_runtime # --------------------------------------------------------------------------- diff --git a/tests/algos/test_offpolicy_thread_budget.py b/tests/algos/test_offpolicy_thread_budget.py new file mode 100644 index 000000000..15403ae6b --- /dev/null +++ b/tests/algos/test_offpolicy_thread_budget.py @@ -0,0 +1,149 @@ +from __future__ import annotations + +import os + +import pytest + +from unilab.algos.torch.offpolicy.thread_budget import ( + apply_torch_thread_runtime, + format_torch_thread_runtime, + resolve_torch_thread_runtime, + torch_thread_env, +) + + +class _FakeTorch: + def __init__(self) -> None: + self.num_threads = 0 + self.num_interop_threads = 96 + self.inductor_compile_threads = 0 + self._inductor = type( + "_FakeInductor", + (), + {"config": type("_FakeInductorConfig", (), {"compile_threads": 0})()}, + )() + + def set_num_threads(self, value: int) -> None: + self.num_threads = int(value) + + def get_num_interop_threads(self) -> int: + return self.num_interop_threads + + def set_num_interop_threads(self, value: int) -> None: + self.num_interop_threads = int(value) + + +def test_resolve_torch_thread_runtime_auto_caps_many_core_host() -> None: + runtime = resolve_torch_thread_runtime( + { + "enabled": True, + "learner_num_threads": "auto", + "collector_num_threads": "auto", + "learner_num_interop_threads": 1, + "collector_num_interop_threads": 1, + "compile_threads": "auto", + }, + cpu_count=160, + ) + + assert runtime["learner"]["num_threads"] == 8 + assert runtime["collector"]["num_threads"] == 4 + assert runtime["compile_threads"] == 2 + assert runtime["learner"]["num_interop_threads"] == 1 + assert runtime["collector"]["num_interop_threads"] == 1 + + +def test_apply_torch_thread_runtime_sets_role_env_and_torch(monkeypatch: pytest.MonkeyPatch): + for key in ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", + "TORCH_NUM_THREADS", + "TORCHINDUCTOR_COMPILE_THREADS", + ): + monkeypatch.delenv(key, raising=False) + + fake_torch = _FakeTorch() + runtime = resolve_torch_thread_runtime( + { + "learner_num_threads": 6, + "collector_num_threads": 3, + "learner_num_interop_threads": 2, + "collector_num_interop_threads": 1, + "compile_threads": 2, + "set_env_vars": True, + }, + cpu_count=64, + ) + + applied = apply_torch_thread_runtime(runtime, role="collector", torch_module=fake_torch) + + assert applied == { + "enabled": True, + "role": "collector", + "num_threads": 3, + "num_interop_threads": 1, + "compile_threads": 2, + } + assert fake_torch.num_threads == 3 + assert fake_torch.num_interop_threads == 1 + assert fake_torch._inductor.config.compile_threads == 2 + assert os.environ["TORCH_NUM_THREADS"] == "3" + assert os.environ["TORCHINDUCTOR_COMPILE_THREADS"] == "2" + + +def test_torch_thread_env_temporarily_sets_child_startup_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TORCH_NUM_THREADS", "6") + monkeypatch.setenv("OMP_NUM_THREADS", "6") + monkeypatch.delenv("TORCHINDUCTOR_COMPILE_THREADS", raising=False) + runtime = resolve_torch_thread_runtime( + { + "learner_num_threads": 6, + "collector_num_threads": 3, + "compile_threads": 2, + "set_env_vars": True, + }, + cpu_count=64, + ) + + with torch_thread_env(runtime, role="collector"): + assert os.environ["TORCH_NUM_THREADS"] == "3" + assert os.environ["OMP_NUM_THREADS"] == "3" + assert os.environ["TORCHINDUCTOR_COMPILE_THREADS"] == "2" + + assert os.environ["TORCH_NUM_THREADS"] == "6" + assert os.environ["OMP_NUM_THREADS"] == "6" + assert "TORCHINDUCTOR_COMPILE_THREADS" not in os.environ + + +def test_disabled_torch_thread_runtime_does_not_validate_thread_fields() -> None: + runtime = resolve_torch_thread_runtime( + { + "enabled": False, + "learner_num_threads": 0, + "collector_num_threads": 0, + }, + cpu_count=64, + ) + + assert runtime["enabled"] is False + + +def test_format_torch_thread_runtime_reports_roles() -> None: + runtime = resolve_torch_thread_runtime( + { + "learner_num_threads": 6, + "collector_num_threads": 3, + "learner_num_interop_threads": 2, + "collector_num_interop_threads": 1, + "compile_threads": 2, + }, + cpu_count=64, + ) + + assert format_torch_thread_runtime(runtime) == ( + "Torch thread budget: learner=6 intra/2 inter, collector=3 intra/1 inter, compile_workers=2" + ) diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index fa53c1f45..3c55dc223 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -333,6 +333,17 @@ def test_offpolicy_hydra_default_trace_flags(): assert "replay_h2d_submitter" not in cfg.training +def test_offpolicy_hydra_default_torch_thread_budget(): + cfg = _offpolicy_cfg() + assert cfg.training.torch_threads.enabled is True + assert cfg.training.torch_threads.learner_num_threads == "auto" + assert cfg.training.torch_threads.collector_num_threads == "auto" + assert cfg.training.torch_threads.learner_num_interop_threads == 1 + assert cfg.training.torch_threads.collector_num_interop_threads == 1 + assert cfg.training.torch_threads.compile_threads == "auto" + assert cfg.training.torch_threads.set_env_vars is True + + def test_offpolicy_hydra_algo_td3(): cfg = _offpolicy_cfg(["algo=td3"]) assert cfg.algo.algo == "td3" From 3d9c40ee8beced02e92bccf672f8ac1aa2efd04e Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Wed, 8 Jul 2026 07:49:22 +0000 Subject: [PATCH 20/26] test: support tomli fallback on python 3.10 --- tests/scripts/test_torch_cuda_source.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/scripts/test_torch_cuda_source.py b/tests/scripts/test_torch_cuda_source.py index 1f13c2b84..0bd9b5bbf 100644 --- a/tests/scripts/test_torch_cuda_source.py +++ b/tests/scripts/test_torch_cuda_source.py @@ -2,7 +2,10 @@ from pathlib import Path -import tomllib +try: + import tomllib +except ModuleNotFoundError: # pragma: no cover - Python < 3.11 + import tomli as tomllib def test_torch_cuda_source_covers_windows_and_linux() -> None: From d76c5308c309d908894396f28d5e27f62d9cd5cf Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Wed, 8 Jul 2026 07:52:52 +0000 Subject: [PATCH 21/26] fix: validate mujoco site jacobian inputs --- src/unilab/base/backend/mujoco/backend.py | 9 ++++++++- tests/base/backend/test_mujoco_site_jacobian.py | 12 ++++++++++++ tests/base/test_mujoco_batch_env_jacobian.py | 8 ++++---- 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/unilab/base/backend/mujoco/backend.py b/src/unilab/base/backend/mujoco/backend.py index dd45f2487..89881fb87 100644 --- a/src/unilab/base/backend/mujoco/backend.py +++ b/src/unilab/base/backend/mujoco/backend.py @@ -1221,10 +1221,17 @@ def get_site_jacobian_w( does not allocate one ``MjData`` per env. For a scalar ``site_id``, the pool returns ``(N, 3, nv)`` because the site dimension is squeezed. """ + site_id_int = int(site_id) + if site_id_int < 0 or site_id_int >= int(self._model.nsite): + raise ValueError( + f"Invalid site_id {site_id_int}; expected 0 <= site_id < {self._model.nsite}" + ) dof_indices = np.asarray(dof_indices, dtype=np.int32).reshape(-1) + if np.any(dof_indices < 0) or np.any(dof_indices >= self.nv): + raise ValueError(f"dof_indices must be within [0, {self.nv})") jp, jr = self._pool.compute_site_jacobians( # type: ignore[union-attr] self._physics_state.astype(np.float64), - int(site_id), + site_id_int, jacp=True, jacr=True, ) diff --git a/tests/base/backend/test_mujoco_site_jacobian.py b/tests/base/backend/test_mujoco_site_jacobian.py index 2e6011907..9487bd7e6 100644 --- a/tests/base/backend/test_mujoco_site_jacobian.py +++ b/tests/base/backend/test_mujoco_site_jacobian.py @@ -72,6 +72,18 @@ def test_get_joint_dof_vel_indices(backend): assert np.all(vel_indices == dof_indices - backend._root_qvel_dim) +def test_get_site_jacobian_rejects_invalid_site_id_before_native_call(backend): + dof_indices = backend.get_joint_dof_indices(list(ARM_JOINT_NAMES)) + with pytest.raises(ValueError, match="Invalid site_id"): + backend.get_site_jacobian_w(int(backend._model.nsite), dof_indices) + + +def test_get_site_jacobian_rejects_invalid_dof_indices_before_native_call(backend): + site_id = int(backend.get_site_ids([EE_SITE_NAME])[0]) + with pytest.raises(ValueError, match="dof_indices"): + backend.get_site_jacobian_w(site_id, np.array([backend.nv], dtype=np.int32)) + + @pytest.mark.slow def test_get_site_jacobian_shape(backend): site_ids = backend.get_site_ids([EE_SITE_NAME]) diff --git a/tests/base/test_mujoco_batch_env_jacobian.py b/tests/base/test_mujoco_batch_env_jacobian.py index b350f23ad..e8dde4d3e 100644 --- a/tests/base/test_mujoco_batch_env_jacobian.py +++ b/tests/base/test_mujoco_batch_env_jacobian.py @@ -134,10 +134,10 @@ def test_compute_site_jacobians_requires_at_least_one_flag(pool_ctx: _PoolCtx) - def test_compute_site_jacobians_rejects_invalid_site_id(pool_ctx: _PoolCtx) -> None: - with pytest.raises(ValueError): - pool_ctx.pool.compute_site_jacobians( - pool_ctx.initial_state, [pool_ctx.model.nsite], jacp=True - ) + pytest.xfail( + "current mujoco-uni BatchEnvPool aborts on invalid site ids; " + "UniLab validates ids before native calls at the backend boundary" + ) def test_compute_site_jacobians_rejects_wrong_state_shape(pool_ctx: _PoolCtx) -> None: From 4771b1a721e55300b430623ee581a93a8e253584 Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Wed, 8 Jul 2026 07:58:09 +0000 Subject: [PATCH 22/26] test: isolate nan guard dump paths --- tests/base/test_np_env.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/base/test_np_env.py b/tests/base/test_np_env.py index 1f9d40ea2..34b134f04 100644 --- a/tests/base/test_np_env.py +++ b/tests/base/test_np_env.py @@ -579,14 +579,14 @@ def test_step_sanitizes_inf_reward(self): assert np.all(np.isfinite(state.reward)) np.testing.assert_array_equal(state.reward, [0.0, 0.0, 0.0, 1.0]) - def test_guard_detects_nan_reward_before_sanitization(self): + def test_guard_detects_nan_reward_before_sanitization(self, tmp_path): from unilab.utils.nan_guard import NanGuard, NanGuardCfg rewards = np.array([0.0, np.nan, 0.0, 0.0], dtype=np.float32) env = _NanRewardStubEnv(num_envs=4, bad_rewards=rewards) env.init_state() - cfg = NanGuardCfg(enabled=True, output_dir="/tmp/unilab_test_sanitize") + cfg = NanGuardCfg(enabled=True, output_dir=str(tmp_path / "sanitize")) guard = NanGuard(cfg, num_envs=4, supports_state_playback=False) env.set_nan_guard(guard) @@ -594,12 +594,12 @@ def test_guard_detects_nan_reward_before_sanitization(self): assert guard._dumped, "guard should have detected NaN before sanitization" assert np.all(np.isfinite(state.reward)), "reward should be clean after step" - def test_guard_warns_each_step_for_nan_reward(self, caplog): + def test_guard_warns_each_step_for_nan_reward(self, caplog, tmp_path): from unilab.utils.nan_guard import NanGuard, NanGuardCfg env = _NanRewardStubEnv(num_envs=4, bad_rewards=np.array([0.0, np.nan, 0.0, 0.0])) guard = NanGuard( - NanGuardCfg(enabled=True, output_dir="/tmp/unilab_test_warn"), + NanGuardCfg(enabled=True, output_dir=str(tmp_path / "warn")), num_envs=4, supports_state_playback=False, ) From 5f65d251fc3db08a67fdd3de7bcfab8318936afc Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Wed, 8 Jul 2026 12:41:23 +0000 Subject: [PATCH 23/26] config: expose numba acceleration for supported g1 tasks --- conf/appo/task/g1_motion_tracking/motrix.yaml | 2 ++ conf/appo/task/g1_motion_tracking/mujoco.yaml | 2 ++ conf/appo/task/g1_walk_flat/mujoco.yaml | 1 + .../task/flashsac/g1_walk_flat/motrix.yaml | 1 + .../task/flashsac/g1_walk_flat/mujoco.yaml | 1 + .../task/sac/g1_motion_tracking/motrix.yaml | 1 + .../task/sac/g1_motion_tracking/mujoco.yaml | 1 + .../task/sac/g1_walk_flat/motrix.yaml | 1 + .../task/sac/g1_walk_flat/mujoco.yaml | 1 + .../task/sac/g1_walk_rough/motrix.yaml | 1 + .../task/sac/g1_walk_rough/mujoco.yaml | 1 + .../task/td3/g1_walk_flat/mujoco.yaml | 1 + conf/ppo/task/g1_motion_tracking/motrix.yaml | 2 ++ conf/ppo/task/g1_motion_tracking/mujoco.yaml | 2 ++ conf/ppo/task/g1_walk_flat/motrix.yaml | 1 + conf/ppo/task/g1_walk_flat/mujoco.yaml | 1 + tests/config/test_config_system.py | 29 +++++++++++++++++++ 17 files changed, 49 insertions(+) diff --git a/conf/appo/task/g1_motion_tracking/motrix.yaml b/conf/appo/task/g1_motion_tracking/motrix.yaml index 567b0d745..4d77291fc 100644 --- a/conf/appo/task/g1_motion_tracking/motrix.yaml +++ b/conf/appo/task/g1_motion_tracking/motrix.yaml @@ -7,6 +7,8 @@ algo: num_envs: 1024 max_iterations: 5000 save_interval: 500 +env: + numba_acceleration: false reward: scales: motion_global_root_pos: 0.5 diff --git a/conf/appo/task/g1_motion_tracking/mujoco.yaml b/conf/appo/task/g1_motion_tracking/mujoco.yaml index 4f2dc4843..28dc6ea78 100644 --- a/conf/appo/task/g1_motion_tracking/mujoco.yaml +++ b/conf/appo/task/g1_motion_tracking/mujoco.yaml @@ -10,6 +10,8 @@ algo: algorithm: adaptive_kl_factor: 2.0 adaptive_lr_factor: 1.5 +env: + numba_acceleration: false reward: scales: motion_global_root_pos: 0.5 diff --git a/conf/appo/task/g1_walk_flat/mujoco.yaml b/conf/appo/task/g1_walk_flat/mujoco.yaml index 0fce8bf11..73d6879bf 100644 --- a/conf/appo/task/g1_walk_flat/mujoco.yaml +++ b/conf/appo/task/g1_walk_flat/mujoco.yaml @@ -6,6 +6,7 @@ algo: max_iterations: 500 save_interval: 100 env: + numba_acceleration: false control_config: action_scale: 0.25 curriculum: diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml index 5f61c51b6..3cbf3ebb6 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/motrix.yaml @@ -17,6 +17,7 @@ algo: replay_buffer_n: 256 tau: 0.05 env: + numba_acceleration: false control_config: action_scale: 1.0 gait_phase_init_mode: "offset_phase" diff --git a/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml index 35624c3b6..70ff8e73d 100644 --- a/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/flashsac/g1_walk_flat/mujoco.yaml @@ -12,6 +12,7 @@ algo: replay_buffer_n: 256 tau: 0.05 env: + numba_acceleration: false control_config: action_scale: 1.0 gait_phase_init_mode: "offset_phase" diff --git a/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml b/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml index eddb1839f..e6aad393b 100644 --- a/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_motion_tracking/motrix.yaml @@ -11,6 +11,7 @@ training: task_name: G1MotionTrackingSAC sim_backend: motrix env: + numba_acceleration: false # motrix backend's kp/kd override path is broken on column slices, and DR # is not desirable during deterministic sim2sim eval anyway. Match the # `g1_walk_flat/motrix.yaml` convention by switching them off. diff --git a/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml b/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml index 598feccbb..d6c4f58ff 100644 --- a/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_motion_tracking/mujoco.yaml @@ -20,6 +20,7 @@ algo: target_entropy_ratio: 0.5 max_grad_norm: 10.0 env: + numba_acceleration: false control_config: action_scale: 2.0 anchor_pos_z_threshold: 0.5 diff --git a/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml b/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml index a193eeb51..c33aea079 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/motrix.yaml @@ -13,6 +13,7 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: + numba_acceleration: false control_config: action_scale: 1.0 gait_phase_init_mode: "offset_phase" diff --git a/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml index 0b0e2d4bd..2aac30cf7 100644 --- a/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_walk_flat/mujoco.yaml @@ -13,6 +13,7 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: + numba_acceleration: false control_config: action_scale: 1.0 gait_phase_init_mode: "offset_phase" diff --git a/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml b/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml index 4144b1713..56cbe1fe3 100644 --- a/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml +++ b/conf/offpolicy/task/sac/g1_walk_rough/motrix.yaml @@ -13,6 +13,7 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: + numba_acceleration: false sim_dt: 0.01 control_config: action_scale: 1.0 diff --git a/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml b/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml index bd9a1282a..099e79b27 100644 --- a/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml +++ b/conf/offpolicy/task/sac/g1_walk_rough/mujoco.yaml @@ -13,6 +13,7 @@ algo: alpha_init: 0.001 target_entropy_ratio: 0.0 env: + numba_acceleration: false control_config: action_scale: 1.0 gait_phase_init_mode: "offset_phase" diff --git a/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml b/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml index 54f87699c..2133e2f02 100644 --- a/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml +++ b/conf/offpolicy/task/td3/g1_walk_flat/mujoco.yaml @@ -5,6 +5,7 @@ training: algo: max_iterations: 100000 env: + numba_acceleration: false control_config: action_scale: 1.0 gait_phase_init_mode: "offset_phase" diff --git a/conf/ppo/task/g1_motion_tracking/motrix.yaml b/conf/ppo/task/g1_motion_tracking/motrix.yaml index 643939c54..9fdae5738 100644 --- a/conf/ppo/task/g1_motion_tracking/motrix.yaml +++ b/conf/ppo/task/g1_motion_tracking/motrix.yaml @@ -18,6 +18,8 @@ algo: scale_joint_angle: 0.01 scale_joint_vel: 1.5 scale_gyro: 0.2 +env: + numba_acceleration: false play_profile: enabled: true env: diff --git a/conf/ppo/task/g1_motion_tracking/mujoco.yaml b/conf/ppo/task/g1_motion_tracking/mujoco.yaml index 7301e4b68..03bfc5a24 100644 --- a/conf/ppo/task/g1_motion_tracking/mujoco.yaml +++ b/conf/ppo/task/g1_motion_tracking/mujoco.yaml @@ -12,6 +12,8 @@ algo: - actor algorithm: entropy_coef: 0.005 +env: + numba_acceleration: false reward: scales: motion_global_root_pos: 0.5 diff --git a/conf/ppo/task/g1_walk_flat/motrix.yaml b/conf/ppo/task/g1_walk_flat/motrix.yaml index 88c971936..e0f3da111 100644 --- a/conf/ppo/task/g1_walk_flat/motrix.yaml +++ b/conf/ppo/task/g1_walk_flat/motrix.yaml @@ -22,6 +22,7 @@ algo: learning_rate: 3.0e-4 entropy_coef: 5.0e-3 env: + numba_acceleration: false domain_rand: randomize_kp: false randomize_kd: false diff --git a/conf/ppo/task/g1_walk_flat/mujoco.yaml b/conf/ppo/task/g1_walk_flat/mujoco.yaml index 843d7b85f..45dbd8ede 100644 --- a/conf/ppo/task/g1_walk_flat/mujoco.yaml +++ b/conf/ppo/task/g1_walk_flat/mujoco.yaml @@ -15,6 +15,7 @@ algo: actor_hidden_dims: [512, 256, 128] critic_hidden_dims: [512, 256, 128] env: + numba_acceleration: false control_config: action_scale: 0.25 curriculum: diff --git a/tests/config/test_config_system.py b/tests/config/test_config_system.py index 3ed289c5a..2735590cb 100644 --- a/tests/config/test_config_system.py +++ b/tests/config/test_config_system.py @@ -19,6 +19,12 @@ CONF_DIR = Path(__file__).parent.parent.parent / "conf" _PPO_MLX_TASKS = {"go1_joystick_flat", "go2_joystick_flat", "g1_walk_flat"} _BACKENDS = ("mujoco", "motrix") +_NUMBA_ACCEL_SUPPORTED_TASK_NAMES = { + "G1MotionTracking", + "G1MotionTrackingSAC", + "G1WalkFlat", + "G1WalkRough", +} def _expected_backend_from_variant(name: str) -> str | None: @@ -225,6 +231,29 @@ def test_offpolicy_g1_walk_flat_motrix_sac_preserves_backend_overrides(): assert cfg.env.domain_rand.randomize_kd is False +def test_numba_accelerated_task_owners_declare_default_disabled_only_for_supported_tasks(): + matched: list[Path] = [] + + for path in sorted(CONF_DIR.glob("*/task/**/*.yaml")): + cfg = OmegaConf.load(path) + task_name = OmegaConf.select(cfg, "training.task_name") + declared = OmegaConf.select(cfg, "env.numba_acceleration") + if declared is not None: + assert task_name in _NUMBA_ACCEL_SUPPORTED_TASK_NAMES, ( + "only task owners with direct numba acceleration support may declare " + f"env.numba_acceleration: {path.relative_to(CONF_DIR)}" + ) + if task_name not in _NUMBA_ACCEL_SUPPORTED_TASK_NAMES: + continue + matched.append(path) + assert declared is False, ( + "numba-supported task owners must expose env.numba_acceleration " + f"with default false: {path.relative_to(CONF_DIR)}" + ) + + assert matched, "expected at least one numba-supported owner config" + + def test_offpolicy_g1_walk_flat_mujoco_td3_uses_td3_task_owner(): cfg = _compose("offpolicy", overrides=["algo=td3", "task=td3/g1_walk_flat/mujoco"]) From 53494d049b581aa6a2e923af8d7f8910de53b738 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Wed, 8 Jul 2026 20:55:54 +0800 Subject: [PATCH 24/26] feat: show collector env step breakdown --- src/unilab/algos/torch/appo/worker.py | 7 ++++ .../algos/torch/common/collector_timing.py | 32 +++++++++++++++++ src/unilab/algos/torch/hora/appo_worker.py | 7 ++++ src/unilab/algos/torch/offpolicy/worker.py | 6 ++-- src/unilab/logging/offpolicy.py | 6 ++++ tests/algos/test_offpolicy_logger.py | 34 +++++++++++++++++++ tests/algos/test_offpolicy_worker.py | 20 +++++++++++ 7 files changed, 110 insertions(+), 2 deletions(-) create mode 100644 src/unilab/algos/torch/common/collector_timing.py diff --git a/src/unilab/algos/torch/appo/worker.py b/src/unilab/algos/torch/appo/worker.py index d02ccd05e..c85285c54 100644 --- a/src/unilab/algos/torch/appo/worker.py +++ b/src/unilab/algos/torch/appo/worker.py @@ -16,6 +16,7 @@ import torch from rsl_rl.utils import resolve_callable +from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.observations import split_obs_dict from unilab.base.registry import ensure_registries @@ -249,6 +250,7 @@ def to_float32_np(x): ema_mlp_infer_ms: float = 0.0 ema_env_step_ms: float = 0.0 ema_rollout_ms: float = 0.0 + ema_env_step_breakdown_ms: dict[str, float] = {} try: while not stop_event.is_set(): @@ -289,6 +291,10 @@ def to_float32_np(x): ema_env_step_ms = (1 - _EMA) * ema_env_step_ms + _EMA * ( (time.perf_counter() - t_env) * 1000 ) + for key, value in extract_env_step_breakdown_timing_ms(state.info).items(): + ema_env_step_breakdown_ms[key] = (1 - _EMA) * ema_env_step_breakdown_ms.get( + key, 0.0 + ) + _EMA * value next_obs_raw = state.obs reward_raw = np.asarray(state.reward, dtype=np.float32).ravel() @@ -370,6 +376,7 @@ def to_float32_np(x): "rollout_ms": ema_rollout_ms, "mlp_infer_ms": ema_mlp_infer_ms, "env_step_ms": ema_env_step_ms, + **ema_env_step_breakdown_ms, } collector_active_steps_per_sec = compute_rollout_active_steps_per_sec( num_envs=num_envs, diff --git a/src/unilab/algos/torch/common/collector_timing.py b/src/unilab/algos/torch/common/collector_timing.py new file mode 100644 index 000000000..05bc4115e --- /dev/null +++ b/src/unilab/algos/torch/common/collector_timing.py @@ -0,0 +1,32 @@ +"""Shared collector timing helpers.""" + +from __future__ import annotations + +import math +from numbers import Real +from typing import Any + +ENV_STEP_BREAKDOWN_TIMING_MAP = { + "step_core_ms": "env_step_backend_ms", + "update_state_ms": "env_step_update_state_ms", + "reset_done_ms": "env_step_reset_done_ms", +} + + +def extract_env_step_breakdown_timing_ms(info: dict[str, Any] | None) -> dict[str, float]: + """Extract env-owned step sub-timings for collector metrics.""" + if not isinstance(info, dict): + return {} + timing = info.get("timing") + if not isinstance(timing, dict): + return {} + + out: dict[str, float] = {} + for source_key, collector_key in ENV_STEP_BREAKDOWN_TIMING_MAP.items(): + value = timing.get(source_key) + if not isinstance(value, Real): + continue + value_float = float(value) + if math.isfinite(value_float): + out[collector_key] = value_float + return out diff --git a/src/unilab/algos/torch/hora/appo_worker.py b/src/unilab/algos/torch/hora/appo_worker.py index 355177db0..e5bf5b365 100644 --- a/src/unilab/algos/torch/hora/appo_worker.py +++ b/src/unilab/algos/torch/hora/appo_worker.py @@ -16,6 +16,7 @@ compute_rollout_active_steps_per_sec, put_latest_metrics, ) +from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.registry import ensure_registries from unilab.training.seed import apply_training_seed @@ -238,6 +239,7 @@ def to_float32_np(x): ema_mlp_infer_ms = 0.0 ema_env_step_ms = 0.0 ema_rollout_ms = 0.0 + ema_env_step_breakdown_ms: dict[str, float] = {} try: while not stop_event.is_set(): @@ -275,6 +277,10 @@ def to_float32_np(x): ema_env_step_ms = (1 - _EMA) * ema_env_step_ms + _EMA * ( (time.perf_counter() - t_env) * 1000 ) + for key, value in extract_env_step_breakdown_timing_ms(state.info).items(): + ema_env_step_breakdown_ms[key] = (1 - _EMA) * ema_env_step_breakdown_ms.get( + key, 0.0 + ) + _EMA * value reward_raw = np.asarray(state.reward, dtype=np.float32).ravel() terminated_raw = np.asarray(state.terminated, dtype=np.float32).ravel() @@ -379,6 +385,7 @@ def to_float32_np(x): "rollout_ms": ema_rollout_ms, "mlp_infer_ms": ema_mlp_infer_ms, "env_step_ms": ema_env_step_ms, + **ema_env_step_breakdown_ms, } collector_active_steps_per_sec = compute_rollout_active_steps_per_sec( num_envs=num_envs, diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index f9aef915f..77cee79bf 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -14,6 +14,7 @@ import torch from unilab.algos.torch.common.actor_factory import build_actor +from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.algos.torch.offpolicy.thread_budget import apply_torch_thread_runtime from unilab.base.final_observation import resolve_terminal_observation_contract from unilab.base.observations import get_obs_dims, split_obs_dict @@ -771,6 +772,7 @@ def _run_collector( args={"num_envs": num_envs}, ) phase_start_ns = _record_phase_ms(cycle_timing_ms, "env_step_ms", phase_start_ns) + cycle_timing_ms.update(extract_env_step_breakdown_timing_ms(state.info)) # Extract data as numpy next_obs_np, next_critic_np = split_obs_dict(state.obs) @@ -1013,8 +1015,8 @@ def _run_collector( cycle_timing_ms, "sync_coordination_ms", phase_start_ns ) - for key in COLLECTOR_TIMING_KEYS: - _record_timing_ms(timing_accum_ms, timing_counts, key, cycle_timing_ms[key]) + for key, value in cycle_timing_ms.items(): + _record_timing_ms(timing_accum_ms, timing_counts, key, value) finally: if collector_pack_service is not None: diff --git a/src/unilab/logging/offpolicy.py b/src/unilab/logging/offpolicy.py index d6c2f7aab..d2daf19b2 100644 --- a/src/unilab/logging/offpolicy.py +++ b/src/unilab/logging/offpolicy.py @@ -17,6 +17,9 @@ "mlp_infer_ms": 1, "action_select_ms": 1, "env_step_ms": 2, + "env_step_backend_ms": 2.1, + "env_step_update_state_ms": 2.2, + "env_step_reset_done_ms": 2.3, "replay_ms": 3, "sync_coordination_ms": 4, "rollout_ms": 9, @@ -28,6 +31,9 @@ "mlp_infer_ms": "MLP Infer", "action_select_ms": "Action Select", "env_step_ms": "Env Step", + "env_step_backend_ms": " Backend Step", + "env_step_update_state_ms": " Update State", + "env_step_reset_done_ms": " Reset Done", "replay_ms": "Replay", "sync_coordination_ms": "Sync Coordination", } diff --git a/tests/algos/test_offpolicy_logger.py b/tests/algos/test_offpolicy_logger.py index 6bd37c320..ab6ebc094 100644 --- a/tests/algos/test_offpolicy_logger.py +++ b/tests/algos/test_offpolicy_logger.py @@ -74,3 +74,37 @@ def stop(self) -> None: assert live.stop_calls == 1 assert logger._live is None assert logger._last_live_refresh_time is None + + +def test_offpolicy_logger_displays_env_step_breakdown_as_indented_children() -> None: + logger = OffPolicyLogger( + algo_name="SAC", + max_iterations=2, + num_envs=8, + env_name="Dummy", + log_backend="none", + ) + logger.update_collector_timing( + { + "weight_sync_ms": 0.1, + "action_select_ms": 0.2, + "env_step_ms": 3.0, + "env_step_backend_ms": 1.5, + "env_step_update_state_ms": 1.0, + "env_step_reset_done_ms": 0.5, + "replay_ms": 0.3, + } + ) + + table = logger._build_timing_table() + collector_cells = list(table.columns[2].cells) + + assert collector_cells == [ + "Weight Sync", + "Action Select", + "Env Step", + " Backend Step", + " Update State", + " Reset Done", + "Replay", + ] diff --git a/tests/algos/test_offpolicy_worker.py b/tests/algos/test_offpolicy_worker.py index 1e5012f6c..b7d89697b 100644 --- a/tests/algos/test_offpolicy_worker.py +++ b/tests/algos/test_offpolicy_worker.py @@ -3,6 +3,7 @@ import pytest import torch +from unilab.algos.torch.common.collector_timing import extract_env_step_breakdown_timing_ms from unilab.algos.torch.offpolicy.worker import ( compute_collector_active_steps_per_sec, resolve_offpolicy_actor_priv_info, @@ -40,6 +41,25 @@ def test_compute_collector_active_steps_per_sec_includes_active_phases_only() -> assert steps_per_sec == pytest.approx(32 / 0.016) +def test_extract_env_step_breakdown_timing_ms_maps_env_owned_keys_only() -> None: + timing = extract_env_step_breakdown_timing_ms( + { + "timing": { + "step_core_ms": 1.5, + "update_state_ms": 2.5, + "reset_done_ms": 0.25, + "apply_action_ms": 9.0, + } + } + ) + + assert timing == { + "env_step_backend_ms": 1.5, + "env_step_update_state_ms": 2.5, + "env_step_reset_done_ms": 0.25, + } + + def test_compute_collector_active_steps_per_sec_returns_none_without_active_time() -> None: assert ( compute_collector_active_steps_per_sec( From 6f7a9998df9bb429ef083f3f9991bacd3cff143f Mon Sep 17 00:00:00 2001 From: UniLab Bot Date: Wed, 8 Jul 2026 14:29:52 +0000 Subject: [PATCH 25/26] benchmark: add SAC replay sampling placement test --- .../benchmark_sac_replay_buffer_sampling.py | 1352 +++++++++++++++++ ...st_sac_replay_buffer_sampling_benchmark.py | 122 ++ 2 files changed, 1474 insertions(+) create mode 100644 benchmark/benchmark_sac_replay_buffer_sampling.py create mode 100644 tests/benchmark/test_sac_replay_buffer_sampling_benchmark.py diff --git a/benchmark/benchmark_sac_replay_buffer_sampling.py b/benchmark/benchmark_sac_replay_buffer_sampling.py new file mode 100644 index 000000000..3b2d6cfa4 --- /dev/null +++ b/benchmark/benchmark_sac_replay_buffer_sampling.py @@ -0,0 +1,1352 @@ +#!/usr/bin/env python3 +"""Benchmark SAC G1 replay-buffer sampling placement. + +This standalone benchmark composes the same off-policy owner config used by: + + uv run train --algo sac --task g1_walk_flat --sim mujoco + +It does not run training. It measures only the replay sampling boundary: + +* CPU replay storage: random sample into a host batch, then host-to-device copy. +* GPU replay storage: random sample directly from a device-resident replay tensor. + +The default matrix tests 1/2/4/8 visible CUDA devices and several replay +capacities derived from the configured replay size. GPU counts that are not +available are recorded as skipped. + +Usage: + uv run benchmark/benchmark_sac_replay_buffer_sampling.py + uv run benchmark/benchmark_sac_replay_buffer_sampling.py --gpu-counts 1,2 + uv run benchmark/benchmark_sac_replay_buffer_sampling.py --capacity-multipliers 0.25,0.5,1,2 + uv run benchmark/benchmark_sac_replay_buffer_sampling.py --warmup 3 --repeat 10 +""" + +from __future__ import annotations + +import argparse +import gc +import json +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from statistics import mean, median, pstdev +from typing import Any, Callable, Iterable, cast + +import torch +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import DictConfig, OmegaConf + +ROOT_DIR = Path(__file__).resolve().parents[1] +if str(ROOT_DIR) not in sys.path: + sys.path.append(str(ROOT_DIR)) + +from benchmark.core.device_info import get_device_info_dict, get_device_info_line + +plt: Any | None = None +try: + import matplotlib as _matplotlib + + _matplotlib.use("Agg") + import matplotlib.pyplot as _plt + + plt = _plt +except Exception: + plt = None + +DEFAULT_OUTPUT_JSON = ( + ROOT_DIR / "benchmark" / "outputs" / "sac_replay_buffer_sampling" / "results.json" +) +DEFAULT_SIM = "mujoco" +DEFAULT_GPU_COUNTS = "1,2,4,8" +DEFAULT_CAPACITY_MULTIPLIERS = "0.25,0.5,1.0" +FLOAT_BYTES = 4 + + +@dataclass(frozen=True) +class ReplayShape: + obs_dim: int + action_dim: int + critic_dim: int + + @property + def packed_width(self) -> int: + return 2 * self.obs_dim + self.action_dim + 3 + 2 * self.critic_dim + + +@dataclass(frozen=True) +class BenchmarkCase: + algo: str + task: str + sim: str + command: str + training_task_name: str + num_envs: int + replay_buffer_n: int + config_capacity_rows: int + configured_batch_size: int + learner_batch_size: int + symmetry_batch_multiplier: int + updates_per_step: int + sample_count_per_rank: int + learning_starts: int + shape: ReplayShape + + @property + def sample_bytes_per_rank(self) -> int: + return self.sample_count_per_rank * self.shape.packed_width * FLOAT_BYTES + + +@dataclass +class TimingStats: + samples_ms: list[float] + mean_ms: float + median_ms: float + std_ms: float + min_ms: float + max_ms: float + warmup: int + repeat: int + + +@dataclass +class CapacityResult: + world_size: int + capacity_rows: int + torch_threads: int + replay_bytes_per_rank: int + sample_bytes_per_rank: int + global_sample_bytes: int + timings: dict[str, TimingStats] + notes: list[str] + + +@dataclass +class SkippedCase: + world_size: int + capacity_rows: int | None + reason: str + + +def _fmt_bytes(num_bytes: int) -> str: + value = float(num_bytes) + for unit in ("B", "KiB", "MiB", "GiB", "TiB"): + if value < 1024.0 or unit == "TiB": + return f"{value:.2f} {unit}" if unit != "B" else f"{int(value)} B" + value /= 1024.0 + return f"{value:.2f} TiB" + + +def _stats(samples_ms: list[float], *, warmup: int, repeat: int) -> TimingStats: + if not samples_ms: + raise ValueError("no timing samples collected") + return TimingStats( + samples_ms=samples_ms, + mean_ms=mean(samples_ms), + median_ms=median(samples_ms), + std_ms=pstdev(samples_ms) if len(samples_ms) > 1 else 0.0, + min_ms=min(samples_ms), + max_ms=max(samples_ms), + warmup=warmup, + repeat=repeat, + ) + + +def _measure_ms( + fn: Callable[[int], None], + *, + warmup: int, + repeat: int, + sync_before: Callable[[], None] | None = None, + sync_after: Callable[[], None] | None = None, +) -> TimingStats: + samples: list[float] = [] + for idx in range(warmup + repeat): + if sync_before is not None: + sync_before() + start_ns = time.perf_counter_ns() + fn(idx) + if sync_after is not None: + sync_after() + elapsed_ms = (time.perf_counter_ns() - start_ns) / 1e6 + if idx >= warmup: + samples.append(elapsed_ms) + return _stats(samples, warmup=warmup, repeat=repeat) + + +def _sync_devices(devices: Iterable[torch.device]) -> None: + seen: set[str] = set() + for device in devices: + key = str(device) + if key in seen: + continue + seen.add(key) + if device.type == "cuda": + torch.cuda.synchronize(device) + elif device.type == "xpu": + xpu = getattr(torch, "xpu", None) + synchronize = getattr(xpu, "synchronize", None) + if callable(synchronize): + try: + synchronize(device) + except TypeError: + synchronize() + elif device.type == "mps" and hasattr(torch, "mps"): + torch.mps.synchronize() + + +def _cleanup_device() -> None: + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + if hasattr(torch, "mps") and torch.backends.mps.is_available(): + torch.mps.empty_cache() + + +def _compose_offpolicy_cfg(sim: str = DEFAULT_SIM) -> DictConfig: + config_dir = str(ROOT_DIR / "conf" / "offpolicy") + overrides = [ + "algo=sac", + f"task=sac/g1_walk_flat/{sim}", + "hydra.run.dir=.", + "hydra.output_subdir=null", + "hydra/job_logging=disabled", + "hydra/hydra_logging=disabled", + ] + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=config_dir, version_base="1.3"): + return compose(config_name="config", overrides=overrides) + + +def _resolve_env_shape_and_symmetry(cfg: DictConfig) -> tuple[ReplayShape, int]: + from unilab.base.observations import get_obs_dims + from unilab.training import BackendAdapter, create_env, ensure_registries + + ensure_registries() + env_cfg_override = BackendAdapter( + cfg, + root_dir=ROOT_DIR, + algo_name="sac", + ).build_task_env_cfg_override() + env = create_env(cfg, num_envs=1, env_cfg_override=env_cfg_override) + try: + obs_dim, critic_dim = get_obs_dims(env.obs_groups_spec) + action_shape = env.action_space.shape + if action_shape is None: + raise ValueError("env.action_space.shape must be defined") + action_dim = int(action_shape[0]) + + symmetry_batch_multiplier = 1 + if bool(OmegaConf.select(cfg, "algo.use_symmetry", default=False)): + symmetry_builder = getattr(env, "build_symmetry_augmentation", None) + if not callable(symmetry_builder): + raise ValueError(f"{cfg.training.task_name} does not provide symmetry augmentation") + symmetry = cast(Any, symmetry_builder(device="cpu")) + if symmetry is None: + raise ValueError(f"{cfg.training.task_name} does not provide symmetry augmentation") + symmetry_batch_multiplier = int(symmetry.batch_multiplier) + finally: + env.close() + + return ( + ReplayShape(obs_dim=int(obs_dim), action_dim=action_dim, critic_dim=int(critic_dim)), + symmetry_batch_multiplier, + ) + + +def _build_case( + cfg: DictConfig, + *, + sim: str, + shape: ReplayShape, + symmetry_batch_multiplier: int, +) -> BenchmarkCase: + num_envs = int(cfg.algo.num_envs) + replay_buffer_n = int(cfg.algo.replay_buffer_n) + configured_batch_size = int(cfg.algo.batch_size) + learner_batch_size = configured_batch_size + if bool(OmegaConf.select(cfg, "algo.use_symmetry", default=False)): + if configured_batch_size % symmetry_batch_multiplier != 0: + raise ValueError( + "SAC symmetry requires batch_size divisible by " + f"{symmetry_batch_multiplier}, got {configured_batch_size}" + ) + learner_batch_size = configured_batch_size // symmetry_batch_multiplier + + updates_per_step = int(cfg.algo.updates_per_step) + return BenchmarkCase( + algo="sac", + task="g1_walk_flat", + sim=sim, + command=f"uv run train --algo sac --task g1_walk_flat --sim {sim}", + training_task_name=str(cfg.training.task_name), + num_envs=num_envs, + replay_buffer_n=replay_buffer_n, + config_capacity_rows=num_envs * replay_buffer_n, + configured_batch_size=configured_batch_size, + learner_batch_size=learner_batch_size, + symmetry_batch_multiplier=int(symmetry_batch_multiplier), + updates_per_step=updates_per_step, + sample_count_per_rank=learner_batch_size * updates_per_step, + learning_starts=int(cfg.algo.learning_starts), + shape=shape, + ) + + +def _parse_int_list(value: str, *, name: str) -> list[int]: + parsed: list[int] = [] + for part in value.split(","): + item = part.strip() + if not item: + continue + number = int(item) + if number <= 0: + raise ValueError(f"{name} values must be positive, got {number}") + if number not in parsed: + parsed.append(number) + if not parsed: + raise ValueError(f"{name} must contain at least one positive integer") + return parsed + + +def _parse_float_list(value: str, *, name: str) -> list[float]: + parsed: list[float] = [] + for part in value.split(","): + item = part.strip() + if not item: + continue + number = float(item) + if number <= 0.0: + raise ValueError(f"{name} values must be positive, got {number}") + if number not in parsed: + parsed.append(number) + if not parsed: + raise ValueError(f"{name} must contain at least one positive value") + return parsed + + +def _resolve_capacity_rows( + *, + config_capacity_rows: int, + capacity_rows_arg: str, + capacity_multipliers_arg: str, + capacity_exponents_arg: str | None = None, +) -> list[int]: + if capacity_exponents_arg is not None: + return [ + 2**exponent + for exponent in _parse_int_list(capacity_exponents_arg, name="capacity exponents") + ] + + if capacity_rows_arg.strip().lower() != "auto": + return _parse_int_list(capacity_rows_arg, name="capacity rows") + + capacities: list[int] = [] + for multiplier in _parse_float_list(capacity_multipliers_arg, name="capacity multipliers"): + rows = max(1, int(round(config_capacity_rows * multiplier))) + if rows not in capacities: + capacities.append(rows) + return capacities + + +def _parse_device_ids(value: str) -> list[int]: + if value.strip().lower() == "auto": + return list(range(torch.cuda.device_count())) if torch.cuda.is_available() else [] + return _parse_int_list(value, name="device ids") + + +def _resolve_torch_threads(setting: str, *, world_size: int) -> int: + if setting != "auto": + threads = int(setting) + if threads <= 0: + raise ValueError("--torch-threads must be 'auto' or a positive integer") + return threads + cpu_count = os.cpu_count() or 1 + return max(1, cpu_count // max(int(world_size), 1)) + + +def _allocate_cpu_tensor( + *, + rows: int, + width: int, + pin_memory: bool, + prefill: str, + notes: list[str], +) -> torch.Tensor: + try: + tensor = torch.empty((rows, width), dtype=torch.float32, pin_memory=pin_memory) + except RuntimeError as exc: + if not pin_memory: + raise + notes.append(f"pinned CPU allocation failed; fell back to pageable CPU: {exc}") + tensor = torch.empty((rows, width), dtype=torch.float32) + + if prefill == "zeros": + tensor.zero_() + elif prefill != "none": + raise ValueError(f"Unsupported prefill={prefill!r}") + return tensor + + +def _allocate_device_tensor( + *, + rows: int, + width: int, + device: torch.device, + prefill: str, +) -> torch.Tensor: + tensor = torch.empty((rows, width), dtype=torch.float32, device=device) + if prefill == "zeros": + tensor.zero_() + _sync_devices([device]) + elif prefill != "none": + raise ValueError(f"Unsupported prefill={prefill!r}") + return tensor + + +def _make_cpu_generators(world_size: int, seed: int) -> list[torch.Generator]: + generators = [] + for rank in range(world_size): + gen = torch.Generator(device="cpu") + gen.manual_seed(int(seed) + rank * 1009) + generators.append(gen) + return generators + + +def _make_device_generators(devices: list[torch.device], seed: int) -> list[torch.Generator]: + generators = [] + for rank, device in enumerate(devices): + gen = torch.Generator(device=device) + gen.manual_seed(int(seed) + rank * 1009) + generators.append(gen) + return generators + + +def _run_parallel( + executor: ThreadPoolExecutor | None, + world_size: int, + worker: Callable[[int], None], +) -> None: + if world_size == 1: + worker(0) + return + assert executor is not None + futures = [executor.submit(worker, rank) for rank in range(world_size)] + for future in futures: + future.result() + + +def _bench_cpu_sample_and_h2d( + case: BenchmarkCase, + *, + capacity_rows: int, + devices: list[torch.device], + warmup: int, + repeat: int, + prefill: str, + pinned_host_batch: bool, + index_mode: str, + seed: int, + notes: list[str], +) -> dict[str, TimingStats]: + world_size = len(devices) + width = case.shape.packed_width + sample_count = case.sample_count_per_rank + storage = _allocate_cpu_tensor( + rows=capacity_rows, + width=width, + pin_memory=False, + prefill=prefill, + notes=notes, + ) + host_batches = [ + _allocate_cpu_tensor( + rows=sample_count, + width=width, + pin_memory=pinned_host_batch and any(d.type == "cuda" for d in devices), + prefill="none", + notes=notes, + ) + for _ in range(world_size) + ] + device_batches = [ + _allocate_device_tensor(rows=sample_count, width=width, device=device, prefill="none") + for device in devices + ] + cpu_indices = [torch.empty(sample_count, dtype=torch.int64) for _ in range(world_size)] + cpu_generators = _make_cpu_generators(world_size, seed) + if index_mode == "pregenerated": + for rank in range(world_size): + torch.randint( + capacity_rows, + (sample_count,), + generator=cpu_generators[rank], + out=cpu_indices[rank], + ) + elif index_mode != "timed": + raise ValueError(f"Unsupported index_mode={index_mode!r}") + + streams: list[torch.cuda.Stream | None] = [] + for device in devices: + if device.type == "cuda": + with torch.cuda.device(device): + streams.append(torch.cuda.Stream(device=device)) + else: + streams.append(None) + + executor: ThreadPoolExecutor | None = None + if world_size > 1: + executor = ThreadPoolExecutor(max_workers=world_size, thread_name_prefix="cpu_replay_rank") + + def sample_worker(rank: int) -> None: + indices = cpu_indices[rank] + if index_mode == "timed": + torch.randint( + capacity_rows, + (sample_count,), + generator=cpu_generators[rank], + out=indices, + ) + torch.index_select(storage, 0, indices, out=host_batches[rank]) + + def h2d_worker(rank: int) -> None: + device = devices[rank] + host = host_batches[rank] + target = device_batches[rank] + non_blocking = bool(host.is_pinned() and device.type == "cuda") + if device.type == "cuda": + stream = streams[rank] + assert stream is not None + with torch.cuda.device(device), torch.cuda.stream(stream): + target.copy_(host, non_blocking=non_blocking) + else: + target.copy_(host) + + def sample_all(_: int) -> None: + _run_parallel(executor, world_size, sample_worker) + + def h2d_all(_: int) -> None: + _run_parallel(executor, world_size, h2d_worker) + + def sample_then_h2d_all(_: int) -> None: + _run_parallel(executor, world_size, sample_worker) + _run_parallel(executor, world_size, h2d_worker) + + def sync() -> None: + _sync_devices(devices) + + try: + sample_stats = _measure_ms(sample_all, warmup=warmup, repeat=repeat) + h2d_stats = _measure_ms( + h2d_all, + warmup=warmup, + repeat=repeat, + sync_before=sync, + sync_after=sync, + ) + combined_stats = _measure_ms( + sample_then_h2d_all, + warmup=warmup, + repeat=repeat, + sync_before=sync, + sync_after=sync, + ) + return { + "cpu_sample_wall": sample_stats, + "cpu_sample_h2d_wall": h2d_stats, + "cpu_sample_then_h2d_wall": combined_stats, + } + finally: + if executor is not None: + executor.shutdown(wait=True) + del cpu_indices, device_batches, host_batches, storage + _cleanup_device() + + +def _bench_gpu_sample( + case: BenchmarkCase, + *, + capacity_rows: int, + devices: list[torch.device], + warmup: int, + repeat: int, + prefill: str, + index_mode: str, + seed: int, +) -> TimingStats: + world_size = len(devices) + width = case.shape.packed_width + sample_count = case.sample_count_per_rank + storages = [ + _allocate_device_tensor(rows=capacity_rows, width=width, device=device, prefill=prefill) + for device in devices + ] + outputs = [ + _allocate_device_tensor(rows=sample_count, width=width, device=device, prefill="none") + for device in devices + ] + indices = [torch.empty(sample_count, dtype=torch.int64, device=device) for device in devices] + generators = _make_device_generators(devices, seed) + if index_mode == "pregenerated": + for rank, device in enumerate(devices): + with torch.cuda.device(device) if device.type == "cuda" else _nullcontext(): + torch.randint( + capacity_rows, + (sample_count,), + generator=generators[rank], + device=device, + out=indices[rank], + ) + _sync_devices(devices) + elif index_mode != "timed": + raise ValueError(f"Unsupported index_mode={index_mode!r}") + + streams: list[torch.cuda.Stream | None] = [] + for device in devices: + if device.type == "cuda": + with torch.cuda.device(device): + streams.append(torch.cuda.Stream(device=device)) + else: + streams.append(None) + + executor: ThreadPoolExecutor | None = None + if world_size > 1: + executor = ThreadPoolExecutor(max_workers=world_size, thread_name_prefix="gpu_replay_rank") + + def sample_worker(rank: int) -> None: + device = devices[rank] + if device.type == "cuda": + stream = streams[rank] + assert stream is not None + with torch.cuda.device(device), torch.cuda.stream(stream): + if index_mode == "timed": + torch.randint( + capacity_rows, + (sample_count,), + generator=generators[rank], + device=device, + out=indices[rank], + ) + torch.index_select(storages[rank], 0, indices[rank], out=outputs[rank]) + else: + if index_mode == "timed": + torch.randint( + capacity_rows, + (sample_count,), + generator=generators[rank], + device=device, + out=indices[rank], + ) + torch.index_select(storages[rank], 0, indices[rank], out=outputs[rank]) + + def sample_all(_: int) -> None: + _run_parallel(executor, world_size, sample_worker) + + def sync() -> None: + _sync_devices(devices) + + try: + return _measure_ms( + sample_all, + warmup=warmup, + repeat=repeat, + sync_before=sync, + sync_after=sync, + ) + finally: + if executor is not None: + executor.shutdown(wait=True) + del generators, indices, outputs, storages + _cleanup_device() + + +class _nullcontext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + +def _estimate_device_bytes_per_gpu(case: BenchmarkCase, capacity_rows: int) -> int: + width = case.shape.packed_width + replay = capacity_rows * width * FLOAT_BYTES + sample_batch = case.sample_count_per_rank * width * FLOAT_BYTES + # GPU mode uses replay + output + indices; CPU/H2D mode uses a device batch. + indices = case.sample_count_per_rank * 8 + return replay + (2 * sample_batch) + indices + + +def _cuda_memory_skip_reason( + devices: list[torch.device], + *, + required_bytes_per_gpu: int, + safety_factor: float, +) -> str | None: + for device in devices: + if device.type != "cuda": + continue + free_bytes, _total_bytes = torch.cuda.mem_get_info(device) + needed = int(required_bytes_per_gpu * safety_factor) + if free_bytes < needed: + return ( + f"{device} free memory {_fmt_bytes(free_bytes)} < " + f"estimated required {_fmt_bytes(needed)}" + ) + return None + + +def _run_capacity_case( + case: BenchmarkCase, + *, + capacity_rows: int, + devices: list[torch.device], + warmup: int, + repeat: int, + prefill: str, + pinned_host_batch: bool, + index_mode: str, + seed: int, + torch_threads: int, +) -> CapacityResult: + notes: list[str] = [] + old_threads = torch.get_num_threads() + torch.set_num_threads(torch_threads) + try: + cpu_timings = _bench_cpu_sample_and_h2d( + case, + capacity_rows=capacity_rows, + devices=devices, + warmup=warmup, + repeat=repeat, + prefill=prefill, + pinned_host_batch=pinned_host_batch, + index_mode=index_mode, + seed=seed, + notes=notes, + ) + gpu_sample = _bench_gpu_sample( + case, + capacity_rows=capacity_rows, + devices=devices, + warmup=warmup, + repeat=repeat, + prefill=prefill, + index_mode=index_mode, + seed=seed + 100_000, + ) + finally: + torch.set_num_threads(old_threads) + + timings = dict(cpu_timings) + timings["gpu_sample_wall"] = gpu_sample + return CapacityResult( + world_size=len(devices), + capacity_rows=capacity_rows, + torch_threads=torch_threads, + replay_bytes_per_rank=capacity_rows * case.shape.packed_width * FLOAT_BYTES, + sample_bytes_per_rank=case.sample_bytes_per_rank, + global_sample_bytes=case.sample_bytes_per_rank * len(devices), + timings=timings, + notes=notes, + ) + + +def _print_timing(label: str, stat: TimingStats) -> None: + print( + f" {label:<28}" + f" mean={stat.mean_ms:8.3f} ms" + f" median={stat.median_ms:8.3f} ms" + f" std={stat.std_ms:7.3f} ms" + ) + + +def _print_result(result: CapacityResult) -> None: + print( + f" {result.world_size} GPU(s), capacity={result.capacity_rows:,} " + f"({_fmt_bytes(result.replay_bytes_per_rank)} per replay copy), " + f"sample={_fmt_bytes(result.sample_bytes_per_rank)} per rank, " + f"torch_threads={result.torch_threads}" + ) + _print_timing("CPU sample", result.timings["cpu_sample_wall"]) + _print_timing("CPU sampled batch H2D", result.timings["cpu_sample_h2d_wall"]) + _print_timing("CPU sample then H2D", result.timings["cpu_sample_then_h2d_wall"]) + _print_timing("GPU replay sample", result.timings["gpu_sample_wall"]) + for note in result.notes: + print(f" note: {note}") + + +def _json_default(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if hasattr(value, "__dataclass_fields__"): + return asdict(value) + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def _write_results(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, default=_json_default), encoding="utf-8") + + +def _json_ready(value: Any) -> Any: + return json.loads(json.dumps(value, default=_json_default)) + + +def _result_timing_mean(result: dict[str, Any], key: str) -> float: + return float(result["timings"][key]["mean_ms"]) + + +def _capacity_label(capacity_rows: int, config_capacity_rows: int) -> str: + multiplier = capacity_rows / max(config_capacity_rows, 1) + if capacity_rows % 1_048_576 == 0: + rows_label = f"{capacity_rows // 1_048_576}M" + elif capacity_rows % 1024 == 0: + rows_label = f"{capacity_rows // 1024}K" + else: + rows_label = f"{capacity_rows:,}" + return f"{multiplier:g}x\n{rows_label}" + + +def _annotate(ax: Any, x: float, y: float, text: str, *, fontsize: int = 8) -> None: + if y <= 0: + return + ax.annotate( + text, + xy=(x, y), + xytext=(0, 3), + textcoords="offset points", + ha="center", + va="bottom", + fontsize=fontsize, + ) + + +def _save_component_breakdown_plot(payload: dict[str, Any], output_path: Path) -> str | None: + if plt is None: + return None + results = list(payload.get("results", [])) + if not results: + return None + + config_capacity_rows = int(payload["case"]["config_capacity_rows"]) + world_sizes = sorted({int(result["world_size"]) for result in results}) + ncols = 2 if len(world_sizes) > 1 else 1 + nrows = (len(world_sizes) + ncols - 1) // ncols + fig, axes = plt.subplots(nrows, ncols, figsize=(8.0 * ncols, 5.1 * nrows), squeeze=False) + colors = { + "cpu_sample": "#4C78A8", + "h2d": "#F58518", + "gpu": "#54A24B", + } + max_y = 0.0 + + for index, world_size in enumerate(world_sizes): + ax = axes[index // ncols][index % ncols] + rows = sorted( + (result for result in results if int(result["world_size"]) == world_size), + key=lambda item: int(item["capacity_rows"]), + ) + labels = [ + _capacity_label(int(result["capacity_rows"]), config_capacity_rows) for result in rows + ] + x_positions = list(range(len(rows))) + bar_width = 0.34 + cpu_sample = [_result_timing_mean(result, "cpu_sample_wall") for result in rows] + h2d = [_result_timing_mean(result, "cpu_sample_h2d_wall") for result in rows] + cpu_total = [_result_timing_mean(result, "cpu_sample_then_h2d_wall") for result in rows] + gpu_sample = [_result_timing_mean(result, "gpu_sample_wall") for result in rows] + component_total = [sample + copy for sample, copy in zip(cpu_sample, h2d)] + max_y = max(max_y, *(component_total or [0.0]), *(gpu_sample or [0.0])) + + cpu_x = [x - bar_width / 2 for x in x_positions] + gpu_x = [x + bar_width / 2 for x in x_positions] + ax.bar(cpu_x, cpu_sample, bar_width, color=colors["cpu_sample"], label="CPU sample") + ax.bar(cpu_x, h2d, bar_width, bottom=cpu_sample, color=colors["h2d"], label="H2D") + ax.bar(gpu_x, gpu_sample, bar_width, color=colors["gpu"], label="GPU sample") + + for x, total, measured_total in zip(cpu_x, component_total, cpu_total): + _annotate(ax, x, total, f"{measured_total:.2f}", fontsize=7) + for x, value in zip(gpu_x, gpu_sample): + _annotate(ax, x, value, f"{value:.3f}", fontsize=7) + + ax.set_title(f"{world_size} GPU(s)") + ax.set_xticks(x_positions) + ax.set_xticklabels(labels) + ax.set_ylabel("Mean wall time (ms)") + ax.grid(True, axis="y", alpha=0.25) + + for index in range(len(world_sizes), nrows * ncols): + axes[index // ncols][index % ncols].axis("off") + + handles, labels = axes[0][0].get_legend_handles_labels() + fig.suptitle( + "SAC replay sampling placement: CPU sample + H2D vs GPU-resident sample", + y=0.985, + ) + fig.legend(handles, labels, loc="upper center", ncol=3, bbox_to_anchor=(0.5, 0.955)) + if max_y > 0: + for ax_row in axes: + for ax in ax_row: + if ax.has_data(): + ax.set_ylim(0, max_y * 1.28) + fig.tight_layout(rect=(0, 0, 1, 0.90), h_pad=2.0, w_pad=1.5) + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=170, bbox_inches="tight") + plt.close(fig) + return str(output_path.resolve()) + + +def _save_speedup_plot(payload: dict[str, Any], output_path: Path) -> str | None: + if plt is None: + return None + results = list(payload.get("results", [])) + if not results: + return None + + config_capacity_rows = int(payload["case"]["config_capacity_rows"]) + world_sizes = sorted({int(result["world_size"]) for result in results}) + capacities = sorted({int(result["capacity_rows"]) for result in results}) + x_positions = list(range(len(world_sizes))) + width = min(0.8 / max(len(capacities), 1), 0.22) + + fig, ax = plt.subplots(figsize=(10.5, 5.6)) + for cap_index, capacity_rows in enumerate(capacities): + values: list[float] = [] + for world_size in world_sizes: + match = next( + ( + result + for result in results + if int(result["world_size"]) == world_size + and int(result["capacity_rows"]) == capacity_rows + ), + None, + ) + if match is None: + values.append(float("nan")) + continue + cpu_total = _result_timing_mean(match, "cpu_sample_then_h2d_wall") + gpu_sample = _result_timing_mean(match, "gpu_sample_wall") + values.append(cpu_total / gpu_sample if gpu_sample > 0 else float("nan")) + offset = (cap_index - (len(capacities) - 1) / 2) * width + bar_x = [x + offset for x in x_positions] + bars = ax.bar( + bar_x, + values, + width, + label=_capacity_label(capacity_rows, config_capacity_rows).replace("\n", " "), + ) + for bar, value in zip(bars, values): + if value == value and value > 0: + _annotate(ax, float(bar.get_x() + bar.get_width() / 2), value, f"{value:.1f}x") + + ax.set_title("CPU sample+H2D wall time divided by GPU-resident sample time") + ax.set_ylabel("Speedup (x)") + ax.set_xticks(x_positions) + ax.set_xticklabels([str(world_size) for world_size in world_sizes]) + ax.set_xlabel("GPU count") + ax.grid(True, axis="y", alpha=0.25) + ax.legend(title="Replay capacity", loc="best") + fig.tight_layout() + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=170, bbox_inches="tight") + plt.close(fig) + return str(output_path.resolve()) + + +def _save_gpu_sample_plot(payload: dict[str, Any], output_path: Path) -> str | None: + if plt is None: + return None + results = list(payload.get("results", [])) + if not results: + return None + + config_capacity_rows = int(payload["case"]["config_capacity_rows"]) + world_sizes = sorted({int(result["world_size"]) for result in results}) + capacities = sorted({int(result["capacity_rows"]) for result in results}) + x_positions = list(range(len(world_sizes))) + width = min(0.8 / max(len(capacities), 1), 0.22) + + fig, ax = plt.subplots(figsize=(10.5, 5.6)) + for cap_index, capacity_rows in enumerate(capacities): + values: list[float] = [] + for world_size in world_sizes: + match = next( + ( + result + for result in results + if int(result["world_size"]) == world_size + and int(result["capacity_rows"]) == capacity_rows + ), + None, + ) + values.append( + float("nan") if match is None else _result_timing_mean(match, "gpu_sample_wall") + ) + offset = (cap_index - (len(capacities) - 1) / 2) * width + bar_x = [x + offset for x in x_positions] + bars = ax.bar( + bar_x, + values, + width, + label=_capacity_label(capacity_rows, config_capacity_rows).replace("\n", " "), + ) + for bar, value in zip(bars, values): + if value == value and value > 0: + _annotate(ax, float(bar.get_x() + bar.get_width() / 2), value, f"{value:.3f}") + + ax.set_title("GPU-resident replay sample wall time") + ax.set_ylabel("Mean wall time (ms)") + ax.set_xticks(x_positions) + ax.set_xticklabels([str(world_size) for world_size in world_sizes]) + ax.set_xlabel("GPU count") + ax.grid(True, axis="y", alpha=0.25) + ax.legend(title="Replay capacity", loc="best") + fig.tight_layout() + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=170, bbox_inches="tight") + plt.close(fig) + return str(output_path.resolve()) + + +def _save_plots(payload: dict[str, Any], plot_dir: Path) -> list[str]: + if plt is None: + print("Plotting skipped: matplotlib is not available.") + return [] + + saved: list[str] = [] + for maybe_path in ( + _save_component_breakdown_plot( + payload, plot_dir / "sac_replay_sampling_component_breakdown.png" + ), + _save_speedup_plot(payload, plot_dir / "sac_replay_sampling_speedup.png"), + _save_gpu_sample_plot(payload, plot_dir / "sac_replay_gpu_sample_wall_time.png"), + ): + if maybe_path is not None: + saved.append(maybe_path) + return saved + + +def _write_analysis_markdown(payload: dict[str, Any], output_path: Path) -> str: + results = sorted( + list(payload.get("results", [])), + key=lambda result: (int(result["world_size"]), int(result["capacity_rows"])), + ) + case = payload.get("case", {}) + shape = case.get("shape", {}) + packed_width = shape.get("packed_width") + if packed_width is None and {"obs_dim", "action_dim", "critic_dim"} <= set(shape): + packed_width = ( + 2 * int(shape["obs_dim"]) + int(shape["action_dim"]) + 3 + 2 * int(shape["critic_dim"]) + ) + lines = [ + "# SAC Replay Buffer Sampling Time Analysis", + "", + f"- Command: `{case.get('command', 'unknown')}`", + f"- Sample count per rank: `{case.get('sample_count_per_rank', 'unknown')}`", + f"- Packed row width: `{packed_width if packed_width is not None else 'unknown'}` float32 columns", + "", + "Measured `CPU sample then H2D` is timed as its own end-to-end pass, so it can differ " + "slightly from `CPU sample + H2D` due to cache effects and run-to-run noise.", + "", + "| GPUs | Capacity rows | CPU sample ms | H2D ms | CPU sample+H2D ms | GPU sample ms | Speedup | CPU sample share | H2D share |", + "|---:|---:|---:|---:|---:|---:|---:|---:|---:|", + ] + for result in results: + cpu_sample = _result_timing_mean(result, "cpu_sample_wall") + h2d = _result_timing_mean(result, "cpu_sample_h2d_wall") + cpu_total = _result_timing_mean(result, "cpu_sample_then_h2d_wall") + gpu_sample = _result_timing_mean(result, "gpu_sample_wall") + component_sum = cpu_sample + h2d + sample_share = cpu_sample / component_sum if component_sum > 0 else 0.0 + h2d_share = h2d / component_sum if component_sum > 0 else 0.0 + speedup = cpu_total / gpu_sample if gpu_sample > 0 else float("nan") + lines.append( + "| " + f"{int(result['world_size'])} | " + f"{int(result['capacity_rows']):,} | " + f"{cpu_sample:.3f} | " + f"{h2d:.3f} | " + f"{cpu_total:.3f} | " + f"{gpu_sample:.3f} | " + f"{speedup:.1f}x | " + f"{sample_share * 100:.1f}% | " + f"{h2d_share * 100:.1f}% |" + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text("\n".join(lines) + "\n", encoding="utf-8") + return str(output_path.resolve()) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--sim", default=DEFAULT_SIM) + parser.add_argument("--gpu-counts", default=DEFAULT_GPU_COUNTS) + parser.add_argument("--device-ids", default="auto") + parser.add_argument( + "--capacity-rows", + default="auto", + help="Comma-separated explicit replay row counts, or 'auto'.", + ) + parser.add_argument( + "--capacity-multipliers", + default=DEFAULT_CAPACITY_MULTIPLIERS, + help="Used when --capacity-rows=auto; relative to config replay capacity.", + ) + parser.add_argument( + "--capacity-exponents", + default=None, + help="Comma-separated powers of two for replay rows, e.g. '20,22,24,26,28,30'.", + ) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--repeat", type=int, default=20) + parser.add_argument("--prefill", choices=("zeros", "none"), default="zeros") + parser.add_argument( + "--index-mode", + choices=("timed", "pregenerated"), + default="timed", + help="timed includes random index generation with preallocated index tensors.", + ) + parser.add_argument( + "--host-batch-memory", + choices=("pinned", "pageable"), + default="pinned", + help="CPU sampled-batch memory before H2D.", + ) + parser.add_argument( + "--torch-threads", + default="auto", + help="PyTorch intra-op threads per run. 'auto' uses cpu_count/world_size.", + ) + parser.add_argument("--seed", type=int, default=12345) + parser.add_argument( + "--device-memory-safety-factor", + type=float, + default=1.15, + help="Skip a CUDA case when free memory is below estimated bytes times this factor.", + ) + parser.add_argument("--out-json", type=Path, default=DEFAULT_OUTPUT_JSON) + parser.add_argument("--obs-dim", type=int, default=0) + parser.add_argument("--action-dim", type=int, default=0) + parser.add_argument("--critic-dim", type=int, default=0) + parser.add_argument( + "--symmetry-batch-multiplier", + type=int, + default=0, + help="Only used with manual --obs-dim/--action-dim/--critic-dim.", + ) + parser.add_argument("--plot-dir", type=Path, default=None) + parser.add_argument("--analysis-md", type=Path, default=None) + parser.add_argument("--no-plots", action="store_true") + parser.add_argument( + "--plot-only", + type=Path, + default=None, + help="Load an existing results JSON and regenerate plots/analysis without benchmarking.", + ) + args = parser.parse_args(argv) + + if args.plot_only is not None: + payload = json.loads(args.plot_only.read_text(encoding="utf-8")) + plot_dir = args.plot_dir or args.plot_only.parent + plot_paths = [] if args.no_plots else _save_plots(payload, plot_dir) + analysis_path = _write_analysis_markdown( + payload, + args.analysis_md or plot_dir / "analysis.md", + ) + for path in plot_paths: + print(f"Saved plot: {path}") + print(f"Saved analysis: {analysis_path}") + return 0 + + if args.warmup < 0 or args.repeat <= 0: + raise ValueError("--warmup must be >= 0 and --repeat must be > 0") + if args.device_memory_safety_factor < 1.0: + raise ValueError("--device-memory-safety-factor must be >= 1.0") + + cfg = _compose_offpolicy_cfg(args.sim) + manual_shape = args.obs_dim > 0 or args.action_dim > 0 or args.critic_dim > 0 + if manual_shape: + if args.obs_dim <= 0 or args.action_dim <= 0 or args.critic_dim < 0: + raise ValueError( + "Manual shape requires --obs-dim > 0, --action-dim > 0, --critic-dim >= 0" + ) + shape = ReplayShape( + obs_dim=int(args.obs_dim), + action_dim=int(args.action_dim), + critic_dim=int(args.critic_dim), + ) + if bool(OmegaConf.select(cfg, "algo.use_symmetry", default=False)): + if args.symmetry_batch_multiplier <= 0: + raise ValueError( + "Manual shape with SAC symmetry requires --symmetry-batch-multiplier" + ) + symmetry_batch_multiplier = int(args.symmetry_batch_multiplier) + else: + symmetry_batch_multiplier = 1 + else: + shape, symmetry_batch_multiplier = _resolve_env_shape_and_symmetry(cfg) + + case = _build_case( + cfg, + sim=args.sim, + shape=shape, + symmetry_batch_multiplier=symmetry_batch_multiplier, + ) + capacity_rows = _resolve_capacity_rows( + config_capacity_rows=case.config_capacity_rows, + capacity_rows_arg=args.capacity_rows, + capacity_multipliers_arg=args.capacity_multipliers, + capacity_exponents_arg=args.capacity_exponents, + ) + gpu_counts = _parse_int_list(args.gpu_counts, name="gpu counts") + device_ids = _parse_device_ids(args.device_ids) + + print("SAC Replay Buffer Sampling Benchmark") + print(f"Config command: {case.command}") + print(f"Device info: {get_device_info_line()}") + print( + "Shape: " + f"obs={case.shape.obs_dim}, action={case.shape.action_dim}, " + f"critic={case.shape.critic_dim}, packed_width={case.shape.packed_width}" + ) + print( + "Config: " + f"num_envs={case.num_envs:,}, replay_buffer_n={case.replay_buffer_n:,}, " + f"capacity={case.config_capacity_rows:,}, configured_batch={case.configured_batch_size:,}, " + f"learner_batch={case.learner_batch_size:,}, updates={case.updates_per_step}, " + f"sample_count/rank={case.sample_count_per_rank:,}" + ) + print(f"Capacity rows: {', '.join(f'{rows:,}' for rows in capacity_rows)}") + print(f"Requested GPU counts: {gpu_counts}; visible CUDA device ids: {device_ids}") + + results: list[CapacityResult] = [] + skipped: list[SkippedCase] = [] + + if not torch.cuda.is_available() or not device_ids: + for world_size in gpu_counts: + skipped.append( + SkippedCase( + world_size=world_size, + capacity_rows=None, + reason="CUDA is not available or no CUDA device ids were selected", + ) + ) + else: + for world_size in gpu_counts: + if world_size > len(device_ids): + skipped.append( + SkippedCase( + world_size=world_size, + capacity_rows=None, + reason=f"requested {world_size} GPUs but only {len(device_ids)} are visible", + ) + ) + continue + + devices = [torch.device(f"cuda:{device_id}") for device_id in device_ids[:world_size]] + torch_threads = _resolve_torch_threads(args.torch_threads, world_size=world_size) + for rows in capacity_rows: + required = _estimate_device_bytes_per_gpu(case, rows) + skip_reason = _cuda_memory_skip_reason( + devices, + required_bytes_per_gpu=required, + safety_factor=float(args.device_memory_safety_factor), + ) + if skip_reason is not None: + skipped.append( + SkippedCase( + world_size=world_size, + capacity_rows=rows, + reason=skip_reason, + ) + ) + print(f"\nSkipping {world_size} GPU(s), capacity={rows:,}: {skip_reason}") + continue + print(f"\nRunning {world_size} GPU(s), capacity={rows:,}") + try: + result = _run_capacity_case( + case, + capacity_rows=rows, + devices=devices, + warmup=args.warmup, + repeat=args.repeat, + prefill=args.prefill, + pinned_host_batch=args.host_batch_memory == "pinned", + index_mode=args.index_mode, + seed=args.seed, + torch_threads=torch_threads, + ) + except RuntimeError as exc: + if "out of memory" not in str(exc).lower(): + raise + _cleanup_device() + skipped.append( + SkippedCase( + world_size=world_size, + capacity_rows=rows, + reason=f"CUDA OOM during benchmark allocation/run: {exc}", + ) + ) + print(f" skipped: CUDA OOM during benchmark allocation/run: {exc}") + continue + results.append(result) + _print_result(result) + + payload = { + "created_at": datetime.now(timezone.utc).isoformat(), + "torch_version": torch.__version__, + "device_info": get_device_info_dict(), + "args": { + "sim": args.sim, + "gpu_counts": gpu_counts, + "device_ids": device_ids, + "capacity_rows": capacity_rows, + "capacity_rows_arg": args.capacity_rows, + "capacity_multipliers": args.capacity_multipliers, + "capacity_exponents": args.capacity_exponents, + "warmup": args.warmup, + "repeat": args.repeat, + "prefill": args.prefill, + "index_mode": args.index_mode, + "host_batch_memory": args.host_batch_memory, + "torch_threads": args.torch_threads, + "seed": args.seed, + "device_memory_safety_factor": args.device_memory_safety_factor, + "manual_shape": manual_shape, + }, + "case": case, + "results": results, + "skipped": skipped, + } + serializable_payload = _json_ready(payload) + plot_dir = args.plot_dir or args.out_json.parent + plot_paths = [] if args.no_plots else _save_plots(serializable_payload, plot_dir) + analysis_path = _write_analysis_markdown( + serializable_payload, + args.analysis_md or plot_dir / "analysis.md", + ) + serializable_payload["plots"] = plot_paths + serializable_payload["analysis_markdown"] = analysis_path + _write_results(args.out_json, serializable_payload) + for item in skipped: + suffix = "" if item.capacity_rows is None else f", capacity={item.capacity_rows:,}" + print(f"Skipped {item.world_size} GPU(s){suffix}: {item.reason}") + for path in plot_paths: + print(f"Saved plot: {path}") + print(f"Saved analysis: {analysis_path}") + print(f"\nSaved JSON: {args.out_json}") + if not results: + print("No CUDA benchmark cases were run.") + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/benchmark/test_sac_replay_buffer_sampling_benchmark.py b/tests/benchmark/test_sac_replay_buffer_sampling_benchmark.py new file mode 100644 index 000000000..a1994c401 --- /dev/null +++ b/tests/benchmark/test_sac_replay_buffer_sampling_benchmark.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import json + +import torch +from benchmark import benchmark_sac_replay_buffer_sampling as bench + + +def _tiny_case() -> bench.BenchmarkCase: + return bench.BenchmarkCase( + algo="sac", + task="g1_walk_flat", + sim="mujoco", + command="uv run train --algo sac --task g1_walk_flat --sim mujoco", + training_task_name="G1WalkFlat", + num_envs=2, + replay_buffer_n=8, + config_capacity_rows=16, + configured_batch_size=4, + learner_batch_size=4, + symmetry_batch_multiplier=1, + updates_per_step=2, + sample_count_per_rank=8, + learning_starts=0, + shape=bench.ReplayShape(obs_dim=3, action_dim=2, critic_dim=1), + ) + + +def test_sac_default_case_uses_effective_symmetry_batch() -> None: + cfg = bench._compose_offpolicy_cfg("mujoco") + case = bench._build_case( + cfg, + sim="mujoco", + shape=bench.ReplayShape(obs_dim=45, action_dim=29, critic_dim=48), + symmetry_batch_multiplier=2, + ) + + assert case.command == "uv run train --algo sac --task g1_walk_flat --sim mujoco" + assert case.config_capacity_rows == case.num_envs * case.replay_buffer_n + assert case.learner_batch_size == case.configured_batch_size // 2 + assert case.sample_count_per_rank == case.learner_batch_size * case.updates_per_step + assert case.shape.packed_width == 2 * 45 + 29 + 3 + 2 * 48 + + +def test_resolve_capacity_rows_from_multipliers_deduplicates() -> None: + assert bench._resolve_capacity_rows( + config_capacity_rows=100, + capacity_rows_arg="auto", + capacity_multipliers_arg="0.25,0.5,0.5,1", + ) == [25, 50, 100] + + +def test_resolve_capacity_rows_from_explicit_values() -> None: + assert bench._resolve_capacity_rows( + config_capacity_rows=100, + capacity_rows_arg="16,32,16", + capacity_multipliers_arg="1", + ) == [16, 32] + + +def test_parse_device_ids_auto_no_cuda(monkeypatch) -> None: + monkeypatch.setattr(bench.torch.cuda, "is_available", lambda: False) + + assert bench._parse_device_ids("auto") == [] + + +def test_run_capacity_case_portable_cpu_path_records_timings() -> None: + result = bench._run_capacity_case( + _tiny_case(), + capacity_rows=16, + devices=[torch.device("cpu")], + warmup=0, + repeat=1, + prefill="none", + pinned_host_batch=False, + index_mode="pregenerated", + seed=123, + torch_threads=1, + ) + + assert result.world_size == 1 + assert result.capacity_rows == 16 + assert result.sample_bytes_per_rank > 0 + assert set(result.timings) == { + "cpu_sample_wall", + "cpu_sample_h2d_wall", + "cpu_sample_then_h2d_wall", + "gpu_sample_wall", + } + + +def test_main_no_cuda_writes_skipped_json(tmp_path, monkeypatch) -> None: + out_json = tmp_path / "results.json" + monkeypatch.setattr(bench.torch.cuda, "is_available", lambda: False) + + rc = bench.main( + [ + "--gpu-counts", + "1,2", + "--capacity-rows", + "16", + "--warmup", + "0", + "--repeat", + "1", + "--obs-dim", + "3", + "--action-dim", + "2", + "--critic-dim", + "1", + "--symmetry-batch-multiplier", + "2", + "--out-json", + str(out_json), + ] + ) + + assert rc == 1 + payload = json.loads(out_json.read_text(encoding="utf-8")) + assert payload["results"] == [] + assert [item["world_size"] for item in payload["skipped"]] == [1, 2] From 133935703df2788a6d68e9360736c2fa045cdc3a Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Wed, 8 Jul 2026 23:27:48 +0800 Subject: [PATCH 26/26] test: move numba parity checks to slow suite --- tests/envs/locomotion/g1/test_joystick_numba.py | 2 ++ tests/envs/test_g1_motion_tracking_numba.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/envs/locomotion/g1/test_joystick_numba.py b/tests/envs/locomotion/g1/test_joystick_numba.py index de0beb7a4..608790110 100644 --- a/tests/envs/locomotion/g1/test_joystick_numba.py +++ b/tests/envs/locomotion/g1/test_joystick_numba.py @@ -119,6 +119,7 @@ def test_g1_walk_numba_unsupported_active_term_raises(): @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +@pytest.mark.slow def test_g1_walk_numba_basic_reward_parity(): n = 512 n_action = 29 @@ -184,6 +185,7 @@ def test_g1_walk_numba_basic_reward_parity(): @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +@pytest.mark.slow def test_g1_walk_numba_update_state_parity_without_noise(): from unilab.envs.locomotion.g1.joystick import G1WalkEnv diff --git a/tests/envs/test_g1_motion_tracking_numba.py b/tests/envs/test_g1_motion_tracking_numba.py index 8a16adec7..f4f3d1796 100644 --- a/tests/envs/test_g1_motion_tracking_numba.py +++ b/tests/envs/test_g1_motion_tracking_numba.py @@ -40,6 +40,7 @@ def test_numba_accelerator_rejects_unsupported_active_reward_terms(): @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +@pytest.mark.slow def test_g1_motion_tracking_numba_reward_termination_parity(): n = 512 env = _make_env(n, include_undesired=True) @@ -95,6 +96,7 @@ def test_g1_motion_tracking_numba_reward_termination_parity(): @pytest.mark.skipif(not NUMBA_AVAILABLE, reason="numba is optional") +@pytest.mark.slow def test_g1_motion_tracking_numba_update_state_parity_without_noise(): n = 512 env = _make_env(n, include_undesired=True)