From a80ad4e15be5f42c0732eb7edab5fe37833d2e2e Mon Sep 17 00:00:00 2001 From: Yves Date: Sat, 18 Apr 2026 23:34:26 +0800 Subject: [PATCH] fix: remove async runner extra kwargs leakage --- src/unilab/algos/torch/appo/runner.py | 6 +- .../algos/torch/offpolicy/multi_gpu_runner.py | 3 +- src/unilab/algos/torch/offpolicy/runner.py | 1 + src/unilab/ipc/async_runner.py | 9 +- tests/algos/test_appo_runner_unit.py | 122 ++++++++++++++++++ tests/algos/test_offpolicy_runner_unit.py | 80 +++++++++++- tests/ipc/test_async_runner.py | 31 +++++ 7 files changed, 239 insertions(+), 13 deletions(-) create mode 100644 tests/algos/test_appo_runner_unit.py diff --git a/src/unilab/algos/torch/appo/runner.py b/src/unilab/algos/torch/appo/runner.py index c33144fcc..7fd5d018f 100644 --- a/src/unilab/algos/torch/appo/runner.py +++ b/src/unilab/algos/torch/appo/runner.py @@ -46,8 +46,8 @@ def __init__( rl_cfg=rl_cfg, device=device, collector_device=collector_device, - num_envs=num_envs, sim_backend=sim_backend, + num_envs=num_envs, ) self.steps_per_env = steps_per_env @@ -95,7 +95,7 @@ def _detect_dims(self): env = registry.make( self.env_name, num_envs=1, - sim_backend=self.extra_kwargs.get("sim_backend", "mujoco"), + sim_backend=self.sim_backend, env_cfg_override=self.env_cfg_overrides if self.env_cfg_overrides else None, ) obs_dim, critic_dim = get_obs_dims(env.obs_groups_spec) @@ -234,7 +234,7 @@ def learn( "critic_weight_param_shapes": critic_weight_param_shapes, "metrics_queue": metrics_queue, "collector_device": self.collector_device, - "sim_backend": self.extra_kwargs.get("sim_backend", "mujoco"), + "sim_backend": self.sim_backend, "env_cfg_override": self.env_cfg_overrides if self.env_cfg_overrides else None, } self._start_collector( diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py index f5f1ab20a..7834fac42 100644 --- a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -380,7 +380,6 @@ def _learn_multi_gpu( # --- Start Collector (CPU, single process, unchanged) --- weight_param_shapes = {k: v.shape for k, v in self.learner.actor.state_dict().items()} - sim_backend: str = self.extra_kwargs.get("sim_backend", "mujoco") collector_kwargs = { "env_name": self.env_name, "num_envs": self.num_envs, @@ -399,7 +398,7 @@ def _learn_multi_gpu( "env_steps_per_sync": self.env_steps_per_sync, "obs_normalization": False, "shared_obs_normalizer_stats": None, - "sim_backend": sim_backend, + "sim_backend": self.sim_backend, "env_cfg_override": self.env_cfg_override, } self._start_collector( diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index 13a23ef0f..51c579105 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -221,6 +221,7 @@ def learn( "env_steps_per_sync": self.env_steps_per_sync, "obs_normalization": self.obs_normalization, "shared_obs_normalizer_stats": shared_obs_normalizer_stats, + "sim_backend": self.sim_backend, "env_cfg_override": self.env_cfg_override, "obs_dim": self.obs_dim, "action_dim": self.action_dim, diff --git a/src/unilab/ipc/async_runner.py b/src/unilab/ipc/async_runner.py index 3352e5007..0b52dd043 100644 --- a/src/unilab/ipc/async_runner.py +++ b/src/unilab/ipc/async_runner.py @@ -26,16 +26,16 @@ def __init__( *, device: str | None = None, collector_device: str | None = None, + sim_backend: str = "mujoco", num_envs: int = 4096, - **kwargs, ): self.env_name = env_name self.env_cfg_overrides = env_cfg_overrides self.rl_cfg = rl_cfg self.device = device or self._get_default_device() self.collector_device = collector_device or self.device + self.sim_backend = sim_backend self.num_envs = num_envs - self.extra_kwargs = kwargs self._collector_process: Any = None self._stop_event = _SPAWN_CTX.Event() @@ -58,10 +58,7 @@ def learn( ) -> None: ... def _start_collector(self, target_fn: Callable, kwargs: dict) -> None: - collector_kwargs = {**self.extra_kwargs, **kwargs} - self._collector_process = _SPAWN_CTX.Process( - target=target_fn, kwargs=collector_kwargs, daemon=True - ) + self._collector_process = _SPAWN_CTX.Process(target=target_fn, kwargs=kwargs, daemon=True) self._collector_process.start() def close(self) -> None: diff --git a/tests/algos/test_appo_runner_unit.py b/tests/algos/test_appo_runner_unit.py new file mode 100644 index 000000000..35a1a0035 --- /dev/null +++ b/tests/algos/test_appo_runner_unit.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import pytest +import torch + +import unilab.algos.torch.appo.runner as appo_runner_module +from unilab.algos.torch.appo.runner import APPORunner + + +class _FakeModule: + def state_dict(self) -> dict[str, torch.Tensor]: + return {"weight": torch.zeros(1)} + + +class _FakeLearner: + def __init__(self) -> None: + self.actor = _FakeModule() + self.critic = _FakeModule() + self.num_learning_epochs = 1 + + def get_state_dict(self) -> dict[str, int]: + return {"iteration": 0} + + +class _FakeSharedOnPolicyStorage: + def __init__( + self, + *, + num_envs: int, + num_steps: int, + obs_dim: int, + action_dim: int, + critic_dim: int, + num_slots: int, + create: bool, + ) -> None: + del num_envs, num_steps, obs_dim, action_dim, critic_dim, num_slots, create + self.name = "fake-storage" + self._write_ptr = object() + self._read_ptr = object() + + def cleanup(self) -> None: + pass + + +class _FakeWeightSync: + def __init__(self) -> None: + self.name = "fake-weight-sync" + + @classmethod + def from_state_dict( + cls, state_dict: dict[str, torch.Tensor], create: bool = True + ) -> "_FakeWeightSync": + del state_dict, create + return cls() + + def cleanup(self) -> None: + pass + + +class _FakeLogger: + def __init__(self, **kwargs) -> None: + del kwargs + self._total_steps = 0 + self._mean_ep_length = 0.0 + + def set_collection_sync(self, enabled: bool, env_steps_per_sync: int) -> None: + del enabled, env_steps_per_sync + + def start(self) -> None: + pass + + def log_status(self, status: str) -> None: + del status + + def log_save(self, ckpt_path: str) -> None: + del ckpt_path + + def finish(self) -> None: + pass + + +def test_appo_runner_uses_explicit_runtime_context( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + captured_detect: dict[str, object] = {} + captured_collector: dict[str, object] = {} + + def fake_detect_dims(self: APPORunner) -> tuple[int, int]: + captured_detect["sim_backend"] = self.sim_backend + self.critic_dim = 7 + self.critic_input_dim = 5 + return (4, 2) + + def capture_start_collector(*, target_fn, kwargs): + del target_fn + captured_collector.update(kwargs) + + monkeypatch.setattr(APPORunner, "_detect_dims", fake_detect_dims) + monkeypatch.setattr(APPORunner, "_build_learner", lambda self: _FakeLearner()) + monkeypatch.setattr(appo_runner_module, "SharedOnPolicyStorage", _FakeSharedOnPolicyStorage) + monkeypatch.setattr(appo_runner_module, "SharedWeightSync", _FakeWeightSync) + monkeypatch.setattr(appo_runner_module, "OffPolicyLogger", _FakeLogger) + monkeypatch.setattr(appo_runner_module.torch, "save", lambda *args, **kwargs: None) + + runner = APPORunner( + env_name="DummyEnv", + env_cfg_overrides={"reward_config": {"scales": {"alive": 1.0}}}, + rl_cfg={"actor": {}, "critic": {}, "algorithm": {}}, + device="cpu", + collector_device="cpu", + sim_backend="motrix", + num_envs=2, + steps_per_env=4, + ) + monkeypatch.setattr(runner, "_start_collector", capture_start_collector) + + runner.learn(max_iterations=0, save_interval=0, log_dir=str(tmp_path)) + + assert captured_detect["sim_backend"] == "motrix" + assert captured_collector["sim_backend"] == "motrix" + assert captured_collector["env_cfg_override"] == {"reward_config": {"scales": {"alive": 1.0}}} diff --git a/tests/algos/test_offpolicy_runner_unit.py b/tests/algos/test_offpolicy_runner_unit.py index fa2118bef..d0d643f52 100644 --- a/tests/algos/test_offpolicy_runner_unit.py +++ b/tests/algos/test_offpolicy_runner_unit.py @@ -5,6 +5,7 @@ import pytest import torch +import unilab.algos.torch.offpolicy.multi_gpu_runner as multi_gpu_runner_module import unilab.algos.torch.offpolicy.runner as runner_module from unilab.algos.torch.offpolicy.runner import ( OffPolicyRunner, @@ -48,8 +49,16 @@ def get_state_dict(self) -> dict[str, int]: class _FakeReplayBuffer: last_instance: "_FakeReplayBuffer | None" = None - def __init__(self, capacity: int, obs_dim: int, action_dim: int, device: str, critic_dim: int): - del capacity, obs_dim, action_dim, device, critic_dim + def __init__( + self, + capacity: int, + obs_dim: int, + action_dim: int, + device: str, + critic_dim: int = 0, + defer_gpu: bool = False, + ): + del capacity, obs_dim, action_dim, device, critic_dim, defer_gpu self.size = torch.zeros(1, dtype=torch.int64) self.ptr = torch.zeros(1, dtype=torch.int64) self.sample_calls = 0 @@ -311,3 +320,70 @@ def fake_sleep(seconds: float) -> None: assert replay_buffer.sample_calls == 1 assert replay_buffer.sample_sizes_at_call == [threshold] assert logger.step_calls and logger.step_calls[0]["iteration"] == 1 + + +def test_offpolicy_runner_passes_explicit_runtime_context_to_collector( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + runner = _make_runner(monkeypatch, sync_collection=False) + runner.sim_backend = "motrix" + runner.env_cfg_override = {"reward_config": {"scales": {"alive": 1.0}}} + captured: dict[str, object] = {} + + def capture_start_collector(*, target_fn, kwargs): + del target_fn + captured.update(kwargs) + + monkeypatch.setattr(runner, "_start_collector", capture_start_collector) + monkeypatch.setattr(runner_module._SPAWN_CTX, "Queue", lambda maxsize=0: queue.Queue()) + monkeypatch.setattr(runner_module.time, "sleep", lambda seconds: None) + + runner.learn(max_iterations=0, save_interval=0, log_dir=str(tmp_path)) + + assert captured["sim_backend"] == "motrix" + assert captured["env_cfg_override"] == {"reward_config": {"scales": {"alive": 1.0}}} + + +def test_multi_gpu_runner_passes_explicit_runtime_context_to_collector( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + monkeypatch.setattr(multi_gpu_runner_module, "ReplayBuffer", _FakeReplayBuffer) + monkeypatch.setattr(multi_gpu_runner_module, "SharedWeightSync", _FakeWeightSync) + monkeypatch.setattr(multi_gpu_runner_module.time, "sleep", lambda seconds: None) + monkeypatch.setattr(runner_module, "get_env_dims", lambda *args, **kwargs: (4, 2, 0)) + monkeypatch.setattr( + multi_gpu_runner_module.tmp, + "spawn", + lambda *args, **kwargs: None, + ) + + learner = _FakeLearner() + runner = multi_gpu_runner_module.MultiGPUOffPolicyRunner( + learner=learner, + env_name="DummyEnv", + algo_type="sac", + learner_kwargs={}, + num_gpus=2, + num_envs=2, + replay_buffer_n=8, + batch_size=8, + learning_starts=6, + updates_per_step=1, + policy_frequency=1, + sync_collection=False, + env_steps_per_sync=1, + device="cpu", + sim_backend="motrix", + env_cfg_override={"reward_config": {"scales": {"alive": 1.0}}}, + ) + captured: dict[str, object] = {} + + def capture_start_collector(*, target_fn, kwargs): + del target_fn + captured.update(kwargs) + + monkeypatch.setattr(runner, "_start_collector", capture_start_collector) + runner.learn(max_iterations=0, save_interval=0, log_dir=str(tmp_path)) + + assert captured["sim_backend"] == "motrix" + assert captured["env_cfg_override"] == {"reward_config": {"scales": {"alive": 1.0}}} diff --git a/tests/ipc/test_async_runner.py b/tests/ipc/test_async_runner.py index bdf995cd5..f44590c27 100644 --- a/tests/ipc/test_async_runner.py +++ b/tests/ipc/test_async_runner.py @@ -82,6 +82,11 @@ def test_init_collector_device_explicit(): assert r.collector_device == "cpu" +def test_init_sim_backend_explicit(): + r = _make_runner(sim_backend="motrix") + assert r.sim_backend == "motrix" + + def test_init_num_envs(): r = _make_runner(num_envs=64) assert r.num_envs == 64 @@ -187,6 +192,16 @@ def _noop_collector(stop_event) -> None: stop_event.wait(timeout=30) +def _collector_report_kwargs( + stop_event, + report_queue, + token: str, + sim_backend: str = "missing", +) -> None: + report_queue.put({"sim_backend": sim_backend, "token": token}) + stop_event.wait(timeout=30) + + def test_start_collector_spawns_process(): """_start_collector() must create and start a subprocess.""" r = _make_runner() @@ -196,6 +211,22 @@ def test_start_collector_spawns_process(): r.close() +def test_start_collector_does_not_merge_runner_runtime_fields(): + r = _make_runner(sim_backend="motrix") + report_queue = _SPAWN_CTX.Queue() + r._start_collector( + target_fn=_collector_report_kwargs, + kwargs={ + "stop_event": r._stop_event, + "report_queue": report_queue, + "token": "ok", + }, + ) + payload = report_queue.get(timeout=5) + assert payload == {"sim_backend": "missing", "token": "ok"} + r.close() + + # --------------------------------------------------------------------------- # __del__ exception handling # ---------------------------------------------------------------------------