Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions docs/sphinx/source/en/2-user_guide/3-backends/7-newton.md
Original file line number Diff line number Diff line change
Expand Up @@ -33,8 +33,10 @@ collector sim processes land on their own physical GPU with no cross-GPU
leakage; single-GPU PPO/SAC regressions pass alongside. Newton/Warp follows
standard CUDA device semantics, so no `CUDA_VISIBLE_DEVICES` pinning (the
Genesis quirk) is needed; the rank-local device reaches spawn collectors as
a `newton_device="cuda:N"` env override (uni_rl 1.0.0's collector-side
process-binding gate only covers mjwarp), and the SAC owner raises the
a `newton_device="cuda:N"` env override, and uni_rl's collector-side
process binding is injection-based — UniLab injects
`bind_backend_process_device_for_backend`, which covers both mjwarp and
newton — while the SAC owner raises the
collector tick-0 timeout to 180 s to cover Warp kernel compilation on the
cold path.

Expand Down
7 changes: 4 additions & 3 deletions docs/sphinx/source/zh_CN/2-user_guide/3-backends/7-newton.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,10 @@ DpRankSupervisor DP=2(`training.devices=[0,1]`)训练冒烟均正常完成
在各自 GPU、无对端泄漏;单卡 PPO/SAC 回归同步通过。Newton/Warp 遵循标准
CUDA 设备语义,无需 Genesis 那样的 `CUDA_VISIBLE_DEVICES` 钉卡;rank
本地设备通过 env override 以 `newton_device="cuda:N"` 传入 spawn
collector(uni_rl 1.0.0 的 collector 进程绑定门只覆盖 mjwarp),SAC
owner 将 collector tick-0 超时提升到 180 s 以覆盖 Warp 内核编译的冷
路径。
collector,且 uni_rl 的 collector 进程绑定为注入式——UniLab 注入的
`bind_backend_process_device_for_backend` 同时覆盖 mjwarp 与
newton;SAC owner 将 collector tick-0 超时提升到 180 s 以覆盖 Warp
内核编译的冷路径。

## 安装

Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,7 @@ viser = ["viser>=1.0.26", "trimesh>=3.21.7"]
# required-environments; elsewhere the extra is empty and the CLI reports a
# targeted runtime diagnostic.
superdex = [
"unisim-core[superdex]>=1.4.0 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
"unisim-core[superdex]>=1.4.2 ; python_version >= '3.12' and sys_platform == 'linux' and platform_machine == 'x86_64'",
]

[dependency-groups]
Expand Down
18 changes: 6 additions & 12 deletions scripts/benchmark/rl/benchmark_replay_buffer_placement.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,33 +191,27 @@ def _xpu_available() -> bool:


def _replay_transfer_manifest(device: torch.device, *, ring_depth: int = 2) -> dict[str, Any]:
"""Describe the uni_rl 1.2.1 replay-transfer backend for ``device``.

uni_rl 1.2.1 builds ``CudaLikeReplayTransferBackend`` for CUDA (ROCm
included) and ``TorchCopyReplayTransferBackend`` otherwise; the async
``submit_h2d`` path and the dedicated XPU backend no longer exist.
"""
if device.type == "cuda":
torch_version = getattr(torch, "version", None)
is_rocm = bool(getattr(torch_version, "hip", None))
return {
"backend": "CudaLikeReplayTransferBackend",
"device_family": "rocm" if is_rocm else "cuda",
"host_memory_kind": "registered_pinned_shared",
"supports_async_submit": True,
"supports_timing_events": True,
"h2d_submitter": "torch_copy_stream" if is_rocm else "pybind11",
"ring_depth": ring_depth,
}
if device.type == "xpu":
return {
"backend": "XpuReplayTransferBackend",
"device_family": "xpu",
"host_memory_kind": "pageable_shared",
"supports_async_submit": True,
"supports_timing_events": False,
"h2d_submitter": "torch_xpu_copy_stream",
"ring_depth": ring_depth,
}
return {
"backend": "TorchCopyReplayTransferBackend",
"device_family": device.type,
"host_memory_kind": "pageable_shared",
"supports_async_submit": False,
"supports_timing_events": False,
"h2d_submitter": "torch_copy",
"ring_depth": ring_depth,
Expand Down
65 changes: 8 additions & 57 deletions src/unilab/base/backend_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,25 +22,6 @@
from unilab.base.scene import SceneCfg


def _legacy_genesis_device_option_error(exc: TypeError) -> bool:
"""Identify an old UniSim adapter rejecting the optional device keyword.

UniSim 1.1 reports unknown backend options from ``GenesisBackend`` while
other compatible releases may expose Python's usual ``unexpected keyword``
wording. Keep the compatibility retry narrowly scoped to those messages;
constructor errors from the actual Genesis runtime must still propagate.
"""

message = str(exc).lower()
mentions_device = "genesis_device_id" in message or "device_id" in message
rejects_keyword = (
"does not accept backend options" in message
or "unexpected keyword argument" in message
or "unexpected keyword" in message
)
return mentions_device and rejects_keyword


def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]:
"""Translate ``EnvCfg`` backend knobs into UniSim adapter options."""
result: dict[str, Any] = {
Expand Down Expand Up @@ -69,18 +50,13 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]:
"isaacsim_render_mode": cfg.isaacsim_render_mode,
"isaacsim_render_width": cfg.isaacsim_render_width,
"isaacsim_render_height": cfg.isaacsim_render_height,
"superdex_execution_mode": cfg.superdex_execution_mode,
}
# Keep the optional key absent for legacy unisim-core releases that do not
# know about Genesis' explicit device argument. Once a rank selects a
# device the key is added below and ``create_backend`` supplies a narrow
# compatibility fallback for those releases.
# Forward the explicit Genesis device id only when a rank selected one;
# when absent, unisim-core's factory default applies and Genesis picks
# its own device.
if cfg.genesis_device_id is not None:
result["genesis_device_id"] = cfg.genesis_device_id
# Keep the default absent so unisim-core releases that predate the
# execution-mode option still accept the SuperDex kwargs; "serial" requires
# the updated adapter.
if cfg.superdex_execution_mode != "batch":
result["superdex_execution_mode"] = cfg.superdex_execution_mode
return result


