fix: remove async runner extra kwargs leakage

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(

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(

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,

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:

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}}}

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}}}

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
 # ---------------------------------------------------------------------------