diff --git a/docs/sphinx/source/adr/ADR-0012-sole-tensor-manager-and-scoped-backends.md b/docs/sphinx/source/adr/ADR-0012-sole-tensor-manager-and-scoped-backends.md index 9cfa294db..40bd20dc0 100644 --- a/docs/sphinx/source/adr/ADR-0012-sole-tensor-manager-and-scoped-backends.md +++ b/docs/sphinx/source/adr/ADR-0012-sole-tensor-manager-and-scoped-backends.md @@ -61,14 +61,16 @@ owns its explicit packed conversion boundary. MuJoCo remains a production backen canonical in-process host-bridge implementation; it is not a test-only fallback. `env.tensor_runtime` and `env.tensor_runtime_device` are removed. Tensor execution is an invariant -of the Manager runtime. Placement is derived from declared backend data plane and rank topology. -No compatibility alias, dual-mode field, or hidden fallback is retained. - -For a `HOST_BRIDGE` backend, CPU-authoritative physics does not imply CPU Manager tensors. -When PyTorch exposes an available GPU and the backend accepts its current `cuda` ordinal, -the Manager uses that GPU, including ROCm PyTorch's `cuda` namespace. CPU-only capabilities -or hosts without an available GPU retain CPU placement. `DEVICE_RESIDENT` adapters keep -their own CUDA-only platform requirements; accepting HIP Torch buffers on a host bridge +of the Manager runtime. Placement is derived from the declared backend data plane and, for a +`HOST_BRIDGE`, the process's explicit carrier request. No NumPy/dual-runtime compatibility alias +or hidden fallback is retained. + +For a `HOST_BRIDGE` backend, CPU-authoritative physics keeps the default Manager/TorchEnv +carriers on CPU. A training process may explicitly request `cuda`/`cuda:` through +`manager_torch_device` (routed from its resolved learner device), including ROCm PyTorch's +`cuda` namespace; UniLab validates the request against the backend's declared Torch devices. +GPU visibility alone is not a placement request. `DEVICE_RESIDENT` adapters keep their own +CUDA-only placement and platform requirements; accepting HIP Torch buffers on a host bridge does not make those physics engines ROCm-compatible. ### Manager-owned RNG diff --git a/docs/sphinx/source/en/1-getting_started/2-installation.md b/docs/sphinx/source/en/1-getting_started/2-installation.md index 0c439f60c..70d663bda 100644 --- a/docs/sphinx/source/en/1-getting_started/2-installation.md +++ b/docs/sphinx/source/en/1-getting_started/2-installation.md @@ -243,10 +243,13 @@ ROCm notes: profile, use `uv run --no-sync` (or `UV_NO_SYNC=1 make ...`) for validation; automatic synchronization would reinstall that profile's CUDA wheels. - With ROCm PyTorch and an available GPU, a `HOST_BRIDGE` backend accepting - `cuda` buffers runs the Manager/TorchEnv tensors on the current GPU. MuJoCo - physics remains on CPU; packed transfers connect it to GPU observations, - actions, rewards, and resets. CUDA-only physics backends remain unsupported - on ROCm. See {doc}`/adr/ADR-0012-sole-tensor-manager-and-scoped-backends`. + `cuda` buffers can use the current GPU for Manager/TorchEnv tensors when the + training process explicitly requests that learner device (which routes + `manager_torch_device`). The direct environment default remains CPU. MuJoCo + physics stays on CPU; packed transfers connect it to explicitly requested GPU + observations, actions, rewards, and resets. CUDA-only physics backends remain + unsupported on ROCm. See + {doc}`/adr/ADR-0012-sole-tensor-manager-and-scoped-backends`. - When installing from PyPI instead of a source checkout, `make sync-rocm` does not apply. Install the torch build validated by the repository from the PyTorch ROCm index first, then `unilab`. The published dependency range is diff --git a/docs/sphinx/source/en/5-reference/5-support_matrix.md b/docs/sphinx/source/en/5-reference/5-support_matrix.md index 69180eaa1..d388ce167 100644 --- a/docs/sphinx/source/en/5-reference/5-support_matrix.md +++ b/docs/sphinx/source/en/5-reference/5-support_matrix.md @@ -74,12 +74,12 @@ This table is derived from UniSim's SDK-free public static inventory. It describ | Backend | Execution / process / data plane | Torch devices | CUDA runtime | Linux+CUDA | macOS | ROCm | Worker | Reset randomization | Fixed variants | Host callbacks | Packed bridge | |---|---|---|---|---|---|---|---|---|---|---|---| -| `mujoco` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | unknown | unknown | Unsupported | Exact | -| `motrix` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | -| `drake` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | +| `mujoco` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | unknown | unknown | Unsupported | Exact | +| `motrix` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | +| `drake` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | | `mjwarp` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Unsupported | | `newton` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Unsupported | -| `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | +| `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Exact | | `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | Unsupported | Unsupported | Unsupported | Unsupported | ### Entrypoint x Task Owner @@ -125,7 +125,8 @@ configuration. If `CUDA_VISIBLE_DEVICES` is set, backend ordinals address that remapped namespace, not host-global physical indices. On macOS and ROCm, use a CPU-authoritative host-bridge backend. On ROCm, -Manager/TorchEnv uses the current GPU when ROCm PyTorch is available and the -backend accepts that `cuda` device; CPU physics does not require CPU Manager -tensors. GPU Torch buffers on a host bridge do not imply GPU physics or a -device-resident backend lifecycle. +Manager/TorchEnv uses the current GPU only when the training process explicitly +requests that learner device and the backend accepts the routed `cuda` +`manager_torch_device`; direct environment construction remains CPU. GPU Torch +buffers on a host bridge do not imply GPU physics or a device-resident backend +lifecycle. diff --git a/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md b/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md index b54ff4645..e560d3ca5 100644 --- a/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md +++ b/docs/sphinx/source/zh_CN/1-getting_started/2-installation.md @@ -218,10 +218,12 @@ ROCm 说明： - 如果手工安装了 ROCm wheel，但项目仍使用默认依赖配置档，验证时使用 `uv run --no-sync`（或 `UV_NO_SYNC=1 make ...`）；自动同步会重新安装默认配置档的 CUDA wheel。 -- 安装 ROCm PyTorch 且 GPU 可用时，接受 `cuda` buffer 的 `HOST_BRIDGE` 后端会让 - Manager/TorchEnv tensor 使用当前 GPU。MuJoCo 物理仿真仍在 CPU 上，通过 packed - 传输连接 GPU 上的 observation、action、reward 与 reset。CUDA-only 物理后端仍不 - 支持 ROCm。见 {doc}`/adr/ADR-0012-sole-tensor-manager-and-scoped-backends`。 +- 安装 ROCm PyTorch 且 GPU 可用时，接受 `cuda` buffer 的 `HOST_BRIDGE` 后端只有在 + 训练进程显式请求该 learner 设备（并路由 `manager_torch_device`）时才会让 + Manager/TorchEnv tensor 使用当前 GPU；直接构造环境的默认仍是 CPU。MuJoCo 物理 + 仿真仍在 CPU 上，通过 packed 传输连接显式请求的 GPU observation、action、reward + 与 reset。CUDA-only 物理后端仍不支持 ROCm。见 + {doc}`/adr/ADR-0012-sole-tensor-manager-and-scoped-backends`。 - 从 PyPI 安装（不克隆仓库）时，`make sync-rocm` 不适用；先从 PyTorch ROCm 索引 安装仓库验证过的 torch build，再安装 `unilab`。发布的依赖范围是 `torch>=2.9,<2.15`，pip 会保留已安装的 ROCm build，不会替换为 CUDA wheel： 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 df54b9164..0407ad665 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 @@ -65,12 +65,12 @@ uv run scripts/generate_support_matrix.py --write | Backend | Execution / process / data plane | Torch devices | CUDA runtime | Linux+CUDA | macOS | ROCm | Worker | Reset randomization | Fixed variants | Host callbacks | Packed bridge | |---|---|---|---|---|---|---|---|---|---|---|---| -| `mujoco` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | unknown | unknown | 不支持 | 支持 | -| `motrix` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | -| `drake` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | +| `mujoco` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | unknown | unknown | 不支持 | 支持 | +| `motrix` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | +| `drake` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | | `mjwarp` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 不支持 | | `newton` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 不支持 | -| `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | +| `superdex` | Host bridge / in-process / host bridge | CPU / CUDA | Required only when the learner requests CUDA state/control buffers | Supported: CPU-authoritative physics with optional CUDA Torch buffers | CPU-authoritative host bridge only; no CUDA physics claim | CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`) | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 支持 | | `genesis` | Device-resident / in-process / direct | CUDA | Required for the entire tensor lifecycle | Supported: Linux CUDA only | Unsupported; no CPU, MPS, or ROCm fallback | Unsupported; no CPU, MPS, or ROCm fallback | In-process; no external Python worker | 不支持 | 不支持 | 不支持 | 不支持 | ### Entrypoint x Task Owner @@ -114,7 +114,8 @@ NVIDIA 驱动与容器运行时不匹配，再调整任务配置。设置了 `CUDA_VISIBLE_DEVICES` 时，后端序号指向重映射后的命名空间，而不是宿主机 全局物理索引。 -在 macOS 与 ROCm 上请使用 CPU-authoritative host-bridge 后端。在 ROCm 上， -ROCm PyTorch 的 GPU 可用且后端接受当前 `cuda` 设备时，Manager/TorchEnv 使用 -当前 GPU；CPU 物理仿真不要求 Manager tensor 使用 CPU。host bridge 使用 GPU -Torch buffer 不代表 GPU physics，也不代表 device-resident 后端生命周期。 +在 macOS 与 ROCm 上请使用 CPU-authoritative host-bridge 后端。在 ROCm 上，只有 +训练进程显式请求该 learner 设备且后端接受被路由的 `cuda` +`manager_torch_device` 时，Manager/TorchEnv 才使用当前 GPU；直接构造环境仍为 +CPU。host bridge 使用 GPU Torch buffer 不代表 GPU physics，也不代表 +device-resident 后端生命周期。 diff --git a/scripts/generate_support_matrix.py b/scripts/generate_support_matrix.py index c92197e6b..d795aeaff 100644 --- a/scripts/generate_support_matrix.py +++ b/scripts/generate_support_matrix.py @@ -25,7 +25,7 @@ _SUPPORT_LEVEL_ENUM = re.compile(r"\bSupportLevel\.([A-Z][A-Z_]+)\b") -_ROCM_HOST_BRIDGE_LABEL = "Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`)" +_ROCM_HOST_BRIDGE_LABEL = "CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`)" def _with_support_level_values(content: str) -> str: @@ -33,8 +33,9 @@ def _with_support_level_values(content: str) -> str: rendered = _SUPPORT_LEVEL_ENUM.sub(lambda match: match.group(1).lower(), content) # UniSim's static inventory is deliberately conservative: it only states # that no ROCm-native CUDA-only fallback exists. Host bridges accepting - # `cuda` Torch buffers also accept ROCm PyTorch's CUDA-namespace API while - # physics remains CPU-authoritative; make that repository guarantee visible. + # `cuda` Torch buffers also accept ROCm PyTorch's CUDA-namespace API, + # but only when the process explicitly requests those Manager carriers; + # make that opt-in repository guarantee visible. return rendered.replace( "CPU-authoritative host bridge only; no ROCm CUDA-only fallback", _ROCM_HOST_BRIDGE_LABEL, diff --git a/src/unilab/base/process_device.py b/src/unilab/base/process_device.py index 7efc17fd5..c46ebba3a 100644 --- a/src/unilab/base/process_device.py +++ b/src/unilab/base/process_device.py @@ -16,8 +16,9 @@ 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. +# their simulator. MuJoCo/Motrix/Drake run CPU-authoritative physics; their +# optional accelerator Manager buffers are selected separately through +# ``manager_torch_device``. BACKEND_ENV_DEVICE_FIELDS: dict[str, str] = { "isaacgym": "isaacgym_device_id", "isaacsim": "isaacsim_device_id", @@ -41,6 +42,11 @@ # integer payload sent to that worker before construction. _EXTERNAL_CUDA_IPC_BACKENDS = {"isaacgym", "isaacsim"} +# HOST_BRIDGE backends accept optional accelerator Torch carriers across their packed +# boundary. A process may explicitly request those carriers on its learner Torch +# device; unset non-accelerator requests retain the owner/default CPU placement. +_HOST_BRIDGE_TORCH_BACKENDS = {"mujoco", "motrix", "drake", "superdex"} + # Set once ``bind_genesis_process_device`` has pinned CUDA_VISIBLE_DEVICES for # this process. Genesis/Quadrants binds its CUDA runtime to the first visible @@ -170,6 +176,35 @@ def apply_backend_env_device_override( return result +def apply_manager_torch_device_override( + env_cfg_override: Mapping[str, Any] | None, + backend_type: str, + *, + learner_device: str | None = None, +) -> dict[str, Any]: + """Return an env override selecting optional HOST_BRIDGE Torch carriers. + + The input mapping is never mutated. Explicit CPU requests force CPU carriers. + Unset and non-CUDA learner requests retain owner/default placement. A CUDA request + is validated against backend capabilities when the environment binds. + DEVICE_RESIDENT backends own placement and never receive this synthetic field. + """ + + result = dict(env_cfg_override) if env_cfg_override is not None else {} + if _normalize_backend(backend_type) not in _HOST_BRIDGE_TORCH_BACKENDS: + return result + if learner_device is None: + return result + device = str(learner_device).strip() + if not device or device.lower() == "cpu": + result["manager_torch_device"] = "cpu" + return result + if device.split(":", 1)[0].lower() != "cuda": + return result + result["manager_torch_device"] = device + return result + + def resolve_backend_process_device(backend_type: str, learner_device: str | None) -> str | None: backend = _normalize_backend(backend_type) if ( @@ -399,6 +434,7 @@ def _reset_genesis_device_pin_for_tests() -> None: __all__ = [ "BACKEND_ENV_DEVICE_FIELDS", "apply_backend_env_device_override", + "apply_manager_torch_device_override", "bind_backend_process_device", "bind_backend_process_device_for_backend", "bind_genesis_process_device", diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index dc5f51db6..562655784 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -25,7 +25,6 @@ TensorExecution, TensorLifecycleCapabilities, TensorProcessTopology, - tensor_device_matches, ) from unilab.base.backend_factory import create_backend, env_backend_kwargs @@ -163,6 +162,11 @@ class ManagerBasedRlEnvCfg(EnvCfg): recorders: dict[str, RecorderTermCfg | None] = _manager_terms_field() reset_owners: dict[str, ResetOwnerCfg | None] = _manager_terms_field() + # HOST_BRIDGE physics remains CPU-authoritative. CPU is therefore the default; + # a training process may request accelerator Torch carriers shared with its + # learner. This selects public-manager carrier placement only; it is neither a + # physics-device claim nor a tensor/NumPy runtime compatibility switch. + manager_torch_device: str | None = None seed: int | None = None is_finite_horizon: bool = False auto_reset: bool = True @@ -171,6 +175,27 @@ class ManagerBasedRlEnvCfg(EnvCfg): critic_observation_group: str | None = None def validate(self) -> None: + if self.manager_torch_device is not None: + if ( + not isinstance(self.manager_torch_device, str) + or not self.manager_torch_device.strip() + ): + raise TypeError( + "ManagerBasedRlEnvCfg manager_torch_device must be 'cpu', " + f"'cuda', or 'cuda:', or None; got {self.manager_torch_device!r}" + ) + try: + requested_device = torch.device(self.manager_torch_device) + except (RuntimeError, TypeError, ValueError) as exc: + raise ValueError( + "ManagerBasedRlEnvCfg manager_torch_device is not a valid Torch " + f"device: {self.manager_torch_device!r}" + ) from exc + if requested_device.type not in {"cpu", "cuda"}: + raise ValueError( + "ManagerBasedRlEnvCfg manager_torch_device must be cpu, cuda, or " + f"None; got {self.manager_torch_device!r}" + ) for name, value in (("sim_dt", self.sim_dt), ("ctrl_dt", self.ctrl_dt)): if isinstance(value, bool) or not isinstance(value, (int, float, np.number)): raise TypeError(f"ManagerBasedRlEnvCfg {name} must be a real number") @@ -371,17 +396,9 @@ def __init__(self, cfg: ManagerBasedRlEnvCfg, backend: SimBackend, num_envs: int ) initial_capabilities = backend.get_tensor_capabilities() - runtime_device = torch.device("cpu") - if initial_capabilities.execution is TensorExecution.DEVICE_RESIDENT: - runtime_device = torch.device("cuda", index=torch.cuda.current_device()) - elif ( - initial_capabilities.execution is TensorExecution.HOST_BRIDGE - and torch.cuda.is_available() - ): - # PyTorch uses the CUDA device namespace for both NVIDIA and ROCm GPUs. - accelerator = torch.device("cuda", index=torch.cuda.current_device()) - if tensor_device_matches(initial_capabilities.torch_devices, accelerator): - runtime_device = accelerator + runtime_device = self._resolve_runtime_device( + initial_capabilities, cfg.manager_torch_device + ) super().__init__(cfg, backend, num_envs, device=runtime_device) actual_seed = cfg.seed if cfg.seed is not None else secrets.randbits(63) cfg.seed = actual_seed @@ -881,6 +898,41 @@ def _validate_manager_tensor_runtime(self) -> None: """Bind the tensor lifecycle after backend materialization.""" self._bind_tensor_runtime() + @staticmethod + def _resolve_runtime_device( + capabilities: TensorLifecycleCapabilities, requested_device: str | None + ) -> torch.device: + """Resolve authoritative Manager carrier placement. + + ``DEVICE_RESIDENT`` owns accelerator storage and has no CPU public + carrier. For ``HOST_BRIDGE``, capability inventory says which device + families are possible, while the caller's request chooses whether this + process pays the packed H2D/D2H boundary. CPU is the default; explicit + ``cuda`` resolves to the current Torch ordinal (including ROCm/HIP's + CUDA namespace). Final capability validation remains fail-closed in + ``TorchEnv._bind_tensor_runtime``. + """ + + if capabilities.execution is TensorExecution.DEVICE_RESIDENT: + return torch.device("cuda", index=torch.cuda.current_device()) + if capabilities.execution is not TensorExecution.HOST_BRIDGE: + raise ValueError( + "TorchEnv got a non-tensor backend capability " + f"{capabilities.execution!r}; expected HOST_BRIDGE or DEVICE_RESIDENT" + ) + if requested_device is None: + return torch.device("cpu") + try: + device = torch.device(requested_device) + except (RuntimeError, TypeError, ValueError) as exc: + raise ValueError( + "Manager HOST_BRIDGE runtime device is not a valid Torch device: " + f"{requested_device!r}" + ) from exc + if device.type == "cuda" and device.index is None: + return torch.device("cuda", index=torch.cuda.current_device()) + return device + def _materialize_backend(self) -> None: """Finalize backend runtime resources before the first reset or step.""" try: diff --git a/src/unilab/scripts/play_interactive.py b/src/unilab/scripts/play_interactive.py index a49616313..ce2088aa4 100644 --- a/src/unilab/scripts/play_interactive.py +++ b/src/unilab/scripts/play_interactive.py @@ -47,6 +47,7 @@ from unilab.base.process_device import ( apply_backend_env_device_override, + apply_manager_torch_device_override, configure_backend_process_device, resolve_backend_env_device_id, ) @@ -846,6 +847,11 @@ def _create_env(num_envs: int): sim_backend, learner_device=device, ) + legacy_env_cfg_override = apply_manager_torch_device_override( + legacy_env_cfg_override, + sim_backend, + learner_device=device, + ) return registry.make( args.task, num_envs=num_envs, @@ -866,6 +872,11 @@ def _create_env(num_envs: int): sim_backend, learner_device=device, ) + env_cfg_override = apply_manager_torch_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 7e1a54aba..97dbcf956 100644 --- a/src/unilab/scripts/train_appo.py +++ b/src/unilab/scripts/train_appo.py @@ -21,6 +21,7 @@ from unilab.base.env_factory import registry_env_factory from unilab.base.process_device import ( apply_backend_env_device_override, + apply_manager_torch_device_override, configure_backend_process_device, pin_genesis_device_before_cuda_init, ) @@ -77,6 +78,15 @@ def build_appo_runner_kwargs( else OmegaConf.select(cfg, "training.device", default=None) ), ) + routed_env_cfg_override = apply_manager_torch_device_override( + routed_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, @@ -258,6 +268,11 @@ def play_appo( str(cfg.training.sim_backend), learner_device=device, ) + play_env_cfg_override = apply_manager_torch_device_override( + 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, @@ -400,6 +415,15 @@ def main(cfg: DictConfig) -> None: else learner_device ), ) + env_cfg_override = apply_manager_torch_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 learner_device + ), + ) seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True) tracker = None diff --git a/src/unilab/scripts/train_offpolicy.py b/src/unilab/scripts/train_offpolicy.py index efa9f69d1..7d3b4eb5b 100644 --- a/src/unilab/scripts/train_offpolicy.py +++ b/src/unilab/scripts/train_offpolicy.py @@ -35,6 +35,7 @@ from unilab.base.env_factory import registry_env_factory from unilab.base.process_device import ( apply_backend_env_device_override, + apply_manager_torch_device_override, bind_backend_process_device_for_backend, configure_backend_process_device, pin_genesis_device_before_cuda_init, @@ -200,11 +201,16 @@ def build_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[st from unilab.utils.device import get_default_device rank_device = resolve_dp_rank_device(rank) or get_default_device() - return apply_backend_env_device_override( + routed = apply_backend_env_device_override( base, str(cfg.training.sim_backend), learner_device=rank_device, ) + return apply_manager_torch_device_override( + routed, + str(cfg.training.sim_backend), + learner_device=rank_device, + ) def build_offpolicy_play_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[str, Any] | None: @@ -213,11 +219,16 @@ def build_offpolicy_play_env_cfg_override(algo_name: str, cfg: DictConfig) -> di from unilab.utils.device import get_default_device rank_device = resolve_dp_rank_device(rank) or get_default_device() - return apply_backend_env_device_override( + routed = apply_backend_env_device_override( base, str(cfg.training.sim_backend), learner_device=rank_device, ) + return apply_manager_torch_device_override( + routed, + str(cfg.training.sim_backend), + learner_device=rank_device, + ) def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): diff --git a/src/unilab/scripts/train_rsl_rl.py b/src/unilab/scripts/train_rsl_rl.py index 3c2a5e925..1a307e5e9 100644 --- a/src/unilab/scripts/train_rsl_rl.py +++ b/src/unilab/scripts/train_rsl_rl.py @@ -16,6 +16,7 @@ from unilab.base.config_adapter import BackendAdapter, create_env from unilab.base.process_device import ( apply_backend_env_device_override, + apply_manager_torch_device_override, configure_backend_process_device, pin_genesis_device_before_cuda_init, resolve_backend_env_device_id, @@ -102,6 +103,11 @@ def build_ppo_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: str(cfg.training.sim_backend), learner_device=learner_device, ) + result = apply_manager_torch_device_override( + result, + str(cfg.training.sim_backend), + 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 @@ -124,11 +130,16 @@ def build_ppo_play_env_cfg_override(cfg: DictConfig) -> dict[str, Any]: local_rank=local_rank, default_device="cpu", ) - return apply_backend_env_device_override( + result = apply_backend_env_device_override( base, str(cfg.training.sim_backend), learner_device=learner_device, ) + return apply_manager_torch_device_override( + result, + str(cfg.training.sim_backend), + learner_device=learner_device, + ) def run_motrix_rsl_play_loop( @@ -320,6 +331,11 @@ def play_rsl_rl(cfg: DictConfig, device: str) -> str | None: str(cfg.training.sim_backend), learner_device=device, ) + play_env_cfg_override = apply_manager_torch_device_override( + play_env_cfg_override, + str(cfg.training.sim_backend), + learner_device=device, + ) session, _policy_obs_mode, _checkpoint_path = create_rsl_rl_playback_session( playback_cfg=playback_cfg, env_factory=lambda n: create_env( @@ -478,6 +494,11 @@ def main(cfg: DictConfig) -> None: str(cfg.training.sim_backend), learner_device=device, ) + env_cfg_override = apply_manager_torch_device_override( + env_cfg_override, + str(cfg.training.sim_backend), + learner_device=device, + ) seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True) # Compute effective max_iterations (supports num_timesteps override) diff --git a/tests/base/backend/test_process_device.py b/tests/base/backend/test_process_device.py index 4652c12d7..9f7d444e0 100644 --- a/tests/base/backend/test_process_device.py +++ b/tests/base/backend/test_process_device.py @@ -13,6 +13,7 @@ import unilab.base.process_device as process_device from unilab.base.process_device import ( apply_backend_env_device_override, + apply_manager_torch_device_override, bind_genesis_process_device, configure_backend_process_device, rank_local_visible_cuda_entries, @@ -150,6 +151,28 @@ def test_backend_env_device_override_routes_rank_local_zero( assert routed["nested"] is owner_override["nested"] +def test_manager_torch_device_override_routes_only_explicit_cuda() -> None: + owner_override = {"manager_torch_device": "owner-cpu", "nested": {"keep": True}} + + routed = apply_manager_torch_device_override(owner_override, "mujoco", learner_device=None) + assert routed is not owner_override + assert routed == owner_override + + routed = apply_manager_torch_device_override(owner_override, "mujoco", learner_device="cpu") + assert routed["manager_torch_device"] == "cpu" + + routed = apply_manager_torch_device_override(owner_override, "mujoco", learner_device="cuda:3") + assert routed["manager_torch_device"] == "cuda:3" + assert routed["nested"] is owner_override["nested"] + assert owner_override["manager_torch_device"] == "owner-cpu" + + routed = apply_manager_torch_device_override(owner_override, "mujoco", learner_device="mps") + assert routed["manager_torch_device"] == "owner-cpu" + + routed = apply_manager_torch_device_override(owner_override, "mjwarp", learner_device="cuda:3") + assert routed["manager_torch_device"] == "owner-cpu" + + def test_non_gpu_backend_is_left_untouched() -> None: owner_override = {"isaacgym_device_id": 0} assert apply_backend_env_device_override(owner_override, "mujoco") == owner_override diff --git a/tests/base/test_genesis_runtime.py b/tests/base/test_genesis_runtime.py index b761a5a8f..5bcb9eb5b 100644 --- a/tests/base/test_genesis_runtime.py +++ b/tests/base/test_genesis_runtime.py @@ -224,7 +224,7 @@ def test_interactive_viewer_renders_frames() -> None: # pos/lookat branch reuses a tilted default up and rolls the camera). expected_pose = genesis_playback.camera_pose_matrix_z_up( *genesis_playback.camera_pose_from_kwargs( - viewer_backend._camera_kwargs, viewer_backend._camera_lookat() + viewer_backend._camera_cfg, viewer_backend._camera_lookat() ) ) actual_pose = np.asarray( diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index e45d07f55..e9dbeb2c7 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -252,19 +252,18 @@ def test_ppo_g1_motion_tracking(): assert cfg.algo.max_iterations == 15000 assert cfg.algo.algorithm.entropy_coef == pytest.approx(0.005) - # Interactive playback defaults to policy mode for this task. - assert cfg.interactive.action_mode == "policy" - assert cfg.env.commands.motion.params.sampling_mode == "start" + assert cfg.interactive.action_mode == "zero" + assert cfg.env.commands.motion.params.sampling_mode == "adaptive" assert cfg.env.commands.motion.params.truncate_on_clip_end is False - assert cfg.env.sim_dt == pytest.approx(0.005) + assert cfg.env.sim_dt == pytest.approx(1.0 / 150.0) assert cfg.env.actions.joint_pos.scale == pytest.approx(X2_ACTION_SCALE[0]) - assert cfg.env.terminations.anchor_pos.params.threshold == pytest.approx(0.5) - assert cfg.env.terminations.ee_body_pos.params.threshold == pytest.approx(0.5) - assert cfg.env.terminations.undesired_contacts is not None - assert cfg.reward.motion_joint_pos.weight == pytest.approx(0.5) - assert cfg.reward.motion_joint_vel.weight == pytest.approx(0.25) - assert cfg.reward.motion_body_pos.weight == pytest.approx(2.0) - assert cfg.reward.motion_body_ori.weight == pytest.approx(1.5) - assert cfg.reward.motion_ee_body_pos_z.weight == pytest.approx(2.0) - assert cfg.reward.action_rate_l2.weight == pytest.approx(-0.005) - assert cfg.reward.undesired_contacts.weight == pytest.approx(-0.1) + assert cfg.env.terminations.anchor_pos.params.threshold == pytest.approx(0.25) + assert cfg.env.terminations.ee_body_pos.params.threshold == pytest.approx(0.25) + assert "undesired_contacts" not in cfg.env.terminations + assert cfg.reward.motion_joint_pos.weight == pytest.approx(0.0) + assert cfg.reward.motion_joint_vel.weight == pytest.approx(0.0) + assert cfg.reward.motion_body_pos.weight == pytest.approx(1.0) + assert cfg.reward.motion_body_ori.weight == pytest.approx(1.0) + assert "motion_ee_body_pos_z" not in cfg.reward + assert cfg.reward.action_rate_l2.weight == pytest.approx(-0.1) + assert "undesired_contacts" not in cfg.reward diff --git a/tests/envs/locomotion/g1/test_g1_owner_contract.py b/tests/envs/locomotion/g1/test_g1_owner_contract.py index ceeeec13a..f66b2d91e 100644 --- a/tests/envs/locomotion/g1/test_g1_owner_contract.py +++ b/tests/envs/locomotion/g1/test_g1_owner_contract.py @@ -657,7 +657,7 @@ def test_g1_walk_profile_runtime_obs_scaling_matches_legacy_layout() -> None: try: env.reset(seed=3) - state = env.step(torch.zeros((2, 29), dtype=torch.float32)) + state = env.step(torch.zeros((2, 29), dtype=torch.float32, device=env.device)) gyro = env._backend.get_sensor_data("torso_gyro") upvector = env._backend.get_sensor_data("torso_upvector") dof_vel = env._backend.get_dof_vel() @@ -693,7 +693,7 @@ def test_g1_legacy_profile_runtime_obs_scaling_matches_legacy_layout() -> None: try: env.reset(seed=3) - state = env.step(torch.zeros((2, 29), dtype=torch.float32)) + state = env.step(torch.zeros((2, 29), dtype=torch.float32, device=env.device)) gyro = env._backend.get_sensor_data("torso_gyro") dof_vel = env._backend.get_dof_vel() linvel = env._backend.get_sensor_data("pelvis_local_linvel") @@ -821,6 +821,7 @@ def _genesis_runtime_available() -> bool: from pathlib import Path import numpy as np +import torch from hydra import compose, initialize_config_dir from hydra.core.global_hydra import GlobalHydra @@ -871,7 +872,7 @@ def _genesis_runtime_available() -> bool: "critic": (2, 101), } for _ in range(12): - state = env.step(torch.zeros((2, 29), dtype=torch.float32)) + state = env.step(torch.zeros((2, 29), dtype=torch.float32, device=env.device)) assert set(state.obs) == {"obs", "critic"} assert state.obs["obs"].shape == (2, 98) assert state.obs["critic"].shape == (2, 101) @@ -931,5 +932,5 @@ def test_g1_walk_flat_genesis_owner_real_runtime_smoke(config_group: str) -> Non timeout=600, ) - assert result.returncode == 0, result.stderr or result.stdout + assert result.returncode == 0, f"stdout:\n{result.stdout}\n--- stderr ---\n{result.stderr}" assert f"[genesis env smoke:{config_group}]" in result.stdout diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 79ce10561..912cd065b 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -1018,10 +1018,19 @@ def test_manager_public_inputs_and_episode_counters_are_tensor_first() -> None: env.close() -def test_backend_capability_derives_declared_host_bridge_device( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize( + ("manager_torch_device", "expected_device"), + [ + (None, torch.device("cpu")), + ("cpu", torch.device("cpu")), + ], +) +def test_backend_capability_defaults_host_bridge_manager_carriers_to_cpu( + manager_torch_device: str | None, + expected_device: torch.device, ) -> None: cfg = _make_cfg() + cfg.manager_torch_device = manager_torch_device env, backend = _make_env(cfg) @@ -1032,19 +1041,31 @@ def test_backend_capability_derives_declared_host_bridge_device( assert all(value.device == env.device for value in env.obs_buf.values()) env.close() - # A CPU-only capability must stay on CPU even when a GPU is visible. This - # exercises the negative branch without requiring CUDA/ROCm hardware. + +def test_host_bridge_ignores_accelerator_availability_without_explicit_request( + monkeypatch: pytest.MonkeyPatch, +) -> None: if not bool(torch.cuda._is_compiled()): pytest.skip("fake CUDA visibility requires a Torch CUDA/ROCm build") original_is_available = torch.cuda.is_available monkeypatch.setattr(torch.cuda, "is_available", lambda: True) try: + cfg = _make_cfg() env, _ = _make_env(cfg) - assert env.device == torch.device("cpu") finally: monkeypatch.setattr(torch.cuda, "is_available", original_is_available) env.close() + assert env.device == torch.device("cpu") + + +def test_manager_torch_device_rejects_non_tensor_device_and_invalid_syntax() -> None: + for value, error in (("mps", ValueError), (12, TypeError), ("", TypeError)): + cfg = _make_cfg() + cfg.manager_torch_device = value + with pytest.raises(error): + cfg.validate() + @pytest.mark.parametrize("device", ["cpu", "cuda"]) def test_mujoco_host_bridge_manager_tensor_lifecycle( @@ -1093,6 +1114,8 @@ def test_mujoco_host_bridge_manager_tensor_lifecycle( terminations={"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True)}, policy_observation_group="actor", ) + if device == "cuda": + cfg.manager_torch_device = device backend = create_backend( "mujoco", scene, 2, cfg.sim_dt, base_name="base", add_body_sensors=True ) diff --git a/tests/scripts/test_support_matrix.py b/tests/scripts/test_support_matrix.py index 8df71d4b4..b71a0c55e 100644 --- a/tests/scripts/test_support_matrix.py +++ b/tests/scripts/test_support_matrix.py @@ -138,7 +138,7 @@ def test_generator_renders_hardened_support_levels_without_enum_leakage(): mujoco_zh = next(line for line in zh_block.splitlines() if line.startswith("| `mujoco` |")) assert mujoco_en.endswith("| unknown | unknown | Unsupported | Exact |") assert mujoco_zh.endswith("| unknown | unknown | 不支持 | 支持 |") - assert "Supported: CPU-authoritative physics with ROCm Torch buffers (`cuda`)" in mujoco_en + assert "CPU-authoritative host bridge with explicit ROCm Torch buffers (`cuda`)" in mujoco_en def test_owner_backend_identity_comes_from_training_sim_backend( diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 841387703..170373cae 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -306,6 +306,23 @@ def test_offpolicy_g1_walk_flat_env_cfg_override_has_rewards_and_events(): assert env_cfg_override["rewards"]["tracking_lin_vel"]["weight"] == pytest.approx(2.0) +def test_host_bridge_env_overrides_route_explicit_manager_torch_device(monkeypatch): + offpolicy_cfg = _offpolicy_cfg(["task=g1_walk_flat/mujoco"]) + ppo_cfg = _ppo_cfg(["task=g1_walk_flat/mujoco", "training.device=cuda:2"]) + ppo_mod = _train_rsl_rl(monkeypatch) + offpolicy_mod = _offpolicy() + + monkeypatch.setattr(offpolicy_mod, "resolve_dp_rank_device", lambda rank: "cuda:2") + monkeypatch.setattr("unilab.utils.device.get_default_device", lambda: "cuda:2", raising=False) + + offpolicy_override = offpolicy_mod.build_offpolicy_env_cfg_override("sac", offpolicy_cfg) + ppo_override = ppo_mod.build_ppo_env_cfg_override(ppo_cfg) + + assert offpolicy_override is not None + assert offpolicy_override["manager_torch_device"] == "cuda:2" + assert ppo_override["manager_torch_device"] == "cuda:2" + + @pytest.mark.parametrize( ("backend", "field"), [ @@ -1529,6 +1546,7 @@ def explore(self, obs, deterministic=True): class FakeEnv: def __init__(self): + self.device = torch.device("cpu") self.obs_groups_spec = {"obs": 4} self.action_space = type("ActionSpace", (), {"shape": (2,)})() self.state = None @@ -1537,12 +1555,12 @@ def init_state(self): self.state = type( "State", (), - {"obs": {"obs": np.zeros((cfg.training.play_env_num, 4), dtype=np.float32)}}, + {"obs": {"obs": torch.zeros((cfg.training.play_env_num, 4), dtype=torch.float32)}}, )() def reset(self, env_ids): batch = len(env_ids) - return ({"obs": np.zeros((batch, 4), dtype=np.float32)}, {}) + return ({"obs": torch.zeros((batch, 4), dtype=torch.float32)}, {}) def step(self, actions): batch = actions.shape[0] @@ -1550,7 +1568,7 @@ def step(self, actions): "State", (), { - "obs": {"obs": np.ones((batch, 4), dtype=np.float32)}, + "obs": {"obs": torch.ones((batch, 4), dtype=torch.float32)}, "info": {}, }, )() @@ -1655,10 +1673,11 @@ def explore(self, obs, deterministic=True): def _make_play_fake_env(captured: dict[str, Any], play_env_num: int): - import numpy as np + import torch class FakeEnv: def __init__(self): + self.device = torch.device("cpu") self.obs_groups_spec = {"obs": 4} self.action_space = type("ActionSpace", (), {"shape": (2,)})() self.state = None @@ -1667,12 +1686,12 @@ def init_state(self): self.state = type( "State", (), - {"obs": {"obs": np.zeros((play_env_num, 4), dtype=np.float32)}}, + {"obs": {"obs": torch.zeros((play_env_num, 4), dtype=torch.float32)}}, )() def reset(self, env_ids): batch = len(env_ids) - return ({"obs": np.zeros((batch, 4), dtype=np.float32)}, {}) + return ({"obs": torch.zeros((batch, 4), dtype=torch.float32)}, {}) def step(self, actions): batch = actions.shape[0] @@ -1680,7 +1699,7 @@ def step(self, actions): "State", (), { - "obs": {"obs": np.ones((batch, 4), dtype=np.float32)}, + "obs": {"obs": torch.ones((batch, 4), dtype=torch.float32)}, "info": {}, }, )() @@ -2651,6 +2670,8 @@ def test_play_interactive_velocity_command_contract_rejects_numpy() -> None: def test_play_interactive_runner_log_dir_uses_algo_log_name(monkeypatch: pytest.MonkeyPatch): import types + from unilab.base.base import EnvPlayCapabilities + mod = _play_interactive() captured: dict[str, object] = {} @@ -2690,10 +2711,18 @@ def sync(self): user_scn = type("Scene", (), {"ngeom": 0})() + class FakePhysicsStateApplier: + def __init__(self, env, model, env_index=0): + del env, model, env_index + + def apply(self, state, data): + del state, data + fake_env = types.SimpleNamespace( obs_groups_spec={"obs": 5}, action_space=types.SimpleNamespace(shape=(3,), low=np.full((3,), -1.0), high=np.ones((3,))), cfg=types.SimpleNamespace(ctrl_dt=0.02), + play_capabilities=EnvPlayCapabilities(supports_physics_state_playback=True), get_playback_model=lambda: object(), get_scene_visual_model_file=lambda: None, get_physics_state_snapshot=lambda: np.zeros((1, 8), dtype=np.float32), @@ -2716,9 +2745,7 @@ def sync(self): ) monkeypatch.setattr(mod, "PPOConfig", lambda: types.SimpleNamespace(to_dict=lambda: {})) monkeypatch.setattr(mod.mujoco, "MjData", lambda model: object()) - monkeypatch.setattr(mod.mujoco, "mj_setState", lambda *args, **kwargs: None) - monkeypatch.setattr(mod.mujoco, "mj_forward", lambda *args, **kwargs: None) - monkeypatch.setattr(mod.mujoco, "mjtState", types.SimpleNamespace(mjSTATE_FULLPHYSICS=0)) + monkeypatch.setattr(mod, "PhysicsStateApplier", FakePhysicsStateApplier) monkeypatch.setattr(mod.mujoco.viewer, "launch_passive", lambda *args, **kwargs: FakeViewer()) args = types.SimpleNamespace( diff --git a/tests/scripts/test_visualization_entrypoints.py b/tests/scripts/test_visualization_entrypoints.py index cde91af80..ecf46a36c 100644 --- a/tests/scripts/test_visualization_entrypoints.py +++ b/tests/scripts/test_visualization_entrypoints.py @@ -7,6 +7,7 @@ import numpy as np import pytest +import torch from unisim.backend.base import CameraCfg # CPU-bound on the single-core CI runner; kept in the slow lane (make test-slow). @@ -136,12 +137,12 @@ def _keyboard_env( obs_contains_command: bool = False, ) -> Any: info: dict[str, Any] = ( - {"commands": np.asarray([[0.37, -0.23, 0.19]], dtype=np.float32)} + {"commands": torch.tensor([[0.37, -0.23, 0.19]], dtype=torch.float32)} if with_commands else {"steps": 0} ) if obs_contains_command: - info["commands"] = np.asarray([[0.37, -0.23, 0.19]], dtype=np.float32) + info["commands"] = torch.tensor([[0.37, -0.23, 0.19]], dtype=torch.float32) commands_cfg = ( type( "Cmds", @@ -159,9 +160,9 @@ def _keyboard_env( cfg = cfg_type() cfg.commands = commands_cfg obs = ( - {"obs": np.asarray([[1.0, 0.37, -0.23, 0.19, 2.0]], dtype=np.float32)} + {"obs": torch.tensor([[1.0, 0.37, -0.23, 0.19, 2.0]], dtype=torch.float32)} if obs_contains_command - else {"obs": np.asarray([[1.0, 2.0, 3.0]], dtype=np.float32)} + else {"obs": torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32)} ) state = type("State", (), {"info": info, "obs": obs})() env_type = type(env_cls_name, (), {"__module__": module}) @@ -297,11 +298,13 @@ def get_scene_visual_model_file(self) -> str: def test_play_interactive_binds_mjwarp_process_device_before_session(monkeypatch): mod = _load_script("play_interactive") - bound: list[tuple[str, str]] = [] + bound: list[tuple[str, str, int | None]] = [] monkeypatch.setattr( mod, "configure_backend_process_device", - lambda backend_type, device: bound.append((backend_type, device)), + lambda backend_type, device, backend_device_id=None: bound.append( + (backend_type, device, backend_device_id) + ), ) monkeypatch.setattr(mod, "_select_playback_device", lambda cfg: "cuda:0") monkeypatch.setattr(mod, "available_backends_for_task", lambda task: ["mjwarp"]) @@ -323,16 +326,18 @@ def _stop_before_session(**kwargs): with pytest.raises(RuntimeError, match="stop-before-session"): mod.play_interactive(args, None) - assert bound == [("mjwarp", "cuda:0")] + assert bound == [("mjwarp", "cuda:0", None)] def test_play_interactive_device_binding_is_noop_for_mujoco(monkeypatch): mod = _load_script("play_interactive") - bound: list[tuple[str, str]] = [] + bound: list[tuple[str, str, int | None]] = [] monkeypatch.setattr( mod, "configure_backend_process_device", - lambda backend_type, device: bound.append((backend_type, device)), + lambda backend_type, device, backend_device_id=None: bound.append( + (backend_type, device, backend_device_id) + ), ) monkeypatch.setattr(mod, "_select_playback_device", lambda cfg: "cuda:0") monkeypatch.setattr(mod, "available_backends_for_task", lambda task: ["mujoco"]) @@ -354,4 +359,4 @@ def _stop_before_session(**kwargs): with pytest.raises(RuntimeError, match="stop-before-session"): mod.play_interactive(args, None) - assert bound == [("mujoco", "cuda:0")] + assert bound == [("mujoco", "cuda:0", None)]