diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index dfd45cf0d..4262549d4 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/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 diff --git a/src/unilab/rl/contracts/env.py b/src/unilab/rl/contracts/env.py index f89a0b142..099545732 100644 --- a/src/unilab/rl/contracts/env.py +++ b/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. diff --git a/src/unilab/visualization/interactive_playback.py b/src/unilab/visualization/interactive_playback.py index 8129d27a2..1843e2c17 100644 --- a/src/unilab/visualization/interactive_playback.py +++ b/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 + # ``/logs/`` 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 - # ``/logs/`` 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( diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index b6748d53e..460126b58 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/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 = {} diff --git a/tests/envs/test_observation_partial_reset.py b/tests/envs/test_observation_partial_reset.py index baf19434d..d08a86ff8 100644 --- a/tests/envs/test_observation_partial_reset.py +++ b/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], diff --git a/tests/test_library_import_boundary.py b/tests/test_library_import_boundary.py index 1efacf403..af031c19f 100644 --- a/tests/test_library_import_boundary.py +++ b/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) diff --git a/tests/visualization/test_interactive_playback.py b/tests/visualization/test_interactive_playback.py index f1bb9d7d6..5369db8c7 100644 --- a/tests/visualization/test_interactive_playback.py +++ b/tests/visualization/test_interactive_playback.py @@ -20,6 +20,7 @@ PlaybackConfig, PlaybackControls, RslRlPlaybackSession, + create_appo_playback_session, create_rsl_rl_playback_session, create_sac_playback_session, make_sim2sim_preflight, @@ -327,7 +328,7 @@ def test_create_rsl_rl_playback_session_runs_sim2sim_preflight(tmp_path: Path) - run_dir.mkdir() checkpoint = run_dir / "model_10.pt" torch.save({"actor": {}}, checkpoint) - preflight_calls: list[str | None] = [] + calls: list[str] = [] class Runner: def __init__(self, wrapped_env, train_cfg, log_dir, device): @@ -339,15 +340,42 @@ def load(self, checkpoint, load_cfg): def get_inference_policy(self, *, device): return lambda obs: torch.ones((1, 2)) + def env_factory(num_envs: int) -> SimpleNamespace: + calls.append("env") + return _rsl_rl_session_test_env() + + def preflight(source_run_dir: str | None) -> None: + assert source_run_dir == str(run_dir) + calls.append("preflight") + kwargs = _rsl_rl_session_kwargs(tmp_path) kwargs["checkpoint_resolver"] = lambda *args: str(checkpoint) kwargs["runner_cls"] = Runner - kwargs["sim2sim_preflight"] = lambda run_dir: preflight_calls.append(run_dir) + kwargs["env_factory"] = env_factory + kwargs["sim2sim_preflight"] = preflight _session, _mode, resolved = create_rsl_rl_playback_session(**kwargs) assert resolved == str(checkpoint) - assert preflight_calls == [str(run_dir)] + assert calls == ["preflight", "env"] + + +def test_rsl_rl_sim2sim_failure_prevents_env_construction(tmp_path: Path) -> None: + checkpoint = tmp_path / "model_10.pt" + checkpoint.touch() + env_calls: list[int] = [] + kwargs = _rsl_rl_session_kwargs(tmp_path) + kwargs["checkpoint_resolver"] = lambda *args: str(checkpoint) + kwargs["env_factory"] = lambda num_envs: env_calls.append(num_envs) + + def reject_checkpoint(source_run_dir: str | None) -> None: + raise RuntimeError("incompatible checkpoint") + + kwargs["sim2sim_preflight"] = reject_checkpoint + with pytest.raises(RuntimeError, match="incompatible checkpoint"): + create_rsl_rl_playback_session(**kwargs) + + assert env_calls == [] def test_create_rsl_rl_playback_session_wraps_load_with_dim_guard(tmp_path: Path) -> None: @@ -376,6 +404,52 @@ def get_inference_policy(self, *, device): create_rsl_rl_playback_session(**kwargs) +def test_appo_sim2sim_failure_prevents_env_construction( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + import unilab.visualization.interactive_playback as interactive_playback + + checkpoint = tmp_path / "model_10.pt" + checkpoint.touch() + env_calls: list[int] = [] + monkeypatch.setattr( + interactive_playback, + "_resolve_appo_checkpoint_from_cfg", + lambda cfg, *, root_dir: (str(checkpoint), str(tmp_path)), + ) + + def reject_checkpoint(source_run_dir, target_cfg, *, algo_name, strict): + raise RuntimeError("incompatible checkpoint") + + monkeypatch.setattr(interactive_playback, "resolve_sim2sim_config", reject_checkpoint) + + def env_factory(num_envs: int) -> SimpleNamespace: + env_calls.append(num_envs) + return _rsl_rl_session_test_env() + + with pytest.raises(RuntimeError, match="incompatible checkpoint"): + create_appo_playback_session( + playback_cfg=PlaybackConfig( + task="Task", + load_run="run", + checkpoint=None, + action_mode="policy", + policy_obs_mode="actor", + algo_log_name="appo", + log_root=None, + ), + cfg=SimpleNamespace(training=SimpleNamespace(sim2sim_strict=True)), + rl_cfg={}, + env_factory=env_factory, + root_dir=tmp_path, + device="cpu", + wrapper_cls=_RslRlTestWrapper, + ) + + assert env_calls == [] + + def test_sac_playback_session_runs_sim2sim_preflight( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, @@ -473,6 +547,61 @@ def fake_preflight(source_run_dir, target_cfg, *, algo_name, strict): assert preflight_calls == [(str(tmp_path), "sac")] +@pytest.mark.parametrize("algo_name", ["sac", "flashsac"]) +def test_offpolicy_sim2sim_failure_prevents_env_construction( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + algo_name: str, +) -> None: + from omegaconf import OmegaConf + + import unilab.utils.checkpoint as checkpoint_utils + import unilab.visualization.interactive_playback as interactive_playback + + checkpoint = tmp_path / "model_10.pt" + checkpoint.touch() + resolved_checkpoint = (str(checkpoint), str(tmp_path)) + env_calls: list[int] = [] + monkeypatch.setattr( + checkpoint_utils, + "resolve_checkpoint_path", + lambda task_log_root, load_run, *, checkpoint=None, suffix=".pt": resolved_checkpoint, + ) + + def reject_checkpoint(source_run_dir, target_cfg, *, algo_name, strict): + raise RuntimeError("incompatible checkpoint") + + monkeypatch.setattr(interactive_playback, "resolve_sim2sim_config", reject_checkpoint) + + def env_factory(num_envs: int) -> None: + env_calls.append(num_envs) + + with pytest.raises(RuntimeError, match="incompatible checkpoint"): + create_sac_playback_session( + playback_cfg=PlaybackConfig( + task="Task", + load_run="run", + checkpoint=None, + action_mode="policy", + policy_obs_mode="actor", + algo_log_name=algo_name, + log_root=None, + ), + cfg=OmegaConf.create( + { + "training": {"sim2sim_strict": True}, + "algo": {"algo_log_name": algo_name}, + } + ), + env_factory=env_factory, + root_dir=tmp_path, + device="cpu", + algo_name=algo_name, + ) + + assert env_calls == [] + + def test_offpolicy_playback_keeps_policy_io_on_env_device() -> None: env = _StrictTensorPlaybackEnv() actor = _ConstantActor(0.5)