diff --git a/docs/sphinx/source/adr/ADR-0000-index.md b/docs/sphinx/source/adr/ADR-0000-index.md index 7a7dc6e8b..1ffecb582 100644 --- a/docs/sphinx/source/adr/ADR-0000-index.md +++ b/docs/sphinx/source/adr/ADR-0000-index.md @@ -21,6 +21,7 @@ orphan: true | [ADR-0005 Unified Obs Critic Env And IPC Contract](ADR-0005-unified-obs-critic-env-and-ipc-contract.md) | Observation / IPC | Accepted | | [ADR-0006 Community Manager API On NumPy Runtime](ADR-0006-community-manager-api-on-numpy-runtime.md) | Manager API / NumPy runtime | Accepted | | [ADR-0007 UniSim Extraction Boundary](ADR-0007-unisim-extraction-boundary.md) | Physics package extraction | Accepted | +| [ADR-0008 SuperDex Native C++ Scene Batch Executor](ADR-0008-superdex-persistent-cpu-workers.md) | Backend CPU scene execution | Accepted | ## ADR Governance diff --git a/docs/sphinx/source/adr/ADR-0008-superdex-persistent-cpu-workers.md b/docs/sphinx/source/adr/ADR-0008-superdex-persistent-cpu-workers.md new file mode 100644 index 000000000..57470f435 --- /dev/null +++ b/docs/sphinx/source/adr/ADR-0008-superdex-persistent-cpu-workers.md @@ -0,0 +1,82 @@ +--- +orphan: true +--- + +# ADR-0008 SuperDex Native C++ Scene Batch Executor + +- Status: Accepted +- Date: 2026-09-07 +- Owners: SuperDex fork / UniSim backend / UniLab task-config maintainers +- Supersedes: None +- Superseded by: None + +## Context + +The initial SuperDex adapter creates one independent native scene for every +environment but advances them serially from Python. `Scene.step()` releases the +GIL, yet a Python thread pool would still leave generalized-force writes and +articulated state reads as separate Python-to-C++ calls. A subprocess/shared +memory design was considered and rejected before implementation because it +would introduce a new runtime protocol beyond the requested integration scope. + +The maintained fork [unilabsim/project_superdex](https://github.com/unilabsim/project_superdex) +can build an extension against the exact Physics and Robotics sources. This +removes the wheel/header ABI mismatch that prevents a downstream native shim. +The roadmap is [UniLab#1533](https://github.com/Motphys/UniLab/issues/1533). + +## Decision + +1. Add `superdex.physics.SceneBatchExecutor` to the fork's `mochi_physics` + pybind extension. It owns persistent C++ worker threads, but does not own + scenes or actors. +2. One executor invocation receives contiguous generalized-force, articulated + pose/velocity, link-state, contact and solver-status arrays. It writes each + actor's forces, advances each distinct scene, and refreshes every runtime + cache before its completion barrier opens. +3. UniSim owns the executor and creates it after cold-path materialization. It + keeps reset, asset conversion, cache-frame conversion and `SimBackend` + ownership in the adapter. `close()` joins the executor + before destroying bots, scenes, and the process-global runtime. +4. `superdex_num_workers=0` selects `min(available physical CPU cores, + num_envs)`. Explicit worker counts are capped at `num_envs`; the SDK stays + single-threaded so the executor is the only physics parallelism layer. + +## Consequences + +The local integration requires Physics and Robotics bindings built from this +fork at the same commit. Older qpos/qvel-only executor builds are rejected. It +neither changes package versions nor publishes a wheel. The executor provides +CPU scene parallelism; it does not claim GPU physics, native rendering, or +dynamics equivalence with MuJoCo. + +Validation must compare serial and parallel trajectories, selected reset, and +complete backend throughput using the same scene, actions, batch, and substep +count. Training evidence must report the actual outer worker count and SDK +thread count separately. + +## Alternatives Considered + +- Python `ThreadPoolExecutor`: useful for a narrow step probe, but does not + fuse force/state binding calls or give the adapter a durable native barrier. +- Subprocess IPC and shared memory: rejected because SuperDex does not require + process isolation and the protocol is outside this roadmap's approved scope. +- Linking a downstream extension to the released wheel: rejected because the + wheel does not publish a stable extension ABI or matching headers. + +## Evidence In Repo + +- `src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml` selects automatic + native workers for the SuperDex owner. +- `src/unilab/scripts/train_offpolicy.py` and + `src/unilab/scripts/train_rsl_rl.py` obtain rank-owned CPU ids from UniRL + before environment construction. +- `tests/algos/test_offpolicy_double_buffer_runner.py` verifies that each rank + receives complete physical-core groups, including SMT siblings. +- `tests/ipc/test_dp_launcher.py` verifies the corresponding resolver contract + through the UniRL dependency. + +## Related Documents + +- {doc}`SuperDex Backend ` +- {doc}`UniSim Extraction Boundary ` +- [UniLab#1533](https://github.com/Motphys/UniLab/issues/1533) diff --git a/docs/sphinx/source/adr/README.md b/docs/sphinx/source/adr/README.md index 772be3814..1ffecb582 100644 --- a/docs/sphinx/source/adr/README.md +++ b/docs/sphinx/source/adr/README.md @@ -20,6 +20,8 @@ orphan: true | [ADR-0004 Registry Bootstrap Contract](ADR-0004-registry-bootstrap-contract.md) | Registry bootstrap | Accepted | | [ADR-0005 Unified Obs Critic Env And IPC Contract](ADR-0005-unified-obs-critic-env-and-ipc-contract.md) | Observation / IPC | Accepted | | [ADR-0006 Community Manager API On NumPy Runtime](ADR-0006-community-manager-api-on-numpy-runtime.md) | Manager API / NumPy runtime | Accepted | +| [ADR-0007 UniSim Extraction Boundary](ADR-0007-unisim-extraction-boundary.md) | Physics package extraction | Accepted | +| [ADR-0008 SuperDex Native C++ Scene Batch Executor](ADR-0008-superdex-persistent-cpu-workers.md) | Backend CPU scene execution | Accepted | ## ADR Governance diff --git a/docs/sphinx/source/conf.py b/docs/sphinx/source/conf.py index 26f7a80db..03ccb3f24 100644 --- a/docs/sphinx/source/conf.py +++ b/docs/sphinx/source/conf.py @@ -277,6 +277,7 @@ _LANGUAGE_PATH_FORWARD: dict[str, str] = { "en/why_unilab": "zh_CN/why_unilab", "en/2-user_guide/3-backends/6-drake": "zh_CN/2-user_guide/3-backends/6-drake", + "en/2-user_guide/3-backends/8-superdex": "zh_CN/2-user_guide/3-backends/8-superdex", "en/1-getting_started/5-faq": "zh_CN/1-getting_started/5-faq", "en/4-developer_guide/1-architecture/6-manager_based_api": ( "zh_CN/4-developer_guide/1-architecture/6-manager_based_api" diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/0-index.md b/docs/sphinx/source/en/2-user_guide/3-backends/0-index.md index 61cf595ea..f6abea40e 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/0-index.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/0-index.md @@ -1,7 +1,7 @@ # Simulation Backends UniLab exposes backend names through registry/config paths, including `mujoco`, -`motrix`, `mjwarp`, `drake`, `isaacgym`, `genesis`, `isaacsim`, and `newton` +`motrix`, `mjwarp`, `drake`, `isaacgym`, `genesis`, `isaacsim`, `newton`, and `superdex` where an owner is registered. User commands select them with `--sim`, which routes to the matching task owner YAML; do not switch a run by overriding `training.sim_backend` alone. @@ -146,4 +146,5 @@ separately authorized issue. 5-genesis 6-drake 7-newton +8-superdex ``` diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md b/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md new file mode 100644 index 000000000..038ee5660 --- /dev/null +++ b/docs/sphinx/source/en/2-user_guide/3-backends/8-superdex.md @@ -0,0 +1,156 @@ +# SuperDex Backend + +SuperDex is an optional CPU physics adapter owned by `unisim.backend.superdex`. +The initial UniLab owner is the fixed-base `FR3JointTarget` task: +`src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml`. It uses seven torque +actions, 21 observation values, joint-state resets and the standard NumPy +manager runtime. Its support level is **Configured**; bounded rollout or short +training checks do not establish full-training performance or platform support. +The implementation is tracked in [#1534](https://github.com/Motphys/UniLab/issues/1534) +under [roadmap #1533](https://github.com/Motphys/UniLab/issues/1533). + +## Local Development Setup + +This development profile uses locally linked UniSim and UniLab checkouts; it +does not require a new published version. SuperDex Physics/Robotics 1.0.0 requires +Python 3.12. CPU physics does not require CUDA. Native rendering and video are +not part of the FR3 owner; training defaults to `no_play=true`. + +In the UniLab checkout, use an existing Python 3.12 virtual environment or create +one, then install the local packages: + +```bash +uv venv --python 3.12 +export UNILAB_LOCAL_UNISIM=/absolute/path/to/unisim +uv pip install -e "${UNILAB_LOCAL_UNISIM}[superdex,mujoco]" -e . --group pyproject.toml:dev +export UV_NO_SYNC=1 +export SUPERDEX_ASSETS_PATH=/absolute/path/to/project_superdex/assets +``` + +`UNILAB_LOCAL_UNISIM` enables the repository's local dependency validation: +tests require an editable installation and check that its metadata and imported +module point to exactly that checkout. Without this variable, the normal +indexed-release requirement remains in force. `UV_NO_SYNC=1` preserves the local +links when running existing Make targets; `uv sync` would re-resolve the locked +release profile. + +Native FR3 assets remain in the upstream SuperDex checkout. The asset hub +registers `bots/arms/fr3_v2/fr3_v2.superdex_bot`, checking its collision SDF, +render files, `LICENSE` and `NOTICE` before constructing physics. No native +robot binaries are bundled or downloaded by UniLab. Set +`env.superdex_assets_root=/absolute/path/to/project_superdex/assets` to override +`SUPERDEX_ASSETS_PATH` for a specific owner invocation. + +## Run the FR3 Task + +```bash +uv run --no-sync train --algo ppo --task fr3_joint_target --sim superdex \ + algo.max_iterations=2 algo.num_steps_per_env=16 \ + algo.algorithm.num_learning_epochs=1 +``` + +The target joint positions, rewards, reset ranges and action scales live in the +task's `base.yaml`. The torque bounds `[20,20,20,20,5,5,5]` Nm are an explicit +research profile, not rated hardware limits. `superdex_effort_limits` declares +the same bounds at the native backend boundary. The SDK remains single-threaded; +the native scene executor below owns all supported CPU parallelism. + +## Default CPU Environment Parallelism + +The backend uses SuperDex's source-built `SceneBatchExecutor`, a persistent C++ +thread pool that owns the barrier across independent scenes. Each substep writes +batched generalized forces, advances scenes, and returns articulated/link state, +contact sensors and solver status without per-environment Python binding calls. +Asset materialization, reset and cache-frame transforms remain owned by the +UniSim adapter. This is CPU thread parallelism, not GPU physics, and it does +not change the PPO/APPO collector, learner or policy contracts. The decision is +recorded in {doc}`/adr/ADR-0008-superdex-persistent-cpu-workers` and +[unisim#41](https://github.com/unilabsim/unisim/issues/41). + +Both task owners select automatic workers by default: + +| Owner option | Meaning | +| --- | --- | +| `env.superdex_num_workers=0` | Automatic: `min(available physical CPU cores, num_envs)` | +| `env.superdex_num_workers=1` | One native C++ scene worker | +| `env.superdex_num_workers=K` | Explicit C++ worker count, capped at `num_envs` | + +With 1024 environments on a 16-core/32-thread host, automatic selection yields +16 workers. Concurrent multi-rank collectors are assigned whole physical-core +groups, including their logical siblings, so ranks do not split an SMT core. +`training.dp_collector_cpu_ids` may instead provide one explicit CPU-id list +per rank. The selected block is applied before SuperDex materializes its native +worker pool. + +For every physics substep, the host runs the pre-step control callback, enters +the native batch barrier, then publishes the refreshed batch before the next +callback. Selected reset preserves caller row order and leaves unselected +environments unchanged. A native worker failure closes the executor and reports +the failure; it does not silently return stale state or switch to serial. +Closing an environment joins its C++ workers before scenes are destroyed. + +Worker count alone is not evidence of speedup. Throughput comparisons must use +the same model, control sequence, batch size and substeps, and report complete +backend/env time, startup, RSS, CPU use and actual worker count. Count environment +control steps, not physics substeps. Existing contact approximations are unchanged. + +The fixed root still has a named entity and readable body state. Reset terms +write joint state; they do not request a floating-root layout. The task does +not require contact sensors, cameras, site Jacobians or runtime material DR. +SuperDex has no native renderer. Its default record playback uses the offline +MuJoCo renderer with the authored MJCF visual model while SuperDex remains the +physics backend. `.superdex_bot` scenes must provide `visual_model_file` for +this path. Selecting playback mode `none` skips playback entirely; it is not +evidence that a checkpoint has executed a rollout. + +`superdex_allow_contact_approximation` defaults to `false`. It is reserved for +explicitly audited MJCF conversion profiles: enabling it accepts a warning about +contact/material approximation, including missing torsional/rolling friction +equivalence. It does not establish arbitrary MJCF task compatibility. +Contact queries report the last completed solver step. Reset clears this state; +it does not provide a fresh geometric overlap test until a positive physics step +has completed. Kinematic body/joint getters are refreshed immediately at reset. + +## Validation and Ownership + +The `go2_joystick_flat/superdex` PPO owner is a research sim2sim profile. It +inherits the MuJoCo owner, preserving 49 actor observations, 52 critic +observations, 12 position-target actions, normalization, network dimensions and +control timing. It disables runtime PD gain randomization and explicitly opts +into contact approximation. The adapter's cold-path MJCF conversion is restricted; +this owner does not imply arbitrary scene support or equivalent walking behavior. + +Create a small source checkpoint with the MuJoCo owner, then pass its path to the +optional checkpoint test: + +```bash +uv run --no-sync train --algo ppo --task go2_joystick_flat --sim mujoco \ + algo.num_envs=2 algo.max_iterations=2 algo.num_steps_per_env=16 \ + algo.algorithm.num_mini_batches=1 algo.algorithm.num_learning_epochs=1 \ + training.device=cpu training.no_play=true +export UNILAB_SUPERDEX_GO2_CHECKPOINT=/absolute/path/to/source/run/model_1.pt +uv run --no-sync pytest tests/envs/test_go2_superdex.py -q +``` + +This validates the source `run_config.json` before environment construction, +checks rejection of changed policy action semantics, loads the actual policy +through the production playback session and executes 64 SuperDex control steps +without a renderer. It checks finite values and interface compatibility; a +two-iteration checkpoint is not expected to walk reliably. + +```bash +uv run --no-sync pytest tests/assets/test_superdex_assets.py \ + tests/envs/test_fr3_superdex.py tests/test_cli_runtime_requirements.py -q +``` + +The optional native tests require the SDK and `SUPERDEX_ASSETS_PATH`; they cover +finite rollout data, selected reset isolation, immediate observation refresh +and a spawned `EnvFactory`. Missing runtime/assets produce an explicit skip; +such a run is not native validation. The base asset/config tests need no native +asset checkout. + +Engine conversion and physics live in UniSim; asset registration, Hydra and task +terms remain in UniLab. See +{doc}`/adr/ADR-0007-unisim-extraction-boundary`, +{doc}`/adr/ADR-0006-community-manager-api-on-numpy-runtime` and +{doc}`/adr/ADR-0002-backend-capability-boundary-for-play-and-snapshot`. diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/0-index.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/0-index.md index c5893368e..6a435f28c 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/0-index.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/0-index.md @@ -2,7 +2,7 @@ UniLab 通过 registry/config 路径暴露后端名称，包括在对应 owner 注册后可用的 `mujoco`、`motrix`、`mjwarp`、`drake`、`isaacgym`、`genesis`、`isaacsim` 和 -`newton`。用户命令通过 +`newton`、`superdex`。用户命令通过 `--sim` 选择后端，该选项会路由到对应的 task owner YAML；不要仅靠 override `training.sim_backend` 来切换一次运行。 @@ -137,4 +137,5 @@ benchmark v1 目前只保留 `BenchmarkCase`、`BenchmarkResult` 和 provenance 5-genesis 6-drake 7-newton +8-superdex ``` diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md new file mode 100644 index 000000000..f7e830743 --- /dev/null +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/8-superdex.md @@ -0,0 +1,128 @@ +# SuperDex 后端 + +SuperDex 是由 `unisim.backend.superdex` 拥有的可选 CPU 物理后端。UniLab 首个 +owner 为固定基 `FR3JointTarget`，配置位于 +`src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml`。任务使用 7 维力矩动作、 +21 维观测、关节状态 reset 和标准 NumPy manager。当前支持等级为 **Configured**； +短 rollout 或少量训练迭代不能证明完整训练效果、性能或跨平台支持。 +实施见 [#1534](https://github.com/Motphys/UniLab/issues/1534)，所属 +roadmap 为 [#1533](https://github.com/Motphys/UniLab/issues/1533)。 + +## 本地开发安装 + +该开发配置使用本地链接的 UniSim 与 UniLab，不要求发布新版本。SuperDex +Physics/Robotics 1.0.0 要求 Python 3.12；CPU 物理不需要 CUDA。FR3 owner 暂不包含 +原生渲染和视频，默认 `no_play=true`。 + +在 UniLab checkout 中使用已有的 Python 3.12 环境，或创建环境后安装本地包： + +```bash +uv venv --python 3.12 +export UNILAB_LOCAL_UNISIM=/absolute/path/to/unisim +uv pip install -e "${UNILAB_LOCAL_UNISIM}[superdex,mujoco]" -e . --group pyproject.toml:dev +export UV_NO_SYNC=1 +export SUPERDEX_ASSETS_PATH=/absolute/path/to/project_superdex/assets +``` + +`UNILAB_LOCAL_UNISIM` 启用严格的本地依赖验证：测试同时检查 editable 安装元数据 +和实际 import 路径确实指向指定 checkout。不设置时保留正常的索引发布包检查。 +`UV_NO_SYNC=1` 让现有 Make 目标保留本地链接；`uv sync` 会重新解析锁定的发布依赖。 + +FR3 原生资产保留在上游 checkout。asset hub 注册 +`bots/arms/fr3_v2/fr3_v2.superdex_bot`，在物理构造前验证 collision SDF、render、 +`LICENSE` 和 `NOTICE`。UniLab 不打包或下载这些二进制。单次运行可通过 +`env.superdex_assets_root=/absolute/path/to/project_superdex/assets` 覆盖环境变量。 + +## 运行 FR3 任务 + +```bash +uv run --no-sync train --algo ppo --task fr3_joint_target --sim superdex \ + algo.max_iterations=2 algo.num_steps_per_env=16 \ + algo.algorithm.num_learning_epochs=1 +``` + +目标关节角、reward、reset 范围和动作缩放由任务 `base.yaml` 声明。力矩上限 +`[20,20,20,20,5,5,5]` Nm 是显式研究配置，不是硬件额定值； +`superdex_effort_limits` 在 native backend 边界声明同样的上限。SDK 固定为单线程； +下面的 native scene executor 是唯一支持的 CPU 并行层。 + +## 默认 CPU 环境并行 + +backend 使用 SuperDex 源码构建的 `SceneBatchExecutor`。它是跨独立 scene 的常驻 +C++ 线程池：每个子步批量写入广义力、推进 scene，并回写 articulation/link state、 +contact sensor 和 solver status，不再逐环境跨越 Python binding。资产物化、reset 和 +cache frame 转换仍由 UniSim adapter 负责。这是 CPU 线程并行，不是 GPU physics；它不改变 +PPO/APPO collector、learner 或 policy contract。决策见 +{doc}`/adr/ADR-0008-superdex-persistent-cpu-workers` 和 +[unisim#41](https://github.com/unilabsim/unisim/issues/41)。 + +两个 task owner 默认选择自动 worker： + +| Owner 选项 | 含义 | +| --- | --- | +| `env.superdex_num_workers=0` | 自动：`min(affinity 内可用物理核心数, num_envs)` | +| `env.superdex_num_workers=1` | 一个 native C++ scene worker | +| `env.superdex_num_workers=K` | 显式 C++ worker 数，最多为 `num_envs` | + +1024 个环境、16 核 32 线程主机上，自动解析为 16 workers。多 rank 并发 collector +按完整物理核心分配，并将同一核心的 logical sibling 放在同一 rank，避免拆分 SMT 核心。 +也可以通过 `training.dp_collector_cpu_ids` 为每个 rank 显式提供 CPU id 列表；该分片在 +SuperDex 创建 native worker pool 前应用。 + +每个物理子步由 host 执行 pre-step control callback，随后进入 native batch barrier， +完成后发布新 batch state，再执行下一 callback。局部 reset 保留请求行顺序，不影响 +未选择环境。native worker 错误会关闭 executor 并报告失败，不静默返回旧状态或回退 +串行。env.close 会在销毁 scene 前 join C++ worker。 + +worker 数不能代替加速证据。吞吐对照应匹配模型、控制序列、batch 和子步，记录完整 +backend/env 时间、startup、RSS、CPU 使用和实际 worker 数。吞吐按 env 控制步统计， +不能把 physics 子步重复计入样本。已有接触近似的物理边界保持不变。 + +固定根具有名称和可读的 body state，但 reset 只写 joint state，不要求 free-root +layout。该任务不需要接触 sensor、相机、site Jacobian 或材料 DR。SuperDex 没有 native +renderer；默认 record 回放使用离线 MuJoCo renderer，物理仍由 SuperDex 执行。 +`.superdex_bot` 场景需要提供 `visual_model_file`。`play_render_mode=none` 会完全跳过 +回放，不能作为 checkpoint rollout 已执行的证据。 + +`superdex_allow_contact_approximation` 默认 `false`，只供经过审核的 MJCF 转换配置 +显式启用。启用后会警告 contact/material 近似，包括 torsional/rolling friction +不等价；这不代表任意 MJCF 任务已兼容。 +接触查询返回最近一次完成求解的结果。reset 清除该结果，需要一次正时间步才会 +产生新接触，不能把 reset 后的 contact 当成即时几何重叠测试。运动学 body/joint +getter 则在 reset 后立即刷新。 + +## 验证与归属 + +`go2_joystick_flat/superdex` PPO owner 是研究性质的 sim2sim 配置。它继承 MuJoCo +owner，保留 49 维 actor 观测、52 维 critic 观测、12 维位置目标动作、归一化、网络 +维度和控制时序；关闭 runtime PD gain DR，并显式接受接触近似。adapter 的冷路径 +MJCF 转换有明确限制，该 owner 不代表任意场景或行走效果等价。 + +先用 MuJoCo owner 创建一个小型来源 checkpoint，再把路径传给可选 checkpoint 测试： + +```bash +uv run --no-sync train --algo ppo --task go2_joystick_flat --sim mujoco \ + algo.num_envs=2 algo.max_iterations=2 algo.num_steps_per_env=16 \ + algo.algorithm.num_mini_batches=1 algo.algorithm.num_learning_epochs=1 \ + training.device=cpu training.no_play=true +export UNILAB_SUPERDEX_GO2_CHECKPOINT=/absolute/path/to/source/run/model_1.pt +uv run --no-sync pytest tests/envs/test_go2_superdex.py -q +``` + +测试在构造 env 前验证来源 `run_config.json`，检查修改动作语义时确实拒绝，随后 +通过 production playback session 加载真实策略，无渲染执行 64 个 SuperDex 控制步。 +它验证有限数值和接口兼容性；只训练两轮的 checkpoint 不以可靠行走为验收标准。 + +```bash +uv run --no-sync pytest tests/assets/test_superdex_assets.py \ + tests/envs/test_fr3_superdex.py tests/test_cli_runtime_requirements.py -q +``` + +原生测试要求 SDK 和 `SUPERDEX_ASSETS_PATH`，覆盖有限数值 rollout、局部 reset +隔离、即时观测刷新及 spawn `EnvFactory`。缺失 SDK/资产会明确 skip，不能将 skip +记为原生验证通过。基础资产和配置测试不依赖原生资产 checkout。 + +引擎转换与物理由 UniSim 拥有；资产注册、Hydra 与任务 term 由 UniLab 拥有。 +相关约束见 {doc}`/adr/ADR-0007-unisim-extraction-boundary`、 +{doc}`/adr/ADR-0006-community-manager-api-on-numpy-runtime` 和 +{doc}`/adr/ADR-0002-backend-capability-boundary-for-play-and-snapshot`。 diff --git a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md index 301fdb91d..984457226 100644 --- a/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/zh_CN/5-reference/5-support_matrix.md @@ -59,76 +59,78 @@ uv run scripts/generate_support_matrix.py --write ### Entrypoint x Task Owner -| Entrypoint | Task owner | MuJoCo | mjwarp | Motrix | IsaacGym | Genesis | IsaacSim | Newton | -|------------|------------|--------|--------|--------|----------|---------|----------|--------| -| PPO (torch) | `go1_joystick_flat` (Go1 joystick) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `go2_joystick_rough` (Go2 joystick rough) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | Tested | Configured | Configured | Configured | Configured | -| PPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `x2_wall_flip_tracking` (X2 wall flip tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `a2_joystick_flat` (a2 joystick flat) | Tested | - | - | - | - | - | - | -| PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_box_tracking` (g1 23dof box tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_motion_tracking_deploy` (g1 23dof motion tracking deploy) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_23dof_walk_rough` (g1 23dof walk rough) | Tested | - | Registered | - | - | - | - | -| PPO (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_box_tracking` (g1 box tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `g1_motion_tracking_deploy` (g1 motion tracking deploy) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `go1_joystick_rough` (go1 joystick rough) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `go2_footstand` (go2 footstand) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `go2w_joystick_flat` (go2w joystick flat) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `go2w_joystick_rough` (go2w joystick rough) | Tested | - | Tested | - | - | - | - | -| PPO (torch) | `stewart_balance` (stewart balance) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `go1_joystick_flat` (Go1 joystick) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | Registered | Registered | Registered | Registered | Registered | -| APPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Registered | - | - | - | - | -| APPO (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Tested | - | - | - | - | -| APPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | - | Tested | - | - | - | - | -| SAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | Tested | Tested | Tested | Configured | Tested | -| SAC (torch) | `g1_walk_rough` (G1 walk rough) | Tested | - | Tested | - | - | - | - | -| SAC (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | Configured | Tested | - | - | - | - | -| SAC (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Registered | - | - | - | - | -| SAC (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Registered | - | - | - | - | -| SAC (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Registered | - | - | - | - | -| SAC (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | -| SAC (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | - | - | - | - | -| SAC (torch) | `g1_23dof_walk_rough` (g1 23dof walk rough) | Tested | - | Tested | - | - | - | - | -| SAC (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Registered | - | - | - | - | -| SAC (torch) | `g1_23dof_wbt_obs` (g1 23dof wbt obs) | Tested | - | Registered | - | - | - | - | -| SAC (torch) | `g1_wbt_obs` (g1 wbt obs) | Tested | - | Registered | - | - | - | - | -| TD3 (torch) | `go1_joystick_flat` (Go1 joystick) | Registered | - | Tested | - | - | - | - | -| TD3 (torch) | `go2_joystick_flat` (Go2 joystick) | Registered | - | Tested | - | - | - | - | -| TD3 (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | Registered | Registered | Registered | Registered | Registered | -| TD3 (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Registered | - | - | - | - | -| FlashSAC (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Registered | - | - | - | - | -| FlashSAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Configured | Tested | Registered | Registered | Registered | Registered | -| FlashSAC (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | - | - | - | - | +| Entrypoint | Task owner | MuJoCo | mjwarp | Motrix | IsaacGym | Genesis | IsaacSim | Newton | SuperDex | +|------------|------------|--------|--------|--------|----------|---------|----------|--------|----------| +| PPO (torch) | `go1_joystick_flat` (Go1 joystick) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Tested | - | - | - | - | Configured | +| PPO (torch) | `go2_joystick_rough` (Go2 joystick rough) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | Tested | Configured | Configured | Configured | Configured | - | +| PPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `x2_wall_flip_tracking` (X2 wall flip tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `sharpa_inhand_grasp` (Sharpa in-hand grasp) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `a2_joystick_flat` (a2 joystick flat) | Tested | - | - | - | - | - | - | - | +| PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `fr3_joint_target` (fr3 joint target) | - | - | - | - | - | - | - | Configured | +| PPO (torch) | `g1_23dof_box_tracking` (g1 23dof box tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_23dof_motion_tracking_deploy` (g1 23dof motion tracking deploy) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_23dof_walk_rough` (g1 23dof walk rough) | Tested | - | Registered | - | - | - | - | - | +| PPO (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_box_tracking` (g1 box tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `g1_motion_tracking_deploy` (g1 motion tracking deploy) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `go1_joystick_rough` (go1 joystick rough) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `go2_footstand` (go2 footstand) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `go2w_joystick_flat` (go2w joystick flat) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `go2w_joystick_rough` (go2w joystick rough) | Tested | - | Tested | - | - | - | - | - | +| PPO (torch) | `stewart_balance` (stewart balance) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `go1_joystick_flat` (Go1 joystick) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Tested | - | - | - | - | Registered | +| APPO (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | Registered | Registered | Registered | Registered | Registered | - | +| APPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `sharpa_inhand` (Sharpa in-hand) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_23dof_climb_tracking` (g1 23dof climb tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Registered | - | - | - | - | - | +| APPO (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Tested | - | - | - | - | - | +| APPO (torch) | `g1_climb_tracking` (g1 climb tracking) | Tested | - | Tested | - | - | - | - | - | +| SAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Tested | Tested | Tested | Tested | Configured | Tested | - | +| SAC (torch) | `g1_walk_rough` (G1 walk rough) | Tested | - | Tested | - | - | - | - | - | +| SAC (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | Configured | Tested | - | - | - | - | - | +| SAC (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | - | Registered | - | - | - | - | - | +| SAC (torch) | `g1_wall_flip_tracking` (G1 wall flip tracking) | Tested | - | Registered | - | - | - | - | - | +| SAC (torch) | `g1_23dof_flip_tracking` (g1 23dof flip tracking) | Tested | - | Registered | - | - | - | - | - | +| SAC (torch) | `g1_23dof_motion_tracking` (g1 23dof motion tracking) | Tested | - | Tested | - | - | - | - | - | +| SAC (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | - | - | - | - | - | +| SAC (torch) | `g1_23dof_walk_rough` (g1 23dof walk rough) | Tested | - | Tested | - | - | - | - | - | +| SAC (torch) | `g1_23dof_wall_flip_tracking` (g1 23dof wall flip tracking) | Tested | - | Registered | - | - | - | - | - | +| SAC (torch) | `g1_23dof_wbt_obs` (g1 23dof wbt obs) | Tested | - | Registered | - | - | - | - | - | +| SAC (torch) | `g1_wbt_obs` (g1 wbt obs) | Tested | - | Registered | - | - | - | - | - | +| TD3 (torch) | `go1_joystick_flat` (Go1 joystick) | Registered | - | Tested | - | - | - | - | - | +| TD3 (torch) | `go2_joystick_flat` (Go2 joystick) | Registered | - | Tested | - | - | - | - | Registered | +| TD3 (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Registered | Registered | Registered | Registered | Registered | Registered | - | +| TD3 (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Registered | - | - | - | - | - | +| FlashSAC (torch) | `go2_joystick_flat` (Go2 joystick) | Tested | - | Registered | - | - | - | - | Registered | +| FlashSAC (torch) | `g1_walk_flat` (G1 walk flat) | Tested | Configured | Tested | Registered | Registered | Registered | Registered | - | +| FlashSAC (torch) | `g1_23dof_walk_flat` (g1 23dof walk flat) | Tested | - | Tested | - | - | - | - | - | ### Source Index - Registry bootstrap: `src/unilab/envs/**` decorators via `unilab.base.registry.ensure_registries()`. - Owner YAML scan: `src/unilab/conf/ppo/task/**`, `src/unilab/conf/appo/task/**`, `src/unilab/conf/sac/task/**`, `src/unilab/conf/td3/task/**`, `src/unilab/conf/flashsac/task/**`. - Generic compose coverage: `tests/config/test_config_system.py::test_supported_task_composes`. +- SuperDex remains `Configured`: FR3 has optional CPU rollout/spawn coverage in `tests/envs/test_fr3_superdex.py`; the Go2 research profile has policy-contract/rollout/checkpoint coverage in `tests/envs/test_go2_superdex.py`. Neither profile claims full-training performance or cross-platform support. - Validated mjwarp entrypoints are explicitly recorded in `_MAINTAINER_VALIDATED_MJWARP_ENTRYPOINT_TASKS`; near-risk coverage lives in `tests/base/test_mjwarp_backend.py`, `tests/base/test_backend_conformance.py`, `tests/base/test_mjwarp_differential.py`, and `tests/base/test_mjwarp_playback.py`. - Validated isaacgym entrypoints are explicitly recorded in `_MAINTAINER_VALIDATED_ISAACGYM_ENTRYPOINT_TASKS` (real hardware via the external Python 3.8 worker runtime; not covered by repo CI). - Validated genesis entrypoints are explicitly recorded in `_MAINTAINER_VALIDATED_GENESIS_ENTRYPOINT_TASKS` (real hardware, genesis-world extra + CUDA; not covered by repo CI); near-risk coverage lives in `tests/base/test_genesis_backend.py` (fake runtime), `tests/base/test_genesis_runtime.py` (real-runtime slow lane), and the genesis env smoke in `tests/envs/locomotion/g1/test_g1_owner_contract.py`. diff --git a/pyproject.toml b/pyproject.toml index 5fd2f7331..759a50b7f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,14 +40,12 @@ classifiers = [ requires-python = ">=3.10,<3.14" dependencies = [ "numpy", - # Physics implementations are provided by the independently released - # unisim-core package from the production PyPI index. - "unisim-core>=1.1.4", - # RL algorithms and async runtimes (PPO/APPO/SAC/TD3/HORA 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.1.1", + # Keep the SuperDex backend contract aligned with the merged UniSim source + # until its next PyPI release. The commit is pinned for reproducible CI. + "unisim-core @ git+https://github.com/unilabsim/unisim.git@037e596", + # Keep the CPU-partition contract aligned with the merged UniRL source + # until its next PyPI release. The commit is pinned for reproducible CI. + "unilab-rl @ git+https://github.com/unilabsim/unilab_rl.git@2cdab3c", "numba>=0.67", "prettytable>=3.10", # torch is a range (not an exact pin) so that published PyPI metadata lets diff --git a/scripts/audit_sim2sim_contracts.py b/scripts/audit_sim2sim_contracts.py index ddf60c46f..b05047ded 100644 --- a/scripts/audit_sim2sim_contracts.py +++ b/scripts/audit_sim2sim_contracts.py @@ -41,6 +41,7 @@ ("mujoco", "genesis"), ("mujoco", "isaacsim"), ("mujoco", "newton"), + ("mujoco", "superdex"), ) diff --git a/scripts/benchmark/superdex_go2_collector.py b/scripts/benchmark/superdex_go2_collector.py new file mode 100644 index 000000000..0f52c3572 --- /dev/null +++ b/scripts/benchmark/superdex_go2_collector.py @@ -0,0 +1,97 @@ +"""Short end-to-end Go2 collector benchmark for the SuperDex roadmap. + +The timed loop includes policy inference, action conversion, env.step (including +physics, observations, reward and termination/reset), and rollout bookkeeping. +It intentionally does not run PPO updates or claim training quality. +""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path + +import numpy as np +import torch +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +from unilab.base import registry +from unilab.base.config_adapter import BackendAdapter + + +def run(mode: str, *, num_envs: int, workers: int, warmup: int, steps: int, repeats: int) -> dict: + root = Path(__file__).resolve().parents[2] + values: list[float] = [] + rewards: list[float] = [] + for repeat in range(repeats): + GlobalHydra.instance().clear() + with initialize_config_dir( + config_dir=str(root / "src/unilab/conf/ppo"), version_base="1.3" + ): + cfg = compose("config", overrides=["task=go2_joystick_flat/superdex"]) + override = BackendAdapter(cfg, root_dir=root).build_task_env_cfg_override() + env = registry.make( + "Go2JoystickFlat", + sim_backend="superdex", + num_envs=num_envs, + env_cfg_override={**override, "superdex_num_workers": workers, "seed": repeat + 1}, + ) + try: + state = env.init_state() + obs_dim = int(env.obs_groups_spec["obs"]) + policy = torch.nn.Sequential( + torch.nn.Linear(obs_dim, 128), + torch.nn.Tanh(), + torch.nn.Linear(128, env.action_space.shape[0]), + torch.nn.Tanh(), + ).eval() + if mode == "unbatched": + env._backend._pre_step_control_fn = lambda _backend, controls: controls + for _ in range(warmup): + with torch.inference_mode(): + action = policy(torch.from_numpy(state.obs["obs"]).float()).numpy() + state = env.step(action) + started = time.perf_counter() + total_reward = 0.0 + total_done = 0 + for _ in range(steps): + with torch.inference_mode(): + action = policy(torch.from_numpy(state.obs["obs"]).float()).numpy() + state = env.step(action) + total_reward += float(np.asarray(state.reward).mean()) + total_done += int(np.asarray(state.terminated).sum()) + values.append(num_envs * steps / (time.perf_counter() - started)) + rewards.append(total_reward / steps) + finally: + env.close() + return { + "mode": mode, + "metric": "collector-env-step/s", + "num_envs": num_envs, + "workers": workers, + "warmup": warmup, + "steps": steps, + "repeats": values, + "median": float(np.median(values)), + "mean_reward_per_step": float(np.mean(rewards)), + "done_count": total_done, + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--mode", choices=("batched", "unbatched"), required=True) + parser.add_argument("--num-envs", type=int, default=128) + parser.add_argument("--workers", type=int, default=16) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--steps", type=int, default=10) + parser.add_argument("--repeats", type=int, default=2) + args = parser.parse_args() + print("SUPERDEX_GO2_COLLECTOR_RESULT=" + json.dumps(run(**vars(args)), sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/scripts/tools/support_matrix.py b/scripts/tools/support_matrix.py index 68a76a8bd..222a5a780 100644 --- a/scripts/tools/support_matrix.py +++ b/scripts/tools/support_matrix.py @@ -22,6 +22,7 @@ "genesis", "isaacsim", "newton", + "superdex", ) # Maintainer-confirmed completed training validations. Keep this mapping narrow: @@ -249,6 +250,9 @@ def _configured_entries(root: Path, spec: EntrypointSpec) -> dict[str, dict[str, def _is_tested(spec: EntrypointSpec, task_slug: str, backend: str, root: Path) -> bool: + if backend == "superdex": + # Bounded native smoke/short training does not establish full training support. + return False if backend == "mjwarp": return ( spec.entrypoint_id, @@ -411,8 +415,8 @@ def render_support_matrix(root: Path | None = None) -> str: "", "### Entrypoint x Task Owner", "", - "| Entrypoint | Task owner | MuJoCo | mjwarp | Motrix | IsaacGym | Genesis | IsaacSim | Newton |", - "|------------|------------|--------|--------|--------|----------|---------|----------|--------|", + "| Entrypoint | Task owner | MuJoCo | mjwarp | Motrix | IsaacGym | Genesis | IsaacSim | Newton | SuperDex |", + "|------------|------------|--------|--------|--------|----------|---------|----------|--------|----------|", ] for row in build_support_rows(resolved_root): @@ -421,7 +425,7 @@ def render_support_matrix(root: Path | None = None) -> str: f"{row.cells['mujoco'].level.label} | {row.cells['mjwarp'].level.label} | " f"{row.cells['motrix'].level.label} | {row.cells['isaacgym'].level.label} | " f"{row.cells['genesis'].level.label} | {row.cells['isaacsim'].level.label} |" - f" {row.cells['newton'].level.label} |" + f" {row.cells['newton'].level.label} | {row.cells['superdex'].level.label} |" ) lines.extend( @@ -432,6 +436,7 @@ def render_support_matrix(root: Path | None = None) -> str: "- Registry bootstrap: `src/unilab/envs/**` decorators via `unilab.base.registry.ensure_registries()`.", "- Owner YAML scan: `src/unilab/conf/ppo/task/**`, `src/unilab/conf/appo/task/**`, `src/unilab/conf/sac/task/**`, `src/unilab/conf/td3/task/**`, `src/unilab/conf/flashsac/task/**`.", "- Generic compose coverage: `tests/config/test_config_system.py::test_supported_task_composes`.", + "- SuperDex remains `Configured`: FR3 has optional CPU rollout/spawn coverage in `tests/envs/test_fr3_superdex.py`; the Go2 research profile has policy-contract/rollout/checkpoint coverage in `tests/envs/test_go2_superdex.py`. Neither profile claims full-training performance or cross-platform support.", "- Validated mjwarp entrypoints are explicitly recorded in `_MAINTAINER_VALIDATED_MJWARP_ENTRYPOINT_TASKS`; near-risk coverage lives in `tests/base/test_mjwarp_backend.py`, `tests/base/test_backend_conformance.py`, `tests/base/test_mjwarp_differential.py`, and `tests/base/test_mjwarp_playback.py`.", "- Validated isaacgym entrypoints are explicitly recorded in `_MAINTAINER_VALIDATED_ISAACGYM_ENTRYPOINT_TASKS` (real hardware via the external Python 3.8 worker runtime; not covered by repo CI).", "- Validated genesis entrypoints are explicitly recorded in `_MAINTAINER_VALIDATED_GENESIS_ENTRYPOINT_TASKS` (real hardware, genesis-world extra + CUDA; not covered by repo CI); near-risk coverage lives in `tests/base/test_genesis_backend.py` (fake runtime), `tests/base/test_genesis_runtime.py` (real-runtime slow lane), and the genesis env smoke in `tests/envs/locomotion/g1/test_g1_owner_contract.py`.", diff --git a/src/unilab/assets/hub.py b/src/unilab/assets/hub.py index 5cd89264b..2713ba639 100644 --- a/src/unilab/assets/hub.py +++ b/src/unilab/assets/hub.py @@ -52,6 +52,54 @@ "x2": (("robots/x2/meshes", "pelvis.STL", "*.STL", "STL"),), } +# Upstream native robots stay in the user's audited SuperDex asset checkout. +# This registry deliberately has no download fallback or SDK import. Relative +# model names in task owners are resolved only beneath the configured root. +SUPERDEX_ROBOT_ASSET_SPECS: dict[str, tuple[str, ...]] = { + "bots/arms/fr3_v2/fr3_v2.superdex_bot": ( + "LICENSE", + "NOTICE", + *(f"collision/fr3_link{i}_collision.mochi.h5" for i in range(8)), + *(f"render/fr3_link{i}_render.glb" for i in range(8)), + ), +} + + +def resolve_superdex_robot_asset(model_file: str, *, assets_root: str | None = None) -> str: + """Resolve a registered native robot in a local SuperDex asset checkout. + + ``assets_root`` or ``SUPERDEX_ASSETS_PATH`` points to the upstream ``assets`` + directory, not the repository root. Required collision, visual and license + files are checked before physics construction. No files are downloaded or + copied into UniLab's package. + """ + root_value = assets_root if assets_root is not None else os.environ.get("SUPERDEX_ASSETS_PATH") + if not root_value or not root_value.strip(): + raise FileNotFoundError( + "SuperDex native assets require env.superdex_assets_root or SUPERDEX_ASSETS_PATH " + "pointing to the project_superdex/assets directory" + ) + root = Path(root_value).expanduser().resolve() + supplied = Path(model_file).expanduser() + resolved = (supplied if supplied.is_absolute() else root / supplied).resolve() + try: + relative = resolved.relative_to(root).as_posix() + except ValueError as exc: + raise ValueError(f"SuperDex robot asset must be inside configured root {root}") from exc + if relative not in SUPERDEX_ROBOT_ASSET_SPECS: + raise ValueError(f"SuperDex robot asset is not registered: {relative}") + required = ( + resolved, + *(resolved.parent / item for item in SUPERDEX_ROBOT_ASSET_SPECS[relative]), + ) + missing = [str(path) for path in required if not path.is_file()] + if missing: + raise FileNotFoundError( + "SuperDex registered robot asset is incomplete; restore the upstream checkout " + f"including collision/render files and license notices: {', '.join(missing)}" + ) + return str(resolved) + def resolve_motion_files( motion_file: str | Sequence[str], diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py index 33f5145d1..02532bc17 100644 --- a/src/unilab/base/backend_factory.py +++ b/src/unilab/base/backend_factory.py @@ -8,12 +8,13 @@ from __future__ import annotations +from dataclasses import replace from typing import TYPE_CHECKING, Any import unisim from unisim.backend.base import SimBackend -from unilab.assets.hub import ensure_robot_assets_for_paths +from unilab.assets.hub import ensure_robot_assets_for_paths, resolve_superdex_robot_asset from unilab.base.process_device import bind_genesis_process_device if TYPE_CHECKING: @@ -44,6 +45,10 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]: """Translate ``EnvCfg`` backend knobs into UniSim adapter options.""" result: dict[str, Any] = { "post_step_forward_sensor": cfg.post_step_forward_sensor, + "superdex_num_workers": cfg.superdex_num_workers, + "superdex_assets_root": cfg.superdex_assets_root, + "superdex_effort_limits": cfg.superdex_effort_limits, + "superdex_allow_contact_approximation": cfg.superdex_allow_contact_approximation, "motrix_max_iterations": cfg.motrix_max_iterations, "chunk_size": cfg.chunk_size, "adaptive_chunk_size": cfg.adaptive_chunk_size, @@ -90,6 +95,18 @@ def create_backend( """Prepare UniLab-owned assets and construct a UniSim backend.""" if scene is None: raise ValueError("SceneCfg must be provided") + superdex_assets_root = kwargs.pop("superdex_assets_root", None) + if backend_type == "superdex" and scene.model_file.endswith(".superdex_bot"): + scene = replace( + scene, + model_file=resolve_superdex_robot_asset( + scene.model_file, assets_root=superdex_assets_root + ), + ) + if backend_type != "superdex": + kwargs.pop("superdex_num_workers", None) + kwargs.pop("superdex_effort_limits", None) + kwargs.pop("superdex_allow_contact_approximation", None) ensure_robot_assets_for_paths( [scene.model_file, scene.visual_model_file, *scene.fragment_files] ) @@ -108,7 +125,10 @@ def create_backend( # not accept MuJoCo's synthetic body-sensor injection. Keep this # capability translation at the owner/backend boundary so env code remains # backend-agnostic. - kwargs["body_state_required"] = body_state_required and backend_type != "newton" + kwargs["body_state_required"] = body_state_required and backend_type not in { + "newton", + "superdex", + } 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 diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py index bd5ea2583..20dd90912 100644 --- a/src/unilab/base/base.py +++ b/src/unilab/base/base.py @@ -35,6 +35,12 @@ class EnvCfg: render_offset_mode: str = "grid" drake_backend_mode: str = "batch" drake_nthread: int = 0 + # SuperDex native robots are resolved through the local asset hub registry. + # 0 selects affinity-aware native C++ scene batching; 1 stays serial. + superdex_num_workers: int = 0 + superdex_assets_root: Optional[str] = None + superdex_effort_limits: Optional[list[float]] = None + superdex_allow_contact_approximation: bool = False motrix_max_iterations: Optional[int] = None post_step_forward_sensor: bool = False adaptive_chunk_size: bool = True @@ -107,6 +113,29 @@ def validate(self): """ if self.sim_dt > self.ctrl_dt: raise ValueError("sim_dt must be less than or equal to ctrl_dt") + if ( + isinstance(self.superdex_num_workers, bool) + or not isinstance(self.superdex_num_workers, int) + or self.superdex_num_workers < 0 + ): + raise ValueError("superdex_num_workers must be an integer >= 0") + if self.superdex_assets_root is not None and ( + not isinstance(self.superdex_assets_root, str) or not self.superdex_assets_root.strip() + ): + raise ValueError("superdex_assets_root must be a non-empty string or None") + if self.superdex_effort_limits is not None: + if not isinstance(self.superdex_effort_limits, list) or not self.superdex_effort_limits: + raise ValueError("superdex_effort_limits must be a non-empty list or None") + if any( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not np.isfinite(value) + or value <= 0 + for value in self.superdex_effort_limits + ): + raise ValueError("superdex_effort_limits must contain finite positive numbers") + if not isinstance(self.superdex_allow_contact_approximation, bool): + raise ValueError("superdex_allow_contact_approximation must be bool") for name, value in ( ("mjwarp_nconmax", self.mjwarp_nconmax), ("mjwarp_njmax", self.mjwarp_njmax), diff --git a/src/unilab/base/registry.py b/src/unilab/base/registry.py index 76c66d03c..533d7c1d6 100644 --- a/src/unilab/base/registry.py +++ b/src/unilab/base/registry.py @@ -49,6 +49,7 @@ def __call__( "genesis", "isaacsim", "newton", + "superdex", ) _DEFAULT_SIM_BACKEND_ORDER: tuple[str, ...] = ("mujoco", "motrix") _REGISTRY_MODULES_ATTR = "__unilab_registry_modules__" diff --git a/src/unilab/cli.py b/src/unilab/cli.py index 2b52ad2c7..cdf063d46 100644 --- a/src/unilab/cli.py +++ b/src/unilab/cli.py @@ -25,6 +25,7 @@ "genesis", "isaacsim", "newton", + "superdex", ) SUPPORTED_RENDER_MODES = ("auto", "interactive", "record", "none") OFFPOLICY_ALGOS = {"sac", "td3", "flashsac"} @@ -128,6 +129,21 @@ def _check_runtime_requirements(algo: str, sim: str) -> None: f"(missing: {joined}). Install it with `uv sync --extra newton` " "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 the locally linked UniSim SuperDex development checkout; " + "the installed unisim-core does not provide that adapter." + ) from exc + + if not superdex_dependencies_available(): + raise SystemExit( + "sim=superdex requires Python 3.12 and the SuperDex Physics/Robotics runtime. " + "Install the locally linked UniSim superdex extra in a Python 3.12 environment; " + "see the SuperDex backend page for local development setup." + ) if sim == "motrix" and find_spec("motrixsim") is None: raise SystemExit( "sim=motrix requires the Motrix extra. Install it with " diff --git a/src/unilab/conf/ppo/config.yaml b/src/unilab/conf/ppo/config.yaml index 96397cf19..16a370b07 100644 --- a/src/unilab/conf/ppo/config.yaml +++ b/src/unilab/conf/ppo/config.yaml @@ -60,6 +60,9 @@ training: # CUDA device, and [d0..dN-1] launches one RSL-RL worker per device. # algo.num_envs is per rank; do not set this together with training.device. devices: null + # list[list[int]] | null; explicit CPU ids per rank. Null automatically + # assigns complete physical-core groups, including logical siblings. + dp_collector_cpu_ids: null device: null logger: tensorboard wandb_project: unilab diff --git a/src/unilab/conf/ppo/task/fr3_joint_target/base.yaml b/src/unilab/conf/ppo/task/fr3_joint_target/base.yaml new file mode 100644 index 000000000..8a9e1e168 --- /dev/null +++ b/src/unilab/conf/ppo/task/fr3_joint_target/base.yaml @@ -0,0 +1,71 @@ +# @package _global_ +# Fixed-base native FR3 task. Targets and research torque limits are task-owned; +# they are not a claim about the hardware's rated actuator limits. +env: + scene: + model_file: bots/arms/fr3_v2/fr3_v2.superdex_bot + entities: + robot: + root_body_name: fr3_link0 + joint_names: &joints [fr3_joint1, fr3_joint2, fr3_joint3, fr3_joint4, fr3_joint5, fr3_joint6, fr3_joint7] + actuator_names: *joints + body_names: [fr3_link0, fr3_link7] + sim_dt: 0.002 + ctrl_dt: 0.01 + max_episode_seconds: 5.0 + seed: 1 + observations: + policy: + terms: + target_error: + func: unilab.tasks.manipulation.fr3.joint_target.JointTargetObservation + params: + entity_name: robot + target: &target [0.1, -0.7, 0.0, -2.2, 0.0, 1.5, 1.57] + joint_vel: + func: unilab.envs.mdp.joint_vel_rel + actions: + func: unilab.envs.mdp.last_action + actions: + effort: + _target_: unilab.envs.mdp.JointEffortActionCfg + entity_name: robot + actuator_names: *joints + scale: + fr3_joint[1-4]: 20.0 + fr3_joint[5-7]: 5.0 + clip: + fr3_joint[1-4]: [-20.0, 20.0] + fr3_joint[5-7]: [-5.0, 5.0] + events: + reset_scene_to_default: + func: unilab.envs.mdp.reset_scene_to_default + mode: reset + reset_joints: + func: unilab.tasks.manipulation.fr3.joint_target.ResetJointOffsets + mode: reset + params: + entity_name: robot + position_range: [-0.05, 0.05] + velocity_range: [-0.01, 0.01] + terminations: + time_out: + func: unilab.envs.mdp.time_out + time_out: true + policy_observation_group: policy + critic_observation_group: null + +reward: + joint_target: + func: unilab.tasks.manipulation.fr3.joint_target.JointTargetReward + weight: 2.0 + params: + entity_name: robot + target: *target + std: 0.5 + joint_velocity: + func: unilab.envs.mdp.joint_vel_l2 + weight: -0.01 + action_rate: + func: unilab.envs.mdp.action_rate_l2 + weight: -0.001 diff --git a/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml b/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml new file mode 100644 index 000000000..2e005373a --- /dev/null +++ b/src/unilab/conf/ppo/task/fr3_joint_target/superdex.yaml @@ -0,0 +1,35 @@ +# @package _global_ +defaults: + - /task/fr3_joint_target/base + - _self_ + +training: + task_name: FR3JointTarget + sim_backend: superdex + device: cpu + no_play: true + play_render_mode: none + +algo: + num_envs: 2 + num_steps_per_env: 32 + max_iterations: 100 + empirical_normalization: false + obs_groups: + actor: [actor] + critic: [actor] + policy: + actor_hidden_dims: [64, 64] + critic_hidden_dims: [64, 64] + init_noise_std: 0.25 + algorithm: + num_mini_batches: 1 + +env: + adaptive_chunk_size: false + superdex_num_workers: 0 + superdex_assets_root: null + superdex_effort_limits: [20.0, 20.0, 20.0, 20.0, 5.0, 5.0, 5.0] + +play_profile: + enabled: false diff --git a/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml b/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml new file mode 100644 index 000000000..461fd2555 --- /dev/null +++ b/src/unilab/conf/ppo/task/go2_joystick_flat/superdex.yaml @@ -0,0 +1,30 @@ +# @package _global_ +# Research sim2sim profile: inherit the exact MuJoCo policy/observation owner. +# SuperDex contact compliance and slide-only friction are not MuJoCo physics +# equivalents. This profile permits that conversion explicitly, not silently. +defaults: + - /task/go2_joystick_flat/mujoco + - _self_ + +training: + task_name: Go2JoystickFlat + sim_backend: superdex + device: cpu + no_play: false + play_render_mode: record + dp_collector_cpu_ids: null + +algo: + num_envs: 2 + algorithm: + num_mini_batches: 1 + +env: + adaptive_chunk_size: false + superdex_num_workers: 0 + superdex_allow_contact_approximation: true + events: + pd_gains: null + +play_profile: + enabled: true diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 8cc12ba86..c64d08576 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -22,6 +22,7 @@ CONFIG_MAPPING_POLICY_KEY, MANAGER_TERM_MAPPING_POLICY, ) +from unilab.base.cpu_runtime import apply_env_cpu_runtime from unilab.base.entity import EntityCfg, EntityScene from unilab.base.np_env import NpEnv, NpEnvState from unilab.base.reset_state import ResetStateTransaction @@ -738,6 +739,9 @@ def make_manager_based_rl_env( ) cfg.validate() + # Constrain the process before backend materialization so native pools size + # themselves from the rank-owned CPU block. + apply_env_cpu_runtime(cfg.cpu_ids) assert cfg.scene is not None base_name, body_state_requested = _resolve_backend_entity_contract(cfg) backend_kwargs = env_backend_kwargs(cfg) diff --git a/src/unilab/scripts/train_offpolicy.py b/src/unilab/scripts/train_offpolicy.py index d02efd727..a5abdb919 100644 --- a/src/unilab/scripts/train_offpolicy.py +++ b/src/unilab/scripts/train_offpolicy.py @@ -190,7 +190,7 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): collector_cpu_ids = resolve_collector_cpu_ids( dp_world_size, dp_rank, - host_cpu_count, + None, explicit=explicit_cpu_ids, ) diff --git a/src/unilab/scripts/train_rsl_rl.py b/src/unilab/scripts/train_rsl_rl.py index 46b2c0cab..0cd216e3f 100644 --- a/src/unilab/scripts/train_rsl_rl.py +++ b/src/unilab/scripts/train_rsl_rl.py @@ -26,6 +26,7 @@ current_torch_distributed_rank, current_torch_distributed_world_size, launch_torchrun_workers, + resolve_collector_cpu_ids, resolve_dp_topology, validate_dp_launchable, ) @@ -90,6 +91,7 @@ def _backend_adapter(cfg: DictConfig) -> BackendAdapter: def build_ppo_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: base = cast(dict[str, Any], _backend_adapter(cfg).build_task_env_cfg_override()) devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None)) + rank = current_torch_distributed_rank() local_rank = current_torch_distributed_local_rank() world_size = current_torch_distributed_world_size() configured_device = OmegaConf.select(cfg, "training.device", default=None) @@ -98,7 +100,7 @@ def build_ppo_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: if world_size > 1 else (f"cuda:{devices[0]}" if devices else configured_device) ) - return apply_backend_env_device_override( + result = apply_backend_env_device_override( base, str(cfg.training.sim_backend), devices=devices, @@ -107,6 +109,16 @@ def build_ppo_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: world_size=world_size, learner_device=learner_device, ) + if world_size > 1: + explicit = OmegaConf.select(cfg, "training.dp_collector_cpu_ids", default=None) + explicit = OmegaConf.to_container(explicit, resolve=True) if explicit is not None else None + result["cpu_ids"] = resolve_collector_cpu_ids( + world_size, + rank, + None, + explicit=explicit, + ) + return result def build_ppo_play_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: diff --git a/src/unilab/tasks/__init__.py b/src/unilab/tasks/__init__.py index 6d59d3958..61757d8a3 100644 --- a/src/unilab/tasks/__init__.py +++ b/src/unilab/tasks/__init__.py @@ -14,6 +14,7 @@ "unilab.tasks.manipulation.allegro_inhand", "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", + "unilab.tasks.manipulation.fr3", "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", ) diff --git a/src/unilab/tasks/locomotion/go2/__init__.py b/src/unilab/tasks/locomotion/go2/__init__.py index 75357ce80..74744878b 100644 --- a/src/unilab/tasks/locomotion/go2/__init__.py +++ b/src/unilab/tasks/locomotion/go2/__init__.py @@ -9,6 +9,7 @@ registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="mujoco") registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="motrix") registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="drake") +registry.register_env("Go2JoystickFlat", make_manager_based_rl_env, sim_backend="superdex") registry.register_env_config("Go2JoystickRough", ManagerBasedRlEnvCfg) registry.register_env("Go2JoystickRough", make_manager_based_rl_env, sim_backend="mujoco") diff --git a/src/unilab/tasks/manipulation/fr3/__init__.py b/src/unilab/tasks/manipulation/fr3/__init__.py new file mode 100644 index 000000000..a4be1bdf7 --- /dev/null +++ b/src/unilab/tasks/manipulation/fr3/__init__.py @@ -0,0 +1,7 @@ +"""Hydra-owned fixed-base FR3 joint-target task registration.""" + +from unilab.base import registry +from unilab.envs import ManagerBasedRlEnvCfg, make_manager_based_rl_env + +registry.register_env_config("FR3JointTarget", ManagerBasedRlEnvCfg) +registry.register_env("FR3JointTarget", make_manager_based_rl_env, sim_backend="superdex") diff --git a/src/unilab/tasks/manipulation/fr3/joint_target.py b/src/unilab/tasks/manipulation/fr3/joint_target.py new file mode 100644 index 000000000..03d98a816 --- /dev/null +++ b/src/unilab/tasks/manipulation/fr3/joint_target.py @@ -0,0 +1,80 @@ +"""Joint-target terms using the public NumPy entity/reset contracts.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np + +if TYPE_CHECKING: + from unilab.base.entity import Entity + from unilab.managers import ManagerTermBaseCfg + from unilab.managers._types import ManagerBasedRlEnv + + +class JointTargetObservation: + """Bind one ordered joint target on the manager construction path.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv) -> None: + self._entity = cast("Entity", env.scene[cfg.params["entity_name"]]) + self._target = np.asarray(cfg.params["target"], dtype=np.float32) + expected = (self._entity.num_joints,) + if self._target.shape != expected or not np.isfinite(self._target).all(): + raise ValueError(f"FR3 joint target must be finite with shape {expected}") + self._target.setflags(write=False) + + def __call__(self, env: ManagerBasedRlEnv, entity_name: str, target: list[float]) -> np.ndarray: + return self._entity.data.joint_pos - self._target + + +class JointTargetReward(JointTargetObservation): + """Reward current joint accuracy with a configured squared-error scale.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv) -> None: + super().__init__(cfg, env) + std = cfg.params["std"] + if ( + isinstance(std, bool) + or not isinstance(std, (int, float)) + or not np.isfinite(std) + or std <= 0 + ): + raise ValueError("FR3 joint target reward std must be finite and positive") + self._variance = float(std) ** 2 + + def __call__( + self, env: ManagerBasedRlEnv, entity_name: str, target: list[float], std: float = 0.5 + ) -> np.ndarray: + error = self._entity.data.joint_pos - self._target + return np.exp(-np.sum(np.square(error), axis=-1) / self._variance) + + +class ResetJointOffsets: + """Stage selected joint resets without requiring a floating root.""" + + def __init__(self, cfg: ManagerTermBaseCfg, env: ManagerBasedRlEnv) -> None: + self._entity = cast("Entity", env.scene[cfg.params["entity_name"]]) + self._ranges = [] + for name in ("position_range", "velocity_range"): + values = np.asarray(cfg.params[name], dtype=np.float64) + if values.shape != (2,) or not np.isfinite(values).all() or values[0] > values[1]: + raise ValueError(f"FR3 reset {name} must be a finite ordered pair") + self._ranges.append((float(values[0]), float(values[1]))) + + def __call__( + self, + env: ManagerBasedRlEnv, + env_ids: np.ndarray | None, + entity_name: str, + position_range: list[float], + velocity_range: list[float], + ) -> None: + if env_ids is None: + raise ValueError("FR3 reset requires explicit environment IDs") + positions = self._entity.data.default_joint_pos[env_ids] + velocities = self._entity.data.default_joint_vel[env_ids] + self._entity.write_joint_state_to_sim( + positions + env.rng.uniform(*self._ranges[0], size=positions.shape), + velocities + env.rng.uniform(*self._ranges[1], size=velocities.shape), + env_ids=env_ids, + ) diff --git a/src/unilab/tasks/migration_matrix.py b/src/unilab/tasks/migration_matrix.py index ef0af5569..8499998f4 100644 --- a/src/unilab/tasks/migration_matrix.py +++ b/src/unilab/tasks/migration_matrix.py @@ -29,6 +29,8 @@ class TaskMigrationRecord: "A2JoystickFlat", "AllegroInhandRotation", "AllegroInhandRotationGrasp", + # #1534 starts directly on the canonical manager runtime; no legacy seam. + "FR3JointTarget", "Go1JoystickFlat", "Go2FootStand", "Go2JoystickFlat", diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index 1d072591c..902f0ea8d 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -377,6 +377,11 @@ def fake_probe_env_factory(num_envs, env_cfg_override): return _FakeEnv() monkeypatch.setattr(module.os, "cpu_count", lambda: cpu_count) + monkeypatch.setattr( + "uni_rl.ipc.dp_launcher.os.sched_getaffinity", + lambda _: set(range(cpu_count)), + raising=False, + ) if backend_binding_calls is not None: monkeypatch.setattr( module, @@ -414,6 +419,10 @@ def test_build_runner_binds_mjwarp_rank_process_to_learner_device( def test_build_runner_partitions_collector_cpus_per_rank(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "uni_rl.ipc.dp_launcher._discover_physical_cpu_groups", + lambda _: [[core, core + 64] for core in range(64)], + ) # Spawned rank: rank comes from the env, world_size from training.devices. monkeypatch.setenv(UNILAB_DP_RANK, "1") monkeypatch.setenv(UNILAB_DP_LOG_DIR, "/tmp/offpolicy_test_run") @@ -422,7 +431,9 @@ def test_build_runner_partitions_collector_cpus_per_rank(monkeypatch: pytest.Mon ["training.devices=[0,1]"], cpu_count=128, ) - assert runner.kwargs["collector_cpu_ids"] == list(range(64, 128)) + assert runner.kwargs["collector_cpu_ids"] == [ + cpu for core in range(32, 64) for cpu in (core, core + 64) + ] assert runner.kwargs["device"] == "cuda:1" # The thread budget is resolved against the rank's CPU share, not the host. assert runner.kwargs["torch_thread_runtime"]["cpu_count"] == 64 @@ -435,6 +446,10 @@ def test_build_runner_partitions_collector_cpus_per_rank(monkeypatch: pytest.Mon def test_build_runner_rank_zero_partitions_without_dp_env(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "uni_rl.ipc.dp_launcher._discover_physical_cpu_groups", + lambda _: [[core, core + 64] for core in range(64)], + ) # Rank 0 carries no UNILAB_DP_* env; world_size must come from the config. monkeypatch.delenv(UNILAB_DP_RANK, raising=False) monkeypatch.delenv(UNILAB_DP_WORLD_SIZE, raising=False) @@ -443,7 +458,9 @@ def test_build_runner_rank_zero_partitions_without_dp_env(monkeypatch: pytest.Mo ["training.devices=[0,1]"], cpu_count=128, ) - assert runner.kwargs["collector_cpu_ids"] == list(range(0, 64)) + assert runner.kwargs["collector_cpu_ids"] == [ + cpu for core in range(32) for cpu in (core, core + 64) + ] assert runner.kwargs["torch_thread_runtime"]["cpu_count"] == 64 diff --git a/tests/algos/test_offpolicy_dp_sync.py b/tests/algos/test_offpolicy_dp_sync.py index 9704cd71e..3dd427dac 100644 --- a/tests/algos/test_offpolicy_dp_sync.py +++ b/tests/algos/test_offpolicy_dp_sync.py @@ -7,6 +7,7 @@ import pytest import torch +import uni_rl.ipc.dp_launcher as dp_launcher from uni_rl.ipc.dp_launcher import UNILAB_DP_LOG_DIR, UNILAB_DP_RANK from uni_rl.ipc.dp_sync import DpParameterSync @@ -481,6 +482,9 @@ def _build_sac_runner_with_dp_fakes(monkeypatch: pytest.MonkeyPatch, overrides: module = _offpolicy() cfg = _offpolicy_cfg(overrides) monkeypatch.setattr(module.os, "cpu_count", lambda: 128) + monkeypatch.setattr( + dp_launcher.os, "sched_getaffinity", lambda _: set(range(128)), raising=False + ) import uni_rl.algos.fast_sac.double_buffer as owner_module @@ -544,6 +548,9 @@ def test_build_runner_multi_gpu_rank0_requires_log_dir(monkeypatch: pytest.Monke monkeypatch.delenv(UNILAB_DP_LOG_DIR, raising=False) monkeypatch.setattr(module, "registry_env_factory", lambda *args, **kwargs: _fake_env_factory) monkeypatch.setattr(module.os, "cpu_count", lambda: 128) + monkeypatch.setattr( + dp_launcher.os, "sched_getaffinity", lambda _: set(range(128)), raising=False + ) with pytest.raises(ValueError, match="log_dir"): module.build_runner("sac", cfg) @@ -701,6 +708,9 @@ def _build_flashsac_runner_with_dp_fakes(monkeypatch: pytest.MonkeyPatch, overri module = _offpolicy() cfg = _offpolicy_cfg(overrides, algo="flashsac") monkeypatch.setattr(module.os, "cpu_count", lambda: 128) + monkeypatch.setattr( + dp_launcher.os, "sched_getaffinity", lambda _: set(range(128)), raising=False + ) import uni_rl.algos.flash_sac.double_buffer as flash_module @@ -726,6 +736,10 @@ def test_build_runner_single_rank_flashsac_keeps_dp_sync_none(monkeypatch: pytes def test_build_runner_multi_gpu_constructs_dp_sync_for_flashsac_rank0( monkeypatch: pytest.MonkeyPatch, ): + monkeypatch.setattr( + "uni_rl.ipc.dp_launcher._discover_physical_cpu_groups", + lambda _: [[core, core + 64] for core in range(64)], + ) monkeypatch.delenv(UNILAB_DP_RANK, raising=False) monkeypatch.delenv(UNILAB_DP_LOG_DIR, raising=False) kwargs = _build_flashsac_runner_with_dp_fakes( @@ -740,11 +754,14 @@ def test_build_runner_multi_gpu_constructs_dp_sync_for_flashsac_rank0( assert dp_sync.rank == 0 assert dp_sync.backend == "nccl" assert dp_sync.rendezvous_path == "/tmp/dp_sync_test_run/.dp_rendezvous" - # Rank 0 collector owns the first contiguous CPU block. - assert kwargs["collector_cpu_ids"] == list(range(64)) + assert kwargs["collector_cpu_ids"] == [cpu for core in range(32) for cpu in (core, core + 64)] def test_build_runner_multi_gpu_flashsac_spawned_rank(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "uni_rl.ipc.dp_launcher._discover_physical_cpu_groups", + lambda _: [[core, core + 64] for core in range(64)], + ) monkeypatch.setenv(UNILAB_DP_RANK, "1") monkeypatch.setenv(UNILAB_DP_LOG_DIR, "/tmp/dp_sync_shared_root") kwargs = _build_flashsac_runner_with_dp_fakes( @@ -756,4 +773,6 @@ def test_build_runner_multi_gpu_flashsac_spawned_rank(monkeypatch: pytest.Monkey assert dp_sync.rank == 1 # Spawned ranks rendezvous on rank 0's run root, not their rank sub-dir. assert dp_sync.rendezvous_path == "/tmp/dp_sync_shared_root/.dp_rendezvous" - assert kwargs["collector_cpu_ids"] == list(range(64, 128)) + assert kwargs["collector_cpu_ids"] == [ + cpu for core in range(32, 64) for cpu in (core, core + 64) + ] diff --git a/tests/assets/test_superdex_assets.py b/tests/assets/test_superdex_assets.py new file mode 100644 index 000000000..62001653a --- /dev/null +++ b/tests/assets/test_superdex_assets.py @@ -0,0 +1,49 @@ +"""Native SuperDex assets resolve only through the audited local registry.""" + +from pathlib import Path + +import pytest + +from unilab.assets.hub import SUPERDEX_ROBOT_ASSET_SPECS, resolve_superdex_robot_asset + +MODEL = "bots/arms/fr3_v2/fr3_v2.superdex_bot" + + +@pytest.fixture +def asset_root(tmp_path: Path) -> Path: + model = tmp_path / MODEL + for path in (model, *(model.parent / item for item in SUPERDEX_ROBOT_ASSET_SPECS[MODEL])): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("fixture", encoding="utf-8") + return tmp_path + + +def test_local_superdex_asset_root_precedence_and_absolute_paths( + asset_root: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("SUPERDEX_ASSETS_PATH", "/missing/other/root") + expected = str(asset_root / MODEL) + assert resolve_superdex_robot_asset(MODEL, assets_root=str(asset_root)) == expected + assert resolve_superdex_robot_asset(expected, assets_root=str(asset_root)) == expected + monkeypatch.setenv("SUPERDEX_ASSETS_PATH", str(asset_root)) + assert resolve_superdex_robot_asset(MODEL) == expected + + +def test_superdex_assets_require_explicit_local_root(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("SUPERDEX_ASSETS_PATH", raising=False) + with pytest.raises(FileNotFoundError, match="SUPERDEX_ASSETS_PATH"): + resolve_superdex_robot_asset(MODEL) + + +def test_superdex_assets_reject_unregistered_and_escaping_paths(asset_root: Path) -> None: + with pytest.raises(ValueError, match="not registered"): + resolve_superdex_robot_asset("unknown.superdex_bot", assets_root=str(asset_root)) + with pytest.raises(ValueError, match="inside configured root"): + resolve_superdex_robot_asset("../outside.superdex_bot", assets_root=str(asset_root)) + + +def test_superdex_assets_fail_before_sdk_on_incomplete_collision(asset_root: Path) -> None: + missing = asset_root / Path(MODEL).parent / "collision/fr3_link4_collision.mochi.h5" + missing.unlink() + with pytest.raises(FileNotFoundError, match="fr3_link4_collision"): + resolve_superdex_robot_asset(MODEL, assets_root=str(asset_root)) diff --git a/tests/base/test_backend_imports.py b/tests/base/test_backend_imports.py index 33ed47ed9..ebd3acead 100644 --- a/tests/base/test_backend_imports.py +++ b/tests/base/test_backend_imports.py @@ -1,11 +1,14 @@ from __future__ import annotations import ast +import json +import os import subprocess import sys import textwrap from importlib.metadata import distribution from pathlib import Path +from urllib.parse import unquote, urlparse _REPO_ROOT = Path(__file__).resolve().parents[2] _MATERIALIZER_CONSUMERS = ( @@ -15,13 +18,36 @@ ) -def test_unisim_dependency_is_installed_from_package_index() -> None: +def test_unisim_dependency_uses_an_approved_source() -> None: direct_url = distribution("unisim-core").read_text("direct_url.json") - assert direct_url is None, ( - "UniLab tests must consume the indexed unisim-core release, not a local, editable, " - "or VCS checkout" - ) + local_checkout = os.environ.get("UNILAB_LOCAL_UNISIM") + if local_checkout: + import unisim + + expected = Path(local_checkout) + assert expected.is_absolute(), "UNILAB_LOCAL_UNISIM must be an absolute checkout path" + expected = expected.resolve(strict=True) + assert direct_url is not None, ( + "Local UniSim profile requires editable installation metadata" + ) + metadata = json.loads(direct_url) + assert metadata.get("dir_info", {}).get("editable") is True + installed_url = urlparse(metadata["url"]) + assert installed_url.scheme == "file" and installed_url.netloc in {"", "localhost"} + assert Path(unquote(installed_url.path)).resolve() == expected + assert unisim.__file__ is not None + assert Path(unisim.__file__).resolve().is_relative_to(expected / "src" / "unisim") + return + + if direct_url is None: + return + + metadata = json.loads(direct_url) + assert metadata.get("url") == "https://github.com/unilabsim/unisim.git" + vcs_info = metadata.get("vcs_info", {}) + assert vcs_info.get("vcs") == "git" + assert vcs_info.get("commit_id"), "Git-sourced UniSim must be pinned to a commit" def test_materializer_consumers_use_unisim_owner_module() -> None: diff --git a/tests/base/test_superdex_backend_options.py b/tests/base/test_superdex_backend_options.py new file mode 100644 index 000000000..821453eee --- /dev/null +++ b/tests/base/test_superdex_backend_options.py @@ -0,0 +1,93 @@ +"""UniLab resolves assets and passes only public SuperDex adapter options.""" + +from typing import Any + +import pytest + +from unilab.base import backend_factory +from unilab.base.base import EnvCfg +from unilab.base.scene import SceneCfg + + +def test_superdex_native_factory_resolves_assets_without_mutating_owner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: dict[str, Any] = {} + + def resolve(model_file: str, *, assets_root: str | None) -> str: + calls["asset"] = (model_file, assets_root) + return "/registered/assets/fr3.superdex_bot" + + def create(name: str, scene: SceneCfg, n: int, dt: float, **kwargs: Any) -> object: + calls["backend"] = (name, scene, n, dt, kwargs) + return object() + + monkeypatch.setattr(backend_factory, "resolve_superdex_robot_asset", resolve) + monkeypatch.setattr(backend_factory, "ensure_robot_assets_for_paths", lambda *_: None) + monkeypatch.setattr(backend_factory.unisim, "create_backend", create) + cfg = EnvCfg(superdex_assets_root="/registered/assets", superdex_effort_limits=[20.0]) + scene = SceneCfg(model_file="bots/arms/fr3_v2/fr3_v2.superdex_bot") + backend_factory.create_backend( + "superdex", + scene, + 2, + 0.002, + body_state_required=True, + **backend_factory.env_backend_kwargs(cfg), + ) + name, resolved_scene, n, dt, kwargs = calls["backend"] + assert (name, n, dt) == ("superdex", 2, 0.002) + assert resolved_scene.model_file == "/registered/assets/fr3.superdex_bot" + assert scene.model_file == "bots/arms/fr3_v2/fr3_v2.superdex_bot" + assert calls["asset"] == (scene.model_file, "/registered/assets") + assert kwargs["superdex_effort_limits"] == [20.0] + assert kwargs["superdex_num_workers"] == 0 + assert kwargs["superdex_allow_contact_approximation"] is False + assert kwargs["body_state_required"] is False + assert "superdex_assets_root" not in kwargs + + +def test_superdex_options_do_not_leak_to_other_adapters(monkeypatch: pytest.MonkeyPatch) -> None: + calls: dict[str, Any] = {} + + def create(*args: Any, **kwargs: Any) -> object: + calls.update(kwargs) + return object() + + monkeypatch.setattr(backend_factory, "ensure_robot_assets_for_paths", lambda *_: None) + monkeypatch.setattr(backend_factory.unisim, "create_backend", create) + backend_factory.create_backend( + "mujoco", + SceneCfg(model_file="scene.xml"), + 1, + 0.01, + **backend_factory.env_backend_kwargs(EnvCfg()), + ) + assert not any(name.startswith("superdex_") for name in calls) + + +@pytest.mark.parametrize("workers", [0, 1, 16]) +def test_large_superdex_batch_delegates_worker_selection_to_unisim( + workers: int, + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + def create(name: str, scene: SceneCfg, num_envs: int, dt: float, **kwargs: Any) -> object: + captured.update(backend=name, num_envs=num_envs, **kwargs) + return object() + + monkeypatch.setattr(backend_factory, "ensure_robot_assets_for_paths", lambda *_: None) + monkeypatch.setattr(backend_factory.unisim, "create_backend", create) + cfg = EnvCfg(superdex_num_workers=workers) + cfg.validate() + backend_factory.create_backend( + "superdex", + SceneCfg(model_file="scene.xml"), + 1024, + 0.01, + **backend_factory.env_backend_kwargs(cfg), + ) + assert captured["backend"] == "superdex" + assert captured["num_envs"] == 1024 + assert captured["superdex_num_workers"] == workers diff --git a/tests/envs/locomotion/go2/test_manager_based_cfg.py b/tests/envs/locomotion/go2/test_manager_based_cfg.py index 907e6dffb..10f975203 100644 --- a/tests/envs/locomotion/go2/test_manager_based_cfg.py +++ b/tests/envs/locomotion/go2/test_manager_based_cfg.py @@ -276,7 +276,7 @@ def test_go2_flat_registry_has_no_legacy_config_fallback() -> None: assert bare_cfg.rewards == {} assert registry.list_registered_envs()["Go2JoystickFlat"] == { "config_factory": "ManagerBasedRlEnvCfg", - "available_backends": ["mujoco", "motrix", "drake"], + "available_backends": ["mujoco", "motrix", "drake", "superdex"], } for legacy_override in ( {"reward_config": {}}, diff --git a/tests/envs/test_fr3_superdex.py b/tests/envs/test_fr3_superdex.py new file mode 100644 index 000000000..3b491d59a --- /dev/null +++ b/tests/envs/test_fr3_superdex.py @@ -0,0 +1,145 @@ +"""FR3 owner composition and optional native CPU rollout integration.""" + +from __future__ import annotations + +import multiprocessing as mp +import os +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra + +from unilab import cli +from unilab.base import registry +from unilab.base.base import EnvCfg +from unilab.base.config_adapter import BackendAdapter +from unilab.base.config_materialization import apply_cfg_overrides +from unilab.base.env_factory import registry_env_factory +from unilab.envs import ManagerBasedRlEnvCfg + +ROOT = Path(__file__).resolve().parents[2] + + +def _owner() -> tuple[Any, dict[str, Any]]: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(ROOT / "src/unilab/conf/ppo"), version_base="1.3"): + owner = compose("config", overrides=["task=fr3_joint_target/superdex"]) + return owner, BackendAdapter(owner, root_dir=ROOT).build_task_env_cfg_override() + + +def test_fr3_owner_has_torque_control_without_free_root_terms( + monkeypatch: pytest.MonkeyPatch, +) -> None: + owner, overrides = _owner() + cfg = registry.materialize_env_config("FR3JointTarget") + assert isinstance(cfg, ManagerBasedRlEnvCfg) + apply_cfg_overrides(cfg, overrides) + cfg.validate() + assert cfg.scene is not None + assert cfg.scene.entities["robot"].root_body_name == "fr3_link0" + assert list(cfg.events) == ["reset_scene_to_default", "reset_joints"] + assert list(cfg.observations["policy"].terms) == ["target_error", "joint_vel", "actions"] + assert cfg.superdex_effort_limits == [20.0] * 4 + [5.0] * 3 + assert cfg.superdex_num_workers == 0 + assert owner.training.no_play and owner.training.device == "cpu" + assert "superdex" in cli.SUPPORTED_SIMS + monkeypatch.setattr(cli, "_check_runtime_requirements", lambda *_: None) + command = cli.build_command( + mode="train", algo="ppo", task="fr3_joint_target", sim="superdex", overrides=[] + ) + assert "task=fr3_joint_target/superdex" in command + + +@pytest.mark.parametrize( + "kwargs", + [ + {"superdex_num_workers": True}, + {"superdex_num_workers": -1}, + {"superdex_num_workers": 1.5}, + {"superdex_assets_root": ""}, + {"superdex_effort_limits": [0.0]}, + {"superdex_effort_limits": [float("nan")]}, + ], +) +def test_superdex_owner_options_reject_invalid_values(kwargs: dict[str, Any]) -> None: + with pytest.raises(ValueError, match="superdex"): + EnvCfg(**kwargs).validate() + + +def _require_runtime() -> None: + if not os.environ.get("SUPERDEX_ASSETS_PATH"): + pytest.skip("native FR3 integration requires SUPERDEX_ASSETS_PATH") + pytest.importorskip("superdex.physics") + pytest.importorskip("superdex.robotics") + + +def test_fr3_native_rollout_and_selected_reset() -> None: + _require_runtime() + _, overrides = _owner() + env = registry.make( + "FR3JointTarget", sim_backend="superdex", num_envs=2, env_cfg_override=overrides + ) + try: + state = env.init_state() + assert env.obs_groups_spec == {"obs": 21} + assert env.action_space.shape == (7,) + for _ in range(32): + state = env.step(np.full((2, 7), 0.05, dtype=np.float32)) + assert np.isfinite(state.obs["obs"]).all() + assert np.isfinite(state.reward).all() + before = env.scene["robot"].data.joint_pos.copy() + counters = state.info["steps"].copy() + obs, _ = env.reset(env_ids=np.array([0], dtype=np.int32)) + assert obs["obs"].shape == (1, 21) + np.testing.assert_array_equal(env.scene["robot"].data.joint_pos[1], before[1]) + assert state.info["steps"][1] == counters[1] + np.testing.assert_allclose( + obs["obs"][0, :7], + env.scene["robot"].data.joint_pos[0] - np.array([0.1, -0.7, 0, -2.2, 0, 1.5, 1.57]), + atol=1e-6, + ) + finally: + env.close() + with pytest.raises(RuntimeError, match="closed"): + env.step(np.zeros((2, 7), dtype=np.float32)) + + +def _spawn_rollout(overrides: dict[str, Any]) -> tuple[int, ...]: + factory = registry_env_factory("FR3JointTarget", "superdex") + env = factory(num_envs=1, env_cfg_override=overrides) + try: + obs, _ = env.reset(env_ids=np.array([0], dtype=np.int32)) + env.step(np.zeros((1, 7), dtype=np.float32)) + return obs["obs"].shape + finally: + env.close() + + +def _spawn_parallel_rollout(overrides: dict[str, Any]) -> tuple[int, ...]: + factory = registry_env_factory("FR3JointTarget", "superdex") + env = factory(num_envs=2, env_cfg_override={**overrides, "superdex_num_workers": 2}) + try: + obs, _ = env.reset(env_ids=np.array([0, 1], dtype=np.int32)) + env.step(np.zeros((2, 7), dtype=np.float32)) + return obs["obs"].shape + finally: + env.close() + + +def test_fr3_factory_survives_spawn() -> None: + _require_runtime() + _, overrides = _owner() + with mp.get_context("spawn").Pool(1) as pool: + result = pool.apply_async(_spawn_rollout, (overrides,)) + assert result.get(timeout=90) == (1, 21) + + +def test_fr3_parallel_backend_can_run_inside_spawn_collector() -> None: + _require_runtime() + _, overrides = _owner() + with mp.get_context("spawn").Pool(1) as pool: + result = pool.apply_async(_spawn_parallel_rollout, (overrides,)) + assert result.get(timeout=150) == (2, 21) diff --git a/tests/envs/test_go2_superdex.py b/tests/envs/test_go2_superdex.py new file mode 100644 index 000000000..5e50674ad --- /dev/null +++ b/tests/envs/test_go2_superdex.py @@ -0,0 +1,176 @@ +"""Go2's research SuperDex owner preserves policy I/O and reset semantics.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from hydra import compose, initialize_config_dir +from hydra.core.global_hydra import GlobalHydra +from omegaconf import OmegaConf + +from unilab.base import registry +from unilab.base.config_adapter import BackendAdapter +from unilab.utils.sim2sim import DENYLIST, CrossBackendIncompatibleError, resolve_sim2sim_config + +ROOT = Path(__file__).resolve().parents[2] + + +def _owner(backend: str) -> tuple[Any, dict[str, Any]]: + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(ROOT / "src/unilab/conf/ppo"), version_base="1.3"): + owner = compose("config", overrides=[f"task=go2_joystick_flat/{backend}"]) + return owner, BackendAdapter(owner, root_dir=ROOT).build_task_env_cfg_override() + + +def test_go2_superdex_preserves_mujoco_policy_contract() -> None: + source, _ = _owner("mujoco") + target, _ = _owner("superdex") + for field in DENYLIST: + assert OmegaConf.select(source, field) == OmegaConf.select(target, field), field + assert target.env.events.pd_gains is None + assert target.env.events.reset_root_state_uniform is not None + assert target.env.superdex_allow_contact_approximation is True + assert target.env.superdex_num_workers == 0 + assert target.env.sim_dt == source.env.sim_dt + assert target.env.ctrl_dt == source.env.ctrl_dt + + +def test_go2_superdex_native_rollout_and_free_root_reset() -> None: + pytest.importorskip("superdex.physics") + if not (ROOT / "src/unilab/assets/robots/go2/assets/base_0.obj").is_file(): + pytest.skip("native Go2 rollout requires the registered Go2 robot assets") + _, override = _owner("superdex") + env = registry.make( + "Go2JoystickFlat", sim_backend="superdex", num_envs=2, env_cfg_override=override + ) + try: + state = env.init_state() + assert env.action_space.shape == (12,) + assert env.obs_groups_spec == {"obs": 49, "critic": 52} + for _ in range(32): + state = env.step(np.zeros((2, 12), dtype=np.float32)) + assert np.isfinite(state.reward).all() + assert all(np.isfinite(value).all() for value in state.obs.values()) + before = env.scene["robot"].data.root_link_pos_w.copy() + obs, _ = env.reset(env_ids=np.array([0], dtype=np.int32)) + assert obs["obs"].shape == (1, 49) + np.testing.assert_array_equal(env.scene["robot"].data.root_link_pos_w[1], before[1]) + np.testing.assert_allclose( + np.linalg.norm(env.scene["robot"].data.root_link_quat_w, axis=-1), 1.0, atol=1e-5 + ) + finally: + env.close() + + +def test_go2_superdex_parallel_matches_serial_env_and_unsorted_reset() -> None: + pytest.importorskip("superdex.physics") + if not (ROOT / "src/unilab/assets/robots/go2/assets/base_0.obj").is_file(): + pytest.skip("native Go2 rollout requires the registered Go2 robot assets") + _, override = _owner("superdex") + envs = [] + try: + for workers in (1, 2): + envs.append( + registry.make( + "Go2JoystickFlat", + sim_backend="superdex", + num_envs=4, + env_cfg_override={**override, "seed": 7, "superdex_num_workers": workers}, + ) + ) + serial, parallel = envs + for env in envs: + env.init_state() + rng = np.random.default_rng(11) + for _ in range(8): + action = rng.uniform(-0.1, 0.1, (4, 12)).astype(np.float32) + reference, result = (env.step(action) for env in envs) + for key in reference.obs: + np.testing.assert_allclose( + result.obs[key], reference.obs[key], atol=2e-5, rtol=2e-5 + ) + np.testing.assert_allclose(result.reward, reference.reward, atol=2e-5, rtol=2e-5) + before = parallel.scene["robot"].data.root_link_pos_w.copy() + # Non-monotonic IDs span both shards and must retain caller row order. + selected = np.array([3, 0], dtype=np.int32) + ref_obs, _ = serial.reset(env_ids=selected) + out_obs, _ = parallel.reset(env_ids=selected) + for key in ref_obs: + np.testing.assert_allclose(out_obs[key], ref_obs[key], atol=2e-5, rtol=2e-5) + np.testing.assert_array_equal( + parallel.scene["robot"].data.root_link_pos_w[[1, 2]], before[[1, 2]] + ) + finally: + for env in reversed(envs): + env.close() + + +def test_go2_superdex_executes_mujoco_checkpoint(tmp_path: Path) -> None: + checkpoint_value = os.environ.get("UNILAB_SUPERDEX_GO2_CHECKPOINT") + if not checkpoint_value: + pytest.skip("set UNILAB_SUPERDEX_GO2_CHECKPOINT to a MuJoCo Go2 PPO checkpoint") + pytest.importorskip("superdex.physics") + import torch + from rsl_rl.runners import OnPolicyRunner + from uni_rl.algos.rsl_rl import RslRlVecEnvWrapper, get_policy_obs_dims, normalize_ppo_train_cfg + + from unilab.training.run import algo_config_dict + from unilab.visualization.interactive_playback import ( + RslRlPlaybackConfig, + create_rsl_rl_playback_session, + infer_checkpoint_actor_input_dim, + make_sim2sim_preflight, + ) + + checkpoint = Path(checkpoint_value).resolve(strict=True) + owner, overrides = _owner("superdex") + # Validate source policy-I/O before any environment is constructed. + resolve_sim2sim_config(str(checkpoint.parent), owner, algo_name="ppo", strict=True) + incompatible = OmegaConf.create(OmegaConf.to_container(owner, resolve=True)) + incompatible.env.actions.joint_pos.scale = 0.5 + with pytest.raises(CrossBackendIncompatibleError, match="env.actions"): + resolve_sim2sim_config(str(checkpoint.parent), incompatible, algo_name="ppo", strict=True) + + env = registry.make( + "Go2JoystickFlat", sim_backend="superdex", num_envs=2, env_cfg_override=overrides + ) + try: + session, _, loaded = create_rsl_rl_playback_session( + playback_cfg=RslRlPlaybackConfig( + task="Go2JoystickFlat", + load_run=str(checkpoint), + checkpoint=None, + action_mode="policy", + policy_obs_mode="flat", + algo_log_name="rsl_rl_ppo", + log_root=str(tmp_path), + num_envs=2, + ), + env_factory=lambda n: env, + algo_config=algo_config_dict(owner), + root_dir=ROOT, + device="cpu", + checkpoint_resolver=lambda *_: str(checkpoint), + checkpoint_input_dim_reader=infer_checkpoint_actor_input_dim, + entrypoint_log_root=lambda *args, **kwargs: tmp_path, + wrapper_cls=RslRlVecEnvWrapper, + runner_cls=OnPolicyRunner, + policy_obs_dims_getter=get_policy_obs_dims, + train_cfg_normalizer=normalize_ppo_train_cfg, + sim2sim_preflight=make_sim2sim_preflight(owner, algo_name="ppo"), + guard_algo_name="ppo", + ) + assert loaded == str(checkpoint) and session.policy is not None + session.reset() + with torch.inference_mode(): + for _ in range(64): + obs = session.step_once() + assert all(torch.isfinite(value).all() for value in obs.values()) + assert np.isfinite(env.state.reward).all() + assert session.step_count == 64 + finally: + env.close() diff --git a/tests/ipc/test_dp_launcher.py b/tests/ipc/test_dp_launcher.py index 507bcaed1..d7f75d2bf 100644 --- a/tests/ipc/test_dp_launcher.py +++ b/tests/ipc/test_dp_launcher.py @@ -480,6 +480,16 @@ def test_resolve_collector_cpu_ids_even_partition(): assert resolve_collector_cpu_ids(2, 1, 128) == list(range(64, 128)) +def test_resolve_collector_cpu_ids_keeps_physical_siblings_together( + monkeypatch: pytest.MonkeyPatch, +): + groups = [[0, 4], [1, 5], [2, 6], [3, 7]] + monkeypatch.setattr(dp_launcher.os, "sched_getaffinity", lambda _: set(range(8)), raising=False) + monkeypatch.setattr(dp_launcher, "_discover_physical_cpu_groups", lambda _: groups) + assert resolve_collector_cpu_ids(2, 0) == [0, 4, 1, 5] + assert resolve_collector_cpu_ids(2, 1) == [2, 6, 3, 7] + + def test_resolve_collector_cpu_ids_remainder_stays_unassigned(): # 129 CPUs / 2 ranks -> 64+64; CPU 128 keeps default OS scheduling. assert resolve_collector_cpu_ids(2, 0, 129) == list(range(0, 64)) diff --git a/tests/scripts/test_audit_sim2sim_contracts.py b/tests/scripts/test_audit_sim2sim_contracts.py index 6012f88fe..20732720c 100644 --- a/tests/scripts/test_audit_sim2sim_contracts.py +++ b/tests/scripts/test_audit_sim2sim_contracts.py @@ -44,6 +44,18 @@ def test_discover_offpolicy_trees_group_by_task() -> None: assert td3["g1_walk_flat"] == ["mujoco"] +def test_go2_superdex_pair_is_audited_and_transferable(monkeypatch: pytest.MonkeyPatch) -> None: + audit = _load_audit_module() + monkeypatch.setattr( + audit, "_discover", lambda tree: {"go2_joystick_flat": ["mujoco", "superdex"]} + ) + rows = audit.audit_tree("ppo") + assert rows[0]["errors"] == {} + assert rows[0]["pairs"][0]["pair"] == "mujoco<->superdex" + assert rows[0]["pairs"][0]["verdict"] == "TRANSFERABLE" + assert rows[0]["pairs"][0]["warn_diffs"] == [] + + @pytest.mark.parametrize( ("tree", "task_variant", "expected_algo"), [ diff --git a/tests/scripts/test_support_matrix.py b/tests/scripts/test_support_matrix.py index 8f0864498..2bb57489c 100644 --- a/tests/scripts/test_support_matrix.py +++ b/tests/scripts/test_support_matrix.py @@ -5,6 +5,14 @@ import pytest from scripts.tools.support_matrix import BACKENDS, EvidenceLevel, build_support_rows + +def test_superdex_fr3_is_configured_without_full_training_claim() -> None: + row = _row("PPO (torch)", "fr3_joint_target") + assert row.cells["superdex"].level == EvidenceLevel.CONFIGURED + go2_row = _row("PPO (torch)", "go2_joystick_flat") + assert go2_row.cells["superdex"].level == EvidenceLevel.CONFIGURED + + # CPU-bound on the single-core CI runner; kept in the slow lane (make test-slow). pytestmark = pytest.mark.slow @@ -37,6 +45,7 @@ def test_support_matrix_marks_validated_g1_mjwarp_entrypoints_as_tested(): "genesis", "isaacsim", "newton", + "superdex", ) assert torch_row.cells["mjwarp"].level == EvidenceLevel.TESTED assert sac_row.cells["mjwarp"].level == EvidenceLevel.TESTED diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 9f4932ecb..531f4c463 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -660,6 +660,25 @@ def test_ppo_gpu_backend_env_uses_torchrun_local_rank( assert override[field] == 1 +def test_ppo_multi_rank_routes_one_cpu_partition_to_the_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + mod = _train_rsl_rl(monkeypatch) + cfg = _ppo_cfg(["task=go2_joystick_flat/superdex", "training.devices=[0,1]"]) + monkeypatch.setenv("RANK", "1") + monkeypatch.setenv("LOCAL_RANK", "1") + monkeypatch.setenv("WORLD_SIZE", "2") + monkeypatch.setattr( + mod, + "resolve_collector_cpu_ids", + lambda world_size, rank, cpu_count, explicit=None: [8, 24], + ) + + override = mod.build_ppo_env_cfg_override(cfg) + + assert override["cpu_ids"] == [8, 24] + + def test_offpolicy_isaacsim_training_and_eval_use_separate_render_overrides(): cfg = _offpolicy_cfg( [ diff --git a/tests/tasks/test_migration_matrix.py b/tests/tasks/test_migration_matrix.py index 849af3f45..9766a43c4 100644 --- a/tests/tasks/test_migration_matrix.py +++ b/tests/tasks/test_migration_matrix.py @@ -22,6 +22,7 @@ def test_registered_tasks_have_explicit_migration_records() -> None: @pytest.mark.parametrize( ("task_name", "family", "target", "status"), [ + ("FR3JointTarget", "manager_based", "complete", "Compatible"), ("SharpaInhandRotation", "sharpa", "compatibility", "Adapted"), ("G1MotionTracking", "motion_tracking", "complete", "Compatible"), ("G1WBTObs", "motion_tracking", "complete", "Compatible"), diff --git a/tests/tasks/test_package_boundary.py b/tests/tasks/test_package_boundary.py index 34d7f5056..6fdcf42e6 100644 --- a/tests/tasks/test_package_boundary.py +++ b/tests/tasks/test_package_boundary.py @@ -21,6 +21,7 @@ "unilab.tasks.manipulation.allegro_inhand", "unilab.tasks.manipulation.sharpa_inhand", "unilab.tasks.manipulation.stewart", + "unilab.tasks.manipulation.fr3", "unilab.tasks.motion_tracking.g1", "unilab.tasks.motion_tracking.x2", ) diff --git a/tests/test_cli_runtime_requirements.py b/tests/test_cli_runtime_requirements.py index 070357511..cb83b19d8 100644 --- a/tests/test_cli_runtime_requirements.py +++ b/tests/test_cli_runtime_requirements.py @@ -41,3 +41,21 @@ def test_check_runtime_requirements_requires_isolated_newton_extra( def test_newton_is_a_supported_sim() -> None: assert "newton" in cli.SUPPORTED_SIMS + + +def test_superdex_missing_runtime_reports_python_and_sdk(monkeypatch: pytest.MonkeyPatch) -> None: + from unisim.backend.superdex import dependencies + + monkeypatch.setattr(dependencies, "superdex_dependencies_available", lambda: False) + with pytest.raises(SystemExit, match="Python 3.12.*Physics/Robotics"): + cli._check_runtime_requirements("ppo", "superdex") + + +def test_superdex_old_unisim_reports_local_link_requirement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import sys + + monkeypatch.setitem(sys.modules, "unisim.backend.superdex.dependencies", None) + with pytest.raises(SystemExit, match="locally linked UniSim SuperDex"): + cli._check_runtime_requirements("ppo", "superdex") diff --git a/uv.lock b/uv.lock index ecd658237..cd4995627 100644 --- a/uv.lock +++ b/uv.lock @@ -5118,8 +5118,8 @@ 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.1.1" }, - { name = "unisim-core", specifier = ">=1.1.4" }, + { name = "unilab-rl", git = "https://github.com/unilabsim/unilab_rl.git?rev=2cdab3c" }, + { name = "unisim-core", git = "https://github.com/unilabsim/unisim.git?rev=037e596" }, { name = "viser", marker = "extra == 'viser'", specifier = ">=1.0.26" }, { name = "wandb" }, { name = "warp-lang", marker = "extra == 'mjwarp'", specifier = "==1.16.0" }, @@ -5140,7 +5140,7 @@ dev = [ [[package]] name = "unilab-rl" version = "1.1.1" -source = { registry = "https://pypi.org/simple" } +source = { git = "https://github.com/unilabsim/unilab_rl.git?rev=2cdab3c#2cdab3c69862ecc16a726b4761bcff263256b6b3" } dependencies = [ { name = "hydra-core" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -5156,20 +5156,15 @@ 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/51/a9/feca1ac03d07248d4870d405e03ae61e683476637f94673b53f154249e36/unilab_rl-1.1.1.tar.gz", hash = "sha256:e75c0a0a414e11d0bd70b0ba7bb45e4e549978e07a76b61429724197d5fa5848", size = 171127, upload-time = "2026-09-08T10:19:53.92Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d9/77/f1c15d42f25455571c29879fca00b9a0197ca69c7aa73a6f4a0980eabbbe/unilab_rl-1.1.1-py3-none-any.whl", hash = "sha256:d156fa47bc886b8b5285ecb6eb0bdd6e0bde302ad3cf9e42b2cb14b91c3270e4", size = 214137, upload-time = "2026-09-08T10:19:52.688Z" }, -] [[package]] name = "unisim-core" version = "1.1.4" -source = { registry = "https://pypi.org/simple" } +source = { git = "https://github.com/unilabsim/unisim.git?rev=037e596#037e5967cdcfc61785b1a685fd05454bbe49c147" } 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/4a/52/6b276891a561c9cdfacd909cbe4e11437efd97064002e92851eaad2002ea/unisim_core-1.1.4.tar.gz", hash = "sha256:9a98c6189489100191ef93e8d5de4a5be051f2a38687ce51422f90ed4d182a19", size = 219821, upload-time = "2026-09-08T10:18:43.113Z" } [[package]] name = "urllib3"