Expand Down Expand Up @@ -112,17 +88,6 @@ def create_backend(
ensure_robot_assets_for_paths(
[scene.model_file, scene.visual_model_file, *scene.fragment_files]
)
if backend_type != "newton":
# Keep the owner translation forward-compatible with unisim-core
# releases that predate the Newton adapter and therefore do not pop
# these optional kwargs in their shared factory.
for key in (
"newton_device",
"newton_nconmax",
"newton_njmax",
"newton_capacity_check_steps",
):
kwargs.pop(key, None)
if backend_type == "drake":
# unisim-core 1.4.2 dropped the Drake-branch filtering of MuJoCo
# root-body options; Drake derives root state from its own plant and
Expand All @@ -139,11 +104,9 @@ def create_backend(
}
if backend_type == "genesis" and kwargs.get("genesis_device_id") is not None:
# Bind before any unisim-core Genesis constructor can call gs.init.
# New unisim-core releases repeat this idempotently; old releases do
# not accept the keyword, so the retry below still gets the correct
# process-wide device. Binding a non-zero id pins
# CUDA_VISIBLE_DEVICES (Quadrants only honors the first visible
# device), so forward the *post-pin* in-process index.
# Binding a non-zero id pins CUDA_VISIBLE_DEVICES (Quadrants only
# honors the first visible device), so forward the *post-pin*
# in-process index.
genesis_device_id = kwargs["genesis_device_id"]
if (
isinstance(genesis_device_id, bool)
Expand All @@ -156,19 +119,7 @@ def create_backend(
)
bound = bind_genesis_process_device(f"cuda:{genesis_device_id}")
kwargs["genesis_device_id"] = int(bound.rsplit(":", 1)[1])
try:
return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **kwargs)
except TypeError as exc:
if backend_type != "genesis" or "genesis_device_id" not in kwargs:
raise
# unisim-core < 1.2 has no Genesis device field and reports the
# unknown option from GenesisBackend. Retry only for that precise
# capability error; unrelated constructor TypeErrors must propagate.
if not _legacy_genesis_device_option_error(exc):
raise
legacy_kwargs = dict(kwargs)
legacy_kwargs.pop("genesis_device_id", None)
return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **legacy_kwargs)
return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **kwargs)


