fix: integrate env and playback correctness fixes

src/unilab/envs/manager_based_rl_env.py

@@ -1539,15 +1539,13 @@ def reset(
         env_ids: torch.Tensor | None = None,
         *,
         seed: int | None = None,
-        options: dict[str, Any] | None = None,
     ) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
         """Execute the selected-row (``env_ids``) Manager reset lifecycle."""
-        del options
+        if seed is not None:
+            self.seed(seed)
         if self._state is None:
             state = self.init_state()
             return state.obs, {"log": state.info.get("log", {})}
-        if seed is not None:
-            self.seed(seed)
         if env_ids is None:
             env_ids = torch.arange(self.num_envs, device=self.device)
             reset_mask = self._all_env_reset_mask

src/unilab/rl/contracts/env.py

@@ -16,9 +16,13 @@
   or ``{"obs": 48, "critic": 101}``. ``"obs"`` is the actor group;
   ``"critic"`` is the optional privileged critic group.
 - ``reset(env_ids)`` accepts a one-dimensional Torch integer tensor on the
-  environment device (or ``None``) and returns selected-row observations.
-  ``init_state()`` performs cold-path initialization and returns the first
-  state; ``state`` is ``None`` before it is called.
+  environment device (or ``None``) and returns full-width ``(num_envs, dim)``
+  tensors for every observation group. Only selected rows are reset and updated;
+  non-reset rows retain their previous observations. History and delay state are
+  reset only for selected rows, and manual reset shares this full-width
+  publication contract with autoreset. ``init_state()`` performs cold-path
+  initialization and returns the first state; ``state`` is ``None`` before it is
+  called.
 - Collectors keep observations, actions, rewards, termination flags, and reset
   indices tensor-native. NumPy is only a host serialization/diagnostic carrier,
   never a trainer boundary.

src/unilab/visualization/interactive_playback.py

