"""End-to-end smoke tests for the ``src/`` diffusion-planner pipeline. Proves the model builds, runs, trains, round-trips through disk, and that every ``main.py`` mode starts and finishes. Nothing here asserts anything about result quality. """ from __future__ import annotations import importlib import pytest import torch from tests.conftest import ( TINY_ENV, assert_cli_ok, discover_modules, requires_cuda, requires_minihack, run_cli, ) SRC_MODULES = discover_modules("src") # ── 1. Imports ─────────────────────────────────────────────────────── def test_src_module_list_is_not_empty(): assert len(SRC_MODULES) > 10, SRC_MODULES @pytest.mark.parametrize("module_name", SRC_MODULES) def test_src_module_imports_cleanly(module_name): importlib.import_module(module_name) # ── 2. Instantiation from the real config ──────────────────────────── def test_model_instantiates_from_real_config(real_cfg): from src.models.denoiser import LocalDiffusionPlannerWithGlobal, make_model model = make_model(real_cfg) assert isinstance(model, LocalDiffusionPlannerWithGlobal) assert sum(p.numel() for p in model.parameters()) > 0 assert model.head.out_features == real_cfg.action_dim def test_local_only_ablation_instantiates_from_real_config(real_cfg): import copy from src.models.denoiser import LocalDiffusionPlanner, make_model cfg = copy.copy(real_cfg) cfg.use_global_stream = False assert isinstance(make_model(cfg), LocalDiffusionPlanner) def test_ema_wraps_real_config_model(real_cfg): from src.models.denoiser import ModelEMA, make_model model = make_model(real_cfg) ema = ModelEMA(model, decay=real_cfg.ema_decay) ema.update(model) assert set(ema.state_dict()) == {n for n, _ in model.named_parameters()} # ── 3. Forward pass ────────────────────────────────────────────────── def test_forward_pass_shape_dtype_and_finiteness(tiny_cfg, tiny_batch): from src.models.denoiser import make_model local, glob, actions = tiny_batch model = make_model(tiny_cfg).eval() with torch.no_grad(): out = model(local, glob, actions, torch.zeros(local.shape[0], dtype=torch.long)) assert set(out) == {"actions", "goal_pred"} assert out["actions"].shape == ( local.shape[0], tiny_cfg.seq_len, tiny_cfg.action_dim, ) assert out["goal_pred"].shape == (local.shape[0], 2) assert out["actions"].dtype is torch.float32 assert out["goal_pred"].dtype is torch.float32 assert torch.isfinite(out["actions"]).all() assert torch.isfinite(out["goal_pred"]).all() def test_forward_pass_accepts_scalar_timestep(tiny_cfg, tiny_batch): from src.models.denoiser import make_model local, glob, actions = tiny_batch model = make_model(tiny_cfg).eval() with torch.no_grad(): out = model(local, glob, actions, 0) assert torch.isfinite(out["actions"]).all() def test_local_only_forward_pass(tiny_cfg, tiny_batch): import copy from src.models.denoiser import make_model cfg = copy.copy(tiny_cfg) cfg.use_global_stream = False local, glob, actions = tiny_batch model = make_model(cfg).eval() with torch.no_grad(): out = model(local, glob, actions, torch.zeros(local.shape[0], dtype=torch.long)) assert out["actions"].shape == (local.shape[0], cfg.seq_len, cfg.action_dim) assert torch.isfinite(out["actions"]).all() def test_forward_masking_and_loss_are_finite(tiny_cfg, tiny_batch): from src.diffusion.forward import q_sample from src.diffusion.loss import mdlm_loss from src.diffusion.schedules import get_schedule from src.models.denoiser import make_model local, glob, actions = tiny_batch schedule_fn = get_schedule(tiny_cfg.noise_schedule) t = torch.rand(actions.shape[0]).clamp(1e-5, 1 - 1e-5) zt = q_sample(actions, t, tiny_cfg.mask_token, tiny_cfg.pad_token, schedule_fn) assert zt.shape == actions.shape assert zt.dtype is torch.int64 model = make_model(tiny_cfg).eval() t_discrete = (t * tiny_cfg.num_diffusion_steps).long().clamp( 0, tiny_cfg.num_diffusion_steps - 1 ) with torch.no_grad(): out = model(local, glob, zt, t_discrete) loss = mdlm_loss( out["actions"], actions, zt, t, tiny_cfg.mask_token, tiny_cfg.pad_token, schedule_fn, ) assert loss.ndim == 0 assert torch.isfinite(loss) @pytest.mark.parametrize("strategy", ["rescale", "cap", "conf"]) def test_remdm_sampler_runs(tiny_cfg, tiny_batch, strategy): import copy from src.diffusion.sampling import remdm_sample from src.models.denoiser import make_model cfg = copy.copy(tiny_cfg) cfg.remask_strategy = strategy local, glob, _ = tiny_batch model = make_model(cfg).eval() seq = remdm_sample(model, local, glob, cfg, "cpu", physics_aware=False) assert seq.shape == (local.shape[0], cfg.seq_len) assert seq.dtype is torch.int64 assert (seq != cfg.mask_token).all() assert (seq >= 0).all() and (seq < cfg.action_dim).all() def test_greedy_sampler_runs(tiny_cfg, tiny_batch): from src.diffusion.sampling import greedy_sample from src.models.denoiser import make_model local, glob, _ = tiny_batch model = make_model(tiny_cfg).eval() seq = greedy_sample(model, local, glob, tiny_cfg, "cpu") assert seq.shape == (local.shape[0], tiny_cfg.seq_len) assert (seq >= 0).all() and (seq < tiny_cfg.action_dim).all() # ── Environment sanity ─────────────────────────────────────────────── @requires_minihack def test_environment_exposes_a_goal_staircase(real_cfg): """Regression: MiniHack's nhdat patch step can fail silently. When it does, every env falls back to the same default level with no staircase, so the BFS oracle has nothing to target and no episode can ever be won. """ from src.envs.minihack_env import make_env env = make_env(TINY_ENV, None, real_cfg) try: env.reset(seed=0) raw = env.last_raw_obs assert (raw["chars"] == ord(">")).sum() >= 1, "no staircase on the map" assert env._get_bfs_distance(raw) is not None finally: env.close() @requires_minihack def test_distinct_envs_produce_distinct_levels(real_cfg): """A silent nhdat failure collapses every env onto one default level.""" from src.envs.minihack_env import make_env sizes = [] for env_id in ("MiniHack-Room-Random-5x5-v0", "MiniHack-Room-Random-15x15-v0"): env = make_env(env_id, None, real_cfg) try: env.reset(seed=0) sizes.append(int((env.last_raw_obs["chars"] != ord(" ")).sum())) finally: env.close() assert sizes[0] != sizes[1], f"both envs rendered {sizes[0]} cells" # ── 4. One training step ───────────────────────────────────────────── def _make_trainer(cfg, trajectory): """Build a real Trainer with the collector/evaluator/logger left out.""" from src.buffer import ReplayBuffer from src.models.denoiser import ModelEMA, make_model from src.planners.online import Trainer buffer = ReplayBuffer(cfg.buffer_capacity, cfg.seq_len, cfg.pad_token) buffer.add(trajectory) model = make_model(cfg) ema = ModelEMA(model, decay=cfg.ema_decay) optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.dagger_lr) trainer = Trainer( model, ema, optimizer, None, buffer, collector=None, evaluator=None, log=None, cfg=cfg, device="cpu", raw_model=model, ) return trainer, model, ema def test_single_training_step_produces_finite_loss(tiny_cfg, tiny_trajectory): trainer, model, ema = _make_trainer(tiny_cfg, tiny_trajectory) model.train() metrics = trainer._train_step() ema.update(model) for key in ("loss", "loss_diff", "loss_aux", "grad_norm"): assert key in metrics assert torch.isfinite(torch.tensor(metrics[key])), (key, metrics[key]) def test_training_step_updates_parameters(tiny_cfg, tiny_trajectory): trainer, model, _ = _make_trainer(tiny_cfg, tiny_trajectory) before = model.head.weight.detach().clone() model.train() trainer._train_step() assert not torch.equal(before, model.head.weight.detach()) def test_training_step_on_empty_buffer_is_a_no_op(tiny_cfg): from src.buffer import ReplayBuffer from src.models.denoiser import ModelEMA, make_model from src.planners.online import Trainer model = make_model(tiny_cfg) trainer = Trainer( model, ModelEMA(model, decay=tiny_cfg.ema_decay), torch.optim.AdamW(model.parameters(), lr=tiny_cfg.dagger_lr), None, ReplayBuffer(tiny_cfg.buffer_capacity, tiny_cfg.seq_len, tiny_cfg.pad_token), collector=None, evaluator=None, log=None, cfg=tiny_cfg, device="cpu", raw_model=model, ) assert trainer._train_step() == { "loss": 0.0, "loss_diff": 0.0, "loss_aux": 0.0, "grad_norm": 0.0, } @requires_cuda def test_amp_training_step_on_cuda(tiny_cfg, tiny_trajectory): import copy cfg = copy.copy(tiny_cfg) cfg.device = "cuda" cfg.use_amp = True from src.buffer import ReplayBuffer from src.models.denoiser import ModelEMA, make_model from src.planners.online import Trainer buffer = ReplayBuffer(cfg.buffer_capacity, cfg.seq_len, cfg.pad_token) buffer.add(tiny_trajectory) model = make_model(cfg).to("cuda") trainer = Trainer( model, ModelEMA(model, decay=cfg.ema_decay), torch.optim.AdamW(model.parameters(), lr=cfg.dagger_lr), None, buffer, collector=None, evaluator=None, log=None, cfg=cfg, device="cuda", raw_model=model, ) model.train() assert torch.isfinite(torch.tensor(trainer._train_step()["loss"])) # ── 5. Save and reload ─────────────────────────────────────────────── def test_checkpoint_roundtrip_preserves_output(tiny_cfg, tiny_batch, tmp_path): from src.models.denoiser import ModelEMA, make_model local, glob, actions = tiny_batch t = torch.zeros(local.shape[0], dtype=torch.long) model = make_model(tiny_cfg).eval() ema = ModelEMA(model, decay=tiny_cfg.ema_decay) with torch.no_grad(): before = model(local, glob, actions, t) path = tmp_path / "checkpoint.pth" torch.save( { "model_state_dict": model.state_dict(), "ema_state_dict": ema.state_dict(), }, path, ) reloaded = make_model(tiny_cfg) ckpt = torch.load(path, map_location="cpu", weights_only=False) reloaded.load_state_dict(ckpt["model_state_dict"]) reloaded.eval() with torch.no_grad(): after = reloaded(local, glob, actions, t) assert torch.equal(before["actions"], after["actions"]) assert torch.equal(before["goal_pred"], after["goal_pred"]) def test_ema_weights_roundtrip(tiny_cfg, tiny_batch, tmp_path): from src.models.denoiser import ModelEMA, make_model local, glob, actions = tiny_batch t = torch.zeros(local.shape[0], dtype=torch.long) model = make_model(tiny_cfg) ema = ModelEMA(model, decay=0.5) ema.update(model) path = tmp_path / "ema.pth" torch.save({"ema_state_dict": ema.state_dict()}, path) eval_model = ema.make_eval_model(model) with torch.no_grad(): before = eval_model(local, glob, actions, t)["actions"] reloaded = make_model(tiny_cfg) reloaded_ema = ModelEMA(reloaded, decay=0.5) reloaded_ema.load_state_dict( torch.load(path, map_location="cpu", weights_only=False)["ema_state_dict"] ) after_model = reloaded_ema.make_eval_model(reloaded) with torch.no_grad(): after = after_model(local, glob, actions, t)["actions"] assert torch.equal(before, after) # ── 6. Entry points ────────────────────────────────────────────────── def test_main_help(): result = run_cli("main.py", "--help") assert_cli_ok(result) assert "--mode" in result.stdout def test_main_rejects_unknown_mode(): assert run_cli("main.py", "--mode", "not-a-mode").returncode != 0 def test_main_inference_requires_a_checkpoint(tiny_config_file): result = run_cli( "main.py", "--mode", "inference", "--config", str(tiny_config_file) ) assert result.returncode != 0 assert "checkpoint" in (result.stdout + result.stderr).lower() @requires_minihack def test_main_smoke_mode_runs(tiny_config_file): result = run_cli("main.py", "--mode", "smoke", "--config", str(tiny_config_file)) assert_cli_ok(result) assert "Smoke Results" in result.stdout @requires_minihack def test_smoke_mode_removes_its_temporary_artefact_directory(tiny_config_file): """A smoke run leaves no `remdm-smoke-*` directory behind (PARITY "Smoke-mode side effects"). `run_smoke` points `cfg.checkpoint_dir` at a fresh `mkdtemp` so checkpoints, config snapshots and eval JSONs stay out of the repository tree -- and it never removed it. Every smoke run since leaked one: 183 of them, 9.3 GB, had accumulated by 2026-08-19, enough to exhaust the 100 GB project quota and fail this suite on `OSError: [Errno 122] Disk quota exceeded`, and 42 more reappeared within a day of the first clear-out. craftax's smoke path has always removed its temporary expert directory in a `finally`. Counting the directories rather than trusting the run to report them is what makes this catch a leak from any path inside the run. """ import glob import tempfile from pathlib import Path pattern = str(Path(tempfile.gettempdir()) / "remdm-smoke-*") def _count() -> int: return len([p for p in glob.glob(pattern) if Path(p).is_dir()]) before = _count() result = run_cli("main.py", "--mode", "smoke", "--config", str(tiny_config_file)) assert_cli_ok(result) assert _count() == before, ( "smoke mode leaked a temporary artefact directory; " f"{_count() - before} left behind under {pattern}" ) @requires_minihack def test_main_collect_mode_runs(tiny_config_file, tmp_path): result = run_cli("main.py", "--mode", "collect", "--config", str(tiny_config_file)) assert_cli_ok(result) assert (tmp_path / "dataset.pt").exists() def test_main_offline_mode_runs(tiny_config_file, tiny_dataset_file, tmp_path): result = run_cli( "main.py", "--mode", "offline", "--config", str(tiny_config_file), "--data", str(tiny_dataset_file), "--override", "total_timesteps=8", ) assert_cli_ok(result) assert list((tmp_path / "checkpoints").glob("offline_*/offline_final.pth")) @requires_minihack def test_main_inference_mode_runs(tiny_config_file, tiny_checkpoint_file, tmp_path): output = tmp_path / "eval.json" result = run_cli( "main.py", "--mode", "inference", "--config", str(tiny_config_file), "--checkpoint", str(tiny_checkpoint_file), "--envs", TINY_ENV, "--episodes", "1", "--output", str(output), ) assert_cli_ok(result) assert output.exists() @requires_minihack def test_main_online_mode_runs(tiny_config_file): result = run_cli( "main.py", "--mode", "online", "--config", str(tiny_config_file), "--no-warm-start", ) assert_cli_ok(result) def test_main_baselines_mode_validates_algo(tiny_config_file): result = run_cli( "main.py", "--mode", "baselines", "--algo", "not-an-algo", "--config", str(tiny_config_file), ) assert result.returncode != 0 @requires_minihack @pytest.mark.slow def test_main_baselines_bc_runs(tiny_config_file): result = run_cli( "main.py", "--mode", "baselines", "--algo", "bc", "--seeds", "0", "--config", str(tiny_config_file), "--override", "baselines_bc_oracle_episodes_per_env=1", "--override", "baselines_bc_epochs=1", "--override", "baselines_bc_batch_size=8", "--override", "baselines_n_envs_per_id=1", "--override", "baselines_eval_episodes_per_env=1", "--override", "baselines_eval_freq_env_steps=1000000", ) assert_cli_ok(result) @requires_minihack @pytest.mark.slow def test_main_baselines_ppo_runs_without_wandb(tiny_config_file): """Regression: SB3 baselines used to die before their first env step. ``WandbCallback`` was built unconditionally, and behind it SB3 refused to start because ``tensorboard_log`` was set without tensorboard installed. """ result = run_cli( "main.py", "--mode", "baselines", "--algo", "ppo", "--seeds", "0", "--config", str(tiny_config_file), "--override", "total_timesteps=64", "--override", "baselines_n_envs_per_id=1", "--override", "baselines_eval_episodes_per_env=1", "--override", "baselines_eval_freq_env_steps=1000000", ) assert_cli_ok(result) def test_offline_builds_one_eval_model_per_eval_point(tiny_cfg, tiny_trajectory): """PERF-O2: the ID and OOD eval blocks share one EMA copy. Both cadences derive from ``offline_eval_every_grad_steps`` (``offline.py:157-161``), so they fire on the same step against the same ``_ema_source``. Each block used to build its own ``copy.deepcopy(model)`` plus EMA apply. """ import copy as _copy import torch from src.buffer import ReplayBuffer from src.models.denoiser import ModelEMA, make_model from src.planners.offline import make_offline_trainer cfg = _copy.deepcopy(tiny_cfg) cfg.offline_total_grad_steps = 3 cfg.offline_eval_every_grad_steps = 1 cfg.offline_checkpoint_every_grad_steps = 10**9 cfg.checkpoint_every_timesteps = 10**9 cfg.offline_log_every = 10**9 cfg.use_amp = False cfg.torch_compile = False buffer = ReplayBuffer(1000, cfg.seq_len, cfg.pad_token) buffer.load_offline_data( {"trajectories": [tiny_trajectory]}, [tiny_trajectory["env_id"]] ) torch.manual_seed(0) model = make_model(cfg) ema = ModelEMA(model, decay=0.9) built: list[int] = [] real_make = ModelEMA.make_eval_model def counting_make(self, src): built.append(1) return real_make(self, src) seen: list[tuple[str, int]] = [] class _StubEvaluator: def evaluate(self, env_ids, eval_model, n_episodes, cfg_, device): seen.append((env_ids[0], id(eval_model))) return {e: {"win_rate": 0.0, "avg_reward": 0.0} for e in env_ids} ModelEMA.make_eval_model = counting_make try: make_offline_trainer(cfg)( model=model, ema_model=ema, buffer=buffer, cfg=cfg, device=torch.device("cpu"), evaluator=_StubEvaluator(), id_envs=["ID_ENV"], ood_envs=["OOD_ENV"], ) finally: ModelEMA.make_eval_model = real_make id_calls = [s for s in seen if s[0] == "ID_ENV"] ood_calls = [s for s in seen if s[0] == "OOD_ENV"] assert id_calls, "the ID eval never fired, so the test proves nothing" assert len(id_calls) == len(ood_calls), "both blocks should fire together" # One EMA copy per eval point, not two. assert len(built) == len(id_calls), ( f"{len(built)} eval models built for {len(id_calls)} eval points" ) # And both blocks were handed the very same model object. for (_, id_obj), (_, ood_obj) in zip(id_calls, ood_calls, strict=True): assert id_obj == ood_obj def test_building_either_model_warns_about_nothing(tiny_cfg): """No warning may originate in this repo's own source (sweep S11-7). `LocalDiffusionPlanner` left `nn.TransformerEncoder` at its default `enable_nested_tensor=True` while its encoder layer sets `norm_first=True`, which makes the nested-tensor path unavailable, so PyTorch warned on every construction. It was the only repo-origin warning in either suite. Filtering it would have hidden the next one; the keyword is passed instead, as the dual-stream sibling always has. """ import warnings from pathlib import Path from src.models.denoiser import ( LocalDiffusionPlanner, LocalDiffusionPlannerWithGlobal, ) root = str(Path(__file__).resolve().parents[1] / "src") with warnings.catch_warnings(record=True) as caught: warnings.simplefilter("always") LocalDiffusionPlannerWithGlobal(tiny_cfg) LocalDiffusionPlanner(tiny_cfg) ours = [w for w in caught if str(w.filename).startswith(root)] assert not ours, [f"{w.filename}:{w.lineno} {w.message}" for w in ours]