diff --git a/.gitignore b/.gitignore index 1e768f36d..d47f65c9b 100644 --- a/.gitignore +++ b/.gitignore @@ -83,9 +83,7 @@ src/unilab/assets/bots/ src/unilab/assets/robots/x2/meshes/*.STL src/unilab/assets/robots/g1/assets/ src/unilab/assets/robots/g1/textures/ -src/unilab/assets/robots/go1/assets/ src/unilab/assets/robots/go2/assets/ -src/unilab/assets/robots/a2/assets/ src/unilab/assets/robots/allegro_hand/assets/ src/unilab/assets/robots/sharpa_wave/meshes/ diff --git a/docs/sphinx/source/api_reference/index.md b/docs/sphinx/source/api_reference/index.md index 5b22144d1..64d62a1c7 100644 --- a/docs/sphinx/source/api_reference/index.md +++ b/docs/sphinx/source/api_reference/index.md @@ -56,7 +56,7 @@ PPO / APPO / SAC / TD3 variants moved to the uni_rl package (issue #1480). :::{grid-item-card} 🏋 `unilab.training` :link: training/index :link-type: doc -Runtime helpers, monitoring, reward bookkeeping, runner orchestration. +Runtime helpers, reward bookkeeping, runner orchestration. ::: :::{grid-item-card} 🔗 Shared-memory runtime → `uni_rl.ipc` diff --git a/docs/sphinx/source/api_reference/training/index.md b/docs/sphinx/source/api_reference/training/index.md index 80d7238f5..3ba11acde 100644 --- a/docs/sphinx/source/api_reference/training/index.md +++ b/docs/sphinx/source/api_reference/training/index.md @@ -2,7 +2,7 @@ Glue between `algos`, `envs` and `ipc`: experiment lifecycle and the top-level `run` helpers invoked by the `train` / `eval` / `demo` CLI -entrypoints. Layer-0 helpers (seeding, monitoring, reward bookkeeping, +entrypoints. Layer-0 helpers (seeding, reward bookkeeping, checkpoint resolution, sim2sim contracts) live in `unilab.utils`; resolved env config adaptation lives in `unilab.base.config_adapter`. diff --git a/docs/sphinx/source/api_reference/utils/index.md b/docs/sphinx/source/api_reference/utils/index.md index 30a8b0133..37f37f9a1 100644 --- a/docs/sphinx/source/api_reference/utils/index.md +++ b/docs/sphinx/source/api_reference/utils/index.md @@ -1,7 +1,7 @@ # `unilab.utils` — Utilities -Device probing, tensor helpers, training seeding, hardware monitoring, -reward bookkeeping, checkpoint resolution, sim2sim contract checks, NaN +Device probing, tensor helpers, training seeding, reward bookkeeping, +checkpoint resolution, sim2sim contract checks, NaN guards, and pure-numpy geometry/rotation helpers shared across envs and scripts. diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md b/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md index 07cbd0d7a..1de946d91 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md @@ -33,8 +33,10 @@ collector sim processes land on their own physical GPU with no cross-GPU leakage; single-GPU PPO/SAC regressions pass alongside. Newton/Warp follows standard CUDA device semantics, so no `CUDA_VISIBLE_DEVICES` pinning (the Genesis quirk) is needed; the rank-local device reaches spawn collectors as -a `newton_device="cuda:N"` env override (uni_rl 1.0.0's collector-side -process-binding gate only covers mjwarp), and the SAC owner raises the +a `newton_device="cuda:N"` env override, and uni_rl's collector-side +process binding is injection-based — UniLab injects +`bind_backend_process_device_for_backend`, which covers both mjwarp and +newton — while the SAC owner raises the collector tick-0 timeout to 180 s to cover Warp kernel compilation on the cold path. diff --git a/docs/sphinx/source/en/2-user_guide/7-tooling/0-index.md b/docs/sphinx/source/en/2-user_guide/7-tooling/0-index.md index 216746f8d..db0ed2c1c 100644 --- a/docs/sphinx/source/en/2-user_guide/7-tooling/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/7-tooling/0-index.md @@ -24,12 +24,6 @@ Configure run logging and experiment metadata. Inspect NaN guard dumps from PPO runs. ::: -:::{grid-item-card} Scene export -:link: 4-scene_export -:link-type: doc -Export MuJoCo scenes and copied assets for inspection. -::: - :::{grid-item-card} Robot import :link: 5-robot_import :link-type: doc @@ -44,6 +38,5 @@ Connect robot assets, model descriptions, control interfaces, and keyframes. 1-onnx_export 2-wandb 3-nan_visualizer -4-scene_export 5-robot_import ``` diff --git a/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md b/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md deleted file mode 100644 index d849ccded..000000000 --- a/docs/sphinx/source/en/2-user_guide/7-tooling/4-scene_export.md +++ /dev/null @@ -1,16 +0,0 @@ -# Scene Export - -Scene export is implemented by `unisim.backend.mujoco.export_scene` and registered -as the `unilab-export-scene` console entry in `pyproject.toml`. It accepts a -MuJoCo XML or MJB model path, writes `scene.xml`, copies mesh assets when they -are discoverable, and can create a zip archive. - -For task-level materialization checks, use the script that constructs an env -from the registry and owner config: - -```bash -uv run scripts/visualize_task_env.py --task G1WalkFlat --backend mujoco --num_envs 4 -``` - -`tests/test_export_scene.py` covers the export helper, including `scene.xml` -creation, reloadability, and zip output. diff --git a/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md b/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md index 89bd1d9dd..245e9e903 100644 --- a/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md +++ b/docs/sphinx/source/en/3-deployment/1-sim_to_real/9-troubleshooting.md @@ -24,8 +24,8 @@ sideways, start here. Almost always one of: 1. **Joint order swapped.** Inspect `policy.onnx` input width and the joint - order in your motor driver. Use `unilab-export-scene` to dump the - training joint order. + order in your motor driver against the training scene XML + (`src/unilab/assets/robots//scene_flat.xml`). 2. **Action scale unit mismatch.** Policy outputs unscaled values; the driver expects rad, but you fed it normalized [-1, 1]. Apply the `env.actions.joint_pos.scale` / default-angle convention from the training diff --git a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md index 7990557d7..8868e3940 100644 --- a/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/en/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -36,7 +36,7 @@ mostly mechanical. 3. Mirror your reward terms; keep the same names so reward parity is diff-able. 4. Translate command sampling — configure `UniformVelocityCommandCfg` under - the owner YAML's `env.commands` (see the Go1 flat owner). + the owner YAML's `env.commands` (see the `go2_joystick_flat` owner). 5. Translate terrain — Legged Gym's heightfield generator has a UniLab counterpart at `unilab.terrains.heightfield_terrains`. diff --git a/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md b/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md index aba440d8d..3419a3ab2 100644 --- a/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md +++ b/docs/sphinx/source/en/4-developer_guide/3-extending/4-new_terrain.md @@ -25,8 +25,6 @@ materialization out of `step()`, `reset()`, and hot domain-randomization loops. - Terrain generator shape and numerical behavior: `tests/terrains/test_terrain_generator.py` - Rough locomotion height-scan and spawn behavior: - `tests/envs/locomotion/test_go2_rough_height_scan.py`, - `tests/envs/locomotion/test_go2_terrain_spawn.py`, `tests/envs/locomotion/test_terrain_spawn.py` - Backend materialization boundaries: `tests/utils/test_xml_utils.py` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md index 6e3f42e1f..6c57cbe58 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md @@ -31,9 +31,10 @@ DpRankSupervisor DP=2（`training.devices=[0,1]`）训练冒烟均正常完成 在各自 GPU、无对端泄漏；单卡 PPO/SAC 回归同步通过。Newton/Warp 遵循标准 CUDA 设备语义，无需 Genesis 那样的 `CUDA_VISIBLE_DEVICES` 钉卡；rank 本地设备通过 env override 以 `newton_device="cuda:N"` 传入 spawn -collector（uni_rl 1.0.0 的 collector 进程绑定门只覆盖 mjwarp），SAC -owner 将 collector tick-0 超时提升到 180 s 以覆盖 Warp 内核编译的冷 -路径。 +collector，且 uni_rl 的 collector 进程绑定为注入式——UniLab 注入的 +`bind_backend_process_device_for_backend` 同时覆盖 mjwarp 与 +newton；SAC owner 将 collector tick-0 超时提升到 180 s 以覆盖 Warp +内核编译的冷路径。 ## 安装 diff --git a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/0-index.md index 235cfb66e..d57a636a4 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/0-index.md @@ -23,12 +23,6 @@ 检查 PPO 运行中的 NaN guard dump。 ::: -:::{grid-item-card} 场景导出 -:link: 4-scene_export -:link-type: doc -导出 MuJoCo 场景及复制的 asset 以供检查。 -::: - :::{grid-item-card} 机器人导入 :link: 5-robot_import :link-type: doc @@ -43,6 +37,5 @@ 1-onnx_export 2-wandb 3-nan_visualizer -4-scene_export 5-robot_import ``` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md b/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md deleted file mode 100644 index 44694a8b6..000000000 --- a/docs/sphinx/source/zh_CN/2-user_guide/7-tooling/4-scene_export.md +++ /dev/null @@ -1,11 +0,0 @@ -# 场景导出 - -场景导出由 `unisim.backend.mujoco.export_scene` 实现，并在 `pyproject.toml` 中注册为 `unilab-export-scene` 控制台入口。它接受一个 MuJoCo XML 或 MJB 模型路径，写出 `scene.xml`，在能够发现 mesh asset 时复制它们，并且可以创建一个 zip 归档。 - -对于 task 级别的实例化检查，请使用从 registry 和 owner config 构造 env 的脚本： - -```bash -uv run scripts/visualize_task_env.py --task G1WalkFlat --backend mujoco --num_envs 4 -``` - -`tests/test_export_scene.py` 覆盖了导出辅助逻辑，包括 `scene.xml` 的创建、可重新加载性以及 zip 输出。 diff --git a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md index 33a9f1077..355b3f93a 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md +++ b/docs/sphinx/source/zh_CN/3-deployment/1-sim_to_real/9-troubleshooting.md @@ -22,8 +22,8 @@ 几乎总是以下之一： -1. **关节顺序被调换。** 检查 `policy.onnx` 的输入宽度与你电机驱动器中的关节顺序。 - 用 `unilab-export-scene` 导出训练时的关节顺序。 +1. **关节顺序被调换。** 检查 `policy.onnx` 的输入宽度，并将你电机驱动器中的关节顺序 + 与训练场景 XML（`src/unilab/assets/robots//scene_flat.xml`）对照。 2. **动作缩放单位不匹配。** 策略输出未缩放的值；驱动器期望的是弧度，而你喂给它的 是归一化的 [-1, 1]。在把目标发送给驱动器之前，应用训练 owner YAML 中的 `env.actions.joint_pos.scale` / 默认角度约定，并原样复现该 owner 解析后的 diff --git a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md index 2b5d3e84b..3486f2ea4 100644 --- a/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md +++ b/docs/sphinx/source/zh_CN/3-deployment/3-framework_migration/2-from_legged_gym.md @@ -32,7 +32,7 @@ Legged Gym 曾是那套 GPU 常驻的 PPO 模板，教会了整个领域如何 2. 在 `src/unilab/tasks/locomotion//` 下创建一个任务模块。 3. 镜像你的 reward 项；保持名称相同，以便 reward 一致性可被 diff。 4. 翻译命令采样 —— 在 owner YAML 的 `env.commands` 下配置 - `UniformVelocityCommandCfg`，参考 Go1 flat owner。 + `UniformVelocityCommandCfg`，参考 `go2_joystick_flat` owner。 5. 翻译地形 —— Legged Gym 的高度场生成器在 UniLab 中有一个对应物，位于 `unilab.terrains.heightfield_terrains`。 diff --git a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md index 47f0d958b..563878a64 100644 --- a/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md +++ b/docs/sphinx/source/zh_CN/4-developer_guide/3-extending/4-new_terrain.md @@ -25,8 +25,6 @@ - 地形生成器形状与数值行为： `tests/terrains/test_terrain_generator.py` - Rough locomotion 高度扫描与 spawn 行为： - `tests/envs/locomotion/test_go2_rough_height_scan.py`、 - `tests/envs/locomotion/test_go2_terrain_spawn.py`、 `tests/envs/locomotion/test_terrain_spawn.py` - 后端 materialization 边界：`tests/utils/test_xml_utils.py` diff --git a/pyproject.rocm.toml b/pyproject.rocm.toml index 10536bdd5..93770addb 100644 --- a/pyproject.rocm.toml +++ b/pyproject.rocm.toml @@ -24,13 +24,13 @@ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", # Physics implementations are provided by the independently released - # unisim-core package. The 1.4.1 release carries the fixed model variant, + # unisim-core package. The 1.4.2 release carries the fixed model variant, # per-world reset-default, per-env gravity, substep body-wrench, and # IsaacGym fixed-variant contracts. - "unisim-core>=1.4.1", + "unisim-core>=1.4.2", # RL algorithms and async runtimes live in the independently released # uni-rl package (distribution name ``unilab-rl``); see pyproject.toml. - "unilab-rl==1.2.0", + "unilab-rl==1.2.1", "torch==2.11.0", "triton-rocm==3.6.0 ; sys_platform == 'linux' and platform_machine == 'x86_64'", "gymnasium", @@ -61,7 +61,6 @@ eval = "unilab.cli:eval_main" demo = "unilab.cli:demo_main" unilab-complete = "unilab.cli_completion:main" unilab-viz-nan = "unilab.utils.nan_viz:main" -unilab-export-scene = "unisim.backend.mujoco.export_scene:main" unilab-render-teaser = "unilab.visualization.teaser:main" unilab-pull-assets = "unilab.assets.pull:main" @@ -163,7 +162,6 @@ venv = ".venv" include = ["src/unilab"] exclude = [ "src/unilab/envs/", # lazy optional `mujoco` import; not installed in the pyright env - "src/unilab/utils/monitoring.py", # optional pynvml/psutil deps "src/unilab/visualization/", # direct mujoco C bindings + optional viser deps ] reportMissingImports = "warning" diff --git a/pyproject.toml b/pyproject.toml index 8e63ba820..5637d1bfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,15 +39,15 @@ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", # Physics implementations are provided by the independently released - # unisim-core package. The 1.4.1 release carries the fixed model variant, + # unisim-core package. The 1.4.2 release carries the fixed model variant, # per-world reset-default, per-env gravity, substep body-wrench, and # IsaacGym fixed-variant contracts. - "unisim-core>=1.4.1", + "unisim-core>=1.4.2", # RL algorithms and async runtimes (PPO/APPO/SAC/TD3 runners, # collectors, IPC, logging) live in the independently released uni-rl # package (distribution name ``unilab-rl``), consumed via the injected # env contract (uni_rl.env_contract.EnvFactory). Published on PyPI. - "unilab-rl==1.2.0", + "unilab-rl==1.2.1", "numba>=0.67", "prettytable>=3.10", # torch is a range (not an exact pin) so that published PyPI metadata lets @@ -89,7 +89,6 @@ eval = "unilab.cli:eval_main" demo = "unilab.cli:demo_main" unilab-complete = "unilab.cli_completion:main" unilab-viz-nan = "unilab.utils.nan_viz:main" -unilab-export-scene = "unisim.backend.mujoco.export_scene:main" unilab-render-teaser = "unilab.visualization.teaser:main" unilab-pull-assets = "unilab.assets.pull:main" @@ -171,7 +170,7 @@ viser = ["viser>=1.0.26", "trimesh>=3.21.7"] # required-environments; elsewhere the extra is empty and the CLI reports a # targeted runtime diagnostic. superdex = [ - "unisim-core[superdex]>=1.4.0 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", + "unisim-core[superdex]>=1.4.2 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'", ] [dependency-groups] @@ -250,7 +249,6 @@ venv = ".venv" include = ["src/unilab"] exclude = [ "src/unilab/envs/", # lazy optional `mujoco` import; not installed in the pyright env - "src/unilab/utils/monitoring.py", # optional pynvml/psutil deps "src/unilab/visualization/", # direct mujoco C bindings + optional viser deps ] reportMissingImports = "warning" diff --git a/scripts/benchmark/env/benchmark_postprocess.py b/scripts/benchmark/env/benchmark_postprocess.py deleted file mode 100644 index 1b5aff3bd..000000000 --- a/scripts/benchmark/env/benchmark_postprocess.py +++ /dev/null @@ -1,1149 +0,0 @@ -import argparse -import json -import math -import sys -import time -from datetime import datetime -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -import matplotlib - -matplotlib.use("Agg") -import matplotlib.pyplot as plt -import mlx.core as mx -import numpy as np -import torch -from mujoco import rollout as mj_rollout - -from unilab.base import registry - - -def ensure_registries() -> None: - """Import locomotion env modules so they are registered.""" - - -from scripts.benchmark.core.device_info import get_device_info_dict, get_device_info_line -from scripts.benchmark.core.task_names import locomotion_env_name - -DEFAULT_ENV_LIST = [256, 512, 1024, 2048, 4096] -DEFAULT_ITERS = 200 -OUTPUT_DIR = Path("scripts/benchmark/outputs/postprocess") -OUTPUT_JSON = OUTPUT_DIR / "latest_postprocess_benchmark.json" -OUTPUT_PNG = OUTPUT_DIR / "latest_postprocess_latency.png" -TORCH_DEVICE = "mps" -OWNER_TASK_ID = "go2_joystick_flat" -ENV_TASK_NAME = locomotion_env_name(OWNER_TASK_ID) - - -def sync_torch_mps(): - if torch.backends.mps.is_available(): - torch.mps.synchronize() - - -def parse_env_list(raw: str) -> list[int]: - if not raw: - return DEFAULT_ENV_LIST - return [int(x.strip()) for x in raw.split(",") if x.strip()] - - -def geomean(values: list[float]) -> float: - vals = [v for v in values if v > 0.0] - if not vals: - return 0.0 - return float(math.exp(sum(math.log(v) for v in vals) / len(vals))) - - -def build_go1_layout() -> dict: - env = registry.make(ENV_TASK_NAME, num_envs=1, sim_backend="mujoco") - layout = { - "sensor_dim": int(env.model.nsensordata), - "physics_dim": int(env.physics_state_dim), - "num_action": int(env.action_space.shape[0]), - "idx_linvel": env.idx_linvel, - "idx_gyro": env.idx_gyro, - "idx_global_linvel": env.idx_global_linvel, - "idx_upvector": env.idx_upvector, - "idx_qpos": int(env._idx_qpos), - "idx_qvel": int(env._idx_qvel), - "nq": int(env.nq), - "nv": int(env.nv), - "default_angles": np.asarray(env.default_angles, dtype=np.float32).copy(), - "tracking_sigma": float(env.cfg.reward_config.tracking_sigma), - "base_height_target": float(env.cfg.reward_config.base_height_target), - "ctrl_dt": float(env.cfg.ctrl_dt), - "reward_scales": dict(env.cfg.reward_config.scales), - "obs_dim": int(env.observation_space.shape[0]), - "command_low": np.asarray(env.cfg.commands.vel_limit[0], dtype=np.float32), - "command_high": np.asarray(env.cfg.commands.vel_limit[1], dtype=np.float32), - } - env.close() - return layout - - -def measure_physics_step_ms(num_envs: int, iters: int) -> float: - env = registry.make(ENV_TASK_NAME, num_envs=num_envs, sim_backend="mujoco") - try: - initial_state, _, _ = env.reset(np.arange(env.num_envs)) - action_low = env.action_space.low.astype(np.float32) - action_high = env.action_space.high.astype(np.float32) - actions = np.random.uniform( - action_low, action_high, size=(env.num_envs, env.action_space.shape[0]) - ).astype(np.float32) - - for _ in range(20): - _ = env.step(actions) - - elapsed = 0.0 - for _ in range(iters): - t0 = time.perf_counter() - _ = env.step(actions) - t1 = time.perf_counter() - elapsed += t1 - t0 - return elapsed / iters * 1000.0 - finally: - env.close() - - -def measure_physics_step_rollout_ms(num_envs: int, iters: int) -> float: - """Measure pure MuJoCo physics stepping via mujoco.rollout.""" - env = registry.make(ENV_TASK_NAME, num_envs=num_envs, sim_backend="mujoco") - try: - initial_state, _, _ = env.reset(np.arange(env.num_envs)) - action_low = env.action_space.low.astype(np.float32) - action_high = env.action_space.high.astype(np.float32) - actions_np = np.random.uniform( - action_low, - action_high, - size=(env.num_envs, env.action_space.shape[0]), - ).astype(np.float32) - control_np = np.broadcast_to( - actions_np[:, None, :], - (env.num_envs, env.cfg.sim_substeps, env.action_space.shape[0]), - ).copy() - state_buf = np.empty( - (env.num_envs, env.cfg.sim_substeps, env.physics_state_dim), dtype=np.float32 - ) - sensor_buf = np.empty( - (env.num_envs, env.cfg.sim_substeps, env.model.nsensordata), dtype=np.float32 - ) - model_batch = [env._model] * env.num_envs - with mj_rollout.Rollout(nthread=env._n_threads) as runner: - for _ in range(20): - runner.rollout( - model_batch, - env._worker_data, - initial_state, - control_np, - nstep=env.cfg.sim_substeps, - state=state_buf, - sensordata=sensor_buf, - ) - initial_state = state_buf[:, -1, :].copy() - - elapsed = 0.0 - for _ in range(iters): - t0 = time.perf_counter() - runner.rollout( - model_batch, - env._worker_data, - initial_state, - control_np, - nstep=env.cfg.sim_substeps, - state=state_buf, - sensordata=sensor_buf, - ) - t1 = time.perf_counter() - elapsed += t1 - t0 - initial_state = state_buf[:, -1, :].copy() - return elapsed / iters * 1000.0 - finally: - env.close() - - -def measure_rollout_bridge_mlx_pipeline_ms( - num_envs: int, - iters: int, - idx_mx: dict, - scalars_mx: dict, -) -> dict: - """End-to-end: MLX action -> MuJoCo rollout -> MLX sensordata/reward.""" - env = registry.make(ENV_TASK_NAME, num_envs=num_envs, sim_backend="mujoco") - try: - initial_state, _, reset_info = env.reset(np.arange(env.num_envs)) - action_low = env.action_space.low.astype(np.float32) - action_high = env.action_space.high.astype(np.float32) - actions_np = np.random.uniform( - action_low, action_high, size=(env.num_envs, env.action_space.shape[0]) - ).astype(np.float32) - actions_mx = mx.array(actions_np, dtype=mx.float32) - last_actions_mx = mx.zeros_like(actions_mx) - commands_mx = mx.array(reset_info["commands"], dtype=mx.float32) - - # Zero-order hold across MuJoCo substeps. - control_mx = mx.broadcast_to( - actions_mx[:, None, :], (env.num_envs, env.cfg.sim_substeps, env.action_space.shape[0]) - ) - control_np = np.array(control_mx) - state_buf = np.empty( - (env.num_envs, env.cfg.sim_substeps, env.physics_state_dim), dtype=np.float32 - ) - sensor_buf = np.empty( - (env.num_envs, env.cfg.sim_substeps, env.model.nsensordata), dtype=np.float32 - ) - mx.eval(control_mx, commands_mx, last_actions_mx) - model_batch = [env._model] * env.num_envs - with mj_rollout.Rollout(nthread=env._n_threads) as runner: - for _ in range(10): - runner.rollout( - model_batch, - env._worker_data, - initial_state=initial_state, - control=control_np, - nstep=env.cfg.sim_substeps, - state=state_buf, - sensordata=sensor_buf, - ) - last_state = state_buf[:, -1, :].copy() - last_sensor = sensor_buf[:, -1, :] - last_state_mx = mx.array(last_state, dtype=mx.float32) - last_sensor_mx = mx.array(last_sensor, dtype=mx.float32) - obs_mx, rew_mx, done_mx = mlx_postprocess_go1( - sensor_mx=last_sensor_mx, - physics_mx=last_state_mx, - current_mx=actions_mx, - last_mx=last_actions_mx, - commands_mx=commands_mx, - idx_mx=idx_mx, - scalars=scalars_mx, - ) - mx.eval(obs_mx, rew_mx, done_mx) - initial_state = last_state - - rollout_elapsed = 0.0 - post_elapsed = 0.0 - for _ in range(iters): - t0 = time.perf_counter() - runner.rollout( - model_batch, - env._worker_data, - initial_state=initial_state, - control=control_np, - nstep=env.cfg.sim_substeps, - state=state_buf, - sensordata=sensor_buf, - ) - t1 = time.perf_counter() - last_state = state_buf[:, -1, :].copy() - last_sensor = sensor_buf[:, -1, :] - last_state_mx = mx.array(last_state, dtype=mx.float32) - last_sensor_mx = mx.array(last_sensor, dtype=mx.float32) - obs_mx, rew_mx, done_mx = mlx_postprocess_go1( - sensor_mx=last_sensor_mx, - physics_mx=last_state_mx, - current_mx=actions_mx, - last_mx=last_actions_mx, - commands_mx=commands_mx, - idx_mx=idx_mx, - scalars=scalars_mx, - ) - mx.eval(obs_mx, rew_mx, done_mx) - t2 = time.perf_counter() - rollout_elapsed += t1 - t0 - post_elapsed += t2 - t1 - initial_state = last_state - - rollout_ms = rollout_elapsed / iters * 1000.0 - post_ms = post_elapsed / iters * 1000.0 - return { - "rollout_with_mlx_action_ms": rollout_ms, - "mlx_postprocess_from_rollout_ms": post_ms, - "total_ms": rollout_ms + post_ms, - } - finally: - env.close() - - -def numpy_postprocess_go1( - sensor_data: np.ndarray, - physics_state: np.ndarray, - current_actions: np.ndarray, - last_actions: np.ndarray, - commands: np.ndarray, - layout: dict, -): - idx_lin = layout["idx_linvel"] - idx_gyro = layout["idx_gyro"] - idx_glin = layout["idx_global_linvel"] - idx_up = layout["idx_upvector"] - idx_qpos = layout["idx_qpos"] - idx_qvel = layout["idx_qvel"] - nq = layout["nq"] - nv = layout["nv"] - default_angles = layout["default_angles"] - scales = layout["reward_scales"] - - linear_vel = sensor_data[:, idx_lin] - gyro = sensor_data[:, idx_gyro] - global_linvel = sensor_data[:, idx_glin] - upvector = sensor_data[:, idx_up] - local_gravity = -upvector - - dof_pos = physics_state[:, idx_qpos + 7 : idx_qpos + nq] - dof_vel = physics_state[:, idx_qvel + 6 : idx_qvel + nv] - diff = dof_pos - default_angles - - obs = np.hstack( - [ - linear_vel, - gyro, - local_gravity, - diff, - dof_vel, - current_actions, - commands, - ] - ).astype(np.float32, copy=False) - - tracking_sigma = layout["tracking_sigma"] - tracking_lin_vel = np.exp( - -np.sum(np.square(commands[:, :2] - linear_vel[:, :2]), axis=1) / tracking_sigma - ) - tracking_ang_vel = np.exp(-np.square(commands[:, 2] - gyro[:, 2]) / tracking_sigma) - lin_vel_z = np.square(global_linvel[:, 2]) - ang_vel_xy = np.sum(np.square(gyro[:, :2]), axis=1) - base_height = np.square(physics_state[:, idx_qpos + 2] - layout["base_height_target"]) - action_rate = np.sum(np.square(current_actions - last_actions), axis=1) - similar_to_default = np.sum(np.abs(diff), axis=1) - - reward = ( - scales.get("tracking_lin_vel", 0.0) * tracking_lin_vel - + scales.get("tracking_ang_vel", 0.0) * tracking_ang_vel - + scales.get("lin_vel_z", 0.0) * lin_vel_z - + scales.get("ang_vel_xy", 0.0) * ang_vel_xy - + scales.get("base_height", 0.0) * base_height - + scales.get("action_rate", 0.0) * action_rate - + scales.get("similar_to_default", 0.0) * similar_to_default - ) - reward = (reward * layout["ctrl_dt"]).astype(np.float32, copy=False) - done = (upvector[:, 2] <= 0.5).astype(np.bool_) - return obs, reward, done - - -def torch_postprocess_go1( - sensor_t: torch.Tensor, - physics_t: torch.Tensor, - current_t: torch.Tensor, - last_t: torch.Tensor, - commands_t: torch.Tensor, - idx_t: dict, - scalars_t: dict, -): - linear_vel = sensor_t[:, idx_t["lin"]] - gyro = sensor_t[:, idx_t["gyro"]] - global_linvel = sensor_t[:, idx_t["glin"]] - upvector = sensor_t[:, idx_t["up"]] - local_gravity = -upvector - - dof_pos = physics_t[:, idx_t["qpos_start"] : idx_t["qpos_end"]] - dof_vel = physics_t[:, idx_t["qvel_start"] : idx_t["qvel_end"]] - diff = dof_pos - idx_t["default_angles"] - - obs = torch.hstack([linear_vel, gyro, local_gravity, diff, dof_vel, current_t, commands_t]) - - tracking_lin_vel = torch.exp( - -torch.sum((commands_t[:, :2] - linear_vel[:, :2]) ** 2, dim=1) - / scalars_t["tracking_sigma"] - ) - tracking_ang_vel = torch.exp( - -((commands_t[:, 2] - gyro[:, 2]) ** 2) / scalars_t["tracking_sigma"] - ) - lin_vel_z = global_linvel[:, 2] ** 2 - ang_vel_xy = torch.sum(gyro[:, :2] ** 2, dim=1) - base_height = (physics_t[:, idx_t["base_height_idx"]] - scalars_t["base_height_target"]) ** 2 - action_rate = torch.sum((current_t - last_t) ** 2, dim=1) - similar_to_default = torch.sum(torch.abs(diff), dim=1) - - reward = ( - scalars_t["tracking_lin_vel"] * tracking_lin_vel - + scalars_t["tracking_ang_vel"] * tracking_ang_vel - + scalars_t["lin_vel_z"] * lin_vel_z - + scalars_t["ang_vel_xy"] * ang_vel_xy - + scalars_t["base_height"] * base_height - + scalars_t["action_rate"] * action_rate - + scalars_t["similar_to_default"] * similar_to_default - ) * scalars_t["ctrl_dt"] - - done = upvector[:, 2] <= 0.5 - return obs, reward, done - - -def mlx_postprocess_go1( - sensor_mx: mx.array, - physics_mx: mx.array, - current_mx: mx.array, - last_mx: mx.array, - commands_mx: mx.array, - idx_mx: dict, - scalars: dict, -): - linear_vel = sensor_mx[:, idx_mx["lin"]] - gyro = sensor_mx[:, idx_mx["gyro"]] - global_linvel = sensor_mx[:, idx_mx["glin"]] - upvector = sensor_mx[:, idx_mx["up"]] - local_gravity = -upvector - - dof_pos = physics_mx[:, idx_mx["qpos_start"] : idx_mx["qpos_end"]] - dof_vel = physics_mx[:, idx_mx["qvel_start"] : idx_mx["qvel_end"]] - diff = dof_pos - idx_mx["default_angles"] - - obs = mx.concatenate( - [linear_vel, gyro, local_gravity, diff, dof_vel, current_mx, commands_mx], axis=1 - ) - - tracking_lin_vel = mx.exp( - -mx.sum((commands_mx[:, :2] - linear_vel[:, :2]) ** 2, axis=1) / scalars["tracking_sigma"] - ) - tracking_ang_vel = mx.exp(-((commands_mx[:, 2] - gyro[:, 2]) ** 2) / scalars["tracking_sigma"]) - lin_vel_z = global_linvel[:, 2] ** 2 - ang_vel_xy = mx.sum(gyro[:, :2] ** 2, axis=1) - base_height = (physics_mx[:, idx_mx["base_height_idx"]] - scalars["base_height_target"]) ** 2 - action_rate = mx.sum((current_mx - last_mx) ** 2, axis=1) - similar_to_default = mx.sum(mx.abs(diff), axis=1) - - reward = ( - scalars["tracking_lin_vel"] * tracking_lin_vel - + scalars["tracking_ang_vel"] * tracking_ang_vel - + scalars["lin_vel_z"] * lin_vel_z - + scalars["ang_vel_xy"] * ang_vel_xy - + scalars["base_height"] * base_height - + scalars["action_rate"] * action_rate - + scalars["similar_to_default"] * similar_to_default - ) * scalars["ctrl_dt"] - done = upvector[:, 2] <= 0.5 - return obs, reward, done - - -def bench_one_envnum(num_envs: int, iters: int, layout: dict): - try: - physics_step_env_wrapper_ms = measure_physics_step_ms(num_envs=num_envs, iters=iters) - except Exception: - physics_step_env_wrapper_ms = 0.0 - physics_step_rollout_ms = measure_physics_step_rollout_ms(num_envs=num_envs, iters=iters) - - sensor_np = np.random.randn(num_envs, layout["sensor_dim"]).astype(np.float32) - physics_np = np.random.randn(num_envs, layout["physics_dim"]).astype(np.float32) - current_actions_np = np.random.uniform(-1.0, 1.0, size=(num_envs, layout["num_action"])).astype( - np.float32 - ) - last_actions_np = np.random.uniform(-1.0, 1.0, size=(num_envs, layout["num_action"])).astype( - np.float32 - ) - commands_np = np.random.uniform( - low=layout["command_low"], - high=layout["command_high"], - size=(num_envs, 3), - ).astype(np.float32) - - idx_t = { - "lin": torch.as_tensor(layout["idx_linvel"], device=TORCH_DEVICE, dtype=torch.long), - "gyro": torch.as_tensor(layout["idx_gyro"], device=TORCH_DEVICE, dtype=torch.long), - "glin": torch.as_tensor(layout["idx_global_linvel"], device=TORCH_DEVICE, dtype=torch.long), - "up": torch.as_tensor(layout["idx_upvector"], device=TORCH_DEVICE, dtype=torch.long), - "qpos_start": int(layout["idx_qpos"] + 7), - "qpos_end": int(layout["idx_qpos"] + layout["nq"]), - "qvel_start": int(layout["idx_qvel"] + 6), - "qvel_end": int(layout["idx_qvel"] + layout["nv"]), - "base_height_idx": int(layout["idx_qpos"] + 2), - "default_angles": torch.as_tensor( - layout["default_angles"], device=TORCH_DEVICE, dtype=torch.float32 - ), - } - scalars_t = { - "tracking_sigma": torch.tensor( - layout["tracking_sigma"], device=TORCH_DEVICE, dtype=torch.float32 - ), - "base_height_target": torch.tensor( - layout["base_height_target"], device=TORCH_DEVICE, dtype=torch.float32 - ), - "ctrl_dt": torch.tensor(layout["ctrl_dt"], device=TORCH_DEVICE, dtype=torch.float32), - "tracking_lin_vel": torch.tensor( - layout["reward_scales"].get("tracking_lin_vel", 0.0), device=TORCH_DEVICE - ), - "tracking_ang_vel": torch.tensor( - layout["reward_scales"].get("tracking_ang_vel", 0.0), device=TORCH_DEVICE - ), - "lin_vel_z": torch.tensor( - layout["reward_scales"].get("lin_vel_z", 0.0), device=TORCH_DEVICE - ), - "ang_vel_xy": torch.tensor( - layout["reward_scales"].get("ang_vel_xy", 0.0), device=TORCH_DEVICE - ), - "base_height": torch.tensor( - layout["reward_scales"].get("base_height", 0.0), device=TORCH_DEVICE - ), - "action_rate": torch.tensor( - layout["reward_scales"].get("action_rate", 0.0), device=TORCH_DEVICE - ), - "similar_to_default": torch.tensor( - layout["reward_scales"].get("similar_to_default", 0.0), - device=TORCH_DEVICE, - ), - } - - idx_mx = { - "lin": mx.array(layout["idx_linvel"], dtype=mx.int32), - "gyro": mx.array(layout["idx_gyro"], dtype=mx.int32), - "glin": mx.array(layout["idx_global_linvel"], dtype=mx.int32), - "up": mx.array(layout["idx_upvector"], dtype=mx.int32), - "qpos_start": int(layout["idx_qpos"] + 7), - "qpos_end": int(layout["idx_qpos"] + layout["nq"]), - "qvel_start": int(layout["idx_qvel"] + 6), - "qvel_end": int(layout["idx_qvel"] + layout["nv"]), - "base_height_idx": int(layout["idx_qpos"] + 2), - "default_angles": mx.array(layout["default_angles"], dtype=mx.float32), - } - scalars_mx = { - "tracking_sigma": mx.array(layout["tracking_sigma"], dtype=mx.float32), - "base_height_target": mx.array(layout["base_height_target"], dtype=mx.float32), - "ctrl_dt": mx.array(layout["ctrl_dt"], dtype=mx.float32), - "tracking_lin_vel": mx.array( - layout["reward_scales"].get("tracking_lin_vel", 0.0), dtype=mx.float32 - ), - "tracking_ang_vel": mx.array( - layout["reward_scales"].get("tracking_ang_vel", 0.0), dtype=mx.float32 - ), - "lin_vel_z": mx.array(layout["reward_scales"].get("lin_vel_z", 0.0), dtype=mx.float32), - "ang_vel_xy": mx.array(layout["reward_scales"].get("ang_vel_xy", 0.0), dtype=mx.float32), - "base_height": mx.array(layout["reward_scales"].get("base_height", 0.0), dtype=mx.float32), - "action_rate": mx.array(layout["reward_scales"].get("action_rate", 0.0), dtype=mx.float32), - "similar_to_default": mx.array( - layout["reward_scales"].get("similar_to_default", 0.0), dtype=mx.float32 - ), - } - try: - bridge_mlx = measure_rollout_bridge_mlx_pipeline_ms( - num_envs=num_envs, - iters=iters, - idx_mx=idx_mx, - scalars_mx=scalars_mx, - ) - except Exception: - bridge_mlx = { - "rollout_with_mlx_action_ms": 0.0, - "mlx_postprocess_from_rollout_ms": 0.0, - "total_ms": 0.0, - } - - # Warmup - for _ in range(20): - obs_np, rew_np, done_np = numpy_postprocess_go1( - sensor_np, physics_np, current_actions_np, last_actions_np, commands_np, layout - ) - _ = torch.as_tensor(obs_np, device=TORCH_DEVICE, dtype=torch.float32) - _ = torch.as_tensor(rew_np, device=TORCH_DEVICE, dtype=torch.float32) - _ = torch.as_tensor(done_np, device=TORCH_DEVICE, dtype=torch.bool) - for _ in range(20): - sensor_t = torch.as_tensor(sensor_np, device=TORCH_DEVICE, dtype=torch.float32) - physics_t = torch.as_tensor(physics_np, device=TORCH_DEVICE, dtype=torch.float32) - cur_t = torch.as_tensor(current_actions_np, device=TORCH_DEVICE, dtype=torch.float32) - last_t = torch.as_tensor(last_actions_np, device=TORCH_DEVICE, dtype=torch.float32) - cmd_t = torch.as_tensor(commands_np, device=TORCH_DEVICE, dtype=torch.float32) - _ = torch_postprocess_go1(sensor_t, physics_t, cur_t, last_t, cmd_t, idx_t, scalars_t) - for _ in range(20): - sensor_mx = mx.array(sensor_np, dtype=mx.float32) - physics_mx = mx.array(physics_np, dtype=mx.float32) - cur_mx = mx.array(current_actions_np, dtype=mx.float32) - last_mx = mx.array(last_actions_np, dtype=mx.float32) - cmd_mx = mx.array(commands_np, dtype=mx.float32) - _ = mlx_postprocess_go1(sensor_mx, physics_mx, cur_mx, last_mx, cmd_mx, idx_mx, scalars_mx) - sync_torch_mps() - - cpu_compute = 0.0 - cpu_transfer = 0.0 - cpu_to_mlx_compute = 0.0 - cpu_to_mlx_transfer = 0.0 - torch_transfer = 0.0 - torch_compute = 0.0 - mlx_transfer = 0.0 - mlx_compute = 0.0 - mlx_torch_transfer = 0.0 - mlx_torch_compute = 0.0 - mlx_to_torch_transfer = 0.0 - - for _ in range(iters): - t0 = time.perf_counter() - obs_np, rew_np, done_np = numpy_postprocess_go1( - sensor_np, physics_np, current_actions_np, last_actions_np, commands_np, layout - ) - t1 = time.perf_counter() - _ = torch.as_tensor(obs_np, device=TORCH_DEVICE, dtype=torch.float32) - _ = torch.as_tensor(rew_np, device=TORCH_DEVICE, dtype=torch.float32) - _ = torch.as_tensor(done_np, device=TORCH_DEVICE, dtype=torch.bool) - sync_torch_mps() - t2 = time.perf_counter() - cpu_compute += t1 - t0 - cpu_transfer += t2 - t1 - - for _ in range(iters): - t0 = time.perf_counter() - obs_np, rew_np, done_np = numpy_postprocess_go1( - sensor_np, physics_np, current_actions_np, last_actions_np, commands_np, layout - ) - t1 = time.perf_counter() - obs_mx = mx.array(obs_np, dtype=mx.float32) - rew_mx = mx.array(rew_np, dtype=mx.float32) - done_mx = mx.array(done_np, dtype=mx.bool_) - mx.eval(obs_mx, rew_mx, done_mx) - t2 = time.perf_counter() - cpu_to_mlx_compute += t1 - t0 - cpu_to_mlx_transfer += t2 - t1 - - for _ in range(iters): - t0 = time.perf_counter() - sensor_t = torch.as_tensor(sensor_np, device=TORCH_DEVICE, dtype=torch.float32) - physics_t = torch.as_tensor(physics_np, device=TORCH_DEVICE, dtype=torch.float32) - cur_t = torch.as_tensor(current_actions_np, device=TORCH_DEVICE, dtype=torch.float32) - last_t = torch.as_tensor(last_actions_np, device=TORCH_DEVICE, dtype=torch.float32) - cmd_t = torch.as_tensor(commands_np, device=TORCH_DEVICE, dtype=torch.float32) - sync_torch_mps() - t1 = time.perf_counter() - obs_t, rew_t, done_t = torch_postprocess_go1( - sensor_t, physics_t, cur_t, last_t, cmd_t, idx_t, scalars_t - ) - _ = obs_t, rew_t, done_t - sync_torch_mps() - t2 = time.perf_counter() - torch_transfer += t1 - t0 - torch_compute += t2 - t1 - - for _ in range(iters): - t0 = time.perf_counter() - sensor_mx = mx.array(sensor_np, dtype=mx.float32) - physics_mx = mx.array(physics_np, dtype=mx.float32) - cur_mx = mx.array(current_actions_np, dtype=mx.float32) - last_mx = mx.array(last_actions_np, dtype=mx.float32) - cmd_mx = mx.array(commands_np, dtype=mx.float32) - mx.eval(sensor_mx, physics_mx, cur_mx, last_mx, cmd_mx) - t1 = time.perf_counter() - obs_mx, rew_mx, done_mx = mlx_postprocess_go1( - sensor_mx, physics_mx, cur_mx, last_mx, cmd_mx, idx_mx, scalars_mx - ) - mx.eval(obs_mx, rew_mx, done_mx) - t2 = time.perf_counter() - mlx_transfer += t1 - t0 - mlx_compute += t2 - t1 - - for _ in range(iters): - t0 = time.perf_counter() - sensor_mx = mx.array(sensor_np, dtype=mx.float32) - physics_mx = mx.array(physics_np, dtype=mx.float32) - cur_mx = mx.array(current_actions_np, dtype=mx.float32) - last_mx = mx.array(last_actions_np, dtype=mx.float32) - cmd_mx = mx.array(commands_np, dtype=mx.float32) - mx.eval(sensor_mx, physics_mx, cur_mx, last_mx, cmd_mx) - t1 = time.perf_counter() - obs_mx, rew_mx, done_mx = mlx_postprocess_go1( - sensor_mx, physics_mx, cur_mx, last_mx, cmd_mx, idx_mx, scalars_mx - ) - mx.eval(obs_mx, rew_mx, done_mx) - t2 = time.perf_counter() - _ = torch.as_tensor(np.array(obs_mx), device=TORCH_DEVICE, dtype=torch.float32) - _ = torch.as_tensor(np.array(rew_mx), device=TORCH_DEVICE, dtype=torch.float32) - _ = torch.as_tensor(np.array(done_mx), device=TORCH_DEVICE, dtype=torch.bool) - sync_torch_mps() - t3 = time.perf_counter() - mlx_torch_transfer += t1 - t0 - mlx_torch_compute += t2 - t1 - mlx_to_torch_transfer += t3 - t2 - - cpu_compute_ms = cpu_compute / iters * 1000.0 - cpu_transfer_ms = cpu_transfer / iters * 1000.0 - cpu_to_mlx_compute_ms = cpu_to_mlx_compute / iters * 1000.0 - cpu_to_mlx_transfer_ms = cpu_to_mlx_transfer / iters * 1000.0 - torch_transfer_ms = torch_transfer / iters * 1000.0 - torch_compute_ms = torch_compute / iters * 1000.0 - mlx_transfer_ms = mlx_transfer / iters * 1000.0 - mlx_compute_ms = mlx_compute / iters * 1000.0 - mlx_torch_transfer_ms = mlx_torch_transfer / iters * 1000.0 - mlx_torch_compute_ms = mlx_torch_compute / iters * 1000.0 - mlx_to_torch_transfer_ms = mlx_to_torch_transfer / iters * 1000.0 - - cpu_total_ms = cpu_compute_ms + cpu_transfer_ms - cpu_to_mlx_total_ms = cpu_to_mlx_compute_ms + cpu_to_mlx_transfer_ms - torch_total_ms = torch_transfer_ms + torch_compute_ms - mlx_total_ms = mlx_transfer_ms + mlx_compute_ms - mlx_to_torch_total_ms = mlx_torch_transfer_ms + mlx_torch_compute_ms + mlx_to_torch_transfer_ms - - return { - "num_envs": num_envs, - "iterations": iters, - "layout": { - "obs_dim": layout["obs_dim"], - "sensor_dim": layout["sensor_dim"], - "physics_dim": layout["physics_dim"], - "num_action": layout["num_action"], - }, - "physics_step_mode": { - "physics_step_rollout_ms": physics_step_rollout_ms, - "physics_step_env_wrapper_ms": physics_step_env_wrapper_ms, - }, - "mlx_rollout_bridge_mode": { - "rollout_with_mlx_action_ms": bridge_mlx["rollout_with_mlx_action_ms"], - "mlx_postprocess_from_rollout_ms": bridge_mlx["mlx_postprocess_from_rollout_ms"], - "total_ms": bridge_mlx["total_ms"], - }, - "cpu_numpy_mode": { - "compute_numpy_ms": cpu_compute_ms, - "transfer_obs_rew_done_to_torch_mps_ms": cpu_transfer_ms, - "total_ms": cpu_total_ms, - "total_with_physics_ms": cpu_total_ms + physics_step_rollout_ms, - }, - "cpu_numpy_to_mlx_mode": { - "compute_numpy_ms": cpu_to_mlx_compute_ms, - "transfer_obs_rew_done_to_mlx_ms": cpu_to_mlx_transfer_ms, - "total_ms": cpu_to_mlx_total_ms, - "total_with_physics_ms": cpu_to_mlx_total_ms + physics_step_rollout_ms, - }, - "torch_mps_mode": { - "transfer_all_numpy_to_torch_mps_ms": torch_transfer_ms, - "compute_postprocess_on_torch_mps_ms": torch_compute_ms, - "total_ms": torch_total_ms, - "total_with_physics_ms": torch_total_ms + physics_step_rollout_ms, - }, - "mlx_mode": { - "transfer_all_numpy_to_mlx_ms": mlx_transfer_ms, - "compute_postprocess_on_mlx_ms": mlx_compute_ms, - "total_ms": mlx_total_ms, - "total_with_physics_ms": mlx_total_ms + physics_step_rollout_ms, - }, - "mlx_to_torch_mps_mode": { - "transfer_all_numpy_to_mlx_ms": mlx_torch_transfer_ms, - "compute_postprocess_on_mlx_ms": mlx_torch_compute_ms, - "transfer_postprocess_result_mlx_to_torch_mps_ms": mlx_to_torch_transfer_ms, - "total_ms": mlx_to_torch_total_ms, - "total_with_physics_ms": mlx_to_torch_total_ms + physics_step_rollout_ms, - }, - "speedup_cpu_div_torch_mps": cpu_total_ms / torch_total_ms if torch_total_ms > 0 else 0.0, - "speedup_cpu_div_mlx": cpu_total_ms / mlx_total_ms if mlx_total_ms > 0 else 0.0, - "speedup_cpu_to_mlx_div_mlx": cpu_to_mlx_total_ms / mlx_total_ms - if mlx_total_ms > 0 - else 0.0, - "speedup_torch_mps_div_mlx": torch_total_ms / mlx_total_ms if mlx_total_ms > 0 else 0.0, - "speedup_cpu_div_mlx_to_torch_mps": cpu_total_ms / mlx_to_torch_total_ms - if mlx_to_torch_total_ms > 0 - else 0.0, - "speedup_torch_mps_div_mlx_to_torch_mps": torch_total_ms / mlx_to_torch_total_ms - if mlx_to_torch_total_ms > 0 - else 0.0, - } - - -def plot_results(results: list[dict], output_png: Path): - envs = [r["num_envs"] for r in results] - x = np.arange(len(envs), dtype=float) - w = 0.15 - - cpu_compute = np.array([r["cpu_numpy_mode"]["compute_numpy_ms"] for r in results]) - cpu_transfer = np.array( - [r["cpu_numpy_mode"]["transfer_obs_rew_done_to_torch_mps_ms"] for r in results] - ) - cpu_to_mlx_compute = np.array([r["cpu_numpy_to_mlx_mode"]["compute_numpy_ms"] for r in results]) - cpu_to_mlx_transfer = np.array( - [r["cpu_numpy_to_mlx_mode"]["transfer_obs_rew_done_to_mlx_ms"] for r in results] - ) - torch_transfer = np.array( - [r["torch_mps_mode"]["transfer_all_numpy_to_torch_mps_ms"] for r in results] - ) - torch_compute = np.array( - [r["torch_mps_mode"]["compute_postprocess_on_torch_mps_ms"] for r in results] - ) - mlx_transfer = np.array([r["mlx_mode"]["transfer_all_numpy_to_mlx_ms"] for r in results]) - mlx_compute = np.array([r["mlx_mode"]["compute_postprocess_on_mlx_ms"] for r in results]) - mlx_to_torch_transfer = np.array( - [ - r["mlx_to_torch_mps_mode"]["transfer_postprocess_result_mlx_to_torch_mps_ms"] - for r in results - ] - ) - physics_step = np.array([r["physics_step_mode"]["physics_step_rollout_ms"] for r in results]) - - fig, (ax_top, ax_bottom) = plt.subplots(2, 1, figsize=(13, 10), sharex=True) - x_cpu = x - 2.0 * w - x_cpu_mlx = x - 1.0 * w - x_torch = x - x_mlx = x + 1.0 * w - x_mlx_torch = x + 2.0 * w - - ax_top.bar(x_cpu, physics_step, w, label="Shared: MuJoCo physics step", color="#9AA0A6") - ax_top.bar(x_cpu_mlx, physics_step, w, label="_nolegend_", color="#9AA0A6") - ax_top.bar(x_torch, physics_step, w, label="_nolegend_", color="#9AA0A6") - ax_top.bar(x_mlx, physics_step, w, label="_nolegend_", color="#9AA0A6") - ax_top.bar(x_mlx_torch, physics_step, w, label="_nolegend_", color="#9AA0A6") - - ax_top.bar( - x_cpu, - cpu_compute, - w, - bottom=physics_step, - label="Path A (CPU->Torch): numpy compute", - color="#D99A9A", - ) - ax_top.bar( - x_cpu, - cpu_transfer, - w, - bottom=physics_step + cpu_compute, - label="Path A (CPU->Torch): obs/rew/done -> torch.mps", - color="#EBCED6", - ) - ax_top.bar( - x_cpu_mlx, - cpu_to_mlx_compute, - w, - bottom=physics_step, - label="Path B (CPU->MLX): numpy compute", - color="#D7B78B", - ) - ax_top.bar( - x_cpu_mlx, - cpu_to_mlx_transfer, - w, - bottom=physics_step + cpu_to_mlx_compute, - label="Path B (CPU->MLX): obs/rew/done -> mlx", - color="#EFD9B7", - ) - ax_top.bar( - x_torch, - torch_transfer, - w, - bottom=physics_step, - label="Path C (Torch native): numpy -> torch.mps", - color="#8FB3CC", - ) - ax_top.bar( - x_torch, - torch_compute, - w, - bottom=physics_step + torch_transfer, - label="Path C (Torch native): torch.mps compute", - color="#A8CFAE", - ) - ax_top.bar( - x_mlx, - mlx_transfer, - w, - bottom=physics_step, - label="Path D (MLX native): numpy -> mlx", - color="#E1B15A", - ) - ax_top.bar( - x_mlx, - mlx_compute, - w, - bottom=physics_step + mlx_transfer, - label="Path D (MLX native): mlx compute", - color="#F3D9A5", - ) - ax_top.bar( - x_mlx_torch, - mlx_transfer, - w, - bottom=physics_step, - label="Path E (MLX->Torch): numpy -> mlx", - color="#B78BD0", - ) - ax_top.bar( - x_mlx_torch, - mlx_compute, - w, - bottom=physics_step + mlx_transfer, - label="Path E (MLX->Torch): mlx compute", - color="#CDA9E4", - ) - ax_top.bar( - x_mlx_torch, - mlx_to_torch_transfer, - w, - bottom=physics_step + mlx_transfer + mlx_compute, - label="Path E (MLX->Torch): mlx result -> torch.mps", - color="#E7D0F4", - ) - - ax_bottom.bar( - x_cpu, cpu_compute, w, label="Path A (CPU->Torch): numpy compute", color="#D99A9A" - ) - ax_bottom.bar( - x_cpu, - cpu_transfer, - w, - bottom=cpu_compute, - label="Path A (CPU->Torch): obs/rew/done -> torch.mps", - color="#EBCED6", - ) - ax_bottom.bar( - x_cpu_mlx, cpu_to_mlx_compute, w, label="Path B (CPU->MLX): numpy compute", color="#D7B78B" - ) - ax_bottom.bar( - x_cpu_mlx, - cpu_to_mlx_transfer, - w, - bottom=cpu_to_mlx_compute, - label="Path B (CPU->MLX): obs/rew/done -> mlx", - color="#EFD9B7", - ) - ax_bottom.bar( - x_torch, - torch_transfer, - w, - label="Path C (Torch native): numpy -> torch.mps", - color="#8FB3CC", - ) - ax_bottom.bar( - x_torch, - torch_compute, - w, - bottom=torch_transfer, - label="Path C (Torch native): torch.mps compute", - color="#A8CFAE", - ) - ax_bottom.bar( - x_mlx, mlx_transfer, w, label="Path D (MLX native): numpy -> mlx", color="#E1B15A" - ) - ax_bottom.bar( - x_mlx, - mlx_compute, - w, - bottom=mlx_transfer, - label="Path D (MLX native): mlx compute", - color="#F3D9A5", - ) - ax_bottom.bar( - x_mlx_torch, mlx_transfer, w, label="Path E (MLX->Torch): numpy -> mlx", color="#B78BD0" - ) - ax_bottom.bar( - x_mlx_torch, - mlx_compute, - w, - bottom=mlx_transfer, - label="Path E (MLX->Torch): mlx compute", - color="#CDA9E4", - ) - ax_bottom.bar( - x_mlx_torch, - mlx_to_torch_transfer, - w, - bottom=mlx_transfer + mlx_compute, - label="Path E (MLX->Torch): mlx result -> torch.mps", - color="#E7D0F4", - ) - - cpu_total = physics_step + cpu_compute + cpu_transfer - cpu_to_mlx_total = physics_step + cpu_to_mlx_compute + cpu_to_mlx_transfer - torch_total = physics_step + torch_transfer + torch_compute - mlx_total = physics_step + mlx_transfer + mlx_compute - mlx_to_torch_total = physics_step + mlx_transfer + mlx_compute + mlx_to_torch_transfer - cpu_post_total = cpu_compute + cpu_transfer - cpu_to_mlx_post_total = cpu_to_mlx_compute + cpu_to_mlx_transfer - torch_post_total = torch_transfer + torch_compute - mlx_post_total = mlx_transfer + mlx_compute - mlx_to_torch_post_total = mlx_transfer + mlx_compute + mlx_to_torch_transfer - - for i in range(len(envs)): - ax_top.text( - x_cpu[i], - cpu_total[i] + 0.03, - f"{cpu_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_top.text( - x_cpu_mlx[i], - cpu_to_mlx_total[i] + 0.03, - f"{cpu_to_mlx_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_top.text( - x_torch[i], - torch_total[i] + 0.03, - f"{torch_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_top.text( - x_mlx[i], - mlx_total[i] + 0.03, - f"{mlx_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_top.text( - x_mlx_torch[i], - mlx_to_torch_total[i] + 0.03, - f"{mlx_to_torch_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - - ax_bottom.text( - x_cpu[i], - cpu_post_total[i] + 0.03, - f"{cpu_post_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_bottom.text( - x_cpu_mlx[i], - cpu_to_mlx_post_total[i] + 0.03, - f"{cpu_to_mlx_post_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_bottom.text( - x_torch[i], - torch_post_total[i] + 0.03, - f"{torch_post_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_bottom.text( - x_mlx[i], - mlx_post_total[i] + 0.03, - f"{mlx_post_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - ax_bottom.text( - x_mlx_torch[i], - mlx_to_torch_post_total[i] + 0.03, - f"{mlx_to_torch_post_total[i]:.2f}", - ha="center", - va="bottom", - fontsize=8, - ) - - ax_top.set_title( - "Go1 latest step+postprocess (Top: with shared physics; Bottom: postprocess only)\n" - f"{get_device_info_line()}" - ) - ax_top.set_ylabel("Time per step (ms)") - ax_top.grid(axis="y", alpha=0.25) - ax_top.legend( - title="Execution paths", - loc="upper left", - bbox_to_anchor=(1.01, 1.0), - borderaxespad=0.0, - fontsize=8, - title_fontsize=9, - ) - - ax_bottom.set_xlabel("num_envs") - ax_bottom.set_ylabel("Postprocess time (ms)") - ax_bottom.set_xticks(x) - ax_bottom.set_xticklabels([str(v) for v in envs]) - ax_bottom.grid(axis="y", alpha=0.25) - fig.tight_layout() - output_png.parent.mkdir(parents=True, exist_ok=True) - fig.savefig(output_png, dpi=180) - plt.close(fig) - - -def main(): - ensure_registries() - parser = argparse.ArgumentParser() - parser.add_argument("--env_list", type=str, default=",".join(str(v) for v in DEFAULT_ENV_LIST)) - parser.add_argument("--iters", type=int, default=DEFAULT_ITERS) - parser.add_argument("--output_json", type=str, default=str(OUTPUT_JSON)) - parser.add_argument("--output_png", type=str, default=str(OUTPUT_PNG)) - args = parser.parse_args() - - if not torch.backends.mps.is_available(): - raise RuntimeError("Torch MPS is not available on this machine.") - - env_list = parse_env_list(args.env_list) - layout = build_go1_layout() - all_results = [] - for nenv in env_list: - one = bench_one_envnum(nenv, args.iters, layout) - all_results.append(one) - print( - f"[{nenv}] Physics(rollout)={one['physics_step_mode']['physics_step_rollout_ms']:.3f} ms, " - f"Physics(env_wrapper)={one['physics_step_mode']['physics_step_env_wrapper_ms']:.3f} ms, " - f"CPU={one['cpu_numpy_mode']['total_with_physics_ms']:.3f} ms, " - f"CPU->MLX={one['cpu_numpy_to_mlx_mode']['total_with_physics_ms']:.3f} ms, " - f"Torch.MPS={one['torch_mps_mode']['total_with_physics_ms']:.3f} ms, " - f"MLX={one['mlx_mode']['total_with_physics_ms']:.3f} ms, " - f"MLX->Torch.MPS={one['mlx_to_torch_mps_mode']['total_with_physics_ms']:.3f} ms, " - f"MLXBridge(rollout+post)={one['mlx_rollout_bridge_mode']['total_ms']:.3f} ms, " - f"cpu/torch={one['speedup_cpu_div_torch_mps']:.3f}, " - f"cpu/mlx={one['speedup_cpu_div_mlx']:.3f}, " - f"cpu->mlx/mlx={one['speedup_cpu_to_mlx_div_mlx']:.3f}, " - f"torch/mlx={one['speedup_torch_mps_div_mlx']:.3f}, " - f"cpu/mlx->torch={one['speedup_cpu_div_mlx_to_torch_mps']:.3f}, " - f"torch/mlx->torch={one['speedup_torch_mps_div_mlx_to_torch_mps']:.3f}" - ) - - output_json = Path(args.output_json) - output_png = Path(args.output_png) - output_json.parent.mkdir(parents=True, exist_ok=True) - payload = { - "meta": { - "timestamp": datetime.now().isoformat(timespec="seconds"), - "device_info": get_device_info_dict(), - "device_torch": TORCH_DEVICE, - "device_mlx": "mlx", - "torch_version": torch.__version__, - "mps_built": bool(torch.backends.mps.is_built()), - "mps_available": bool(torch.backends.mps.is_available()), - "env_list": env_list, - "iters": args.iters, - "task": OWNER_TASK_ID, - "env_task_name": ENV_TASK_NAME, - "note": ( - "Postprocess logic synchronized with latest " - f"{OWNER_TASK_ID} (env: {ENV_TASK_NAME}); physics step uses mujoco.rollout." - ), - }, - "summary": { - "geomean_speedup_cpu_over_mlx_postprocess": geomean( - [r["speedup_cpu_div_mlx"] for r in all_results] - ), - "geomean_speedup_cpu_to_mlx_over_mlx_postprocess": geomean( - [r["speedup_cpu_to_mlx_div_mlx"] for r in all_results] - ), - "geomean_speedup_torch_over_mlx_postprocess": geomean( - [r["speedup_torch_mps_div_mlx"] for r in all_results] - ), - "geomean_env_wrapper_over_rollout_physics": geomean( - [ - r["physics_step_mode"]["physics_step_env_wrapper_ms"] - / max(r["physics_step_mode"]["physics_step_rollout_ms"], 1e-12) - for r in all_results - ] - ), - }, - "results": all_results, - } - with output_json.open("w", encoding="utf-8") as f: - json.dump(payload, f, indent=2) - - plot_results(all_results, output_png) - print(f"Saved JSON: {output_json}") - print(f"Saved PNG: {output_png}") - - -if __name__ == "__main__": - main() diff --git a/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py b/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py index 61e607b64..00f54da2f 100644 --- a/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py +++ b/scripts/benchmark/physics/benchmark_mujoco_rollout_model_sharing.py @@ -670,7 +670,7 @@ def _plot_summary_plots( def main() -> None: parser = argparse.ArgumentParser( description=( - "Benchmark MuJoCo rollout on go1/go2/g1: shared single model vs " + "Benchmark MuJoCo rollout on go2/g1: shared single model vs " "per-env distinct models, optionally before/after discardvisual" ) ) @@ -683,7 +683,7 @@ def main() -> None: parser.add_argument( "--robots", type=str, - default="go1,go2,g1", + default="go2,g1", help="Comma separated robot names to benchmark from src/unilab/assets/robots.", ) parser.add_argument( diff --git a/scripts/benchmark/rl/benchmark_replay_buffer_placement.py b/scripts/benchmark/rl/benchmark_replay_buffer_placement.py index e98fcc82c..4010c0592 100644 --- a/scripts/benchmark/rl/benchmark_replay_buffer_placement.py +++ b/scripts/benchmark/rl/benchmark_replay_buffer_placement.py @@ -191,6 +191,12 @@ def _xpu_available() -> bool: def _replay_transfer_manifest(device: torch.device, *, ring_depth: int = 2) -> dict[str, Any]: + """Describe the uni_rl 1.2.1 replay-transfer backend for ``device``. + + uni_rl 1.2.1 builds ``CudaLikeReplayTransferBackend`` for CUDA (ROCm + included) and ``TorchCopyReplayTransferBackend`` otherwise; the async + ``submit_h2d`` path and the dedicated XPU backend no longer exist. + """ if device.type == "cuda": torch_version = getattr(torch, "version", None) is_rocm = bool(getattr(torch_version, "hip", None)) @@ -198,26 +204,14 @@ def _replay_transfer_manifest(device: torch.device, *, ring_depth: int = 2) -> d "backend": "CudaLikeReplayTransferBackend", "device_family": "rocm" if is_rocm else "cuda", "host_memory_kind": "registered_pinned_shared", - "supports_async_submit": True, "supports_timing_events": True, "h2d_submitter": "torch_copy_stream" if is_rocm else "pybind11", "ring_depth": ring_depth, } - if device.type == "xpu": - return { - "backend": "XpuReplayTransferBackend", - "device_family": "xpu", - "host_memory_kind": "pageable_shared", - "supports_async_submit": True, - "supports_timing_events": False, - "h2d_submitter": "torch_xpu_copy_stream", - "ring_depth": ring_depth, - } return { "backend": "TorchCopyReplayTransferBackend", "device_family": device.type, "host_memory_kind": "pageable_shared", - "supports_async_submit": False, "supports_timing_events": False, "h2d_submitter": "torch_copy", "ring_depth": ring_depth, diff --git a/scripts/benchmark/xmls/humanoid/humanoid100.xml b/scripts/benchmark/xmls/humanoid/humanoid100.xml deleted file mode 100644 index e7fef05bb..000000000 --- a/scripts/benchmark/xmls/humanoid/humanoid100.xml +++ /dev/null @@ -1,150 +0,0 @@ - - - - - diff --git a/scripts/convert_motion_29dof_to_23dof.py b/scripts/convert_motion_29dof_to_23dof.py deleted file mode 100644 index f90ee6f82..000000000 --- a/scripts/convert_motion_29dof_to_23dof.py +++ /dev/null @@ -1,167 +0,0 @@ -#!/usr/bin/env python3 -"""Convert 29-DoF G1 motion NPZ files to 23-DoF. - -Removes the 6 RM (redundant manipulator) joints: - - waist_roll_joint, waist_pitch_joint - - left_wrist_pitch_joint, left_wrist_yaw_joint - - right_wrist_pitch_joint, right_wrist_yaw_joint - -And maps body data from 31-body 29-DoF XML order to 24-body 23-DoF XML order. - -29-DoF joint indices to remove (0-indexed): [13, 14, 20, 21, 27, 28] -29-DoF body indices to keep (0-indexed): [0..12, 15..20, 23..27] - → maps to 24 bodies in 23-DoF XML order (pelvis + 23 links) -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import numpy as np - -# 29-DoF joint indices to REMOVE (0-indexed) -# waist_roll=13, waist_pitch=14, left_wrist_pitch=20, left_wrist_yaw=21, -# right_wrist_pitch=27, right_wrist_yaw=28 -JOINT_REMOVE_IDX = [13, 14, 20, 21, 27, 28] - -# 29-DoF body indices to KEEP → 25 bodies in 23-DoF NPZ order -# -# The exported NPZ includes the MuJoCo world body (id=0) at index 0, -# so NPZ index = MuJoCo body ID for all child bodies. -# Keeping: world(0), pelvis→right_ankle_roll(1-13), -# torso_link→left_wrist_roll_link(16-21), -# right_shoulder_pitch_link→right_wrist_roll_link(24-28) -BODY_KEEP_IDX = [ - 0, # world body (MuJoCo id=0) - 1, - 2, - 3, - 4, - 5, - 6, - 7, - 8, - 9, - 10, - 11, - 12, - 13, # pelvis → right_ankle_roll_link - 16, - 17, - 18, - 19, - 20, - 21, # torso_link → left_wrist_roll_link - 24, - 25, - 26, - 27, - 28, # right_shoulder_pitch_link → right_wrist_roll_link -] -# Result: 25 bodies (world + 24 child bodies matching 23-DoF XML body order) -# MuJoCo body ID directly indexes into this array (NPZ index = MuJoCo body ID). - - -def convert_motion_npz(src_path: Path, dst_path: Path, is_box: bool = False) -> None: - """Convert a 29-DoF motion NPZ to 23-DoF and save.""" - data = np.load(src_path, allow_pickle=True) - - # Build output dict - out: dict[str, np.ndarray] = {} - - # Copy fps (no change) - out["fps"] = data["fps"] - - # Joint data: remove 6 RM joint columns - joint_keep = [i for i in range(data["joint_pos"].shape[1]) if i not in JOINT_REMOVE_IDX] - assert len(joint_keep) == 23, f"Expected 23 kept joints, got {len(joint_keep)}" - out["joint_pos"] = data["joint_pos"][:, joint_keep] - out["joint_vel"] = data["joint_vel"][:, joint_keep] - - # Body data: keep only mapped body indices - assert len(BODY_KEEP_IDX) == 25, ( - f"Expected 25 kept bodies (world + 24 child), got {len(BODY_KEEP_IDX)}" - ) - for key in ("body_pos_w", "body_quat_w", "body_lin_vel_w", "body_ang_vel_w"): - out[key] = data[key][:, BODY_KEEP_IDX] - - # Object data (box tracking): pass through unchanged - object_keys = ("object_pos_w", "object_quat_w", "object_lin_vel_w", "object_ang_vel_w") - for key in object_keys: - if key in data: - out[key] = data[key] - - # Save - dst_path.parent.mkdir(parents=True, exist_ok=True) - np.savez_compressed(dst_path, **out) - - # Print summary - src_shapes = {k: data[k].shape for k in data.keys()} - dst_shapes = {k: out[k].shape for k in out.keys()} - print(f"Converted: {src_path.name}") - print(f" Source shapes: {src_shapes}") - print(f" Output shapes: {dst_shapes}") - print(f" Saved to: {dst_path}") - - -def main() -> None: - parser = argparse.ArgumentParser(description="Convert 29-DoF G1 motion NPZ files to 23-DoF") - parser.add_argument( - "input", - nargs="*", - help="Input NPZ files (default: predefined G1 motion files)", - ) - parser.add_argument( - "--output-dir", - "-o", - type=Path, - default=None, - help="Output directory (default: same as input dir + '_23dof' suffix)", - ) - parser.add_argument( - "--suffix", - default="_23dof", - help="Suffix for output files (default: _23dof)", - ) - parser.add_argument( - "--copy-box-object-data", - action="store_true", - help="Copy object pose data for box tracking files", - ) - args = parser.parse_args() - - # Default: convert the known 29-DoF motion files - motion_root = ( - Path(__file__).resolve().parent.parent / "src" / "unilab" / "assets" / "motions" / "g1" - ) - - default_inputs = [ - motion_root / "flip_360_001__A304.npz", - motion_root / "flip_from_wall_104__A304.npz", - motion_root / "sub3_largebox_003_boxconverted.npz", - ] - - inputs = [Path(p) for p in args.input] if args.input else default_inputs - - for src_path in inputs: - if not src_path.exists(): - print(f"WARNING: {src_path} not found, skipping") - continue - - if args.output_dir: - dst_dir = args.output_dir - else: - # Default: same directory - dst_dir = src_path.parent - - stem = src_path.stem # e.g. "flip_360_001__A304" - dst_name = f"{stem}{args.suffix}.npz" - dst_path = dst_dir / dst_name - - is_box = "box" in stem.lower() or "largebox" in stem.lower() - convert_motion_npz(src_path, dst_path, is_box=is_box) - - -if __name__ == "__main__": - main() diff --git a/src/unilab/assets/motions/g1/23dof_amp/.gitkeep b/src/unilab/assets/motions/g1/23dof_amp/.gitkeep deleted file mode 100644 index e69de29bb..000000000 diff --git a/src/unilab/assets/robots/hfields/hfield.png b/src/unilab/assets/robots/hfields/hfield.png deleted file mode 100644 index 7b48e7205..000000000 Binary files a/src/unilab/assets/robots/hfields/hfield.png and /dev/null differ diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index 55dcee972..f4ffa1676 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -22,25 +22,6 @@ from unilab.base.scene import SceneCfg -def _legacy_genesis_device_option_error(exc: TypeError) -> bool: - """Identify an old UniSim adapter rejecting the optional device keyword. - - UniSim 1.1 reports unknown backend options from ``GenesisBackend`` while - other compatible releases may expose Python's usual ``unexpected keyword`` - wording. Keep the compatibility retry narrowly scoped to those messages; - constructor errors from the actual Genesis runtime must still propagate. - """ - - message = str(exc).lower() - mentions_device = "genesis_device_id" in message or "device_id" in message - rejects_keyword = ( - "does not accept backend options" in message - or "unexpected keyword argument" in message - or "unexpected keyword" in message - ) - return mentions_device and rejects_keyword - - def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: """Translate ``EnvCfg`` backend knobs into UniSim adapter options.""" result: dict[str, Any] = { @@ -69,18 +50,13 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: "isaacsim_render_mode": cfg.isaacsim_render_mode, "isaacsim_render_width": cfg.isaacsim_render_width, "isaacsim_render_height": cfg.isaacsim_render_height, + "superdex_execution_mode": cfg.superdex_execution_mode, } - # Keep the optional key absent for legacy unisim-core releases that do not - # know about Genesis' explicit device argument. Once a rank selects a - # device the key is added below and ``create_backend`` supplies a narrow - # compatibility fallback for those releases. + # Forward the explicit Genesis device id only when a rank selected one; + # when absent, unisim-core's factory default applies and Genesis picks + # its own device. if cfg.genesis_device_id is not None: result["genesis_device_id"] = cfg.genesis_device_id - # Keep the default absent so unisim-core releases that predate the - # execution-mode option still accept the SuperDex kwargs; "serial" requires - # the updated adapter. - if cfg.superdex_execution_mode != "batch": - result["superdex_execution_mode"] = cfg.superdex_execution_mode return result @@ -112,17 +88,12 @@ def create_backend( ensure_robot_assets_for_paths( [scene.model_file, scene.visual_model_file, *scene.fragment_files] ) - if backend_type != "newton": - # Keep the owner translation forward-compatible with unisim-core - # releases that predate the Newton adapter and therefore do not pop - # these optional kwargs in their shared factory. - for key in ( - "newton_device", - "newton_nconmax", - "newton_njmax", - "newton_capacity_check_steps", - ): - kwargs.pop(key, None) + if backend_type == "drake": + # unisim-core 1.4.2 dropped the Drake-branch filtering of MuJoCo + # root-body options; Drake derives root state from its own plant and + # DrakeBackend rejects the keywords. Pop them at the owner boundary. + kwargs.pop("base_name", None) + kwargs.pop("push_body_name", None) # Newton reconstructs body state from its compiled articulation and does # not accept MuJoCo's synthetic body-sensor injection. Keep this # capability translation at the owner/backend boundary so env code remains @@ -133,11 +104,9 @@ def create_backend( } if backend_type == "genesis" and kwargs.get("genesis_device_id") is not None: # Bind before any unisim-core Genesis constructor can call gs.init. - # New unisim-core releases repeat this idempotently; old releases do - # not accept the keyword, so the retry below still gets the correct - # process-wide device. Binding a non-zero id pins - # CUDA_VISIBLE_DEVICES (Quadrants only honors the first visible - # device), so forward the *post-pin* in-process index. + # Binding a non-zero id pins CUDA_VISIBLE_DEVICES (Quadrants only + # honors the first visible device), so forward the *post-pin* + # in-process index. genesis_device_id = kwargs["genesis_device_id"] if ( isinstance(genesis_device_id, bool) @@ -150,19 +119,7 @@ def create_backend( ) bound = bind_genesis_process_device(f"cuda:{genesis_device_id}") kwargs["genesis_device_id"] = int(bound.rsplit(":", 1)[1]) - try: - return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **kwargs) - except TypeError as exc: - if backend_type != "genesis" or "genesis_device_id" not in kwargs: - raise - # unisim-core < 1.2 has no Genesis device field and reports the - # unknown option from GenesisBackend. Retry only for that precise - # capability error; unrelated constructor TypeErrors must propagate. - if not _legacy_genesis_device_option_error(exc): - raise - legacy_kwargs = dict(kwargs) - legacy_kwargs.pop("genesis_device_id", None) - return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **legacy_kwargs) + return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **kwargs) __all__ = ["SimBackend", "create_backend", "env_backend_kwargs"] diff --git a/src/unilab/base/env_factory.py b/src/unilab/base/env_factory.py index 46c256714..20f8f70b5 100644 --- a/src/unilab/base/env_factory.py +++ b/src/unilab/base/env_factory.py @@ -42,9 +42,11 @@ def make_registry_env( # carry the explicit cold-path id in the opaque override and bind # immediately before registry construction. Binding a non-zero id pins # CUDA_VISIBLE_DEVICES for this process, so forward the post-pin - # in-process index downstream. Newer unisim-core versions repeat this - # check in GenesisBackend itself, making this compatibility guard - # idempotent. + # in-process index downstream. This binding is required, not a + # compatibility shim: spawn collectors are fresh interpreters that cannot + # inherit the parent process's binding, and only UniLab's + # ``bind_genesis_process_device`` sets the pinned-namespace flag its + # resolution helpers consult. if sim_backend == "genesis" and env_cfg_override is not None: genesis_device_id = env_cfg_override.get("genesis_device_id") if genesis_device_id is not None: diff --git a/src/unilab/base/process_device.py b/src/unilab/base/process_device.py index ee30357c1..37ab1e0d4 100644 --- a/src/unilab/base/process_device.py +++ b/src/unilab/base/process_device.py @@ -26,9 +26,12 @@ # Newton consumes an explicit ``cuda:N`` device string (``newton_device``) -# instead of an integer id. uni_rl's collector-side binder gate only knows -# mjwarp, so the rank-local device must reach spawn collectors through the -# env override rather than through process binding. +# instead of an integer id. uni_rl's collector-side process binding is +# injection-based: it binds through the caller-supplied ``bind_device`` +# callable and fails closed when none is injected. UniLab injects +# ``bind_backend_process_device_for_backend`` (which covers mjwarp and +# newton); the env override additionally forwards the rank-local device +# string so spawn collectors can pass it to the Newton adapter. BACKEND_ENV_DEVICE_STR_FIELDS: dict[str, str] = { "newton": "newton_device", } diff --git a/src/unilab/cli.py b/src/unilab/cli.py index 926eb2bae..e7fd7944d 100644 --- a/src/unilab/cli.py +++ b/src/unilab/cli.py @@ -133,13 +133,7 @@ def _check_runtime_requirements(algo: str, sim: str) -> None: "in a source checkout (or `pip install unilab[newton]`)." ) if sim == "superdex": - try: - from unisim.backend.superdex.dependencies import superdex_dependencies_available - except ImportError as exc: - raise SystemExit( - "sim=superdex requires unisim-core>=1.1.5 with the SuperDex adapter; " - "the installed unisim-core does not provide that adapter." - ) from exc + from unisim.backend.superdex.dependencies import superdex_dependencies_available if not superdex_dependencies_available(): raise SystemExit( diff --git a/src/unilab/training/__init__.py b/src/unilab/training/__init__.py index 7dda18b39..ac96b36a7 100644 --- a/src/unilab/training/__init__.py +++ b/src/unilab/training/__init__.py @@ -35,19 +35,16 @@ resolve_offpolicy_checkpoint_path, resolve_task_checkpoint_path, ) -from unilab.utils.monitoring import HardwareMonitor from unilab.utils.seed import ( TrainingSeedInfo, apply_configured_training_seed, apply_training_seed, - derive_worker_seed, resolve_training_seed, ) __all__ = [ "BackendAdapter", "ExperimentTracker", - "HardwareMonitor", "algo_config_dict", "apply_env_nan_guard", "assert_offpolicy_task_choice_matches_algo", @@ -69,7 +66,6 @@ "TrainingSeedInfo", "apply_configured_training_seed", "apply_training_seed", - "derive_worker_seed", "resolve_appo_checkpoint_path", "resolve_offpolicy_checkpoint_path", "resolve_training_seed", diff --git a/src/unilab/utils/device.py b/src/unilab/utils/device.py index 004d64b37..85242b254 100644 --- a/src/unilab/utils/device.py +++ b/src/unilab/utils/device.py @@ -4,7 +4,6 @@ import re import subprocess from functools import lru_cache -from typing import Callable, cast import torch @@ -26,106 +25,6 @@ def get_default_device() -> str: return "cpu" -def _device_count(device_type: str) -> int | None: - if device_type == "cuda": - return int(torch.cuda.device_count()) - if device_type == "xpu": - xpu = getattr(torch, "xpu", None) - device_count = getattr(xpu, "device_count", None) - if callable(device_count): - return int(cast(Callable[[], int], device_count)()) - return None - - -def _mps_available() -> bool: - mps = getattr(torch.backends, "mps", None) - is_available = getattr(mps, "is_available", None) - return bool(callable(is_available) and is_available()) - - -def _parse_device_alias(value: str) -> tuple[str, int | None]: - raw = value.strip().lower() - if not raw: - raise ValueError("Device alias must not be empty") - if ":" not in raw: - return raw, None - base, index_text = raw.split(":", 1) - if not index_text: - raise ValueError(f"Device alias {value!r} has an empty index") - try: - index = int(index_text) - except ValueError as exc: - raise ValueError(f"Device alias {value!r} has a non-integer index") from exc - if index < 0: - raise ValueError(f"Device alias {value!r} has a negative index") - return base, index - - -def _resolve_indexed_device(device_type: str, index: int | None, original: str) -> str: - count = _device_count(device_type) - if count is not None and index is not None and index >= count: - raise ValueError( - f"Requested device {original!r} resolves to {device_type}:{index}, " - f"but only {count} {device_type} device(s) are available" - ) - return device_type if index is None else f"{device_type}:{index}" - - -def _resolve_mps_alias(index: int | None, original: str) -> str: - if not _mps_available(): - raise ValueError(f"Requested device {original!r} requires MPS, but MPS is unavailable") - if index not in (None, 0): - raise ValueError( - f"Requested device {original!r} cannot be mapped to MPS; only index 0 is valid" - ) - return "mps" - - -def resolve_torch_device_alias(device: str | None, *, default: str = "cpu") -> str: - """Resolve a cross-platform torch device alias to a concrete device string. - - ``gpu`` is an abstract accelerator alias. ``cuda`` is also accepted on - macOS/MPS for config portability and maps to ``mps`` when CUDA is absent. - The function validates the resolved device and never silently falls back to - CPU for unavailable accelerators. - """ - original = default if device is None else str(device) - base, index = _parse_device_alias(original) - - if base == "cpu": - if index is not None: - raise ValueError(f"CPU device {original!r} must not include an index") - return "cpu" - - if base == "mps": - return _resolve_mps_alias(index, original) - - if base == "xpu": - if not _xpu_available(): - raise ValueError(f"Requested device {original!r} requires XPU, but XPU is unavailable") - return _resolve_indexed_device("xpu", index, original) - - if base == "cuda": - if torch.cuda.is_available(): - return _resolve_indexed_device("cuda", index, original) - if _mps_available(): - return _resolve_mps_alias(index, original) - raise ValueError(f"Requested device {original!r} requires CUDA, but CUDA is unavailable") - - if base == "gpu": - if torch.cuda.is_available(): - return _resolve_indexed_device("cuda", index, original) - if _xpu_available(): - return _resolve_indexed_device("xpu", index, original) - if _mps_available(): - return _resolve_mps_alias(index, original) - raise ValueError( - f"Requested device {original!r} requires an accelerator, but none is available" - ) - - raise ValueError(f"Unsupported device alias {original!r}; expected cpu, gpu, cuda, mps, or xpu") - - def _is_macos() -> bool: return platform.system() == "Darwin" diff --git a/src/unilab/utils/geometry.py b/src/unilab/utils/geometry.py index 4c0779390..c07c5cc85 100644 --- a/src/unilab/utils/geometry.py +++ b/src/unilab/utils/geometry.py @@ -20,16 +20,6 @@ ) -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_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) @@ -57,38 +47,9 @@ def np_roll_pitch_from_quat(quat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: 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_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) diff --git a/src/unilab/utils/monitoring.py b/src/unilab/utils/monitoring.py deleted file mode 100644 index cafcd38a0..000000000 --- a/src/unilab/utils/monitoring.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Hardware monitoring utilities for performance profiling.""" - -from typing import Dict - -import torch - -try: - import psutil - - HAS_PSUTIL = True -except ImportError: - HAS_PSUTIL = False - - -class HardwareMonitor: - """Monitor CPU, GPU, memory usage.""" - - def __init__(self): - self.has_psutil = HAS_PSUTIL - if self.has_psutil: - self.process = psutil.Process() - - self.has_cuda = torch.cuda.is_available() - if self.has_cuda: - try: - import pynvml - - pynvml.nvmlInit() - self.nvml_handle = pynvml.nvmlDeviceGetHandleByIndex(0) - self.has_nvml = True - except Exception: - self.has_nvml = False - else: - self.has_nvml = False - - def get_metrics(self) -> Dict[str, float]: - """Get current hardware metrics.""" - metrics = {} - - # CPU & Memory (requires psutil) - if self.has_psutil: - metrics["cpu_percent"] = self.process.cpu_percent() - metrics["cpu_count"] = psutil.cpu_count() - mem = self.process.memory_info() - metrics["memory_rss_mb"] = mem.rss / 1024 / 1024 - metrics["memory_percent"] = self.process.memory_percent() - - # GPU - if self.has_cuda: - metrics["gpu_memory_allocated_mb"] = torch.cuda.memory_allocated() / 1024 / 1024 - metrics["gpu_memory_reserved_mb"] = torch.cuda.memory_reserved() / 1024 / 1024 - - if self.has_nvml: - import pynvml - - util = pynvml.nvmlDeviceGetUtilizationRates(self.nvml_handle) - metrics["gpu_utilization"] = util.gpu - metrics["gpu_memory_utilization"] = util.memory - - return metrics diff --git a/src/unilab/utils/seed.py b/src/unilab/utils/seed.py index 87bb40fd5..9dff5f52f 100644 --- a/src/unilab/utils/seed.py +++ b/src/unilab/utils/seed.py @@ -61,15 +61,6 @@ def resolve_training_seed(cfg: Any) -> TrainingSeedInfo: return TrainingSeedInfo(configured_seed=None, configured_seed_source=None, effective_seed=None) -def derive_worker_seed(base_seed: int | None, worker_index: int = 0) -> int | None: - """Derive deterministic subprocess seeds from the effective run seed.""" - if base_seed is None: - return None - if worker_index < 0: - raise ValueError(f"worker_index must be non-negative, got {worker_index}") - return int(base_seed) + int(worker_index) + 1 - - def apply_training_seed( seed: int | None, *, diff --git a/src/unilab/utils/sim2sim.py b/src/unilab/utils/sim2sim.py index 536b685d9..bfe7e11bd 100644 --- a/src/unilab/utils/sim2sim.py +++ b/src/unilab/utils/sim2sim.py @@ -249,42 +249,3 @@ def policy_load_dim_guard( "`uv run scripts/audit_sim2sim_contracts.py`.\n" f"Original load error:\n{exc}" ) from exc - - -class Sim2SimConfigResolver: - """Object facade over the module-level sim2sim contract API.""" - - ALLOWLIST = ALLOWLIST - WARNING_LIST = WARNING_LIST - DENYLIST = DENYLIST - ENV_STRUCTURAL_DENYLIST = ENV_STRUCTURAL_DENYLIST - - @staticmethod - def extract_snapshot(full_cfg: DictConfig) -> dict[str, Any]: - """See :func:`extract_contract_snapshot`.""" - return extract_contract_snapshot(full_cfg) - - @staticmethod - def resolve( - source_run_dir: str | Path | None, - target_cfg: DictConfig, - *, - algo_name: str | None = None, - strict: bool = True, - ) -> DictConfig | None: - """See :func:`resolve_sim2sim_config`.""" - return resolve_sim2sim_config( - source_run_dir, target_cfg, algo_name=algo_name, strict=strict - ) - - @staticmethod - def load_dim_guard( - *, - env_obs_dim: int | None = None, - env_action_dim: int | None = None, - algo_name: str | None = None, - ): - """See :func:`policy_load_dim_guard`.""" - return policy_load_dim_guard( - env_obs_dim=env_obs_dim, env_action_dim=env_action_dim, algo_name=algo_name - ) diff --git a/tests/assets/test_hub.py b/tests/assets/test_hub.py index a4695e72f..664adf326 100644 --- a/tests/assets/test_hub.py +++ b/tests/assets/test_hub.py @@ -255,7 +255,6 @@ def test_robot_asset_specs_cover_hf_hosted_robots(): for robot, specs in ROBOT_ASSET_SPECS.items(): assert specs, robot for directory, marker, pattern, label in specs: - # go2w additionally reference the shared go2 mesh dir. assert directory.startswith("robots/") assert marker and pattern and label diff --git a/tests/base/backend/test_drake_batch_pool.py b/tests/base/backend/test_drake_batch_pool.py index 59f07bc90..0a359a729 100644 --- a/tests/base/backend/test_drake_batch_pool.py +++ b/tests/base/backend/test_drake_batch_pool.py @@ -39,7 +39,7 @@ def _run_clean_python(code: str) -> str: return result.stdout -_GO1_POOL_HELPER = """ +_GO2_POOL_HELPER = """ import xml.etree.ElementTree as ET import numpy as np @@ -53,8 +53,8 @@ def _run_clean_python(code: str) -> str: ) -def make_go1_pool(nbatch, nthread): - source_model = ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml" +def make_go2_pool(nbatch, nthread): + source_model = ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml" drake_model = materialize_drake_compatible_mjcf(source_model) contract = parse_mjcf_model_contract(drake_model.model_file) ( @@ -123,7 +123,7 @@ def test_batch_backend_mode_rejects_existing_pydrake_module() -> None: try: create_backend( "drake", - SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml")), + SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml")), 1, 0.01, drake_backend_mode="batch", @@ -148,7 +148,7 @@ def test_create_backend_rejects_pydrake_mode() -> None: with pytest.raises(ValueError, match="drake_backend_mode='batch'"): create_backend( "drake", - SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml")), + SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml")), 1, 0.01, drake_backend_mode="pydrake", @@ -166,7 +166,7 @@ def test_drake_backend_constructs_without_task_base_name() -> None: backend = create_backend( "drake", - SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml")), + SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml")), 1, 0.01, drake_backend_mode="batch", @@ -251,14 +251,14 @@ def test_drake_uni_runtime_import_is_lazy_and_pydrake_free() -> None: not _batch_extension_built(), reason="optional Drake batch extension has not been built", ) -def test_drake_batch_pool_go1_smoke_shapes_and_time() -> None: +def test_drake_batch_pool_go2_smoke_shapes_and_time() -> None: output = _run_clean_python( - _GO1_POOL_HELPER + _GO2_POOL_HELPER + textwrap.dedent( """ import json - pool, qpos, _, state = make_go1_pool(2, 1) + pool, qpos, _, state = make_go2_pool(2, 1) control = np.tile(qpos[7:], (2, 1)) state_only = pool.step(state, 2, control, None, False) output = pool.step(state, 2, control, None, True) @@ -285,7 +285,7 @@ def test_drake_batch_pool_go1_smoke_shapes_and_time() -> None: "has_sensor_data": True, "nthread": 1, "sensor_finite": True, - "sensor_shape": [2, 42], + "sensor_shape": [2, 68], "state_finite": True, "state_only_has_sensor_data": False, "state_shape": [2, 38], @@ -300,12 +300,12 @@ def test_drake_batch_pool_go1_smoke_shapes_and_time() -> None: ) def test_drake_batch_pool_uses_thread_workspaces_not_env_workspaces() -> None: output = _run_clean_python( - _GO1_POOL_HELPER + _GO2_POOL_HELPER + textwrap.dedent( """ import json - _, qpos, qvel, state = make_go1_pool(4, 1) + _, qpos, qvel, state = make_go2_pool(4, 1) for env_index in range(4): row_qpos = qpos.copy() row_qpos[0] += 0.05 * env_index @@ -313,7 +313,7 @@ def test_drake_batch_pool_uses_thread_workspaces_not_env_workspaces() -> None: state[env_index, 1 + qpos.size :] = qvel def make_pool(nthread): - return make_go1_pool(4, nthread)[0] + return make_go2_pool(4, nthread)[0] control = np.tile(qpos[7:], (4, 1)) serial_pool = make_pool(1) @@ -349,7 +349,7 @@ def make_pool(nthread): "parity_sensor": True, "parity_state": True, "reset_sensor_finite": True, - "reset_sensor_shape": [1, 42], + "reset_sensor_shape": [1, 68], "reset_times": [0.0], "reset_x": [1.23], "serial_workspace_count": 1, @@ -365,12 +365,12 @@ def make_pool(nthread): ) def test_drake_batch_pool_worker_exception_reaches_python() -> None: output = _run_clean_python( - _GO1_POOL_HELPER + _GO2_POOL_HELPER + textwrap.dedent( """ import json - pool, qpos, _, state = make_go1_pool(4, 2) + pool, qpos, _, state = make_go2_pool(4, 2) state[2, 0] = np.nan control = np.tile(qpos[7:], (4, 1)) try: @@ -491,7 +491,7 @@ def test_create_backend_batch_mode_avoids_pydrake_and_steps() -> None: assert "pydrake" not in sys.modules backend = create_backend( "drake", - SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml")), + SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml")), 2, 0.01, drake_backend_mode="batch", @@ -506,7 +506,7 @@ def test_create_backend_batch_mode_avoids_pydrake_and_steps() -> None: backend.step(backend.get_dof_pos(), nsteps=1) diagnostics = backend.diagnostics() foot_contact = backend.get_sensor_data("FL_foot_contact") - body_ids = backend.get_body_ids(["trunk", "FR_calf"]) + body_ids = backend.get_body_ids(["base", "FR_calf"]) body_pos = backend.get_body_pos_w(body_ids) summary = { "cls": type(backend).__name__, @@ -532,7 +532,7 @@ def test_create_backend_batch_mode_avoids_pydrake_and_steps() -> None: "body_finite": True, "body_shape": [2, 2, 3], "cls": "DrakeBackend", - "contact_shape": [2, 3], + "contact_shape": [2, 1], "contact_nonzero": True, "diagnostic_mode": "batch", "foot_shape": [2, 3], @@ -561,7 +561,7 @@ def test_drake_backend_pre_step_hook_refreshes_between_substeps() -> None: backend = create_backend( "drake", - SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml")), + SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml")), 1, 0.01, drake_backend_mode="batch", @@ -608,13 +608,13 @@ def test_drake_backend_body_frame_getters_use_compact_root_frame() -> None: backend = create_backend( "drake", - SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go1/scene_flat.xml")), + SceneCfg(model_file=str(ASSETS_ROOT_PATH / "robots/go2/scene_flat.xml")), 1, 0.01, drake_backend_mode="batch", drake_nthread=1, ) - base_id = backend.get_body_ids(["trunk"]) + base_id = backend.get_body_ids(["base"]) other_id = backend.get_body_ids(["FR_calf"]) base_pos = backend.get_body_pos_b(base_id) base_quat = backend.get_body_quat_b(base_id) diff --git a/tests/base/backend/test_process_device.py b/tests/base/backend/test_process_device.py index 8c51b61f3..7c1b693f9 100644 --- a/tests/base/backend/test_process_device.py +++ b/tests/base/backend/test_process_device.py @@ -170,8 +170,8 @@ def test_non_gpu_backend_is_left_untouched() -> None: def test_newton_override_carries_cuda_device_string() -> None: - # uni_rl's collector binder gate only knows mjwarp, so newton's rank-local - # device must reach spawn collectors as a ``cuda:N`` override string. + # Newton consumes an explicit ``cuda:N`` string, so its rank-local device + # reaches spawn collectors as an override string rather than an integer id. owner_override: dict[str, object] = {"newton_device": None, "nested": {"keep": True}} routed = apply_backend_env_device_override( owner_override, diff --git a/tests/base/test_sim_backend.py b/tests/base/test_sim_backend.py index df1bbddf5..2a7fe29dc 100644 --- a/tests/base/test_sim_backend.py +++ b/tests/base/test_sim_backend.py @@ -24,7 +24,6 @@ def _xml(robot: str, scene: str = "scene_flat.xml") -> str: BASIC_ROBOTS = [ pytest.param(dict(model_file=_xml("g1"), base_name="pelvis"), id="g1"), - pytest.param(dict(model_file=_xml("go1"), base_name="trunk"), id="go1"), pytest.param(dict(model_file=_xml("go2"), base_name="base"), id="go2"), ] diff --git a/tests/base/test_superdex_backend_options.py b/tests/base/test_superdex_backend_options.py index 8adc2df69..7d15f306f 100644 --- a/tests/base/test_superdex_backend_options.py +++ b/tests/base/test_superdex_backend_options.py @@ -93,7 +93,7 @@ def create(name: str, scene: SceneCfg, num_envs: int, dt: float, **kwargs: Any) assert captured["superdex_num_workers"] == workers -def test_superdex_execution_mode_is_forwarded_only_when_serial( +def test_superdex_execution_mode_is_always_forwarded( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, Any] = {} @@ -111,13 +111,11 @@ def create(name: str, scene: SceneCfg, num_envs: int, dt: float, **kwargs: Any) backend_factory.create_backend( "superdex", scene, 1, 0.002, **backend_factory.env_backend_kwargs(serial) ) - # Legacy unisim-core releases predate the option; the default stays absent - # so they still accept the SuperDex kwargs. assert captured["superdex_execution_mode"] == "serial" backend_factory.create_backend( "superdex", scene, 1, 0.002, **backend_factory.env_backend_kwargs(EnvCfg()) ) - assert "superdex_execution_mode" not in captured + assert captured["superdex_execution_mode"] == "batch" def test_superdex_execution_mode_validates_and_does_not_leak( diff --git a/tests/benchmark/test_replay_buffer_placement_benchmark.py b/tests/benchmark/test_replay_buffer_placement_benchmark.py index 4b2ce02c7..6ba406cdc 100644 --- a/tests/benchmark/test_replay_buffer_placement_benchmark.py +++ b/tests/benchmark/test_replay_buffer_placement_benchmark.py @@ -133,7 +133,7 @@ def test_replay_transfer_manifest_records_backend_fields() -> None: assert manifest["backend"] == "TorchCopyReplayTransferBackend" assert manifest["device_family"] == "cpu" assert manifest["host_memory_kind"] == "pageable_shared" - assert manifest["supports_async_submit"] is False + assert "supports_async_submit" not in manifest assert manifest["ring_depth"] == 2 diff --git a/tests/test_cli_runtime_requirements.py b/tests/test_cli_runtime_requirements.py index de8fbb7af..b4c79390e 100644 --- a/tests/test_cli_runtime_requirements.py +++ b/tests/test_cli_runtime_requirements.py @@ -62,13 +62,3 @@ def test_superdex_missing_runtime_reports_python_and_sdk(monkeypatch: pytest.Mon SystemExit, match=r"Python 3\.12.*Physics/Robotics.*uv sync --extra superdex" ): cli._check_runtime_requirements("ppo", "superdex") - - -def test_superdex_old_unisim_reports_adapter_requirement( - monkeypatch: pytest.MonkeyPatch, -) -> None: - import sys - - monkeypatch.setitem(sys.modules, "unisim.backend.superdex.dependencies", None) - with pytest.raises(SystemExit, match="unisim-core>=1.1.5 with the SuperDex adapter"): - cli._check_runtime_requirements("ppo", "superdex") diff --git a/tests/test_export_scene.py b/tests/test_export_scene.py deleted file mode 100644 index 044e3fd91..000000000 --- a/tests/test_export_scene.py +++ /dev/null @@ -1,66 +0,0 @@ -"""Tests for scene export tool.""" - -from __future__ import annotations - -from pathlib import Path -from zipfile import ZipFile - -import pytest - -pytest.importorskip("mujoco") - -from unisim.backend.mujoco.export_scene import export_scene, main - -MINIMAL_XML = """\ - - - - - - - -""" - - -@pytest.fixture() -def model_file(tmp_path: Path) -> str: - p = tmp_path / "test_model.xml" - p.write_text(MINIMAL_XML) - return str(p) - - -def test_export_creates_xml(model_file: str, tmp_path: Path): - out = tmp_path / "export_out" - export_scene(model_file, str(out)) - assert (out / "scene.xml").is_file() - - -def test_export_directory_structure(model_file: str, tmp_path: Path): - out = tmp_path / "export_out2" - export_scene(model_file, str(out)) - assert out.is_dir() - assert (out / "scene.xml").is_file() - - -def test_exported_scene_reloadable(model_file: str, tmp_path: Path): - import mujoco - - out = tmp_path / "export_out3" - export_scene(model_file, str(out)) - model = mujoco.MjModel.from_xml_path(str(out / "scene.xml")) - assert model.ngeom > 0 - - -def test_export_scene_zip_contains_xml(model_file: str, tmp_path: Path): - out = tmp_path / "export_zip" - zip_path = Path(export_scene(model_file, str(out), as_zip=True)) - assert zip_path.is_file() - assert zip_path.suffix == ".zip" - with ZipFile(zip_path) as zf: - assert "scene.xml" in zf.namelist() - - -def test_export_scene_cli_creates_xml(model_file: str, tmp_path: Path): - out = tmp_path / "export_cli" - assert main([model_file, "-o", str(out)]) == 0 - assert (out / "scene.xml").is_file() diff --git a/tests/training/test_seed_contract.py b/tests/training/test_seed_contract.py index 437acb012..e58953354 100644 --- a/tests/training/test_seed_contract.py +++ b/tests/training/test_seed_contract.py @@ -12,7 +12,6 @@ from unilab.utils.seed import ( apply_training_seed, - derive_worker_seed, resolve_training_seed, ) @@ -53,12 +52,6 @@ def test_apply_training_seed_rejects_negative_seed(): apply_training_seed(-1) -def test_derive_worker_seed_is_deterministic_and_distinct_from_base_seed(): - assert derive_worker_seed(10, worker_index=0) == 11 - assert derive_worker_seed(10, worker_index=3) == 14 - assert derive_worker_seed(None, worker_index=3) is None - - @pytest.mark.parametrize( ("config_dir", "overrides"), [ diff --git a/tests/training/test_sim2sim_resolver.py b/tests/training/test_sim2sim_resolver.py index d45104454..a734e8487 100644 --- a/tests/training/test_sim2sim_resolver.py +++ b/tests/training/test_sim2sim_resolver.py @@ -20,7 +20,6 @@ ENV_STRUCTURAL_DENYLIST, WARNING_LIST, CrossBackendIncompatibleError, - Sim2SimConfigResolver, extract_contract_snapshot, policy_load_dim_guard, resolve_sim2sim_config, @@ -397,42 +396,6 @@ def test_dim_guard_does_not_swallow_keyerror(): raise KeyError("actor") -# --- Sim2SimConfigResolver class facade + user-level bypass ------------------------ - - -def test_resolver_class_exposes_field_lists(): - assert Sim2SimConfigResolver.DENYLIST is DENYLIST - assert Sim2SimConfigResolver.WARNING_LIST is WARNING_LIST - assert Sim2SimConfigResolver.ALLOWLIST is ALLOWLIST - assert Sim2SimConfigResolver.ENV_STRUCTURAL_DENYLIST is ENV_STRUCTURAL_DENYLIST - - -def test_resolver_class_resolve_delegates_and_raises(tmp_path): - _write_sidecar(tmp_path, extract_contract_snapshot(_mujoco_cfg())) - target = _mujoco_cfg() - target.env.control_config.action_scale = 0.5 - with pytest.raises(CrossBackendIncompatibleError): - Sim2SimConfigResolver.resolve(tmp_path, target) - - -def test_resolver_class_strict_false_is_user_bypass(tmp_path, capsys): - # training.sim2sim_strict=false maps to strict=False: a DENYLIST denial becomes a - # warning and play proceeds with the target cfg (the load-time dim guard still bites). - _write_sidecar(tmp_path, extract_contract_snapshot(_mujoco_cfg())) - target = _mujoco_cfg() - target.env.control_config.action_scale = 0.5 - assert Sim2SimConfigResolver.resolve(tmp_path, target, strict=False) is target - assert "action_scale" in capsys.readouterr().out - - -def test_resolver_class_extract_and_dim_guard_delegate(): - snap = Sim2SimConfigResolver.extract_snapshot(_mujoco_cfg()) - assert snap["env.control_config.action_scale"] == 0.25 - with pytest.raises(CrossBackendIncompatibleError): - with Sim2SimConfigResolver.load_dim_guard(env_obs_dim=5, env_action_dim=2): - raise RuntimeError("size mismatch for actor.0.weight") - - def _compose_task(task: str) -> Any: conf_dir = str(Path(__file__).resolve().parents[2] / "src" / "unilab" / "conf" / "ppo") GlobalHydra.instance().clear() diff --git a/tests/utils/test_utils_package_policy.py b/tests/utils/test_utils_package_policy.py index f026b3e88..d35cc516c 100644 --- a/tests/utils/test_utils_package_policy.py +++ b/tests/utils/test_utils_package_policy.py @@ -8,7 +8,6 @@ "checkpoint", "device", "geometry", - "monitoring", "nan_guard", "nan_viz", "reward", diff --git a/uv.lock b/uv.lock index 211b3f755..d6919a5a9 100644 --- a/uv.lock +++ b/uv.lock @@ -5239,9 +5239,9 @@ requires-dist = [ { name = "trimesh", marker = "extra == 'newton'", specifier = ">=3.21.7" }, { name = "trimesh", marker = "extra == 'viser'", specifier = ">=3.21.7" }, { name = "typing-extensions" }, - { name = "unilab-rl", specifier = "==1.2.0" }, - { name = "unisim-core", specifier = ">=1.4.1" }, - { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = ">=1.4.0" }, + { name = "unilab-rl", specifier = "==1.2.1" }, + { name = "unisim-core", specifier = ">=1.4.2" }, + { name = "unisim-core", extras = ["superdex"], marker = "python_full_version >= '3.12' and platform_machine == 'x86_64' and sys_platform == 'linux' and extra == 'superdex'", specifier = ">=1.4.2" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, @@ -5260,7 +5260,7 @@ dev = [ [[package]] name = "unilab-rl" -version = "1.2.0" +version = "1.2.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "hydra-core" }, @@ -5277,20 +5277,20 @@ dependencies = [ { name = "torch", version = "2.9.0+cu130", source = { registry = "https://download-r2.pytorch.org/whl/cu130" }, marker = "platform_machine == 'aarch64' and sys_platform == 'linux'" }, { name = "wandb" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f5/9d/7a646bc044cadb48cbd1d0c58832faf59e037a45f22432cff2a7ad5a7b01/unilab_rl-1.2.0.tar.gz", hash = "sha256:bb8b1c713b961ea4b24ea72c39d443e057531e71882fa058502af947b98768da", size = 148289, upload-time = "2026-09-10T05:20:52.549Z" } +sdist = { url = "https://files.pythonhosted.org/packages/65/44/2280a1b68c604b3e4b8bd350e34cde0b9d55e91f43266bfa23c7d571a7cb/unilab_rl-1.2.1.tar.gz", hash = "sha256:e9749709e55ceeb9321cd0e2aa1b655cbd276901398003702dc13fb8d71cda85", size = 144847, upload-time = "2026-09-15T14:59:48.995Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/67/0ea1fb5f0c9a7e8604b0415b018deb3286024191a84d357f1b3221d2a41d/unilab_rl-1.2.0-py3-none-any.whl", hash = "sha256:617317dcca581d1af2daec36d50de7214067dd34f5d85a4f02d7cd38de1df611", size = 183372, upload-time = "2026-09-10T05:20:51.254Z" }, + { url = "https://files.pythonhosted.org/packages/43/d0/384be1d5ed60dfbec4620340784a84042227f396a62cc1a6b965744042eb/unilab_rl-1.2.1-py3-none-any.whl", hash = "sha256:516874fca1649bcfd080c7226e039fd646f7acdfff4c1e97b6f2a27c3264c152", size = 178664, upload-time = "2026-09-15T14:59:47.945Z" }, ] [[package]] name = "unisim-core" -version = "1.4.1" +version = "1.4.2" 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'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f9/57/4da896cef1f76e9902cfe4deb0f5dcf678a2217507dd39fe286a7edb6984/unisim_core-1.4.1.tar.gz", hash = "sha256:aa71cc1f90615929be2e75851d636271675666caa140629c042e6a46fb45d657", size = 269776, upload-time = "2026-09-15T08:21:16.014Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/e7/5954d490d66facf3449f1032d5cc13288fb4fd1166eaba4eea8987c4f4ca/unisim_core-1.4.2.tar.gz", hash = "sha256:e27d89292083cdff646d2dbeb0ac787ee4351e7ef811c20c6fed2a40578785b7", size = 269413, upload-time = "2026-09-16T04:22:39.708Z" } [package.optional-dependencies] superdex = [ diff --git a/uv.rocm.lock b/uv.rocm.lock index 8a620d3ce..7725f2723 100644 --- a/uv.rocm.lock +++ b/uv.rocm.lock @@ -3789,8 +3789,8 @@ requires-dist = [ { name = "trimesh", marker = "extra == 'viser'", specifier = ">=3.21.7" }, { name = "triton-rocm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==3.6.0", index = "https://download.pytorch.org/whl/rocm7.2" }, { name = "typing-extensions" }, - { name = "unilab-rl", specifier = "==1.2.0" }, - { name = "unisim-core", specifier = ">=1.4.1" }, + { name = "unilab-rl", specifier = "==1.2.1" }, + { name = "unisim-core", specifier = ">=1.4.2" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, ] @@ -3807,7 +3807,7 @@ dev = [ [[package]] name = "unilab-rl" -version = "1.2.0" +version = "1.2.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "hydra-core" }, @@ -3823,20 +3823,20 @@ dependencies = [ { name = "torch", version = "2.11.0+rocm7.2", source = { registry = "https://download.pytorch.org/whl/rocm7.2" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "wandb" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f5/9d/7a646bc044cadb48cbd1d0c58832faf59e037a45f22432cff2a7ad5a7b01/unilab_rl-1.2.0.tar.gz", hash = "sha256:bb8b1c713b961ea4b24ea72c39d443e057531e71882fa058502af947b98768da", size = 148289, upload-time = "2026-09-10T05:20:52.549Z" } +sdist = { url = "https://files.pythonhosted.org/packages/65/44/2280a1b68c604b3e4b8bd350e34cde0b9d55e91f43266bfa23c7d571a7cb/unilab_rl-1.2.1.tar.gz", hash = "sha256:e9749709e55ceeb9321cd0e2aa1b655cbd276901398003702dc13fb8d71cda85", size = 144847, upload-time = "2026-09-15T14:59:48.995Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/67/0ea1fb5f0c9a7e8604b0415b018deb3286024191a84d357f1b3221d2a41d/unilab_rl-1.2.0-py3-none-any.whl", hash = "sha256:617317dcca581d1af2daec36d50de7214067dd34f5d85a4f02d7cd38de1df611", size = 183372, upload-time = "2026-09-10T05:20:51.254Z" }, + { url = "https://files.pythonhosted.org/packages/43/d0/384be1d5ed60dfbec4620340784a84042227f396a62cc1a6b965744042eb/unilab_rl-1.2.1-py3-none-any.whl", hash = "sha256:516874fca1649bcfd080c7226e039fd646f7acdfff4c1e97b6f2a27c3264c152", size = 178664, upload-time = "2026-09-15T14:59:47.945Z" }, ] [[package]] name = "unisim-core" -version = "1.4.1" +version = "1.4.2" 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'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f9/57/4da896cef1f76e9902cfe4deb0f5dcf678a2217507dd39fe286a7edb6984/unisim_core-1.4.1.tar.gz", hash = "sha256:aa71cc1f90615929be2e75851d636271675666caa140629c042e6a46fb45d657", size = 269776, upload-time = "2026-09-15T08:21:16.014Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b9/e7/5954d490d66facf3449f1032d5cc13288fb4fd1166eaba4eea8987c4f4ca/unisim_core-1.4.2.tar.gz", hash = "sha256:e27d89292083cdff646d2dbeb0ac787ee4351e7ef811c20c6fed2a40578785b7", size = 269413, upload-time = "2026-09-16T04:22:39.708Z" } [[package]] name = "urllib3"