__all__ = ["SimBackend", "create_backend", "env_backend_kwargs"]
8 changes: 5 additions & 3 deletions src/unilab/base/env_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,11 @@ def make_registry_env(
# carry the explicit cold-path id in the opaque override and bind
# immediately before registry construction. Binding a non-zero id pins
# CUDA_VISIBLE_DEVICES for this process, so forward the post-pin
# in-process index downstream. Newer unisim-core versions repeat this
# check in GenesisBackend itself, making this compatibility guard
# idempotent.
# in-process index downstream. This binding is required, not a
# compatibility shim: spawn collectors are fresh interpreters that cannot
# inherit the parent process's binding, and only UniLab's
# ``bind_genesis_process_device`` sets the pinned-namespace flag its
# resolution helpers consult.
if sim_backend == "genesis" and env_cfg_override is not None:
genesis_device_id = env_cfg_override.get("genesis_device_id")
if genesis_device_id is not None:
Expand Down
9 changes: 6 additions & 3 deletions src/unilab/base/process_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,9 +26,12 @@


# Newton consumes an explicit ``cuda:N`` device string (``newton_device``)
# instead of an integer id. uni_rl's collector-side binder gate only knows
# mjwarp, so the rank-local device must reach spawn collectors through the
# env override rather than through process binding.
# instead of an integer id. uni_rl's collector-side process binding is
# injection-based: it binds through the caller-supplied ``bind_device``
# callable and fails closed when none is injected. UniLab injects
# ``bind_backend_process_device_for_backend`` (which covers mjwarp and
# newton); the env override additionally forwards the rank-local device
# string so spawn collectors can pass it to the Newton adapter.
BACKEND_ENV_DEVICE_STR_FIELDS: dict[str, str] = {
"newton": "newton_device",
}
Expand Down
8 changes: 1 addition & 7 deletions src/unilab/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,13 +133,7 @@ def _check_runtime_requirements(algo: str, sim: str) -> None:
"in a source checkout (or `pip install unilab[newton]`)."
)
if sim == "superdex":
try:
from unisim.backend.superdex.dependencies import superdex_dependencies_available
except ImportError as exc:
raise SystemExit(
"sim=superdex requires unisim-core>=1.1.5 with the SuperDex adapter; "
"the installed unisim-core does not provide that adapter."
) from exc
from unisim.backend.superdex.dependencies import superdex_dependencies_available

if not superdex_dependencies_available():
raise SystemExit(
Expand Down
4 changes: 2 additions & 2 deletions tests/base/backend/test_process_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,8 +170,8 @@ def test_non_gpu_backend_is_left_untouched() -> None:


def test_newton_override_carries_cuda_device_string() -> None:
# uni_rl's collector binder gate only knows mjwarp, so newton's rank-local
# device must reach spawn collectors as a ``cuda:N`` override string.
# Newton consumes an explicit ``cuda:N`` string, so its rank-local device
# reaches spawn collectors as an override string rather than an integer id.
owner_override: dict[str, object] = {"newton_device": None, "nested": {"keep": True}}
routed = apply_backend_env_device_override(
owner_override,
Expand Down
6 changes: 2 additions & 4 deletions tests/base/test_superdex_backend_options.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ def create(name: str, scene: SceneCfg, num_envs: int, dt: float, **kwargs: Any)
assert captured["superdex_num_workers"] == workers


def test_superdex_execution_mode_is_forwarded_only_when_serial(
def test_superdex_execution_mode_is_always_forwarded(
monkeypatch: pytest.MonkeyPatch,
) -> None:
captured: dict[str, Any] = {}
Expand All @@ -111,13 +111,11 @@ def create(name: str, scene: SceneCfg, num_envs: int, dt: float, **kwargs: Any)
backend_factory.create_backend(
"superdex", scene, 1, 0.002, **backend_factory.env_backend_kwargs(serial)
)
# Legacy unisim-core releases predate the option; the default stays absent
# so they still accept the SuperDex kwargs.
assert captured["superdex_execution_mode"] == "serial"
backend_factory.create_backend(
"superdex", scene, 1, 0.002, **backend_factory.env_backend_kwargs(EnvCfg())
)
assert "superdex_execution_mode" not in captured
assert captured["superdex_execution_mode"] == "batch"


def test_superdex_execution_mode_validates_and_does_not_leak(
Expand Down
2 changes: 1 addition & 1 deletion tests/benchmark/test_replay_buffer_placement_benchmark.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,7 +133,7 @@ def test_replay_transfer_manifest_records_backend_fields() -> None:
assert manifest["backend"] == "TorchCopyReplayTransferBackend"
assert manifest["device_family"] == "cpu"
assert manifest["host_memory_kind"] == "pageable_shared"
assert manifest["supports_async_submit"] is False
assert "supports_async_submit" not in manifest
assert manifest["ring_depth"] == 2


Expand Down
10 changes: 0 additions & 10 deletions tests/test_cli_runtime_requirements.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,13 +62,3 @@ def test_superdex_missing_runtime_reports_python_and_sdk(monkeypatch: pytest.Mon
SystemExit, match=r"Python 3\.12.*Physics/Robotics.*uv sync --extra superdex"
):
cli._check_runtime_requirements("ppo", "superdex")


def test_superdex_old_unisim_reports_adapter_requirement(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import sys

monkeypatch.setitem(sys.modules, "unisim.backend.superdex.dependencies", None)
with pytest.raises(SystemExit, match="unisim-core>=1.1.5 with the SuperDex adapter"):
cli._check_runtime_requirements("ppo", "superdex")
2 changes: 1 addition & 1 deletion uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.