From 02b4f1468940687803ae4a29c2146afea302af7c Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Fri, 9 Oct 2026 00:11:16 +0800 Subject: [PATCH 1/2] fix(training): stabilize APPO and MJWarp tests --- src/unilab/scripts/train_appo.py | 50 +++++++++++++--------- tests/base/test_mjwarp_dr.py | 4 +- tests/scripts/test_train_script_configs.py | 3 ++ tests/scripts/test_train_scripts.py | 31 ++++++++++++++ 4 files changed, 66 insertions(+), 22 deletions(-) diff --git a/src/unilab/scripts/train_appo.py b/src/unilab/scripts/train_appo.py index 97dbcf956..1a3181872 100644 --- a/src/unilab/scripts/train_appo.py +++ b/src/unilab/scripts/train_appo.py @@ -60,6 +60,32 @@ def _is_cuda_device(device: str | None) -> bool: return device is not None and str(device).strip().lower().split(":", 1)[0] == "cuda" +def _env_carrier_device( + cfg: DictConfig, + collector_device: str | None, + *, + runner_device: str | None = None, +) -> str | None: + """Resolve the process device used by Manager Torch carriers. + + The APPO collector owns the environment even when its actor runs on another + accelerator family. An explicit CUDA collector request therefore wins; + otherwise fall back to the runner request. Device auto-detection is handled + separately and must not opt CPU-authoritative HOST_BRIDGE managers into CUDA. + """ + if collector_device is not None and _is_cuda_device(collector_device): + return collector_device + return ( + runner_device + if runner_device is not None + else OmegaConf.select( + cfg, + "training.device", + default=None, + ) + ) + + def build_appo_runner_kwargs( cfg: DictConfig, env_cfg_override: dict | None, @@ -72,20 +98,12 @@ def build_appo_runner_kwargs( 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) - ), + learner_device=_env_carrier_device(cfg, collector_device), ) 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) - ), + learner_device=_env_carrier_device(cfg, collector_device), ) runner_kwargs = { @@ -409,20 +427,12 @@ def main(cfg: DictConfig) -> None: 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 - ), + learner_device=_env_carrier_device(cfg, collector_device, runner_device=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 - ), + learner_device=_env_carrier_device(cfg, collector_device, runner_device=learner_device), ) seed_info = apply_configured_training_seed(cfg, torch_runtime=True, cuda=True) diff --git a/tests/base/test_mjwarp_dr.py b/tests/base/test_mjwarp_dr.py index c66ea4812..e4344bd5e 100644 --- a/tests/base/test_mjwarp_dr.py +++ b/tests/base/test_mjwarp_dr.py @@ -385,7 +385,7 @@ def test_low_friction_world_changes_kick_response(backend: Any) -> None: rows = np.asarray([1, 3], dtype=np.int32) qpos, qvel = _stand_state(backend, rows.size) geom_friction = _broadcast_rows(backend.get_geom_friction(), rows.size) - geom_friction[..., 0] = 0.01 + geom_friction[..., 0] = 0.1 backend.set_state( rows, qpos, @@ -405,7 +405,7 @@ def test_low_friction_world_changes_kick_response(backend: Any) -> None: backend.step(_stand_ctrl(backend), nsteps=3) assert np.isfinite(backend.get_base_pos()).all() displacement = np.linalg.norm(backend.get_base_pos()[:, :2] - start, axis=1) - # Near-zero slide friction changes the kick response of the mutated worlds + # Low slide friction changes the kick response of the mutated worlds # while the default-friction worlds stay pairwise consistent. assert abs(displacement[1] - displacement[0]) > 0.003 # Same-condition worlds agree up to contact-solver reduction noise. diff --git a/tests/scripts/test_train_script_configs.py b/tests/scripts/test_train_script_configs.py index 9a58cf224..97e4aa21a 100644 --- a/tests/scripts/test_train_script_configs.py +++ b/tests/scripts/test_train_script_configs.py @@ -92,6 +92,9 @@ def test_appo_task_configs_load(task, tmp_path): f"task={task}", "algo.max_iterations=1", "training.no_play=true", + # HOST_BRIDGE Torch carriers are opt-in; do not let the smoke test + # vary with the host's accelerator availability. + "training.device=cpu", *motion_overrides, ], capture_output=True, diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 170373cae..636d0f726 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -933,6 +933,37 @@ def test_build_appo_runner_kwargs_forwards_sim_backend(): assert runner_kwargs["env_cfg_overrides"]["rewards"] == {} +def test_appo_host_bridge_carriers_require_explicit_cuda(monkeypatch): + mod = _train_appo() + cfg = _appo_cfg(["task=g1_motion_tracking/mujoco"]) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + + auto = mod.build_appo_runner_kwargs(cfg, env_cfg_override={}, collector_device=None) + assert "manager_torch_device" not in auto["env_cfg_overrides"] + + collector_cpu = mod.build_appo_runner_kwargs( + cfg, + env_cfg_override={}, + collector_device="cpu", + ) + assert "manager_torch_device" not in collector_cpu["env_cfg_overrides"] + + explicit_cpu_cfg = _appo_cfg(["task=g1_motion_tracking/mujoco", "training.device=cpu"]) + explicit_cpu = mod.build_appo_runner_kwargs( + explicit_cpu_cfg, + env_cfg_override={}, + collector_device=None, + ) + assert explicit_cpu["env_cfg_overrides"]["manager_torch_device"] == "cpu" + + collector_cuda = mod.build_appo_runner_kwargs( + cfg, + env_cfg_override={}, + collector_device="cuda:1", + ) + assert collector_cuda["env_cfg_overrides"]["manager_torch_device"] == "cuda:1" + + def test_run_motrix_play_loop_runs_without_physics_state(): mod = _train_appo() From 9ba306a9eacfa3b22eef9d28681fe7fce53e80c6 Mon Sep 17 00:00:00 2001 From: TATP-233 Date: Fri, 9 Oct 2026 00:26:47 +0800 Subject: [PATCH 2/2] chore: trigger current-head CI