diff --git a/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md b/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md index eefa41df7..49f30d7b1 100644 --- a/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md +++ b/docs/sphinx/source/en/2-user_guide/3-backends/5-genesis.md @@ -87,6 +87,12 @@ uv run train --algo ppo --task g1_walk_flat --sim genesis \ algo.num_envs=64 algo.max_iterations=3 ``` +During off-policy runner construction, UniLab temporarily sets +`GS_PARA_LEVEL=2` only for the one-environment dimension probe. This makes the +probe warm the same Genesis disk-cache lane as the scaled collector, after +which the previous environment value (or unset state) is restored before +training starts. + ## Playback and Rendering Genesis native rendering is a declared capability and attaches lazily after diff --git a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md index 11c12657c..4d792ef79 100644 --- a/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md +++ b/docs/sphinx/source/zh_CN/2-user_guide/3-backends/5-genesis.md @@ -81,6 +81,10 @@ uv run train --algo ppo --task g1_walk_flat --sim genesis \ algo.num_envs=64 algo.max_iterations=3 ``` +构建 off-policy runner 时，UniLab 只在单环境维度探测期间临时设置 +`GS_PARA_LEVEL=2`，让探测预热与规模化 collector 相同的 Genesis 磁盘缓存 +lane；训练开始前会恢复原有环境变量值或未设置状态。 + ## Playback 与渲染 Genesis 原生渲染是已声明能力，且在 `scene.build` 之后惰性挂载（训练热 diff --git a/src/unilab/scripts/train_offpolicy.py b/src/unilab/scripts/train_offpolicy.py index 4a58bbb0b..4b1aed752 100644 --- a/src/unilab/scripts/train_offpolicy.py +++ b/src/unilab/scripts/train_offpolicy.py @@ -10,7 +10,7 @@ import datetime import os import sys -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from functools import partial from pathlib import Path from typing import Any, cast @@ -92,6 +92,24 @@ def build_failure_summary(exc: BaseException, run_summary: Any | None = None) -> return summary +@contextmanager +def _genesis_probe_para_level(sim_backend: str): + """Keep the one-env Genesis probe on the collector's kernel cache lane.""" + if sim_backend != "genesis": + yield + return + + previous_value = os.environ.get("GS_PARA_LEVEL") + os.environ["GS_PARA_LEVEL"] = "2" + try: + yield + finally: + if previous_value is None: + os.environ.pop("GS_PARA_LEVEL", None) + else: + os.environ["GS_PARA_LEVEL"] = previous_value + + def build_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[str, Any] | None: base = _build_offpolicy_env_cfg_override(algo_name, cfg, root_dir=Path.cwd()) devices = resolve_dp_topology(OmegaConf.select(cfg, "training.devices", default=None)) @@ -248,26 +266,27 @@ def build_runner(algo_name: str, cfg: DictConfig, log_dir: str | None = None): str(cfg.training.sim_backend), ), } - if algo_name == "sac": - from uni_rl.algos.fast_sac.double_buffer import ( - build_sac_double_buffer_runner, - ) + with _genesis_probe_para_level(str(cfg.training.sim_backend)): + if algo_name == "sac": + from uni_rl.algos.fast_sac.double_buffer import ( + build_sac_double_buffer_runner, + ) - runner = build_sac_double_buffer_runner(cfg, **builder_kwargs) - elif algo_name == "td3": - from uni_rl.algos.fast_td3.double_buffer import ( - build_td3_double_buffer_runner, - ) + runner = build_sac_double_buffer_runner(cfg, **builder_kwargs) + elif algo_name == "td3": + from uni_rl.algos.fast_td3.double_buffer import ( + build_td3_double_buffer_runner, + ) - runner = build_td3_double_buffer_runner(cfg, **builder_kwargs) - elif algo_name == "flashsac": - from uni_rl.algos.flash_sac.double_buffer import ( - build_flashsac_double_buffer_runner, - ) + runner = build_td3_double_buffer_runner(cfg, **builder_kwargs) + elif algo_name == "flashsac": + from uni_rl.algos.flash_sac.double_buffer import ( + build_flashsac_double_buffer_runner, + ) - runner = build_flashsac_double_buffer_runner(cfg, **builder_kwargs) - else: - raise ValueError(f"Unsupported algo: {algo_name}") + runner = build_flashsac_double_buffer_runner(cfg, **builder_kwargs) + else: + raise ValueError(f"Unsupported algo: {algo_name}") return runner diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index e28e579df..ce96960ab 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -9,6 +9,7 @@ import importlib.util import json +import os import sys import types from pathlib import Path @@ -1369,6 +1370,42 @@ def test_offpolicy_default_device_xpu_before_mps(): assert _offpolicy().default_device(mock_torch) == "xpu" +def test_offpolicy_genesis_probe_para_level_restores_previous_value( + monkeypatch: pytest.MonkeyPatch, +): + mod = _offpolicy() + monkeypatch.setenv("GS_PARA_LEVEL", "0") + + with mod._genesis_probe_para_level("genesis"): + assert os.environ["GS_PARA_LEVEL"] == "2" + + assert os.environ["GS_PARA_LEVEL"] == "0" + + +def test_offpolicy_genesis_probe_para_level_restores_unset_variable( + monkeypatch: pytest.MonkeyPatch, +): + mod = _offpolicy() + monkeypatch.delenv("GS_PARA_LEVEL", raising=False) + + with mod._genesis_probe_para_level("genesis"): + assert os.environ["GS_PARA_LEVEL"] == "2" + + assert "GS_PARA_LEVEL" not in os.environ + + +def test_offpolicy_genesis_probe_para_level_ignores_other_backends( + monkeypatch: pytest.MonkeyPatch, +): + mod = _offpolicy() + monkeypatch.setenv("GS_PARA_LEVEL", "0") + + with mod._genesis_probe_para_level("mujoco"): + assert os.environ["GS_PARA_LEVEL"] == "0" + + assert os.environ["GS_PARA_LEVEL"] == "0" + + def test_offpolicy_default_device_cpu_fallback(): mock_torch = MagicMock() mock_torch.cuda.is_available.return_value = False