@@ -482,11 +482,6 @@ def create_rsl_rl_playback_session(
         policy_obs_dims_getter = get_policy_obs_dims
 
     device_name = default_device(torch) if device is None else str(device)
-    env = env_factory(int(playback_cfg.num_envs))
-    if env is None:
-        raise RuntimeError("Playback env factory did not return an environment.")
-    actor_obs_dim, flat_obs_dim = policy_obs_dims_getter(env.obs_groups_spec)
-
     policy_obs_mode = playback_cfg.policy_obs_mode
     checkpoint_path: str | None = None
     if playback_cfg.action_mode == "policy":
@@ -497,6 +492,14 @@ def create_rsl_rl_playback_session(
             playback_cfg.algo_log_name,
             playback_cfg.log_root,
         )
+        if checkpoint_path is not None and sim2sim_preflight is not None:
+            sim2sim_preflight(str(Path(checkpoint_path).parent))
+
+    env = env_factory(int(playback_cfg.num_envs))
+    if env is None:
+        raise RuntimeError("Playback env factory did not return an environment.")
+    actor_obs_dim, flat_obs_dim = policy_obs_dims_getter(env.obs_groups_spec)
+    if playback_cfg.action_mode == "policy":
         if policy_obs_mode == "auto" and checkpoint_path is not None:
             ckpt_dim = checkpoint_input_dim_reader(checkpoint_path)
             if ckpt_dim == actor_obs_dim:
@@ -526,8 +529,6 @@ def create_rsl_rl_playback_session(
         if checkpoint_path is None:
             log("WARNING: no checkpoint found - falling back to zero actions.")
         else:
-            if sim2sim_preflight is not None:
-                sim2sim_preflight(str(Path(checkpoint_path).parent))
             log_dir = str(
                 entrypoint_log_root(
                     Path(root_dir),
@@ -663,14 +664,6 @@ def create_appo_playback_session(
     """Create an APPO interactive playback session."""
 
     device_name = default_device(torch) if device is None else str(device)
-    env = env_factory(int(playback_cfg.num_envs))
-    if env is None:
-        raise RuntimeError("Playback env factory did not return an environment.")
-
-    policy_obs_mode = playback_cfg.policy_obs_mode
-    wrapped_env = wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode)
-    policy = None
-    actor = None
     checkpoint_path: str | None = None
     if playback_cfg.action_mode == "policy":
         checkpoint_path, checkpoint_dir = _resolve_appo_checkpoint_from_cfg(cfg, root_dir=root_dir)
@@ -686,22 +679,36 @@ def create_appo_playback_session(
                 algo_name="appo",
                 strict=bool(getattr(cfg.training, "sim2sim_strict", True)),
             )
-            actor, actor_obs_dim, action_dim = _build_appo_actor(
-                env=env,
-                wrapped_env=wrapped_env,
-                cfg=cfg,
-                rl_cfg=rl_cfg,
-                device=device_name,
-            )
-            checkpoint = torch.load(checkpoint_path, map_location=device_name, weights_only=True)
-            with policy_load_dim_guard(
-                env_obs_dim=actor_obs_dim,
-                env_action_dim=action_dim,
-                algo_name="appo",
-            ):
-                actor.load_state_dict(checkpoint["actor"])
-            policy = actor
-            log(f"Loading APPO checkpoint: {checkpoint_path}")
+
+    env = env_factory(int(playback_cfg.num_envs))
+    if env is None:
+        raise RuntimeError("Playback env factory did not return an environment.")
+
+    policy_obs_mode = playback_cfg.policy_obs_mode
+    wrapped_env = wrapper_cls(env, device=device_name, policy_obs_mode=policy_obs_mode)
+    policy = None
+    actor = None
+    if (
+        playback_cfg.action_mode == "policy"
+        and checkpoint_path is not None
+        and Path(checkpoint_path).exists()
+    ):
+        actor, actor_obs_dim, action_dim = _build_appo_actor(
+            env=env,
+            wrapped_env=wrapped_env,
+            cfg=cfg,
+            rl_cfg=rl_cfg,
+            device=device_name,
+        )
+        checkpoint = torch.load(checkpoint_path, map_location=device_name, weights_only=True)
+        with policy_load_dim_guard(
+            env_obs_dim=actor_obs_dim,
+            env_action_dim=action_dim,
+            algo_name="appo",
+        ):
+            actor.load_state_dict(checkpoint["actor"])
+        policy = actor
+        log(f"Loading APPO checkpoint: {checkpoint_path}")
 
     log(f"Action mode: {playback_cfg.action_mode}")
     return (
@@ -976,6 +983,46 @@ def create_sac_playback_session(
     from unilab.utils.checkpoint import resolve_checkpoint_path
 
     device_name = default_device(torch, str(device) if device is not None else None)
+    checkpoint_path: str | None = None
+    checkpoint_available = False
+    if playback_cfg.action_mode == "policy":
+        algo_log_name = str(cfg.algo.algo_log_name)
+        if playback_cfg.log_root is not None:
+            task_log_root = (
+                get_log_root(
+                    Path(root_dir),
+                    algo_log_name=algo_log_name,
+                    log_root=playback_cfg.log_root,
+                )
+                / playback_cfg.task
+            )
+        else:
+            # Without an explicit log_root the lookup stays on the repo-rooted
+            # ``<root_dir>/logs/<algo>`` tree; the test-only
+            # ``UNILAB_TEST_LOG_ROOT`` override is reserved for cfg-driven
+            # resolution (``training.resolve_configured_checkpoint_path``).
+            task_log_root = (Path(root_dir) / "logs" / algo_log_name).absolute() / playback_cfg.task
+        resolved_path, resolved_dir = resolve_checkpoint_path(
+            task_log_root,
+            playback_cfg.load_run,
+            checkpoint=normalize_checkpoint_value(playback_cfg.checkpoint),
+        )
+        checkpoint_path = str(resolved_path) if resolved_path is not None else None
+        checkpoint_dir = str(resolved_dir) if resolved_dir is not None else None
+        checkpoint_available = checkpoint_path is not None and os.path.exists(checkpoint_path)
+        if not checkpoint_available:
+            log(
+                f"WARNING: no {algo_name} checkpoint found for "
+                f"load_run={playback_cfg.load_run} - falling back to zero actions."
+            )
+        else:
+            resolve_sim2sim_config(
+                checkpoint_dir,
+                cfg,
+                algo_name=algo_name,
+                strict=bool(getattr(cfg.training, "sim2sim_strict", True)),
+            )
+
     env = env_factory(int(playback_cfg.num_envs))
     if env is None:
         raise RuntimeError("Playback env factory did not return an environment.")
@@ -996,7 +1043,6 @@ def create_sac_playback_session(
         actor_kwargs.update(_flashsac_actor_kwargs(cfg))
 
     actor = None
-    checkpoint_path: str | None = None
     normalizer = None
     if bool(getattr(cfg.algo, "obs_normalization", False)):
         from unilab.rl.algos.common.normalization import EmpiricalNormalization
@@ -1013,42 +1059,9 @@ def create_sac_playback_session(
             **actor_kwargs,
         )
         actor.eval()
-        algo_log_name = str(cfg.algo.algo_log_name)
-        if playback_cfg.log_root is not None:
-            task_log_root = (
-                get_log_root(
-                    Path(root_dir),
-                    algo_log_name=algo_log_name,
-                    log_root=playback_cfg.log_root,
-                )
-                / playback_cfg.task
-            )
-        else:
-            # Without an explicit log_root the lookup stays on the repo-rooted
-            # ``<root_dir>/logs/<algo>`` tree; the test-only
-            # ``UNILAB_TEST_LOG_ROOT`` override is reserved for cfg-driven
-            # resolution (``training.resolve_configured_checkpoint_path``).
-            task_log_root = (Path(root_dir) / "logs" / algo_log_name).absolute() / playback_cfg.task
-        resolved_path, resolved_dir = resolve_checkpoint_path(
-            task_log_root,
-            playback_cfg.load_run,
-            checkpoint=normalize_checkpoint_value(playback_cfg.checkpoint),
-        )
-        checkpoint_path = str(resolved_path) if resolved_path is not None else None
-        checkpoint_dir = str(resolved_dir) if resolved_dir is not None else None
-        if checkpoint_path is None or not os.path.exists(checkpoint_path):
-            log(
-                f"WARNING: no {algo_name} checkpoint found for "
-                f"load_run={playback_cfg.load_run} - falling back to zero actions."
-            )
+        if checkpoint_path is None or not checkpoint_available:
             actor = None
         else:
-            resolve_sim2sim_config(
-                checkpoint_dir,
-                cfg,
-                algo_name=algo_name,
-                strict=bool(getattr(cfg.training, "sim2sim_strict", True)),
-            )
             checkpoint = torch.load(checkpoint_path, map_location=session_device, weights_only=True)
             expected_actor_semantics = (
                 flashsac_play_actor_semantics(

tests/envs/test_manager_based_rl_env.py

@@ -3533,6 +3533,29 @@ def test_manual_reset_is_required_when_autoreset_is_disabled() -> None:
     env.step(torch.zeros((2, 1), dtype=torch.float32))
 
 
+def test_cold_reset_applies_seed_before_initialization() -> None:
+    first, _ = _make_env()
+    second, _ = _make_env()
+    different, _ = _make_env()
+    try:
+        first_obs, _ = first.reset(seed=123)
+        second_obs, _ = second.reset(seed=123)
+        different_obs, _ = different.reset(seed=124)
+
+        assert first.state is not None
+        assert second.state is not None
+        assert different.state is not None
+        for group in first_obs:
+            torch.testing.assert_close(first_obs[group], second_obs[group])
+        torch.testing.assert_close(first.state.info["steps"], second.state.info["steps"])
+        assert torch.any(first.state.info["steps"] != different.state.info["steps"])
+        assert set(different_obs) == set(first_obs)
+    finally:
+        first.close()
+        second.close()
+        different.close()
+
+
 def test_reset_seed_updates_the_shared_generator_in_place() -> None:
     cfg = _make_cfg(include_optional_managers=False)
     cfg.events = {}

tests/envs/test_observation_partial_reset.py

@@ -77,7 +77,7 @@ def test_observation_mask_reset_contract() -> None:
         for group, values in reset_obs.items():
             assert values.shape == obs_before[group].shape
             assert np.isfinite(values).all()
-            np.testing.assert_array_equal(env.state.obs[group][reset_ids], values)
+            np.testing.assert_array_equal(env.state.obs[group][reset_ids], values[reset_ids])
             np.testing.assert_array_equal(
                 env.state.obs[group][keep_ids],
                 obs_before[group][keep_ids],

tests/test_library_import_boundary.py

@@ -63,9 +63,16 @@ def _imports(path: Path) -> set[str]:
     return modules
 
 
+def _display_path(path: Path) -> str:
+    try:
+        return path.relative_to(_REPO_ROOT).as_posix()
+    except ValueError:
+        return path.as_posix()
+
+
 def _layer_violations(root: Path, forbidden: tuple[str, ...]) -> list[tuple[str, str]]:
     return [
-        (path.relative_to(_REPO_ROOT).as_posix(), module)
+        (_display_path(path), module)
         for path in sorted(root.rglob("*.py"))
         for module in sorted(_imports(path))
         if any(module == layer or module.startswith(f"{layer}.") for layer in forbidden)
@@ -175,7 +182,12 @@ def test_utils_does_not_import_higher_layers() -> None:
     "adapters": (
         "unilab.training",
         "unilab.base",
-        "unilab.algos",
+        "unilab.rl.algos",
+        "unilab.envs",
+        "unilab.tasks",
+        "unilab.managers",
+        "unilab.scripts",
+        "unisim",
     ),
 }
 
@@ -187,6 +199,40 @@ def test_rl_layer_respects_internal_dag(layer: str) -> None:
     assert violations == [], f"src/unilab/rl/{layer} violates the RL-internal DAG: {violations}"
 
 
+@pytest.mark.parametrize(
+    "forbidden",
+    (
+        "unilab.training",
+        "unilab.base",
+        "unilab.rl.algos",
+        "unilab.envs",
+        "unilab.tasks",
+        "unilab.managers",
+        "unilab.scripts",
+        "unisim",
+    ),
+)
+def test_layer_violations_catches_forbidden_adapter_imports(tmp_path: Path, forbidden: str) -> None:
+    path = tmp_path / "forbidden.py"
+    path.write_text(f"from {forbidden} import marker\n", encoding="utf-8")
+
+    assert _layer_violations(tmp_path, _RL_DAG_FORBIDDEN["adapters"]) == [
+        (path.as_posix(), forbidden)
+    ]
+
+
+def test_layer_violations_allows_adapter_contracts_typing_import(tmp_path: Path) -> None:
+    path = tmp_path / "allowed.py"
+    path.write_text(
+        "from typing import TYPE_CHECKING\n"
+        "if TYPE_CHECKING:\n"
+        "    from unilab.rl.contracts import observations\n",
+        encoding="utf-8",
+    )
+
+    assert _layer_violations(tmp_path, _RL_DAG_FORBIDDEN["adapters"]) == []
+
+
 def test_rl_logging_does_not_import_other_unilab_layers() -> None:
     violations = [
         (path.relative_to(_REPO_ROOT).as_posix(), module)