diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md
index 4277eff75..358f813bc 100644
--- a/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md
+++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/1-ppo.md
@@ -50,6 +50,13 @@ device. Do not set `training.device` and `training.devices` together. The
configured order is preserved, including when the parent already has
`CUDA_VISIBLE_DEVICES` set.
+For the IsaacGym, IsaacSim, and Genesis owners, the same topology is also
+applied to the simulator environment. Torchrun workers receive the local
+index inside their remapped `CUDA_VISIBLE_DEVICES` list (for example, host
+device 5 is sent as `device_id=1` when the worker sees `[4,5]`); off-policy
+workers keep the parent process's visible index namespace. Genesis selects
+its process-wide session before `gs.init`.
+
`algo.num_envs` is a **per-rank** count, not a global budget. For `W` ranks,
`N` configured envs, and rollout length `T`:
diff --git a/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md b/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md
index d45cdf62f..1d945c93f 100644
--- a/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md
+++ b/docs/sphinx/source/en/2-user_guide/2-algorithms/3-sac.md
@@ -54,6 +54,11 @@ materialization. The collector therefore does not fall back to Warp's fresh-proc
of `cuda:0`. The local binding is recorded as `collector_backend_device` in the runtime
manifest.
+IsaacGym, IsaacSim, and Genesis receive the rank-selected simulator device
+through the environment override as well. Off-policy collectors use the
+parent's visible CUDA indices; Genesis binds its process-wide session before
+initialization.
+
MuJoCo has a committed multi-GPU scaling benchmark. The mjwarp per-rank placement contract is
covered by `tests/base/backend/test_process_device.py` and the off-policy runner/worker unit
tests; the repository does not currently contain an mjwarp multi-GPU throughput or convergence
diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md
index 26094744c..718f33115 100644
--- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md
+++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/1-ppo.md
@@ -48,6 +48,11 @@ uv run train --algo ppo --task g1_motion_tracking --sim mujoco \
`training.devices`。父进程已有 `CUDA_VISIBLE_DEVICES` 时，配置索引仍按父进程可见
设备解释，并保留用户给定顺序。
+对 IsaacGym、IsaacSim 和 Genesis owner，同一拓扑也会传给环境仿真器。torchrun
+worker 继承重映射后的 `CUDA_VISIBLE_DEVICES`，因此传给 worker 的是本地索引（例如
+worker 看到 `[4,5]` 时，主机设备 5 传为 `device_id=1`）；off-policy worker 保持父进程
+可见设备索引。Genesis 会在 `gs.init` 前选择每个进程的 session 设备。
+
`algo.num_envs` 是**每个 rank** 的环境数，不是全局预算。设 rank 数为 `W`、配置
环境数为 `N`、rollout 长度为 `T`：
diff --git a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md
index c0e2705b5..734044002 100644
--- a/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md
+++ b/docs/sphinx/source/zh_CN/2-user_guide/2-algorithms/3-sac.md
@@ -97,6 +97,8 @@ MuJoCo worker 线程逐核绑定外，collector 进程本身（含 Numba 并行
- MuJoCo 有已提交的多卡 scaling benchmark；mjwarp 的 per-rank device placement 有
`tests/base/backend/test_process_device.py` 与 off-policy runner/worker 单测覆盖，但仓库中
尚无 mjwarp 多卡吞吐或收敛 benchmark。
+- IsaacGym、IsaacSim 和 Genesis 的 collector 环境也会收到 rank 对应的仿真设备。off-policy
+ 使用父进程可见的 CUDA 索引；Genesis 在初始化 `gs.init` 前绑定进程级 session 设备。
- 仅单节点：rank 之间通过 run 目录里的 FileStore rendezvous，NCCL 走 TCP
loopback（默认 `NCCL_P2P_DISABLE=1` / `NCCL_SHM_DISABLE=1`，环境变量显式设置
时优先）——部分机型（如 RTX 6000D）的 NCCL P2P/SHM peer transport 不可靠，
diff --git a/src/unilab/base/backend_factory.py b/src/unilab/base/backend_factory.py
index e9bca6e02..40df7af90 100644
--- a/src/unilab/base/backend_factory.py
+++ b/src/unilab/base/backend_factory.py
@@ -14,15 +14,35 @@
from unisim.backend.base import SimBackend
from unilab.assets.hub import ensure_robot_assets_for_paths
+from unilab.base.process_device import bind_genesis_process_device
if TYPE_CHECKING:
from unilab.base.base import EnvCfg
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."""
- return {
+ result: dict[str, Any] = {
"post_step_forward_sensor": cfg.post_step_forward_sensor,
"motrix_max_iterations": cfg.motrix_max_iterations,
"chunk_size": cfg.chunk_size,
@@ -45,6 +65,13 @@ def env_backend_kwargs(cfg: "EnvCfg") -> dict[str, Any]:
"isaacsim_render_width": cfg.isaacsim_render_width,
"isaacsim_render_height": cfg.isaacsim_render_height,
}
+ # 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.
+ if cfg.genesis_device_id is not None:
+ result["genesis_device_id"] = cfg.genesis_device_id
+ return result
def create_backend(
@@ -63,7 +90,38 @@ def create_backend(
[scene.model_file, scene.visual_model_file, *scene.fragment_files]
)
kwargs["body_state_required"] = body_state_required
- return unisim.create_backend(backend_type, scene, num_envs, sim_dt, **kwargs)
+ 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.
+ genesis_device_id = kwargs["genesis_device_id"]
+ if (
+ isinstance(genesis_device_id, bool)
+ or not isinstance(genesis_device_id, int)
+ or genesis_device_id < 0
+ ):
+ raise ValueError(
+ "genesis_device_id must be a non-negative integer or None, "
+ f"got {genesis_device_id!r}"
+ )
+ 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)
__all__ = ["SimBackend", "create_backend", "env_backend_kwargs"]
diff --git a/src/unilab/base/base.py b/src/unilab/base/base.py
index 59b6dd82e..8f386d291 100644
--- a/src/unilab/base/base.py
+++ b/src/unilab/base/base.py
@@ -56,6 +56,10 @@ class EnvCfg:
# backend defaults (device 0, generous handshake/step timeout).
isaacgym_device_id: Optional[int] = None
isaacgym_worker_timeout_s: Optional[float] = None
+ # ``genesis`` owns one process-wide GPU session. The explicit device id
+ # must be selected before ``gs.init`` so each data-parallel rank gets its
+ # own simulator device; ``None`` keeps Genesis' current/default device.
+ genesis_device_id: Optional[int] = None
# ``genesis`` drops the MJCF global block at import (REPORT #1372
# §3.3), so integrator / constraint solver / friction cone / solver
# iterations must be explicit owner fields. ``None`` keeps the Genesis
@@ -123,6 +127,15 @@ def validate(self):
"isaacgym_worker_timeout_s must be a positive number or None, "
f"got {self.isaacgym_worker_timeout_s!r}"
)
+ if self.genesis_device_id is not None and (
+ isinstance(self.genesis_device_id, bool)
+ or not isinstance(self.genesis_device_id, int)
+ or self.genesis_device_id < 0
+ ):
+ raise ValueError(
+ "genesis_device_id must be a non-negative integer or None, "
+ f"got {self.genesis_device_id!r}"
+ )
for name, value in (
("genesis_integrator", self.genesis_integrator),
("genesis_constraint_solver", self.genesis_constraint_solver),
diff --git a/src/unilab/base/env_factory.py b/src/unilab/base/env_factory.py
index 72ea44dcf..46c256714 100644
--- a/src/unilab/base/env_factory.py
+++ b/src/unilab/base/env_factory.py
@@ -16,6 +16,8 @@
from uni_rl.env_contract import EnvFactory, EnvProtocol
+from unilab.base.process_device import bind_genesis_process_device
+
def make_registry_env(
task_name: str,
@@ -34,6 +36,32 @@ def make_registry_env(
from unilab.base.registry import ensure_registries
ensure_registries()
+ # Genesis owns a process-wide session whose Quadrants runtime binds the
+ # first entry of CUDA_VISIBLE_DEVICES. Off-policy/APPO collectors are
+ # fresh spawn processes, so the parent-side binding cannot reach them;
+ # 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.
+ 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:
+ if (
+ isinstance(genesis_device_id, bool)
+ or not isinstance(genesis_device_id, int)
+ or genesis_device_id < 0
+ ):
+ raise ValueError(
+ "genesis_device_id must be a non-negative integer or None, "
+ f"got {genesis_device_id!r}"
+ )
+ bound = bind_genesis_process_device(f"cuda:{genesis_device_id}")
+ env_cfg_override = {
+ **env_cfg_override,
+ "genesis_device_id": int(bound.rsplit(":", 1)[1]),
+ }
# ABEnv satisfies EnvProtocol at runtime (reset/set_nan_guard live on
# NpEnv); the declared ABEnv type predates the uni_rl protocol.
return cast(
diff --git a/src/unilab/base/process_device.py b/src/unilab/base/process_device.py
index 1302c7ff9..b2278f649 100644
--- a/src/unilab/base/process_device.py
+++ b/src/unilab/base/process_device.py
@@ -1,19 +1,258 @@
-"""Training-worker device routing for UniSim backends."""
+"""Training-worker device routing for UniSim backends.
+
+The learner device and the device consumed by a simulator are related, but
+they do not always use the same index namespace. In particular, the
+off-policy launcher keeps the host-visible CUDA namespace while the PPO
+``torchrun`` launcher remaps ``CUDA_VISIBLE_DEVICES`` and therefore exposes a
+rank-local index to child processes. The helpers in this module keep that
+translation on the cold path, next to the backend process binding contract.
+"""
from __future__ import annotations
-from typing import cast
+import os
+import warnings
+from collections.abc import Mapping, Sequence
+from typing import Any, cast
+
+# These backends consume an explicit integer device id while materializing
+# their simulator. MuJoCo/Motrix/Drake either run on the host or own their
+# device selection internally and must not receive a synthetic override.
+BACKEND_ENV_DEVICE_FIELDS: dict[str, str] = {
+ "isaacgym": "isaacgym_device_id",
+ "isaacsim": "isaacsim_device_id",
+ "genesis": "genesis_device_id",
+}
+
+
+# Set once ``bind_genesis_process_device`` has pinned CUDA_VISIBLE_DEVICES for
+# this process. Genesis/Quadrants binds its CUDA runtime to the first visible
+# device regardless of torch's current device (verified on genesis_world
+# 1.3.3 / Quadrants 1.3.0, issue #1508), so a non-zero request is honored by
+# shrinking visibility to the target GPU and using the in-process index 0.
+# The flag flips the resolution helpers into the pinned namespace.
+_genesis_device_pinned = False
+
+
+def _normalize_backend(backend_type: str) -> str:
+ if not isinstance(backend_type, str) or not backend_type.strip():
+ raise ValueError(f"backend_type must be a non-empty string, got {backend_type!r}")
+ return backend_type.strip().lower()
+
+
+def _normalize_device_indices(devices: Sequence[int] | None) -> tuple[int, ...] | None:
+ if devices is None:
+ return None
+ normalized: list[int] = []
+ for entry in devices:
+ if isinstance(entry, bool) or not isinstance(entry, int):
+ raise ValueError(
+ f"training.devices entries must be integer CUDA indices, got {entry!r}"
+ )
+ if entry < 0:
+ raise ValueError(f"training.devices entries must be non-negative, got {entry}")
+ normalized.append(int(entry))
+ if len(set(normalized)) != len(normalized):
+ raise ValueError(f"training.devices must not contain duplicates, got {normalized}")
+ return tuple(normalized)
+
+
+def _cuda_device_index(device: str | None) -> int | None:
+ """Extract an integer CUDA index from a device string.
+
+ ``cuda`` without an explicit suffix is resolved through the current CUDA
+ device when possible. This is deliberately a cold-path helper; it is not
+ used from environment ``step``/``reset`` loops.
+ """
+
+ if device is None:
+ return None
+ value = str(device).strip().lower()
+ if value == "cuda":
+ try:
+ import torch
+
+ if torch.cuda.is_available():
+ return int(torch.cuda.current_device())
+ except Exception:
+ # Device discovery is only a fallback for an unindexed alias. A
+ # configured topology below remains authoritative if available.
+ pass
+ return 0
+ if not value.startswith("cuda:"):
+ return None
+ index_text = value.split(":", 1)[1].strip()
+ if not index_text:
+ raise ValueError(f"CUDA device alias {device!r} has an empty index")
+ try:
+ index = int(index_text)
+ except ValueError as exc:
+ raise ValueError(f"CUDA device alias {device!r} has a non-integer index") from exc
+ if index < 0:
+ raise ValueError(f"CUDA device alias {device!r} has a negative index")
+ return index
+
+
+def resolve_backend_env_device_id(
+ backend_type: str,
+ *,
+ devices: Sequence[int] | None = None,
+ rank: int = 0,
+ local_rank: int | None = None,
+ world_size: int = 1,
+ learner_device: str | None = None,
+) -> int | None:
+ """Resolve the integer simulator device id for a training rank.
+
+ Args:
+ backend_type: Selected UniSim backend.
+ devices: ``training.devices`` in the host-visible namespace. This is
+ used by the off-policy launcher and by single-process PPO.
+ rank: Off-policy data-parallel rank (rank zero by default).
+ local_rank: ``LOCAL_RANK`` from torchrun. In a distributed PPO worker
+ this is the logical index inside the launcher's remapped
+ ``CUDA_VISIBLE_DEVICES`` list.
+ world_size: Torchrun world size. Values greater than one select the
+ ``local_rank`` namespace; values of one select ``devices[rank]``.
+ learner_device: Explicit learner device fallback when no topology was
+ configured (for example APPO or a single-device play command).
+
+ Returns ``None`` for backends without an explicit simulator device field.
+ For a torchrun worker the returned id is intentionally *local* (rather
+ than the host index in ``devices``), because the worker subprocess inherits
+ the remapped ``CUDA_VISIBLE_DEVICES`` environment.
+ """
+
+ backend = _normalize_backend(backend_type)
+ field = BACKEND_ENV_DEVICE_FIELDS.get(backend)
+ if field is None:
+ return None
+
+ if backend == "genesis" and _genesis_device_pinned:
+ # Post-pin the process sees exactly one CUDA device; every rank-local
+ # consumer (learner probe, spawn collector, playback) must use the
+ # in-process index 0 regardless of the original host/rank topology.
+ return 0
+
+ normalized_devices = _normalize_device_indices(devices)
+ world_size = int(world_size)
+ if world_size < 1:
+ raise ValueError(f"world_size must be positive, got {world_size}")
+
+ if world_size > 1:
+ resolved_local_rank = int(rank if local_rank is None else local_rank)
+ if resolved_local_rank < 0 or resolved_local_rank >= world_size:
+ raise ValueError(
+ f"local_rank={resolved_local_rank} is out of range for world_size={world_size}"
+ )
+ if normalized_devices is not None and len(normalized_devices) != world_size:
+ raise ValueError(
+ f"training.devices has {len(normalized_devices)} entries but "
+ f"WORLD_SIZE={world_size}"
+ )
+ # torchrun launch_torchrun_workers remaps CVD to the selected physical
+ # devices. Isaac workers inherit that environment, so LOCAL_RANK is
+ # the correct payload index.
+ return resolved_local_rank
+
+ if normalized_devices:
+ resolved_rank = int(rank)
+ if resolved_rank < 0 or resolved_rank >= len(normalized_devices):
+ raise ValueError(
+ f"rank={resolved_rank} is out of range for training.devices="
+ f"{list(normalized_devices)}"
+ )
+ return normalized_devices[resolved_rank]
+
+ return _cuda_device_index(learner_device)
+
+
+def apply_backend_env_device_override(
+ env_cfg_override: Mapping[str, Any] | None,
+ backend_type: str,
+ *,
+ devices: Sequence[int] | None = None,
+ rank: int = 0,
+ local_rank: int | None = None,
+ world_size: int = 1,
+ learner_device: str | None = None,
+) -> dict[str, Any]:
+ """Return an env override carrying the rank-selected simulator device.
+
+ The input mapping is never mutated. If no topology/device can be
+ resolved, the owner-configured value is preserved. This lets one helper
+ serve training, playback, and custom entrypoints while retaining the
+ historical default (device zero) for single-process calls.
+ """
+
+ result = dict(env_cfg_override) if env_cfg_override is not None else {}
+ backend = _normalize_backend(backend_type)
+ field = BACKEND_ENV_DEVICE_FIELDS.get(backend)
+ if field is None:
+ return result
+ device_id = resolve_backend_env_device_id(
+ backend,
+ devices=devices,
+ rank=rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=learner_device,
+ )
+ if device_id is not None:
+ result[field] = int(device_id)
+ return result
+
+
+def warn_if_backend_device_collision(
+ backend_type: str,
+ *,
+ devices: Sequence[int] | None,
+ rank: int,
+ device_id: int | None,
+ source: str = "environment",
+) -> None:
+ """Warn when a multi-rank simulator still resolves to device zero.
+
+ This is a transition guard for older adapters/configuration paths. Rank
+ zero legitimately owns device zero; only a non-zero rank resolving to zero
+ is a collision. The warning is intentionally emitted at construction
+ time, never from a hot simulation path.
+ """
+
+ backend = _normalize_backend(backend_type)
+ if backend not in BACKEND_ENV_DEVICE_FIELDS:
+ return
+ if backend == "genesis" and _genesis_device_pinned:
+ # A successful pin places every rank on its own physical GPU; the
+ # in-process index 0 that follows is not a collision.
+ return
+ normalized_devices = _normalize_device_indices(devices)
+ if (
+ normalized_devices is None
+ or len(normalized_devices) <= 1
+ or int(rank) <= 0
+ or device_id != 0
+ ):
+ return
+ warnings.warn(
+ f"{backend} rank {int(rank)} resolved its {source} device to 0 while "
+ f"training.devices={list(normalized_devices)} requests multiple devices; "
+ "all simulator workers may be sharing GPU 0",
+ RuntimeWarning,
+ stacklevel=2,
+ )
def resolve_backend_process_device(backend_type: str, learner_device: str | None) -> str | None:
- if backend_type != "mjwarp":
+ backend = _normalize_backend(backend_type)
+ if backend not in {"mjwarp", "genesis"}:
return None
if learner_device is None:
- raise ValueError("mjwarp requires an explicit CUDA process device")
+ raise ValueError(f"{backend} requires an explicit CUDA process device")
resolved = str(learner_device).strip()
if resolved.split(":", 1)[0].lower() != "cuda":
raise ValueError(
- f"mjwarp requires a CUDA process device shared with its learner; got {resolved!r}"
+ f"{backend} requires a CUDA process device shared with its learner; got {resolved!r}"
)
return resolved
@@ -22,6 +261,8 @@ def configure_backend_process_device(backend_type: str, learner_device: str | No
resolved = resolve_backend_process_device(backend_type, learner_device)
if resolved is None:
return None
+ if _normalize_backend(backend_type) == "genesis":
+ return bind_genesis_process_device(resolved)
return bind_backend_process_device(resolved)
@@ -38,8 +279,139 @@ def bind_backend_process_device(resolved: str) -> str | None:
return cast(str | None, bind_mjwarp_process_device(resolved))
+def _pin_cuda_visible_devices(index: int) -> None:
+ """Shrink ``CUDA_VISIBLE_DEVICES`` to the single entry at ``index``.
+
+ The index addresses the *current* visibility namespace: with no variable
+ set it is the host index, otherwise it indexes into the existing entries
+ (which may be physical indices or UUIDs). This only works before the
+ first CUDA context exists, so an already-initialized torch runtime fails
+ closed with an actionable error instead of crashing inside the engine.
+ """
+
+ import torch
+
+ if torch.cuda.is_initialized():
+ raise RuntimeError(
+ "genesis device routing must pin CUDA_VISIBLE_DEVICES before any CUDA "
+ "context is created in this process, but torch CUDA is already "
+ "initialized; move configure_backend_process_device earlier in the "
+ "entrypoint (before seeding/learner construction)"
+ )
+ raw = os.environ.get("CUDA_VISIBLE_DEVICES")
+ entries = [entry.strip() for entry in raw.split(",") if entry.strip()] if raw else None
+ if entries is None:
+ count = int(torch.cuda.device_count())
+ if index >= count:
+ raise ValueError(
+ f"genesis device index {index} is out of range; torch.cuda.device_count()={count}"
+ )
+ target = str(index)
+ else:
+ if index >= len(entries):
+ raise ValueError(
+ f"genesis device index {index} is out of range for "
+ f"CUDA_VISIBLE_DEVICES={raw!r} ({len(entries)} entr(ies))"
+ )
+ target = entries[index]
+ os.environ["CUDA_VISIBLE_DEVICES"] = target
+ # ``device_count`` is lru-cached; drop any pre-pin host-wide count so
+ # later callers observe the pinned single-device namespace.
+ cache_clear = getattr(torch.cuda.device_count, "cache_clear", None)
+ if callable(cache_clear):
+ cache_clear()
+
+
+def pin_genesis_device_before_cuda_init(
+ backend_type: str,
+ *,
+ devices: Sequence[int] | None = None,
+ rank: int = 0,
+ local_rank: int | None = None,
+ world_size: int = 1,
+ learner_device: str | None = None,
+) -> str | None:
+ """Pin Genesis to its rank device before the first torch CUDA call.
+
+ ``torch.cuda.is_available()`` already latches ``CUDA_VISIBLE_DEVICES`` in
+ the CUDA runtime, so the pin must run ahead of *any* torch CUDA query —
+ entrypoints should call this before registry/bootstrap/device detection.
+ Pure config topology (``training.devices`` / ``LOCAL_RANK`` / an explicit
+ ``cuda:N`` learner device) resolves without touching torch. Returns the
+ in-process device the caller must use when a pin happened, else ``None``.
+ """
+
+ if _normalize_backend(backend_type) != "genesis":
+ return None
+ device_id = resolve_backend_env_device_id(
+ backend_type,
+ devices=devices,
+ rank=rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=learner_device,
+ )
+ if not device_id:
+ return None
+ return bind_genesis_process_device(f"cuda:{device_id}")
+
+
+def bind_genesis_process_device(resolved: str) -> str:
+ """Select the CUDA device used by an in-process Genesis worker.
+
+ Genesis initializes a process-wide session whose Quadrants CUDA runtime
+ always binds the first entry of ``CUDA_VISIBLE_DEVICES``; torch's current
+ device alone is *not* honored (issue #1508). A non-zero request is
+ therefore honored by pinning visibility to the target GPU before any CUDA
+ context exists, after which the in-process device is ``cuda:0``. The
+ returned string is the device the rest of this process must actually use —
+ callers that computed a pre-pin device (learner, probes, collectors) have
+ to adopt the returned value. Binding must happen before constructing the
+ backend (and before ``gs.init``), including in spawn-based collector
+ processes, and remains in effect for the lifetime of the process.
+ """
+
+ global _genesis_device_pinned
+ device = str(resolved).strip()
+ index = _cuda_device_index(device)
+ if index is None:
+ raise ValueError(f"genesis requires a CUDA process device; got {resolved!r}")
+ import torch
+
+ # Pin *before* any torch CUDA query: even ``torch.cuda.is_available()``
+ # latches CUDA_VISIBLE_DEVICES in the runtime, after which rewriting it
+ # would silently keep the process on the first previously visible GPU.
+ if index > 0:
+ if not _genesis_device_pinned:
+ _pin_cuda_visible_devices(index)
+ _genesis_device_pinned = True
+ # Already pinned (or just pinned): the only valid in-process device is
+ # index 0. A stale pre-pin index from the same rank maps onto it.
+ index = 0
+ if not torch.cuda.is_available():
+ raise ValueError(
+ f"genesis requires CUDA device {device!r}, but CUDA is unavailable in this process"
+ )
+ torch.cuda.set_device(index)
+ return f"cuda:{index}"
+
+
+def _reset_genesis_device_pin_for_tests() -> None:
+ """Clear the process pin latch; test-only seam (the CVD rewrite itself is
+ reverted via ``monkeypatch.setitem``/``delitem`` on ``os.environ``)."""
+
+ global _genesis_device_pinned
+ _genesis_device_pinned = False
+
+
__all__ = [
+ "BACKEND_ENV_DEVICE_FIELDS",
+ "apply_backend_env_device_override",
"bind_backend_process_device",
+ "bind_genesis_process_device",
"configure_backend_process_device",
+ "pin_genesis_device_before_cuda_init",
+ "resolve_backend_env_device_id",
"resolve_backend_process_device",
+ "warn_if_backend_device_collision",
]
diff --git a/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml b/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml
index e204a3ad7..d2b94efc5 100644
--- a/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml
+++ b/src/unilab/conf/ppo/task/g1_walk_flat/genesis.yaml
@@ -33,6 +33,12 @@ algo:
actor_hidden_dims: [512, 256, 128]
critic_hidden_dims: [512, 256, 128]
env:
+ # GPU device id for the process-wide Genesis session. Multi-GPU training
+ # overrides this cold-path field from the rank topology; zero preserves the
+ # single-device owner default. A non-zero id pins CUDA_VISIBLE_DEVICES for
+ # the whole process before the first CUDA context exists (Quadrants only
+ # honors the first visible device), after which the in-process index is 0.
+ genesis_device_id: 0
# Re-declares the MJCF that Genesis drops
# at import; the remaining global options keep the Genesis defaults.
genesis_integrator: implicitfast
diff --git a/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml b/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml
index 911663387..8a974fe9c 100644
--- a/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml
+++ b/src/unilab/conf/sac/task/g1_walk_flat/genesis.yaml
@@ -36,6 +36,12 @@ algo:
alpha_init: 0.001
target_entropy_ratio: 0.0
env:
+ # GPU device id for the process-wide Genesis session. Multi-GPU training
+ # overrides this cold-path field from the rank topology; zero preserves the
+ # single-device owner default. A non-zero id pins CUDA_VISIBLE_DEVICES for
+ # the whole process before the first CUDA context exists (Quadrants only
+ # honors the first visible device), after which the in-process index is 0.
+ genesis_device_id: 0
# Re-declares the MJCF that Genesis drops
# at import; the remaining global options keep the Genesis defaults.
genesis_integrator: implicitfast
diff --git a/src/unilab/scripts/play_hora_appo.py b/src/unilab/scripts/play_hora_appo.py
index d4d336650..26bb5f69c 100644
--- a/src/unilab/scripts/play_hora_appo.py
+++ b/src/unilab/scripts/play_hora_appo.py
@@ -17,7 +17,7 @@
from typing import Any, cast
import torch
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
from uni_rl.algos.hora.appo_runner import HoraAPPORunner
from uni_rl.algos.hora.models import build_hora_shared_actor_critic
from uni_rl.algos.hora.observations import (
@@ -34,6 +34,10 @@
from unisim.backend.base import log_playback_plan
from unilab.base.config_adapter import BackendAdapter, create_env
+from unilab.base.process_device import (
+ apply_backend_env_device_override,
+ configure_backend_process_device,
+)
from unilab.utils.sim2sim import policy_load_dim_guard, resolve_sim2sim_config
@@ -107,12 +111,6 @@ def play_hora_appo(
from rsl_rl.utils import resolve_callable
from tensordict import TensorDict
- env_cfg_override = BackendAdapter(
- cfg,
- root_dir=root_dir,
- algo_name="appo",
- ).build_task_env_cfg_override()
-
device = cfg.training.device or (
"cuda"
if torch.cuda.is_available()
@@ -120,6 +118,26 @@ def play_hora_appo(
if torch.backends.mps.is_available()
else "cpu"
)
+ # Genesis owns a process-wide session and must select its CUDA device
+ # before the first backend construction. Keep the rank/device routing in
+ # the shared owner-layer helper so HORA playback follows the same contract
+ # as the generic APPO/PPO/off-policy play paths. A non-zero Genesis
+ # request pins CUDA_VISIBLE_DEVICES; the bound in-process device replaces
+ # the requested one for the policy and the env override.
+ sim_backend = str(OmegaConf.select(cfg, "training.sim_backend", default="mujoco"))
+ if str(device).strip().lower().startswith("cuda"):
+ bound_device = configure_backend_process_device(sim_backend, device)
+ if bound_device is not None:
+ device = bound_device
+ env_cfg_override = apply_backend_env_device_override(
+ BackendAdapter(
+ cfg,
+ root_dir=root_dir,
+ algo_name="appo",
+ ).build_task_env_cfg_override(),
+ sim_backend,
+ learner_device=device,
+ )
print(f"Using device for play: {device}")
env = cast(
diff --git a/src/unilab/scripts/play_interactive.py b/src/unilab/scripts/play_interactive.py
index 1f746db0b..e141c876b 100644
--- a/src/unilab/scripts/play_interactive.py
+++ b/src/unilab/scripts/play_interactive.py
@@ -50,7 +50,10 @@
normalize_ppo_train_cfg,
)
-from unilab.base.process_device import configure_backend_process_device
+from unilab.base.process_device import (
+ apply_backend_env_device_override,
+ configure_backend_process_device,
+)
from unilab.training import (
algo_config_dict,
ensure_registries,
@@ -910,20 +913,44 @@ def play_interactive(args, cfg: DictConfig | None = None, *, algo: str | None =
# mjwarp requires an active CUDA Warp device; bind it process-wide before
# any env is constructed (same pattern as the offpolicy train entrypoint).
- # No-op for backends without a device binding requirement.
- configure_backend_process_device(sim_backend, device)
+ # No-op for backends without a device binding requirement. A non-zero
+ # Genesis request pins CUDA_VISIBLE_DEVICES; the bound in-process device
+ # replaces the requested one for the policy and the env overrides below.
+ if str(device).strip().lower().startswith("cuda"):
+ bound_device = configure_backend_process_device(sim_backend, device)
+ if bound_device is not None:
+ device = bound_device
def _create_env(num_envs: int):
if cfg is None:
- return registry.make(args.task, num_envs=num_envs, sim_backend=sim_backend)
+ # Legacy programmatic callers do not carry a composed Hydra
+ # config, but the selected playback device still has to reach
+ # GPU-backed backend adapters (notably Isaac and Genesis).
+ legacy_env_cfg_override = apply_backend_env_device_override(
+ None,
+ sim_backend,
+ learner_device=device,
+ )
+ return registry.make(
+ args.task,
+ num_envs=num_envs,
+ sim_backend=sim_backend,
+ env_cfg_override=legacy_env_cfg_override or None,
+ )
from unilab.base.config_adapter import create_env
+ env_cfg_override: dict[str, Any] | None
if algo in _OFFPOLICY_INTERACTIVE_ALGOS:
env_cfg_override = build_offpolicy_env_cfg_override(algo, cfg, root_dir=Path.cwd())
else:
env_cfg_override = build_play_backend_adapter(
cfg, root_dir=Path.cwd(), algo_name=algo
).build_task_env_cfg_override()
+ env_cfg_override = apply_backend_env_device_override(
+ env_cfg_override,
+ sim_backend,
+ learner_device=device,
+ )
try:
return create_env(
cfg,
diff --git a/src/unilab/scripts/train_appo.py b/src/unilab/scripts/train_appo.py
index 28df757e8..56182ff5b 100644
--- a/src/unilab/scripts/train_appo.py
+++ b/src/unilab/scripts/train_appo.py
@@ -20,6 +20,11 @@
create_env,
)
from unilab.base.env_factory import registry_env_factory
+from unilab.base.process_device import (
+ apply_backend_env_device_override,
+ configure_backend_process_device,
+ pin_genesis_device_before_cuda_init,
+)
from unilab.training import (
algo_config_dict,
build_run_dir_name,
@@ -45,6 +50,12 @@ def _training_resume_requested(load_run: Any) -> bool:
return str(load_run) not in {"", "-1"}
+def _is_cuda_device(device: str | None) -> bool:
+ """Return whether a device alias names CUDA (indexed or unindexed)."""
+
+ return device is not None and str(device).strip().lower().split(":", 1)[0] == "cuda"
+
+
def build_appo_runner_kwargs(
cfg: DictConfig,
env_cfg_override: dict | None,
@@ -54,12 +65,22 @@ def build_appo_runner_kwargs(
if rl_cfg is None:
rl_cfg = algo_config_dict(cfg)
+ routed_env_cfg_override = apply_backend_env_device_override(
+ env_cfg_override,
+ str(cfg.training.sim_backend),
+ learner_device=(
+ collector_device
+ if collector_device is not None and _is_cuda_device(collector_device)
+ else OmegaConf.select(cfg, "training.device", default=None)
+ ),
+ )
+
runner_kwargs = {
"env_name": cfg.training.task_name,
"env_factory": registry_env_factory(
str(cfg.training.task_name), str(cfg.training.sim_backend)
),
- "env_cfg_overrides": env_cfg_override,
+ "env_cfg_overrides": routed_env_cfg_override,
"rl_cfg": rl_cfg,
"device": cfg.training.device,
"collector_device": collector_device,
@@ -195,6 +216,17 @@ def play_appo(
log_root=None,
num_envs=cfg.training.play_env_num,
)
+ if _is_cuda_device(device):
+ # Genesis pins CUDA_VISIBLE_DEVICES for a non-zero request; adopt the
+ # bound in-process device for the policy and the env override.
+ bound_device = configure_backend_process_device(str(cfg.training.sim_backend), device)
+ if bound_device is not None:
+ device = bound_device
+ play_env_cfg_override = apply_backend_env_device_override(
+ BackendAdapter(cfg, root_dir=Path.cwd(), algo_name="appo").build_play_env_cfg_override(),
+ str(cfg.training.sim_backend),
+ learner_device=device,
+ )
session, _policy_obs_mode, _checkpoint_path = create_appo_playback_session(
playback_cfg=playback_cfg,
cfg=cfg,
@@ -202,9 +234,7 @@ def play_appo(
env_factory=lambda n: create_env(
cfg,
num_envs=n,
- env_cfg_override=BackendAdapter(
- cfg, root_dir=Path.cwd(), algo_name="appo"
- ).build_play_env_cfg_override(),
+ env_cfg_override=play_env_cfg_override,
),
root_dir=Path.cwd(),
device=device,
@@ -266,12 +296,17 @@ def forward(self, obs: torch.Tensor) -> torch.Tensor:
@hydra.main(version_base="1.3", config_path="../conf/appo", config_name="config")
def main(cfg: DictConfig) -> None:
- ensure_registries()
+ # Genesis/Quadrants binds the first CUDA_VISIBLE_DEVICES entry, and even
+ # torch.cuda.is_available() latches the variable in the CUDA runtime, so
+ # the pin must precede registry bootstrap and device auto-detection
+ # (issue #1508). APPO has no device topology; only an explicit learner
+ # device can require the pin.
+ pinned_device = pin_genesis_device_before_cuda_init(
+ str(cfg.training.sim_backend),
+ learner_device=cfg.training.device,
+ )
- seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True)
- env_cfg_override = BackendAdapter(
- cfg, root_dir=Path.cwd(), algo_name="appo"
- ).build_task_env_cfg_override()
+ ensure_registries()
# Convert algo config to plain dict for APPORunner / RSL-RL internals
rl_cfg = algo_config_dict(cfg)
@@ -300,6 +335,35 @@ def main(cfg: DictConfig) -> None:
if torch.backends.mps.is_available()
else "cpu"
)
+ if pinned_device is not None:
+ # The process was pinned to its requested GPU; use the in-process
+ # index for both learner and collector.
+ learner_device = pinned_device
+ if collector_device is not None and _is_cuda_device(collector_device):
+ collector_device = pinned_device
+
+ if _is_cuda_device(learner_device):
+ # A non-zero Genesis request pins CUDA_VISIBLE_DEVICES process-wide
+ # (Quadrants only honors the first visible device); the bound
+ # in-process device replaces both the learner and collector device.
+ bound_device = configure_backend_process_device(
+ str(cfg.training.sim_backend), learner_device
+ )
+ if bound_device is not None:
+ learner_device = bound_device
+ if collector_device is not None and _is_cuda_device(collector_device):
+ collector_device = bound_device
+
+ env_cfg_override = apply_backend_env_device_override(
+ BackendAdapter(cfg, root_dir=Path.cwd(), algo_name="appo").build_task_env_cfg_override(),
+ str(cfg.training.sim_backend),
+ learner_device=(
+ collector_device
+ if collector_device is not None and _is_cuda_device(collector_device)
+ else learner_device
+ ),
+ )
+ seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True)
tracker = None
if not cfg.training.play_only:
diff --git a/src/unilab/scripts/train_offpolicy.py b/src/unilab/scripts/train_offpolicy.py
index 4ffe3658a..524d5cc54 100644
--- a/src/unilab/scripts/train_offpolicy.py
+++ b/src/unilab/scripts/train_offpolicy.py
@@ -30,7 +30,14 @@
from unilab.base.config_adapter import create_env
from unilab.base.env_factory import registry_env_factory
-from unilab.base.process_device import bind_backend_process_device, configure_backend_process_device
+from unilab.base.process_device import (
+ apply_backend_env_device_override,
+ bind_backend_process_device,
+ configure_backend_process_device,
+ pin_genesis_device_before_cuda_init,
+ resolve_backend_env_device_id,
+ warn_if_backend_device_collision,
+)
from unilab.training import (
assert_offpolicy_task_choice_matches_algo,
build_run_dir_name,
@@ -85,16 +92,41 @@ def build_failure_summary(exc: BaseException, run_summary: Any | None = None) ->
def build_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[str, Any] | None:
- return _build_offpolicy_env_cfg_override(algo_name, cfg, root_dir=Path.cwd())
+ base = _build_offpolicy_env_cfg_override(algo_name, cfg, root_dir=Path.cwd())
+ devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None))
+ rank = current_dp_rank()
+ from unilab.utils.device import get_default_device
+
+ rank_device = resolve_dp_rank_device(devices, rank) or get_default_device()
+ return apply_backend_env_device_override(
+ base,
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=rank,
+ world_size=1,
+ learner_device=rank_device,
+ )
def build_offpolicy_play_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[str, Any] | None:
- return _build_offpolicy_play_env_cfg_override(algo_name, cfg, root_dir=Path.cwd())
+ base = _build_offpolicy_play_env_cfg_override(algo_name, cfg, root_dir=Path.cwd())
+ devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None))
+ rank = current_dp_rank()
+ from unilab.utils.device import get_default_device
+
+ rank_device = resolve_dp_rank_device(devices, rank) or get_default_device()
+ return apply_backend_env_device_override(
+ base,
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=rank,
+ world_size=1,
+ learner_device=rank_device,
+ )
def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None):
"""Build algorithm runner from unified Hydra config."""
- env_cfg_override = build_offpolicy_env_cfg_override(algo_name, cfg)
env_factory = registry_env_factory(str(cfg.training.task_name), str(cfg.training.sim_backend))
from uni_rl.offpolicy.thread_budget import (
apply_torch_thread_runtime,
@@ -113,10 +145,43 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None):
from unilab.utils.device import get_default_device
rank_device = resolve_dp_rank_device(dp_devices, dp_rank) or get_default_device()
+ routed_device_id = resolve_backend_env_device_id(
+ str(cfg.training.sim_backend),
+ devices=dp_devices,
+ rank=dp_rank,
+ world_size=1,
+ learner_device=rank_device,
+ )
+ warn_if_backend_device_collision(
+ str(cfg.training.sim_backend),
+ devices=dp_devices,
+ rank=dp_rank,
+ device_id=routed_device_id,
+ source="collector",
+ )
# Bind backend-global device state before algorithm builders materialize
# their probe envs. The spawned collector repeats this binding in its own
- # process using the same rank-local device.
- configure_backend_process_device(str(cfg.training.sim_backend), rank_device)
+ # process using the same rank-local device. A non-zero Genesis request
+ # pins CUDA_VISIBLE_DEVICES for the whole rank process (Quadrants only
+ # honors the first visible device), so the bound in-process device
+ # replaces rank_device for the learner, the probe, and the override below.
+ if str(rank_device).strip().lower().startswith("cuda"):
+ bound_device = configure_backend_process_device(str(cfg.training.sim_backend), rank_device)
+ if bound_device is not None:
+ rank_device = bound_device
+ # ``training.devices`` is host-visible for the off-policy supervisor. The
+ # collector subprocess inherits that namespace, so pass the physical index
+ # through the owner EnvCfg rather than leaving Isaac/Genesis at YAML's
+ # historical device-0 default. The same override reaches the learner-side
+ # dimension probe and the collector env.
+ env_cfg_override = apply_backend_env_device_override(
+ build_offpolicy_env_cfg_override(algo_name, cfg),
+ str(cfg.training.sim_backend),
+ devices=dp_devices,
+ rank=dp_rank,
+ world_size=1,
+ learner_device=rank_device,
+ )
host_cpu_count = os.cpu_count() or 1
explicit_cpu_ids = getattr(cfg.training, "dp_collector_cpu_ids", None)
if explicit_cpu_ids is not None:
@@ -214,7 +279,36 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None:
return None
devices = resolve_dp_topology(cfg.training.devices)
- device = default_device(torch, resolve_dp_rank_device(devices, current_dp_rank()))
+ dp_rank = current_dp_rank()
+ device = default_device(torch, resolve_dp_rank_device(devices, dp_rank))
+ play_device_id = resolve_backend_env_device_id(
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=dp_rank,
+ world_size=1,
+ learner_device=device,
+ )
+ warn_if_backend_device_collision(
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=dp_rank,
+ device_id=play_device_id,
+ source="playback",
+ )
+ if str(device).strip().lower().startswith("cuda"):
+ # Genesis pins CUDA_VISIBLE_DEVICES for a non-zero request; adopt the
+ # bound in-process device for the policy and the env override.
+ bound_device = configure_backend_process_device(str(cfg.training.sim_backend), device)
+ if bound_device is not None:
+ device = bound_device
+ play_env_cfg_override = apply_backend_env_device_override(
+ build_offpolicy_play_env_cfg_override(algo_name, cfg),
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=dp_rank,
+ world_size=1,
+ learner_device=device,
+ )
print(f"Using device for play: {device}")
playback_cfg = RslRlPlaybackConfig(
@@ -233,7 +327,7 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None:
env_factory=lambda n: create_env(
cfg,
num_envs=n,
- env_cfg_override=build_offpolicy_play_env_cfg_override(algo_name, cfg),
+ env_cfg_override=play_env_cfg_override,
),
root_dir=Path.cwd(),
device=device,
@@ -316,11 +410,37 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None:
def main(cfg: DictConfig) -> None:
enable_faulthandler()
- ensure_registries()
devices = resolve_dp_topology(cfg.training.devices)
rank = current_dp_rank()
+ # Genesis/Quadrants binds the first CUDA_VISIBLE_DEVICES entry, and even
+ # torch.cuda.is_available() latches the variable in the CUDA runtime, so
+ # the pin must precede registry bootstrap and device auto-detection
+ # (issue #1508). Pure config topology resolves without touching torch.
+ pinned_device = pin_genesis_device_before_cuda_init(
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=rank,
+ world_size=1,
+ learner_device=OmegaConf.select(cfg, "training.device", default=None),
+ )
+
+ ensure_registries()
+
rank_device = apply_dp_rank_config(cfg, devices, rank)
+ if pinned_device is not None:
+ # The process was pinned to its rank GPU; use the in-process index.
+ rank_device = pinned_device
+
+ # Bind before seed initialization and before any rank-local env/probe is
+ # materialized. ``build_runner`` repeats the binding defensively because
+ # it is also a public assembly seam used by tests and custom callers. A
+ # non-zero Genesis request pins CUDA_VISIBLE_DEVICES here; the bound
+ # in-process device replaces rank_device for the tracker and runner.
+ if rank_device is not None and str(rank_device).strip().lower().startswith("cuda"):
+ bound_device = configure_backend_process_device(str(cfg.training.sim_backend), rank_device)
+ if bound_device is not None:
+ rank_device = bound_device
seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True)
algo_name = cfg.algo.algo
diff --git a/src/unilab/scripts/train_rsl_rl.py b/src/unilab/scripts/train_rsl_rl.py
index 3d5cb785b..77e1aba39 100644
--- a/src/unilab/scripts/train_rsl_rl.py
+++ b/src/unilab/scripts/train_rsl_rl.py
@@ -33,6 +33,13 @@
from unisim.backend.mujoco.xml import materialize_scene_visual_override
from unilab.base.config_adapter import BackendAdapter, create_env
+from unilab.base.process_device import (
+ apply_backend_env_device_override,
+ configure_backend_process_device,
+ pin_genesis_device_before_cuda_init,
+ resolve_backend_env_device_id,
+ warn_if_backend_device_collision,
+)
from unilab.base.run_control import RunComplete
from unilab.training import (
algo_config_dict,
@@ -78,11 +85,47 @@ def _backend_adapter(cfg: DictConfig) -> BackendAdapter:
def build_ppo_env_cfg_override(cfg: DictConfig) -> dict[str, Any]:
- return cast(dict[str, Any], _backend_adapter(cfg).build_task_env_cfg_override())
+ base = cast(dict[str, Any], _backend_adapter(cfg).build_task_env_cfg_override())
+ devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None))
+ local_rank = current_torch_distributed_local_rank()
+ world_size = current_torch_distributed_world_size()
+ configured_device = OmegaConf.select(cfg, "training.device", default=None)
+ learner_device = (
+ f"cuda:{local_rank}"
+ if world_size > 1
+ else (f"cuda:{devices[0]}" if devices else configured_device)
+ )
+ return apply_backend_env_device_override(
+ base,
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=local_rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=learner_device,
+ )
def build_ppo_play_env_cfg_override(cfg: DictConfig) -> dict[str, Any]:
- return cast(dict[str, Any], _backend_adapter(cfg).build_play_env_cfg_override())
+ base = cast(dict[str, Any], _backend_adapter(cfg).build_play_env_cfg_override())
+ devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None))
+ local_rank = current_torch_distributed_local_rank()
+ world_size = current_torch_distributed_world_size()
+ configured_device = OmegaConf.select(cfg, "training.device", default=None)
+ learner_device = (
+ f"cuda:{local_rank}"
+ if world_size > 1
+ else (f"cuda:{devices[0]}" if devices else configured_device)
+ )
+ return apply_backend_env_device_override(
+ base,
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=local_rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=learner_device,
+ )
def run_motrix_rsl_play_loop(
@@ -235,12 +278,45 @@ def _normalize_play_train_cfg(train_cfg: dict[str, Any]) -> dict[str, Any]:
log_root=None,
num_envs=cfg.training.play_env_num,
)
+ play_devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None))
+ play_world_size = current_torch_distributed_world_size()
+ play_local_rank = current_torch_distributed_local_rank()
+ play_device_id = resolve_backend_env_device_id(
+ str(cfg.training.sim_backend),
+ devices=play_devices,
+ rank=play_local_rank,
+ local_rank=play_local_rank,
+ world_size=play_world_size,
+ learner_device=device,
+ )
+ warn_if_backend_device_collision(
+ str(cfg.training.sim_backend),
+ devices=play_devices,
+ rank=play_local_rank,
+ device_id=play_device_id,
+ source="playback",
+ )
+ if str(device).strip().lower().startswith("cuda"):
+ # A non-zero Genesis request pins CUDA_VISIBLE_DEVICES here; adopt the
+ # bound in-process device for both the policy and the env override.
+ bound_device = configure_backend_process_device(str(cfg.training.sim_backend), device)
+ if bound_device is not None:
+ device = bound_device
+ play_env_cfg_override = apply_backend_env_device_override(
+ build_ppo_play_env_cfg_override(cfg),
+ str(cfg.training.sim_backend),
+ devices=play_devices,
+ rank=play_local_rank,
+ local_rank=play_local_rank,
+ world_size=play_world_size,
+ learner_device=device,
+ )
session, _policy_obs_mode, _checkpoint_path = create_rsl_rl_playback_session(
playback_cfg=playback_cfg,
env_factory=lambda n: create_env(
cfg,
num_envs=n,
- env_cfg_override=build_ppo_play_env_cfg_override(cfg),
+ env_cfg_override=play_env_cfg_override,
),
algo_config=rl_cfg,
root_dir=Path.cwd(),
@@ -352,11 +428,29 @@ def main(cfg: DictConfig) -> None:
"training.devices or set training.log_dir explicitly"
)
+ # Genesis/Quadrants binds the first CUDA_VISIBLE_DEVICES entry, and even
+ # torch.cuda.is_available() latches the variable in the CUDA runtime, so
+ # the pin must precede registry bootstrap and device auto-detection
+ # (issue #1508). Pure config topology resolves without touching torch.
+ pinned_device = pin_genesis_device_before_cuda_init(
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=local_rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=str(configured_device) if configured_device is not None else None,
+ )
+
+ if pinned_device is not None and world_size > 1:
+ # rsl_rl's multi-GPU guard expects device == cuda:{LOCAL_RANK}. Inside
+ # the pinned single-device namespace the rank-local index is 0, so
+ # adopt it for downstream consumers (parameter sync keeps using the
+ # unchanged global RANK/WORLD_SIZE).
+ os.environ["LOCAL_RANK"] = "0"
+ local_rank = 0
+
ensure_registries()
apply_rsl_rl_rank_seed(cfg, rank)
- seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True)
- env_cfg_override = build_ppo_env_cfg_override(cfg)
-
device = resolve_rsl_rl_device(
configured_device=str(configured_device) if configured_device is not None else None,
devices=devices,
@@ -364,7 +458,47 @@ def main(cfg: DictConfig) -> None:
local_rank=local_rank,
default_device=get_default_device(),
)
+ if pinned_device is not None:
+ # The process was pinned to its rank GPU; use the in-process index.
+ device = pinned_device
+ # PPO workers launched by torchrun inherit the launcher's remapped
+ # CUDA_VISIBLE_DEVICES, so LOCAL_RANK (not the host index in
+ # training.devices) is the simulator payload id. Single-process workers
+ # retain the configured host-visible index. Route this before env
+ # construction and before Genesis/torch global initialization.
+ env_device_id = resolve_backend_env_device_id(
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=local_rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=device,
+ )
+ warn_if_backend_device_collision(
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=local_rank,
+ device_id=env_device_id,
+ source="training",
+ )
+ if str(device).strip().lower().startswith("cuda"):
+ # Genesis pins CUDA_VISIBLE_DEVICES for a non-zero request (Quadrants
+ # only honors the first visible device); the bound value is the device
+ # this process must actually use afterwards.
+ bound_device = configure_backend_process_device(str(cfg.training.sim_backend), device)
+ if bound_device is not None:
+ device = bound_device
print(f"[rank {rank}/{world_size}] Using device: {device}")
+ env_cfg_override = apply_backend_env_device_override(
+ build_ppo_env_cfg_override(cfg),
+ str(cfg.training.sim_backend),
+ devices=devices,
+ rank=local_rank,
+ local_rank=local_rank,
+ world_size=world_size,
+ learner_device=device,
+ )
+ seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True)
# Compute effective max_iterations (supports num_timesteps override)
max_iterations = cfg.algo.max_iterations
diff --git a/tests/base/backend/test_process_device.py b/tests/base/backend/test_process_device.py
index e4f68c2de..d4ee3a7ec 100644
--- a/tests/base/backend/test_process_device.py
+++ b/tests/base/backend/test_process_device.py
@@ -2,14 +2,22 @@
from __future__ import annotations
+import os
+import warnings
from types import SimpleNamespace
import pytest
+import torch
from unisim.backend.mjwarp import runtime as mjwarp_runtime
+import unilab.base.process_device as process_device
from unilab.base.process_device import (
+ apply_backend_env_device_override,
+ bind_genesis_process_device,
configure_backend_process_device,
+ resolve_backend_env_device_id,
resolve_backend_process_device,
+ warn_if_backend_device_collision,
)
@@ -71,3 +79,178 @@ def test_mjwarp_binding_rejects_non_cuda_warp_resolution(
with pytest.raises(RuntimeError, match="active CUDA Warp device"):
mjwarp_runtime.bind_mjwarp_process_device("cuda:1")
+
+
+@pytest.mark.parametrize(
+ "backend_type",
+ ["isaacgym", "isaacsim", "genesis"],
+)
+def test_offpolicy_rank_routes_host_visible_backend_device(backend_type: str) -> None:
+ assert (
+ resolve_backend_env_device_id(
+ backend_type,
+ devices=(0, 1),
+ rank=1,
+ world_size=1,
+ learner_device="cuda:1",
+ )
+ == 1
+ )
+
+
+@pytest.mark.parametrize(
+ "backend_type",
+ ["isaacgym", "isaacsim", "genesis"],
+)
+def test_torchrun_rank_routes_local_backend_device(backend_type: str) -> None:
+ # torchrun remaps CUDA_VISIBLE_DEVICES to [4, 5], so rank 1 must send
+ # local index 1 to the worker rather than host-visible index 5.
+ assert (
+ resolve_backend_env_device_id(
+ backend_type,
+ devices=(4, 5),
+ rank=1,
+ local_rank=1,
+ world_size=2,
+ )
+ == 1
+ )
+
+
+def test_backend_env_device_override_does_not_mutate_owner_mapping() -> None:
+ owner_override = {"isaacgym_device_id": 0, "nested": {"keep": True}}
+ routed = apply_backend_env_device_override(
+ owner_override,
+ "isaacgym",
+ devices=(0, 1),
+ rank=1,
+ world_size=1,
+ )
+
+ assert routed["isaacgym_device_id"] == 1
+ assert owner_override["isaacgym_device_id"] == 0
+ assert routed["nested"] is owner_override["nested"]
+
+
+def test_nonzero_rank_device_zero_emits_collision_warning() -> None:
+ with pytest.warns(RuntimeWarning, match=r"training\.devices=\[0, 1\]"):
+ warn_if_backend_device_collision(
+ "genesis",
+ devices=(0, 1),
+ rank=1,
+ device_id=0,
+ )
+
+
+def test_non_gpu_backend_is_left_untouched() -> None:
+ owner_override = {"isaacgym_device_id": 0}
+ assert (
+ apply_backend_env_device_override(
+ owner_override,
+ "mujoco",
+ devices=(0, 1),
+ rank=1,
+ world_size=1,
+ )
+ == owner_override
+ )
+
+
+@pytest.fixture
+def genesis_pin_state(monkeypatch: pytest.MonkeyPatch):
+ """Host-free Genesis pin lane: stub torch CUDA state and the pin latch."""
+
+ # monkeypatch.delitem on a missing key records no undo, so the pin's
+ # direct os.environ write would leak into later tests; restore manually.
+ saved_cvd = os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ set_calls: list[int] = []
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+ monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False)
+ monkeypatch.setattr(torch.cuda, "device_count", lambda: 2)
+ monkeypatch.setattr(torch.cuda, "set_device", set_calls.append)
+ process_device._reset_genesis_device_pin_for_tests()
+ try:
+ yield set_calls
+ finally:
+ process_device._reset_genesis_device_pin_for_tests()
+ if saved_cvd is None:
+ os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+ else:
+ os.environ["CUDA_VISIBLE_DEVICES"] = saved_cvd
+
+
+def test_genesis_nonzero_device_pins_visible_devices(genesis_pin_state: list[int]) -> None:
+ # Quadrants only honors the first visible device (issue #1508), so a
+ # non-zero request shrinks CUDA_VISIBLE_DEVICES and reports the
+ # in-process index 0 to the rest of the process.
+ assert bind_genesis_process_device("cuda:1") == "cuda:0"
+ assert os.environ["CUDA_VISIBLE_DEVICES"] == "1"
+ assert genesis_pin_state == [0]
+
+
+def test_genesis_pin_translates_existing_visible_devices(
+ genesis_pin_state: list[int],
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setitem(os.environ, "CUDA_VISIBLE_DEVICES", "4,5")
+
+ assert bind_genesis_process_device("cuda:1") == "cuda:0"
+ assert os.environ["CUDA_VISIBLE_DEVICES"] == "5"
+
+
+def test_genesis_device_zero_binds_without_pin(genesis_pin_state: list[int]) -> None:
+ assert bind_genesis_process_device("cuda:0") == "cuda:0"
+ assert "CUDA_VISIBLE_DEVICES" not in os.environ
+ assert genesis_pin_state == [0]
+
+
+def test_genesis_pin_fails_closed_after_cuda_init(
+ genesis_pin_state: list[int],
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True)
+
+ with pytest.raises(RuntimeError, match="before any CUDA context"):
+ bind_genesis_process_device("cuda:1")
+ assert "CUDA_VISIBLE_DEVICES" not in os.environ
+
+
+def test_genesis_pin_is_idempotent_for_repeated_rank_binding(
+ genesis_pin_state: list[int],
+) -> None:
+ # Entrypoints bind twice (main + defensive runner rebind); a stale pre-pin
+ # index from the same rank maps onto the pinned in-process device.
+ assert bind_genesis_process_device("cuda:1") == "cuda:0"
+ assert bind_genesis_process_device("cuda:1") == "cuda:0"
+ assert os.environ["CUDA_VISIBLE_DEVICES"] == "1"
+ assert genesis_pin_state == [0, 0]
+
+
+def test_genesis_resolution_uses_pinned_namespace(genesis_pin_state: list[int]) -> None:
+ bind_genesis_process_device("cuda:1")
+
+ # After the pin every in-process consumer resolves to index 0, and the
+ # transition collision guard stays quiet because each rank owns its GPU.
+ assert (
+ resolve_backend_env_device_id(
+ "genesis",
+ devices=(0, 1),
+ rank=1,
+ world_size=1,
+ learner_device="cuda:1",
+ )
+ == 0
+ )
+ with warnings.catch_warnings():
+ warnings.simplefilter("error")
+ warn_if_backend_device_collision("genesis", devices=(0, 1), rank=1, device_id=0)
+
+
+def test_genesis_pin_rejects_index_beyond_visible_devices(
+ genesis_pin_state: list[int],
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ monkeypatch.setitem(os.environ, "CUDA_VISIBLE_DEVICES", "3")
+
+ with pytest.raises(ValueError, match="CUDA_VISIBLE_DEVICES"):
+ bind_genesis_process_device("cuda:1")
diff --git a/tests/base/test_genesis_backend.py b/tests/base/test_genesis_backend.py
index b7fe99fa7..b4dc9e697 100644
--- a/tests/base/test_genesis_backend.py
+++ b/tests/base/test_genesis_backend.py
@@ -745,6 +745,7 @@ def test_constructor_validation_and_factory_wiring(fake_genesis, tiny_model_file
def test_env_cfg_genesis_fields_validate_and_reach_factory() -> None:
cfg = EnvCfg(
+ genesis_device_id=1,
genesis_integrator="implicitfast",
genesis_constraint_solver="cg",
genesis_friction_cone="pyramidal",
@@ -753,11 +754,14 @@ def test_env_cfg_genesis_fields_validate_and_reach_factory() -> None:
cfg.validate()
kwargs = env_backend_kwargs(cfg)
assert (
+ kwargs["genesis_device_id"],
kwargs["genesis_integrator"],
kwargs["genesis_constraint_solver"],
kwargs["genesis_friction_cone"],
kwargs["genesis_solver_iterations"],
- ) == ("implicitfast", "cg", "pyramidal", 30)
+ ) == (1, "implicitfast", "cg", "pyramidal", 30)
+ with pytest.raises(ValueError, match="genesis_device_id must be a non-negative integer"):
+ EnvCfg(genesis_device_id=-1).validate()
with pytest.raises(ValueError, match="genesis_integrator must be a non-empty string"):
EnvCfg(genesis_integrator="").validate()
with pytest.raises(ValueError, match="genesis_solver_iterations must be a positive integer"):
diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py
index ae102d288..1c2d0a7e7 100644
--- a/tests/envs/test_env_configs.py
+++ b/tests/envs/test_env_configs.py
@@ -213,6 +213,7 @@ def test_g1_walk_flat_genesis_owner_composes_and_materializes(config_group):
assert env_cfg.scene.fragment_files == []
assert env_cfg.scene.terrain is None
assert env_cfg.scene.default_keyframe_name == "stand"
+ assert env_cfg.genesis_device_id == 0
# The owner re-declares the MJCF that
# Genesis drops at import; the other global options stay at Genesis
# defaults (None).
diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py
index dde7e3fd3..e03cc37d6 100644
--- a/tests/scripts/test_train_scripts.py
+++ b/tests/scripts/test_train_scripts.py
@@ -614,6 +614,52 @@ def test_offpolicy_g1_walk_flat_env_cfg_override_has_rewards_and_events():
assert env_cfg_override["events"]["pd_gains"] is None
+@pytest.mark.parametrize(
+ ("backend", "field"),
+ [
+ ("isaacgym", "isaacgym_device_id"),
+ ("isaacsim", "isaacsim_device_id"),
+ ("genesis", "genesis_device_id"),
+ ],
+)
+def test_offpolicy_gpu_backend_env_follows_dp_rank(
+ monkeypatch: pytest.MonkeyPatch, backend: str, field: str
+) -> None:
+ """Off-policy collectors receive the host-visible rank device."""
+
+ mod = _offpolicy()
+ cfg = _offpolicy_cfg([f"task=g1_walk_flat/{backend}", "training.devices=[0,1]"])
+ monkeypatch.setenv("UNILAB_DP_RANK", "1")
+
+ override = mod.build_offpolicy_env_cfg_override("sac", cfg)
+
+ assert override is not None
+ assert override[field] == 1
+
+
+@pytest.mark.parametrize(
+ ("backend", "field"),
+ [
+ ("isaacgym", "isaacgym_device_id"),
+ ("isaacsim", "isaacsim_device_id"),
+ ("genesis", "genesis_device_id"),
+ ],
+)
+def test_ppo_gpu_backend_env_uses_torchrun_local_rank(
+ monkeypatch: pytest.MonkeyPatch, backend: str, field: str
+) -> None:
+ """PPO workers pass a local index after torchrun remaps CUDA visibility."""
+
+ mod = _train_rsl_rl(monkeypatch)
+ cfg = _ppo_cfg([f"task=g1_walk_flat/{backend}", "training.devices=[4,5]"])
+ monkeypatch.setenv("LOCAL_RANK", "1")
+ monkeypatch.setenv("WORLD_SIZE", "2")
+
+ override = mod.build_ppo_env_cfg_override(cfg)
+
+ assert override[field] == 1
+
+
def test_offpolicy_isaacsim_training_and_eval_use_separate_render_overrides():
cfg = _offpolicy_cfg(
[